aegis_sim.submodels.genetics.composite.architecture

  1import numpy as np
  2from aegis_sim import constants
  3from aegis_sim import variables
  4
  5from aegis_sim.submodels.genetics.composite.interpreter import Interpreter
  6from aegis_sim import parameterization
  7from aegis_sim.parameterization import parametermanager
  8from aegis_sim.submodels.genetics import ploider
  9
 10
 11class CompositeArchitecture:
 12    """
 13
 14    GUI
 15    - when pleiotropy is not needed;
 16    - it is quick, easy to analyze, delivers a diversity of phenotypes
 17    - every trait (surv repr muta neut) can be evolvable or not
 18    - if not evolvable, the value is set by !!!
 19    - if evolvable, it can be agespecific or age-independent
 20    - probability of a trait at each age is determined by a BITS_PER_LOCUS adjacent bits forming a "locus" / gene
 21    - the method by which these loci are converted into a phenotypic value is the Interpreter type
 22
 23    """
 24
 25    def __init__(self, BITS_PER_LOCUS, AGE_LIMIT, THRESHOLD):
 26        self.BITS_PER_LOCUS = BITS_PER_LOCUS
 27        self.n_loci = sum(trait.length for trait in parameterization.traits.values())
 28        self.length = self.n_loci * BITS_PER_LOCUS
 29        self.AGE_LIMIT = AGE_LIMIT
 30
 31        self.evolvable = [trait for trait in parameterization.traits.values() if trait.evolvable]
 32
 33        self.interpreter = Interpreter(
 34            self.BITS_PER_LOCUS,
 35            THRESHOLD,
 36        )
 37
 38        # Fixed seed=0 so all populations (including hybridizing ones with different RANDOM_SEEDs)
 39        # share an identical physical genome layout; locus_permutation[i] = physical position of logical locus i
 40        self.locus_permutation = np.random.default_rng(0).permutation(self.n_loci)
 41
 42        # Per-locus dominance coefficient h, indexed by *physical* locus position
 43        # (because diploid_to_haploid operates on the physical layout, before reorder).
 44        # Built from each trait's G_<trait>_dominance value (default 0.5 = codominant).
 45        self.dominance_per_locus = np.full(self.n_loci, 0.5, dtype=np.float32)
 46        for trait in parameterization.traits.values():
 47            if trait.length == 0:
 48                continue
 49            phys_pos = self.locus_permutation[trait.start:trait.end]
 50            self.dominance_per_locus[phys_pos] = np.float32(trait.dominance)
 51
 52        # Optional pleiotropy. None unless PHENOMAP_SPECS is given.
 53        self.phenomap_matrix = self._build_phenomap(parametermanager.parameters.PHENOMAP_SPECS)
 54
 55    def _build_phenomap(self, PHENOMAP_SPECS):
 56        """Build a genotype-phenotype matrix from PHENOMAP_SPECS (aegis v1 semantics).
 57
 58        The matrix is the identity plus off-diagonal weights: the diagonal means every
 59        locus keeps encoding its own trait at its own age (so age-specific survival is
 60        preserved), and each spec adds a *pleiotropic* effect of one locus on another
 61        trait/age on top. This is what makes antagonistic pleiotropy expressible --
 62        one locus raising surv at an early age and lowering it at a late one.
 63
 64        Each spec is ``[source_trait, source_index, target_trait, target_age, weight]``
 65        with ``source_index`` (1..trait length) and ``target_age`` (1..AGE_LIMIT) 1-based,
 66        as exported by v1. Indices are resolved in *logical* (trait x age) order, which is
 67        the order compute() reorders genomes into.
 68
 69        Returns None when no specs are given, in which case compute() is unchanged.
 70        """
 71        if not PHENOMAP_SPECS:
 72            return None
 73
 74        map_ = np.diag(np.ones(self.n_loci, dtype=np.float32))
 75        for spec in PHENOMAP_SPECS:
 76            source_trait, source_index, target_trait, target_age, weight = spec
 77            source = parameterization.traits[source_trait]
 78            target = parameterization.traits[target_trait]
 79            geno_i = source.start + (int(source_index) - 1)
 80            pheno_i = target.start + (int(target_age) - 1)
 81            assert source.start <= geno_i < source.end, (
 82                f"PHENOMAP_SPECS source index {source_index} is out of range for trait "
 83                f"'{source_trait}' which has {source.length} loci; "
 84                f"set G_{source_trait}_agespecific to at least {source_index}"
 85            )
 86            assert target.start <= pheno_i < target.end, (
 87                f"PHENOMAP_SPECS target age {target_age} is out of range for trait "
 88                f"'{target_trait}' which has {target.length} loci"
 89            )
 90            map_[geno_i, pheno_i] = np.float32(weight)
 91        return map_
 92
 93    def get_number_of_bits(self):
 94        return ploider.ploider.y * self.n_loci * self.BITS_PER_LOCUS
 95
 96    def get_shape(self):
 97        return (ploider.ploider.y, self.n_loci, self.BITS_PER_LOCUS)
 98
 99    def init_genome_array(self, popsize):
100        # TODO enable agespecific False
101        array = variables.rng.random(size=(popsize, *self.get_shape()))
102
103        for trait in parameterization.traits.values():
104            phys_pos = self.locus_permutation[trait.start:trait.end]
105            array[:, :, phys_pos, :] = array[:, :, phys_pos, :] < trait.initgeno
106
107        return array
108
109    def compute(self, genomes):
110
111        if genomes.shape[1] == 1:  # Do not calculate mean if genomes are haploid
112            genomes = genomes[:, 0]
113        else:
114            genomes = ploider.ploider.diploid_to_haploid(genomes, dominance_per_locus=self.dominance_per_locus)
115
116        # Reorder from physical storage order to logical (trait×age) order
117        genomes = genomes[:, self.locus_permutation, :]
118
119        interpretome = np.zeros(shape=(genomes.shape[0], genomes.shape[1]), dtype=np.float32)
120        for trait in parameterization.traits.values():
121            loci = genomes[:, trait.slice]  # fetch
122            probs = self.interpreter.call(loci, trait.interpreter)  # interpret
123            interpretome[:, trait.slice] += probs  # add back
124
125        # Apply pleiotropy, if any, to the raw [0, 1] interpreter output -- before the
126        # lo/hi mapping below, so that spec weights are expressed in interpreter units
127        # (aegis v1 order: interpret -> phenomap -> lo/hi).
128        if self.phenomap_matrix is not None:
129            interpretome = np.clip(interpretome.dot(self.phenomap_matrix), 0, 1).astype(np.float32)
130
131        # Map the [0, 1] interpreter output onto the trait's [lo, hi] phenotypic range.
132        # This is what makes G_<trait>_lo / G_<trait>_hi do anything — without this scaling
133        # the lo/hi parameters are parsed but discarded. Defaults: G_surv_lo=0.7, G_surv_hi=1.0
134        # (surv never goes to 0); G_repr_lo=0, G_repr_hi=0.5; others lo=0, hi=1.
135        for trait in parameterization.traits.values():
136            interpretome[:, trait.slice] = trait.lo + (trait.hi - trait.lo) * interpretome[:, trait.slice]
137
138        return interpretome
139
140    # def diffuse(self, probs):
141    #     window_size = parametermanager.parameters.DIFFUSION_FACTOR * 2 + 1
142    #     p = np.empty(shape=(probs.shape[0], probs.shape[1] + window_size - 1))
143    #     p[:, :window_size] = np.repeat(probs[:, 0], window_size).reshape(-1, window_size)
144    #     p[:, window_size - 1 :] = probs[:]
145    #     diffusome = np.convolve(p[0], np.ones(window_size) / window_size, mode="valid")
146
147    def get_map(self):
148        pass