base_constructors.py

# -*- coding: utf-8 -*-
""" Creating (base) mixtures of categorical distributions. """

import numpy as np
import tensorflow as tf
from itertools import product
import re

import logging
logger = logging.getLogger(__name__)

from base_categorical import MultivariateFactorizedCategorical
from base_mixtures import FactorizedCategoricalMixture, FactorizedIndependentCategoricalMixture


# To avoid NaNs due, e.g., log(0)
EPS_PROBABILITY = 1e-15


def yield_probabilities_delta_diagonal(N, K, eps_prob=EPS_PROBABILITY):
    while True:
        for k in range(K):        
            probs = np.ones( (N, K) ) * eps_prob
            probs[:, k] = 1.
            probs /= np.sum(probs, -1, keepdims=True)
            yield probs        


def yield_probabilities_delta_random(N, K, eps_prob=EPS_PROBABILITY):
    while True:
        probs = np.ones( (N,K) ) * eps_prob
        ks = np.random.choice(range(K), N)
        probs[list(range(N)), ks] = 1.
        probs /= np.sum(probs, -1, keepdims=True)
        yield probs      


def yield_probabilities_delta(N, K, eps_prob=EPS_PROBABILITY):
    if K**N>1000000:
        logger.warn("""[yield_probabilities_delta] WARNING: 
                       There is too many combinations of variables and categories! 
                       Generating diagonal and then random instead of all possible!""")
        diagonal = yield_probabilities_delta_diagonal(N, K, eps_prob)
        for k in range(K):        
            yield next(diagonal)

        random = yield_probabilities_delta_random(N, K, eps_prob)
        while True:
            yield next(random)    

    while True:
        for probability_concentrantion in np.array(list(product(*[range(i) for i in [K]*N]))):
            probs = tf.one_hot(probability_concentrantion, K).numpy()
            probs += eps_prob
            probs /= np.sum(probs, -1, keepdims=True)
            yield probs


def yield_probabilities_uniform(N, K):
    while True:
        yield np.ones( (N, K) )*1/K    


def yield_probabilities_random(N, K):
    while True:
        b = np.random.random( (N,K) )
        b /= np.sum(b, -1, keepdims=True) 
        yield b


def sample_probabilities_from_dirichlet(N, K, alpha, eps_prob=EPS_PROBABILITY):
    logger.debug("[sample_probabilities_from_dirichlet] N=%s K=%s alpha=%s" % (N, K, alpha) )
    probs = np.random.dirichlet(np.ones(K)*alpha, N) + eps_prob
    probs /= np.sum(probs, -1, keepdims=True)
    return probs  


def create_base_mixture(N, K, B=None, specification="p", independent_variables=False,
                              mixing_logits=None, dtype=tf.keras.backend.floatx(), 
                              eps_prob=EPS_PROBABILITY):
    """ Returns a factorized categorical mixture distribution constructed according to specification.

        Args:
            N  number of dimensions (or variables).
            K  cardinality
            B  max number of components in a mixture (default K^N)
            specification:  Selects (in a loop) consecutive probability components: 
                            u-uniform, r-random, 
                            p-delta, d-delta on diagonal, 1-delta in a random position. 
                            For example `pppuur` means three deltas (in concecutive positons),
                            two uniform distributions and one random distribution
                            (then repeat if necessary).
                            Positions of peaks are generated by enumerating 
                            all variables and all categories.
            independent_variables:  Assume independence between dimensions (variables), i.e.,
                                    return FactorizedIndependentCategoricalMixture instead of 
                                    FactorizedCategoricalMixture.
    """
    if B is None or B<=0: B = K**N
    if mixing_logits is None:
        mixing_logits = tf.math.log([1./B]*B)
    assert len(mixing_logits.shape)==1
    assert B==mixing_logits.shape[0]

    mixing_distribution = MultivariateFactorizedCategorical(logits=mixing_logits, dtype=dtype) 

    uniform = yield_probabilities_uniform(N, K)
    random = yield_probabilities_random(N, K)

    peak = yield_probabilities_delta(N, K, eps_prob=eps_prob)
    diagonal = yield_probabilities_delta_diagonal(N, K, eps_prob=eps_prob)
    peak_random = yield_probabilities_delta_random(N, K, eps_prob=eps_prob)


    components = []
    specification = specification*B
    specification = re.findall("\w\([^\)]*\)|\w", specification)[:B]
    logger.debug("[create_base_mixture] specification=%s" % (specification) )
    while len(components)<B:
        b = len(components)
        component_specs = specification[b]
        component_type = component_specs[0]

        if component_type=='u': 
            component = MultivariateFactorizedCategorical( probs=next(uniform) )
        elif component_type=='r':
            component = MultivariateFactorizedCategorical( probs=next(random) )
        elif component_type=='p':
            component = MultivariateFactorizedCategorical( probs=next(peak) )
        elif component_type=='d':
            component = MultivariateFactorizedCategorical( probs=next(diagonal) )
        elif component_type=='1':
            component = MultivariateFactorizedCategorical( probs=next(peak_random) )
        elif component_type=='D':
            alpha = float(re.findall("\(.*\)", component_specs)[0][1:-1])
            component = MultivariateFactorizedCategorical( probs=sample_probabilities_from_dirichlet(N, K, alpha, eps_prob=eps_prob) )
        else:
          raise ValueError("Unknown probability block type! Given: %s. Supported u/r/p." % component_type)
        components.append(component)
    
    if independent_variables:
        return FactorizedIndependentCategoricalMixture(mixing_distribution=mixing_distribution, components=components)
    else:
        return FactorizedCategoricalMixture(mixing_distribution=mixing_distribution, components=components)


create_categorical_peaks = create_base_mixture
create_categorical_blocks = create_base_mixture