| """DGP with discrete heterogeneity aligned to coarse simplex groups.""" |
| import numpy as np |
|
|
| from .base import BaseDGP, DGPSample |
| from .pure_scale import PureScaleDGP |
| from ..utils.simplex import aitchison_dist, ilr, ilr_inv |
|
|
|
|
| class DiscreteGroupsDGP(PureScaleDGP): |
| """Step-function scale heterogeneity based on the predicted top class.""" |
|
|
| def __init__( |
| self, |
| K: int = 10, |
| sigma_low: float = 0.08, |
| sigma_high: float = 0.30, |
| d_x: int = 5, |
| easy_classes: int = 5, |
| ): |
| super().__init__(K=K, sigma_min=sigma_low, c=sigma_high - sigma_low, d_x=d_x) |
| self.sigma_low = sigma_low |
| self.sigma_high = sigma_high |
| self.easy_classes = easy_classes |
|
|
| def _sigma(self, u: np.ndarray) -> np.ndarray: |
| top_class = np.argmax(u, axis=1) |
| is_easy = top_class < self.easy_classes |
| return np.where(is_easy, self.sigma_low, self.sigma_high) |
|
|
| def sample(self, n: int, rng: np.random.Generator) -> DGPSample: |
| self._init_weights(rng) |
| X = rng.standard_normal((n, self.d_x)) |
| mu = self._mu(X) |
| sigma = self._sigma(mu) |
|
|
| Z_mu = ilr(mu) |
| eps = rng.standard_normal((n, self.K - 1)) |
| Y = ilr_inv(Z_mu + sigma[:, None] * eps, K=self.K) |
| U = mu |
| R = aitchison_dist(Y, U) |
| return DGPSample(X=X, Y=Y, U=U, R=R, sigma_true=sigma) |
|
|
