<reponame>bilgelm/NiMARE<gh_stars>0
"""
Topic modeling with generalized correspondence latent Dirichlet allocation.
"""
import logging
import os.path as op

import numpy as np
import pandas as pd
import nibabel as nib
from scipy.stats import multivariate_normal

from ...base import AnnotationModel
from ...due import due, Doi
from ...utils import get_template

LGR = logging.getLogger(__name__)

@due.dcite(Doi('10.1371/journal.pcbi.1005649'),
           description='Introduces GC-LDA decoding.')
class GCLDAModel(AnnotationModel):
    """
    Generate a GCLDA topic model.

    Parameters
    ----------
    count_df : :obj:`pandas.DataFrame`
        A DataFrame with feature counts for the model. The index is 'id',
        used for identifying studies. Other columns are features (e.g.,
        unigrams and bigrams from Neurosynth), where each value is the number
        of times the feature is found in a given article.
    coordinates_df : :obj:`pandas.DataFrame`
        A DataFrame with a list of foci in the dataset. The index is 'id',
        used for identifying studies. Additional columns include 'i', 'j' and
        'k' (the matrix indices of the foci in standard space).
    n_topics : :obj:`int`, optional
        Number of topics to generate in model. The default is 100.
    n_regions : :obj:`int`, optional
        Number of subregions per topic (>=1). The default is 2.
    alpha : :obj:`float`, optional
        Prior count on topics for each document. The default is 0.1.
    beta : :obj:`float`, optional
        Prior count on word-types for each topic. The default is 0.01.
    gamma : :obj:`float`, optional
        Prior count added to y-counts when sampling z assignments. The
        default is 0.01.
    delta : :obj:`float`, optional
        Prior count on subregions for each topic. The default is 1.0.
    dobs : :obj:`int`, optional
        Spatial region 'default observations' (# observations weighting
        Sigma estimates in direction of default 'roi_size' value). The
        default is 25.
    roi_size : :obj:`float`, optional
        Default spatial 'region of interest' size (default value of
        diagonals in covariance matrix for spatial distribution, which the
        distributions are biased towards). The default is 50.0.
    symmetric : :obj:`bool`, optional
        Whether or not to use symmetry constraint on subregions. Symmetry
        requires n_regions = 2. The default is False.
    seed_init : :obj:`int`, optional
        Initial value of random seed. The default is 1.
    name : :obj:`str`, optional
        Name of model.
    """
    def __init__(self, count_df, coordinates_df, mask='Mni152_2mm',
                 n_topics=100, n_regions=2, symmetric=True, alpha=.1,
                 beta=.01, gamma=.01, delta=1.0, dobs=25, roi_size=50.0,
                 seed_init=1, name='gclda'):
        LGR.info('Constructing/Initializing GCLDA Model')

        # --- Checking to make sure parameters are valid
        if (symmetric is True) and (n_regions != 2):
            # symmetric model only valid if R = 2
            raise ValueError('Cannot run a symmetric model unless #Subregions '
                             '(n_regions) == 2 !')

        # Initialize sampling parameters
        self.iter = 0  # Tracks the global sampling iteration of the model
        self.seed = 0  # Tracks current random seed to use (gets incremented
                       # after initialization and each sampling update)

        # Set up model hyperparameters
        # Pseudo-count hyperparams need to be floats so that when sampling
        # distributions are computed the count matrices/vectors are converted
        # to floats
        self.params = {
            'n_topics': n_topics,  # Number of topics (T)
            'n_regions': n_regions,  # Number of subregions (R)
            'alpha': alpha,  # Prior count on topics for each doc
            'beta': beta,  # Prior count on word-types for each topic
            'gamma': gamma,  # Prior count added to y-counts when sampling z assignments
            'delta': delta,  # Prior count on subregions for each topic
            'roi_size': roi_size,  # Default ROI (default covariance spatial
                                   # region we regularize towards) (not in paper)
            'dobs': dobs,  # Sample constant (# observations weighting
                           # sigma in direction of default covariance)
                           # (not in paper)
            'symmetric': symmetric,  # Use constrained symmetry on subregions?
                                     # (only for n_regions = 2)
            'seed_init': seed_init,  # Random seed for initializing model
            }
        self.model_name = ('{0}_{1}T_{2}R_alpha{3:.3f}_beta{4:.3f}_'
                           'gamma{5:.3f}_delta{6:.3f}_{7}dobs_{8:.1f}roi_{9}'
                           'symmetric_{10}').format(
            name, self.params['n_topics'], self.params['n_regions'],
            self.params['alpha'], self.params['beta'],
            self.params['gamma'], self.params['delta'],
            self.params['dobs'], self.params['roi_size'],
            self.params['symmetric'], self.params['seed_init'])

        # Prepare data
        if isinstance(mask, str) and not op.isfile(mask):
            self.mask = get_template(mask, mask='brain')
        elif isinstance(mask, str) and op.isfile(mask):
            self.mask = nib.load(mask)
        elif isinstance(mask, nib.Nifti1Image):
            self.mask = mask
        else:
            raise Exception('Input "mask" could not be figured out.')

        # Import all word-labels into a list
        # List of word-strings (wtoken_word_idx values are indices into this list)
        self.vocabulary = count_df.columns.tolist()

        # Extract document and word indices from count_df
        count_df.index = count_df.index.astype(str)
        ids = count_df.index.tolist()
        docidx_mapper = {id_: i for (i, id_) in enumerate(ids)}
        self.ids = ids

        # Create docidx column
        count_df['id'] = count_df.index
        count_df['docidx'] = count_df['id'].map(docidx_mapper)
        count_df = count_df.dropna(subset=['docidx'])
        count_df = count_df.drop('id', 1)

        # Remove words not found anywhere in the corpus
        count_df = count_df.loc[:, (count_df != 0).any(axis=0)]

        # Get updated vocabulary
        word_labels = count_df.columns.tolist()
        word_labels.remove('docidx')
        self.word_labels = word_labels
        widx_mapper = {word: i for (i, word) in enumerate(self.word_labels)}

        # Melt dataframe and create widx column
        widx_df = pd.melt(count_df, id_vars=['docidx'], var_name='word',
                          value_name='count')
        widx_df['widx'] = widx_df['word'].map(widx_mapper)

        # Replicate rows based on count
        widx_df = widx_df.loc[np.repeat(widx_df.index.values, widx_df['count'])]
        widx_df = widx_df[['docidx', 'widx']].astype(int)
        widx_df.sort_values(by=['docidx', 'widx'], inplace=True)
        # List of document-indices for word-tokens
        self.wtoken_doc_idx = widx_df['docidx'].tolist()
        # List of word-indices for word-tokens
        self.wtoken_word_idx = widx_df['widx'].tolist()

        # Import all peak-indices into lists
        if 'id' not in coordinates_df.columns:
            coordinates_df['id'] = coordinates_df.index
        coordinates_df['docidx'] = coordinates_df['id'].astype(str).map(docidx_mapper)
        coordinates_df = coordinates_df.dropna(subset=['docidx'])
        coordinates_df = coordinates_df[['docidx', 'x', 'y', 'z']]
        coordinates_df['docidx'] = coordinates_df['docidx'].astype(int)
        # List of document-indices for peak-tokens x
        self.ptoken_doc_idx = coordinates_df['docidx'].tolist()
        self.peak_vals = coordinates_df[['x', 'y', 'z']].values

        # Seed random number generator
        np.random.seed(self.params['seed_init'])  # pylint: disable=no-member

        # Preallocate vectors of assignment indices
        self.wtoken_topic_idx = np.zeros(
            len(self.wtoken_word_idx), dtype=int)  # word->topic assignments

        # Randomly initialize peak->topic assignments (y) ~ unif(1...n_topics)
        self.peak_topic_idx = np.random.randint(
            self.params['n_topics'],  # pylint: disable=no-member
            size=(len(self.ptoken_doc_idx)))

        self.peak_region_idx = np.zeros(
            len(self.ptoken_doc_idx), dtype=int)  # peak->region assignments

        # Preallocate count matrices
        # Peaks: D x T: Number of peak-tokens assigned to each topic per document
        self.n_peak_tokens_doc_by_topic = np.zeros(
            (len(self.ids), self.params['n_topics']), dtype=int)

        # Peaks: R x T: Number of peak-tokens assigned to each subregion per topic
        self.n_peak_tokens_region_by_topic = np.zeros(
            (self.params['n_regions'], self.params['n_topics']), dtype=int)

        # Words: W x T: Number of word-tokens assigned to each topic per word-type
        self.n_word_tokens_word_by_topic = np.zeros(
            (len(self.word_labels), self.params['n_topics']), dtype=int)

        # Words: D x T: Number of word-tokens assigned to each topic per document
        self.n_word_tokens_doc_by_topic = np.zeros(
            (len(self.ids), self.params['n_topics']), dtype=int)

        # Words: 1 x T: Total number of word-tokens assigned to each topic (across all docs)
        self.total_n_word_tokens_by_topic = np.zeros(
            (1, self.params['n_topics']), dtype=int)

        # Preallocate Gaussians for all subregions
        # Regions_Mu & Regions_Sigma: Gaussian mean and covariance for all
        # subregions of all topics
        # Formed using lists (over topics) of lists (over subregions) of numpy
        # arrays
        #   regions_mu = (n_topics, n_regions, 1, n_peak_dims)
        #   regions_sigma = (n_topics, n_regions, n_peak_dims, n_peak_dims)
        self.regions_mu = []
        self.regions_sigma = []
        for i_topic in range(self.params['n_topics']):
            topic_mu = []
            topic_sigma = []
            for j_region in range(self.params['n_regions']):
                topic_mu.append(np.zeros((1, self.peak_vals.shape[1])))
                topic_sigma.append(np.zeros(
                    (self.peak_vals.shape[1], self.peak_vals.shape[1])))
            self.regions_mu.append(topic_mu)  # (\mu^{(t)}_r)
            self.regions_sigma.append(topic_sigma)  # (\sigma^{(t)}_r)

        # Initialize lists for tracking log-likelihood of data over sampling iterations
        self.loglikely_iter = []  # Tracks iteration we compute each loglikelihood at
        self.loglikely_x = []  # Tracks log-likelihood of peak tokens
        self.loglikely_w = []  # Tracks log-likelihood of word tokens
        self.loglikely_tot = []  # Tracks log-likelihood of peak + word tokens

        # Initialize peak->subregion assignments (r)
        if not self.params['symmetric']:
            # if symmetric model use deterministic assignment :
            #     if peak_val[0] > 0, r = 1, else r = 0
            self.peak_region_idx[:] = np.random.randint(
                self.params['n_regions'],  # pylint: disable=no-member
                size=(len(self.ptoken_doc_idx)))
        else:
            # if asymmetric model, randomly sample r ~ unif(1...n_regions)
            self.peak_region_idx[:] = (self.peak_vals[:, 0] > 0).astype(int)

        # Update model vectors and count matrices to reflect y and r assignments
        for i_ptoken in range(len(self.ptoken_doc_idx)):
            # document -idx (d)
            doc = self.ptoken_doc_idx[i_ptoken]
            topic = self.peak_topic_idx[i_ptoken]  # peak-token -> topic assignment (y_i)
            region = self.peak_region_idx[i_ptoken]  # peak-token -> subregion assignment (c_i)
            self.n_peak_tokens_doc_by_topic[doc, topic] += 1  # Increment document-by-topic counts
            self.n_peak_tokens_region_by_topic[region, topic] += 1  # Increment region-by-topic

        # Randomly Initialize Word->Topic Assignments (z) for each word
        # token w_i: sample z_i proportional to p(topic|doc_i)
        for i_wtoken in range(len(self.wtoken_word_idx)):
            # w_i word-type
            word = self.wtoken_word_idx[i_wtoken]

            # w_i doc-index
            doc = self.wtoken_doc_idx[i_wtoken]

            # Estimate p(t|d) for current doc
            p_topic_g_doc = self.n_peak_tokens_doc_by_topic[doc] + self.params['gamma']

            # Sample a topic from p(t|d) for the z-assignment
            probs = np.cumsum(p_topic_g_doc)  # Compute a cdf of the sampling
                                              # distribution for z
            # Which elements of cdf are less than random sample?
            sample_locs = probs < np.random.rand() * probs[-1]  # pylint: disable=no-member
            sample_locs = np.where(sample_locs)  # How many elements of cdf are
                                                 # less than sample
            topic = len(sample_locs[0])  # z = # elements of cdf less than
                                         # rand-sample

            # Update model assignment vectors and count-matrices to reflect z
            self.wtoken_topic_idx[i_wtoken] = topic  # Word-token -> topic assignment (z_i)
            self.n_word_tokens_word_by_topic[word, topic] += 1
            self.total_n_word_tokens_by_topic[0, topic] += 1
            self.n_word_tokens_doc_by_topic[doc, topic] += 1

        # Get Initial Spatial Parameter Estimates
        self._update_regions()

        # Get Log-Likelihood of data for Initialized model and save to
        # variables tracking loglikely
        self.compute_log_likelihood()

    def update(self, loglikely_freq=1, verbose=2):
        """
        Run a complete update cycle (sample z, sample y&r, update regions).

        Parameters
        ----------
        loglikely_freq : :obj:`int`, optional
            The frequency with which log-likelihood is updated. Default value
            is 1 (log-likelihood is updated every iteration).
        verbose : {0, 1, 2}, optional
            Determines how much info is printed to console. 0 = none,
            1 = a little, 2 = a lot. Default value is 2.
        """
        self.iter += 1  # Update total iteration count

        if verbose == 2:
            LGR.info('Iter {0:04d}: Sampling z'.format(self.iter))
        self.seed += 1
        self._update_word_topic_assignments(self.seed)  # Update z-assignments

        if verbose == 2:
            LGR.info('Iter {0:04d}: Sampling y|r'.format(self.iter))
        self.seed += 1
        self._update_peak_assignments(self.seed)  # Update y-assignments

        if verbose == 2:
            LGR.info('Iter {0:04d}: Updating spatial params'.format(self.iter))
        self._update_regions()  # Update gaussian estimates for all subregions

        # Only update loglikelihood every 'loglikely_freq' iterations
        # (Computing log-likelihood isn't necessary and slows things down a bit)
        if self.iter % loglikely_freq == 0:
            if verbose == 2:
                LGR.info('Iter {0:04d}: Computing log-likelihood'.format(self.iter))
            self.compute_log_likelihood()  # Compute log-likelihood of
                                           # model in current state
            if verbose > 0:
                LGR.info('Iter {0:04d} Log-likely: x = {1:10.1f}, w = {2:10.1f}, '
                         'tot = {3:10.1f}'.format(self.iter, self.loglikely_x[-1],
                                                  self.loglikely_w[-1],
                                                  self.loglikely_tot[-1]))

    def run(n_iters=10, loglikely_freq=10, verbose=1):
        """
        Run multiple iterations.
        """
        for i in range(model.iter, n_iters):
            model.update(loglikely_freq=loglikely_freq, verbose=verbose)

    def _update_word_topic_assignments(self, randseed):
        """
        Update wtoken_topic_idx (z) indicator variables assigning words->topics.

        Parameters
        ----------
        randseed : :obj:`int`
            Random seed for this iteration.
        """
        # --- Seed random number generator
        np.random.seed(randseed)  # pylint: disable=no-member

        # Loop over all word tokens
        for i_wtoken in range(len(self.wtoken_word_idx)):
            # Get indices for current token
            word = self.wtoken_word_idx[i_wtoken]  # w_i word-type
            doc = self.wtoken_doc_idx[i_wtoken]  # w_i doc-index
            topic = self.wtoken_topic_idx[i_wtoken]  # current topic assignment for
                                                     # word token w_i

            # Decrement count-matrices to remove current wtoken_topic_idx
            self.n_word_tokens_word_by_topic[word, topic] -= 1
            self.total_n_word_tokens_by_topic[0, topic] -= 1
            self.n_word_tokens_doc_by_topic[doc, topic] -= 1

            # Get sampling distribution:
            #    p(z_i|z,d,w) ~ p(w|t) * p(t|d)
            #                 ~ p_w_t * p_topic_g_doc
            p_word_g_topic = (self.n_word_tokens_word_by_topic[word, :] + self.params['beta']) /\
                             (self.total_n_word_tokens_by_topic +
                              self.params['beta'] * len(self.word_labels))
            p_topic_g_doc = self.n_peak_tokens_doc_by_topic[doc, :] + self.params['gamma']
            probs = p_word_g_topic * p_topic_g_doc  # The unnormalized sampling distribution

            # Sample a z_i assignment for the current word-token from the sampling distribution
            probs = np.squeeze(probs) / np.sum(probs)  # Normalize the sampling
                                                       # distribution
            # Numpy returns a [1 x T] vector with a '1' in the index of sampled topic
            vec = np.random.multinomial(1, probs)  # pylint: disable=no-member
            topic = np.where(vec)[0][0]  # Extract selected topic from vector

            # Update the indices and the count matrices using the sampled z assignment
            self.wtoken_topic_idx[i_wtoken] = topic  # Update w_i topic-assignment
            self.n_word_tokens_word_by_topic[word, topic] += 1
            self.total_n_word_tokens_by_topic[0, topic] += 1
            self.n_word_tokens_doc_by_topic[doc, topic] += 1

    def _update_peak_assignments(self, randseed):
        """
        Update y / r indicator variables assigning peaks->topics/subregions.

        Parameters
        ----------
        randseed : :obj:`int`
            Random seed for this iteration.
        """
        # Seed random number generator
        np.random.seed(randseed)  # pylint: disable=no-member

        # Retrieve p(x|r,y) for all subregions
        peak_probs = self._get_peak_probs(self)

        # Iterate over all peaks x, and sample a new y and r assignment for each
        for i_ptoken in range(len(self.ptoken_doc_idx)):
            doc = self.ptoken_doc_idx[i_ptoken]
            topic = self.peak_topic_idx[i_ptoken]
            region = self.peak_region_idx[i_ptoken]

            # Decrement count in Subregion x Topic count matrix
            self.n_peak_tokens_region_by_topic[region, topic] -= 1

            # Decrement count in Document x Topic count matrix
            self.n_peak_tokens_doc_by_topic[doc, topic] -= 1

            # Retrieve the probability of generating current x from all
            # subregions: [R x T] array of probs
            p_x_subregions = (peak_probs[i_ptoken, :, :]).transpose()

            # Compute the probabilities of all subregions given doc
            #     p(r|d) ~ p(r|t) * p(t|d)
            # Counts of subregions per topic + prior: p(r|t)
            p_region_g_topic = self.n_peak_tokens_region_by_topic + self.params['delta']

            # Normalize the columns such that each topic's distribution over
            # subregions sums to 1
            p_region_g_topic = p_region_g_topic / np.sum(p_region_g_topic, axis=0)

            # Counts of topics per document + prior: p(t|d)
            p_topic_g_doc = self.n_peak_tokens_doc_by_topic[doc, :] + self.params['alpha']

            # Reshape from (ntopics,) to (nregions, ntopics) with duplicated rows
            p_topic_g_doc = np.array([p_topic_g_doc] * self.params['n_regions'])

            # Compute p(subregion | document): p(r|d) ~ p(r|t) * p(t|d)
            # [R x T] array of probs
            p_region_g_doc = p_topic_g_doc * p_region_g_topic

            # Compute the multinomial probability: p(z|y)
            # Need the current vector of all z and y assignments for current doc
            # The multinomial from which z is sampled is proportional to number
            # of y assigned to each topic, plus constant \gamma
            doc_y_counts = self.n_peak_tokens_doc_by_topic[doc, :] + self.params['gamma']
            doc_z_counts = self.n_word_tokens_doc_by_topic[doc, :]
            p_peak_g_topic = self._compute_prop_multinomial_from_zy_vectors(
                doc_z_counts, doc_y_counts)

            # Reshape from (ntopics,) to (nregions, ntopics) with duplicated rows
            p_peak_g_topic = np.array([p_peak_g_topic] * self.params['n_regions'])

            # Get the full sampling distribution:
            # [R x T] array containing the proportional probability of all y/r combinations
            probs_pdf = p_x_subregions * p_region_g_doc * p_peak_g_topic

            # Convert from a [R x T] matrix into a [R*T x 1] array we can sample from
            probs_pdf = probs_pdf.transpose().ravel()

            # Normalize the sampling distribution
            probs_pdf = probs_pdf / np.sum(probs_pdf)

            # Sample a single element (corresponding to a y_i and c_i assignment
            # for the peak token) from the sampling distribution
            # Returns a [1 x R*T] vector with a '1' in location that was sampled
            vec = np.random.multinomial(1, probs_pdf)  # pylint: disable=no-member
            sample_idx = np.where(vec)[0][0]  # Extract linear index value from vector

            # Transform the linear index of the sampled element into the
            # subregion/topic (r/y) assignment indices
            # Subregion sampled (r)
            region = np.remainder(sample_idx, self.params['n_regions'])  # pylint: disable=no-member
            topic = int(np.floor(sample_idx / self.params['n_regions']))  # Topic sampled (y)

            # Update the indices and the count matrices using the sampled y/r assignments
            self.n_peak_tokens_region_by_topic[region, topic] += 1  # Increment count in
                                                                    # Subregion x Topic count
                                                                    # matrix
            self.n_peak_tokens_doc_by_topic[doc, topic] += 1  # Increment count in
                                                              # Document x Topic count matrix
            self.peak_topic_idx[i_ptoken] = topic  # Update y->topic assignment
            self.peak_region_idx[i_ptoken] = region  # Update y->subregion assignment

    def _update_regions(self):
        """
        Update spatial distribution parameters (Gaussians params for all
        subregions).
        Updates regions_mu and regions_sigma, indicating location and
        distribution of each subregion.
        """
        # Generate default ROI based on default_width
        default_roi = self.params['roi_size'] * np.eye(self.peak_vals.shape[1])

        if not self.params['symmetric']:
            # For each region, compute a mean and a regularized covariance matrix
            for i_topic in range(self.params['n_topics']):
                for j_region in range(self.params['n_regions']):
                    # Get all peaks assigned to current topic & subregion
                    idx = (self.peak_topic_idx == i_topic) & (self.peak_region_idx == j_region)
                    vals = self.peak_vals[idx]
                    n_obs = self.n_peak_tokens_region_by_topic[j_region, i_topic]

                    # Estimate mean
                    # If there are no observations, we set mean equal to zeros,
                    # otherwise take MLE
                    if n_obs == 0:
                        mu = np.zeros([self.peak_vals.shape[1]])
                    else:
                        mu = np.mean(vals, axis=0)

                    # Estimate covariance
                    # if there are 1 or fewer observations, we set sigma_hat
                    # equal to default ROI, otherwise take MLE
                    if n_obs <= 1:
                        c_hat = default_roi
                    else:
                        c_hat = np.cov(np.transpose(vals))

                    # Regularize the covariance, using the ratio of observations
                    # to dobs (default constant # observations)
                    d_c = n_obs / (n_obs + self.params['dobs'])
                    sigma = d_c * c_hat + (1-d_c) * default_roi

                    # Store estimates in model object
                    self.regions_mu[i_topic][j_region][:] = mu
                    self.regions_sigma[i_topic][j_region][:] = sigma
        else:
            # With symmetric subregions, we jointly compute all estimates for
            # subregions 1 & 2, constraining the means to be symmetric w.r.t.
            # the origin along x-dimension
            for i_topic in range(self.params['n_topics']):
                # Get all peaks assigned to current topic & subregion 1
                idx1 = (self.peak_topic_idx == i_topic) & (self.peak_region_idx == 0)
                vals1 = self.peak_vals[idx1]
                n_obs1 = self.n_peak_tokens_region_by_topic[0, i_topic]

                # Get all peaks assigned to current topic & subregion 2
                idx2 = (self.peak_topic_idx == i_topic) & (self.peak_region_idx == 1)
                vals2 = self.peak_vals[idx2]
                n_obs2 = self.n_peak_tokens_region_by_topic[1, i_topic]

                # Get all peaks assigned to current topic & either subregion
                allvals = self.peak_vals[idx1 | idx2]

                # Estimate means
                # If there are no observations, we set mean equal to zeros,
                # otherwise take MLE

                # Estimate independent mean for subregion 1
                if n_obs1 == 0:
                    m = np.zeros([self.peak_vals.shape[1]])
                else:
                    m = np.mean(vals1, axis=0)

                # Estimate independent mean for subregion 2
                if n_obs2 == 0:
                    n = np.zeros([self.peak_vals.shape[1]])
                else:
                    n = np.mean(vals2, axis=0)

                # Estimate the weighted means of all dims, where for dim1 we
                # compute the mean w.r.t. absolute distance from the origin
                weighted_mean_dim1 = (-m[0]*n_obs1 + n[0]*n_obs2) / (n_obs1 + n_obs2)
                weighted_mean_otherdims = np.mean(allvals[:, 1:], axis=0)

                # Store weighted mean estimates
                mu1 = np.zeros([1, self.peak_vals.shape[1]])
                mu2 = np.zeros([1, self.peak_vals.shape[1]])
                mu1[0, 0] = -weighted_mean_dim1
                mu1[0, 1:] = weighted_mean_otherdims
                mu2[0, 0] = weighted_mean_dim1
                mu2[0, 1:] = weighted_mean_otherdims

                # Store estimates in model object
                self.regions_mu[i_topic][0][:] = mu1
                self.regions_mu[i_topic][1][:] = mu2

                # Estimate Covariances
                # Covariances are estimated independently
                # Cov for subregion 1
                if n_obs1 <= 1:
                    c_hat1 = default_roi
                else:
                    c_hat1 = np.cov(np.transpose(vals1))

                # Cov for subregion 2
                if n_obs2 <= 1:
                    c_hat2 = default_roi
                else:
                    c_hat2 = np.cov(np.transpose(vals2))

                # Regularize the covariances, using the ratio of observations
                # to sample_constant
                d_c_1 = (n_obs1) / (n_obs1 + self.params['dobs'])
                d_c_2 = (n_obs2) / (n_obs2 + self.params['dobs'])
                sigma1 = d_c_1 * c_hat1 + (1 - d_c_1) * default_roi
                sigma2 = d_c_2 * c_hat2 + (1 - d_c_2) * default_roi

                # Store estimates in model object
                self.regions_sigma[i_topic][0][:] = sigma1
                self.regions_sigma[i_topic][1][:] = sigma2

    @due.dcite(Doi('10.1145/1577069.1755845'),
               description='Describes method for computing log-likelihood '
                           'used in model.')
    def compute_log_likelihood(self, model=None, update_vectors=True):
        """
        Compute Log-likelihood of a model object given current model.
        Computes the log-likelihood of data in any model object (either train
        or test) given the posterior predictive distributions over peaks and
        word-types for the model. Note that this is not computing the joint
        log-likelihood of model parameters and data.

        Parameters
        ----------
        model : :obj:`gclda.Model`, optional
            The model for which log-likelihoods will be calculated.
            If not provided, log-likelihood will be calculated for the current
            model (self).
        update_vectors : :obj:`bool`, optional
            Whether to update model's log-likelihood vectors or not.

        Returns
        -------
        x_loglikely : :obj:`float`
            Total log-likelihood of all peak tokens.
        w_loglikely : :obj:`float`
            Total log-likelihood of all word tokens.
        tot_loglikely : :obj:`float`
            Total log-likelihood of peak + word tokens.

        References
        ----------
        [1] <NAME>., <NAME>., <NAME>., & <NAME>. (2009).
        Distributed algorithms for topic models. Journal of Machine Learning
        Research, 10(Aug), 1801-1828.
        """
        if model is None:
            model = self
        elif update_vectors:
            LGR.info('External model detected: Disabling update_vectors')
            update_vectors = False

        # Pre-compute all probabilities from count matrices that are needed
        # for loglikelihood computations
        # Compute docprobs for y = ND x NT: p( y_i=t | d )
        doccounts = self.n_peak_tokens_doc_by_topic + self.params['alpha']
        doccounts_sum = np.sum(doccounts, axis=1)
        docprobs_y = np.transpose(np.transpose(doccounts) / doccounts_sum)

        # Compute docprobs for z = ND x NT: p( z_i=t | y^(d) )
        doccounts = self.n_peak_tokens_doc_by_topic + self.params['gamma']
        doccounts_sum = np.sum(doccounts, axis=1)
        docprobs_z = np.transpose(np.transpose(doccounts) / doccounts_sum)

        # Compute regionprobs = NR x NT: p( r | t )
        regioncounts = (self.n_peak_tokens_region_by_topic) + self.params['delta']
        regioncounts_sum = np.sum(regioncounts, axis=0)
        regionprobs = regioncounts / regioncounts_sum

        # Compute wordprobs = NW x NT: p( w | t )
        wordcounts = self.n_word_tokens_word_by_topic + self.params['beta']
        wordcounts_sum = np.sum(wordcounts, axis=0)
        wordprobs = wordcounts / wordcounts_sum

        # Get the matrix giving p(x_i|r,t) for all x:
        #    NY x NT x NR matrix of probabilities of all peaks given all
        #    topic/subregion spatial distributions
        peak_probs = self._get_peak_probs(model)

        # Compute observed peaks (x) Loglikelihood:
        # p(x|model, doc) = p(topic|doc) * p(subregion|topic) * p(x|subregion)
        #                 = p_topic_g_doc * p_region_g_topic * p_x_r
        # Initialize variable tracking total loglikelihood of all x tokens
        x_loglikely = 0

        # Go over all observed peaks and add p(x|model) to running total
        for i_ptoken in range(len(self.ptoken_doc_idx)):
            doc = self.ptoken_doc_idx[i_ptoken] - 1  # convert didx from 1-idx to 0-idx
            p_x = 0  # Running total for p(x|d) across subregions:
                     # Compute p(x_i|d) for each subregion separately and then
                     # sum across the subregions
            for j_region in range(self.params['n_regions']):
                # p(t|d) - p(topic|doc)
                p_topic_g_doc = docprobs_y[doc]

                # p(r|t) - p(subregion|topic)
                p_region_g_topic = regionprobs[j_region]

                # p(r|d) - p(subregion|document) = p(topic|doc)*p(subregion|topic)
                p_region_g_doc = p_topic_g_doc * p_region_g_topic

                # p(x|r) - p(x|subregion)
                p_x_r = peak_probs[i_ptoken, :, j_region]

                # p(x|subregion,doc) = sum_topics ( p(subregion|doc) * p(x|subregion) )
                p_x_rd = np.dot(p_region_g_doc, p_x_r)
                p_x += p_x_rd  # Add probability for current subregion to total
                               # probability for token across subregions
            # Add probability for current token to running total for all x tokens
            x_loglikely += np.log(p_x)  # pylint: disable=no-member

        # Compute observed words (w) Loglikelihoods:
        # p(w|model, doc) = p(topic|doc) * p(word|topic)
        #                 = p_topic_g_doc * p_w_t
        w_loglikely = 0  # Initialize variable tracking total loglikelihood of all w tokens

        # Compute a matrix of posterior predictives over words:
        # = ND x NW p(w|d) = sum_t ( p(t|d) * p(w|t) )
        p_wtoken_g_doc = np.dot(docprobs_z, np.transpose(wordprobs))

        # Go over all observed word tokens and add p(w|model) to running total
        for i_wtoken in range(len(self.wtoken_word_idx)):
            # convert wtoken_word_idx from 1-idx to 0-idx
            word_token = self.wtoken_word_idx[i_wtoken] - 1
            # convert wtoken_doc_idx from 1-idx to 0-idx
            doc = self.wtoken_doc_idx[i_wtoken] - 1
            # Probability of sampling current w token from d
            p_wtoken = p_wtoken_g_doc[doc, word_token]
            # Add log-probability of current token to running total for all w tokens
            w_loglikely += np.log(p_wtoken)  # pylint: disable=no-member
        tot_loglikely = x_loglikely + w_loglikely

        # Update model log-likelihood history vector (if update_vectors == True)
        if update_vectors:
            self.loglikely_iter.append(self.iter)
            self.loglikely_x.append(x_loglikely)
            self.loglikely_w.append(w_loglikely)
            self.loglikely_tot.append(tot_loglikely)

        # Return loglikely values (used when computing log-likelihood for a
        # model-object containing hold-out data)
        return (x_loglikely, w_loglikely, tot_loglikely)

    def _get_peak_probs(self, model):
        """
        Compute a matrix giving p(x|r,t), using all x values in a model
        object, and each topic's spatial parameters.

        Returns
        -------
        peak_probs : :obj:`numpy.ndarray` of :obj:`numpy.64`
            nPeaks x nTopics x nRegions matrix of probabilities, giving
            probability of sampling each peak (x) from all subregions.
        """
        peak_probs = np.zeros(
            (len(model.ptoken_doc_idx), self.params['n_topics'],
             self.params['n_regions']), dtype=float)
        for i_topic in range(self.params['n_topics']):
            for j_region in range(self.params['n_regions']):
                pdf = multivariate_normal.pdf(
                    model.peak_vals,
                    mean=self.regions_mu[i_topic][j_region][0],
                    cov=self.regions_sigma[i_topic][j_region])
                peak_probs[:, i_topic, j_region] = pdf
        return peak_probs

    def _compute_prop_multinomial_from_zy_vectors(self, z, y):
        """
        Compute proportional multinomial probabilities of current x vector
        given current y vector, for all proposed y_i values.
        Note that this only returns values proportional to the relative
        probabilities of all proposals for y_i.

        Parameters
        ----------
        z : :obj:`numpy.ndarray` of :obj:`numpy.int64`
            A 1-by-T vector of current z counts for document d.
        y : :obj:`numpy.ndarray` of :obj:`numpy.float64`
            A 1-by-T vector of current y counts (plus gamma) for document d.

        Returns
        -------
        p : :obj:`numpy.ndarray` of :obj:`numpy.float64`
            A 1-by-T vector giving the proportional probability of z, given
            that topic t was incremented.
        """
        # Compute the proportional probabilities in log-space
        logp = z * np.log((y+1) / y)  # pylint: disable=no-member
        p = np.exp(logp - np.max(logp))  # Add a constant before exponentiating
                                         # to avoid any underflow issues
        return p

    def get_spatial_probs(self):
        """
        Get conditional probability of selecting each voxel in the brain mask
        given each topic.

        Returns
        -------
        p_voxel_g_topic : :obj:`numpy.ndarray` of :obj:`numpy.float64`
            A voxel-by-topic array of conditional probabilities: p(voxel|topic).
            For cell ij, the value is the probability of voxel i being selected
            given topic j has already been selected.
        p_topic_g_voxel : :obj:`numpy.ndarray` of :obj:`numpy.float64`
            A voxel-by-topic array of conditional probabilities: p(topic|voxel).
            For cell ij, the value is the probability of topic j being selected
            given voxel i is active.
        """
        affine = self.mask_img.affine
        mask_ijk = np.vstack(np.where(self.mask_img.get_data())).T
        mask_xyz = nib.affines.apply_affine(affine, mask_ijk)

        spatial_dists = np.zeros((mask_xyz.shape[0], self.params['n_topics']), float)
        for i_topic in range(self.params['n_topics']):
            for j_region in range(self.params['n_regions']):
                pdf = multivariate_normal.pdf(
                    mask_xyz,
                    mean=self.regions_mu[i_topic][j_region][0],
                    cov=self.regions_sigma[i_topic][j_region])
                spatial_dists[:, i_topic] += pdf
        p_topic_g_voxel = spatial_dists / np.sum(spatial_dists, axis=1)[:, None]
        p_topic_g_voxel = np.nan_to_num(p_topic_g_voxel, 0)  # might be unnecessary

        p_voxel_g_topic = spatial_dists / np.sum(spatial_dists, axis=0)[None, :]
        p_voxel_g_topic = np.nan_to_num(p_voxel_g_topic, 0)  # might be unnecessary

        return p_topic_g_voxel, p_voxel_g_topic

    def save_model_params(self, out_dir, n_top_words=15):
        """
        Run all export-methods: calls all save-methods to export parameters to
        files.

        Parameters
        ----------
        out_dir : :obj:`str`
            The name of the output directory.
        n_top_words : :obj:`int`, optional
            The number of words associated with each topic to report in topic
            word probabilities file.
        """
        # If output directory doesn't exist, make it
        if not isdir(out_dir):
            mkdir(out_dir)

        # print topic-word distributions for top-K words in easy-to-read format
        out_file = join(out_dir, 'Topic_X_Word_Probs.csv')
        self._save_topic_word_probs(out_file, n_top_words=n_top_words)

        # print topic x word count matrix: m.n_word_tokens_word_by_topic
        out_file = join(out_dir, 'Topic_X_Word_CountMatrix.csv')
        self._save_topic_word_counts(out_file)

        # print activation-assignments to topics and subregions:
        # Peak_x, Peak_y, Peak_z, peak_topic_idx, peak_region_idx
        out_file = join(out_dir, 'ActivationAssignments.csv')
        self._save_activation_assignments(out_file)

    def _save_activation_assignments(self, out_file):
        """
        Save Peak->Topic and Peak->Subregion assignments for all x-tokens in
        model to file.

        Parameters
        ----------
        out_file : :obj:`str`
            The name of the output file.
        """
        ptopic_idx = self.peak_topic_idx[i_ptoken]+1
        pregion_idx = self.peak_region_idx[i_ptoken]+1
        data = np.hstack((self.peak_vals[i_ptoken, :3],
                          ptopic_idx[:, None],
                          pregion_idx[:, None]))

        df = pd.DataFrame(data=data,
                          columns=['Peak_X', 'Peak_Y', 'Peak_Z',
                                   'Topic_Assignment', 'Subregion_Assignment'])
        df.to_csv(out_file, index=False)

    def _save_topic_word_counts(self, out_file):
        """
        Save Topic->Word counts for all topics and words to file.

        Parameters
        ----------
        out_file : :obj:`str`
            The name of the output file.
        """
        with open(out_file, 'w+') as fid:
            # Print the topic-headers
            fid.write('WordLabel,')
            for i_topic in range(self.params['n_topics']):
                fid.write('Topic_{0:02d},'.format(i_topic+1))
            fid.write('\n')

            # For each row / wlabel: wlabel-string and its count under each
            # topic (the \phi matrix before adding \beta and normalizing)
            for i_word in range(len(self.word_labels)):
                fid.write('{0},'.format(self.word_labels[i_word]))

                # Print counts under all topics
                for j_topic in range(self.params['n_topics']):
                    fid.write('{0},'.format(self.n_word_tokens_word_by_topic[i_word,
                                                                             j_topic]))
                # Newline for next wlabel row
                fid.write('\n')

    def _save_topic_word_probs(self, out_file, n_top_words=15):
        """
        Save Topic->Word probability distributions for top K words to file.

        Parameters
        ----------
        out_file : :obj:`str`
            The name of the output file.
        n_top_words : :obj:`int`, optional
            The number of top words to be written out for each topic.
        """
        with open(out_file, 'w+') as fid:
            # Compute topic->word probs and marginal topic-probs
            wprobs = self.n_word_tokens_word_by_topic + self.params['beta']

            # Marginal topicprobs
            topic_probs = np.sum(wprobs, axis=0) / np.sum(wprobs)
            wprobs = wprobs / np.sum(wprobs, axis=0)  # Normalized word-probs

            # Get the sorted probabilities and indices of words under each topic
            rnk_vals = np.sort(wprobs, axis=0)
            rnk_vals = rnk_vals[::-1]
            rnk_idx = np.argsort(wprobs, axis=0)
            rnk_idx = rnk_idx[::-1]

            # Print the topic-headers
            for i_topic in range(self.params['n_topics']):
                # Print each topic and its marginal probability to columns
                fid.write('Topic_{0:02d},{1:.4f},'.format(i_topic+1,
                                                          topic_probs[i_topic]))
            fid.write('\n')

            # Print the top K word-strings and word-probs for each topic
            for i in range(n_top_words):
                for j_topic in range(self.params['n_topics']):
                    # Print the kth word in topic t and its probability
                    fid.write('{0},{1:.4f},'.format(self.word_labels[rnk_idx[i, j_topic]],
                                                    rnk_vals[i, j_topic]))
                fid.write('\n')


<reponame>zmlabe/ExtremeEvents<gh_stars>1-10
"""
Script plots LRP maps for correct and incorrect cases

Author    : <NAME>
Date      : 19 January 2021
"""

### Import modules
import numpy as np
import scipy.stats as sts
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, addcyclic, shiftgrid
import calc_Utilities as UT
import calc_dataFunctions as df
import palettable.cubehelix as cm
import palettable.scientific.sequential as ssss
import palettable.scientific.diverging as dddd
import calc_Stats as dSS
import cmocean
from sklearn.metrics import accuracy_score
from netCDF4 import Dataset

### Set preliminaries
directoryfigure = '/Users/zlabe/Desktop/ExtremeEvents_NewANN_v1/LENS/'
directorydata = '/Users/zlabe/Documents/Research/ExtremeEvents/Data/'
reg_name = 'Globe'
dataset = 'LENS'
dataset_obs = '20CRv3'
rm_ensemble_mean = True
variq = ['T2M']
seasons = ['annual']
land_only = False
ocean_only = False
rm_merid_mean = False
rm_annual_mean = False
rm_ensemble_mean = True
ensnum = 40
num_of_class = 2
iterations = 100
letters = ["a","b","c","d","e","f","g","h","i","j","k","l","m"]
labely = [r'1920-2009',r'1920-2009',r'1920-2009',
          r'2009-2099',r'2009-2099',r'2009-2099']
headers = [r'\textbf{GOOD}',r'\textbf{BAD}',r'\textbf{DIFF}']

### Create sample class labels for 1920-2099
if num_of_class == 3:
    yearlabels = np.arange(1920,2099+1,1)
    lengthlabels = yearlabels.shape[0]//num_of_class
    array1 = np.asarray([0]*lengthlabels)
    array2 = np.asarray([1]*lengthlabels)
    array3 = np.asarray([2]*lengthlabels)
    classesl = np.concatenate([array1,array2,array3],axis=None)
elif num_of_class == 2:
    yearlabels = np.arange(1920,2099+1,1)
    lengthlabels = yearlabels.shape[0]//num_of_class
    array1 = np.asarray([0]*lengthlabels)
    array2 = np.asarray([1]*lengthlabels)
    classesl = np.concatenate([array1,array2],axis=None)

### Read in prediction data
trainq = np.genfromtxt(directorydata + 'training_Century_%s_%s_%s_%s_iterations%s_v3.txt' % (variq[0],seasons[0],reg_name,dataset,iterations))
testq = np.genfromtxt(directorydata + 'testing_Century_%s_%s_%s_%s_iterations%s_v3.txt' % (variq[0],seasons[0],reg_name,dataset,iterations))
obsq = np.genfromtxt(directorydata + 'obsout_Century_%s_%s_%s_%s-%s_iterations%s_v3.txt' % (variq[0],seasons[0],reg_name,dataset_obs,dataset,iterations))

### Reshape
train = np.reshape(trainq,(trainq.shape[0]//yearlabels.shape[0],yearlabels.shape[0],trainq.shape[1]))
test = np.reshape(testq,(testq.shape[0]//yearlabels.shape[0],yearlabels.shape[0],testq.shape[1]))
obs = obsq

### Read in LRP training data
data1 = Dataset(directorydata + 'LRP_Maps_train_AnnualT2M_LENS.nc')
lat1 = data1.variables['lat'][:]
lon1 = data1.variables['lon'][:]
lrptrain = data1.variables['LRP'][:].reshape(train.shape[0],yearlabels.shape[0],lat1.shape[0]*lon1.shape[0])
data1.close()

### Read in LRP testing data
data2 = Dataset(directorydata + 'LRP_Maps_test_AnnualT2M_LENS.nc')
lrptest = data2.variables['LRP'][:].reshape(test.shape[0],yearlabels.shape[0],lat1.shape[0]*lon1.shape[0])
data2.close()

### Meshgrid
lon2,lat2 = np.meshgrid(lon1,lat1)

###############################################################################
###############################################################################
###############################################################################
### Calculate accuracy
def truelabel(data):
    """
    Calculate argmax
    """
    maxindexdata = np.empty((data.shape[0],data.shape[1]))
    for i in range(data.shape[0]):
        maxindexdata[i,:] = np.argmax(data[i,:,:],axis=1)
    meanmaxindexdata= np.nanmean(maxindexdata,axis=0)

    return maxindexdata,meanmaxindexdata

def accuracyTotalTime(data_pred,data_true):
    """
    Compute accuracy for the entire time series
    """
    accdata_pred = np.empty((data_pred.shape[0]))
    for i in range(data_pred.shape[0]):
        accdata_pred[i] = accuracy_score(data_true,data_pred[i,:])

    return accdata_pred

def accuracyTPeriodTime(data_pred,data_true):
    """
    Compute accuracy for the three periods
    """
    time = data_true.shape[0]
    period = int(time//2)

    accdata_pred = np.empty((data_pred.shape[0],2))
    for i in range(data_pred.shape[0]):
        for save,j in enumerate(range(0,time,period)):
            accdata_pred[i,save] = accuracy_score(data_true[j:j+period],
                                                data_pred[i,j:j+period])

    return accdata_pred

### Calculate statistics
indextrain,meanindextrain = truelabel(train)
indextest,meanindextest = truelabel(test)

acctrain = accuracyTotalTime(indextrain,classesl)
acctest = accuracyTotalTime(indextest,classesl)

periodtrain = accuracyTPeriodTime(indextrain,classesl)
periodtest = accuracyTPeriodTime(indextest,classesl)

### Save good lrp maps
def lrpType1(indexdata,lrpdata,classesl,lat1,lon1):
    lrpdata_good = []
    lrpdata_bad = []
    for i in range(indexdata.shape[0]):
        for j in range(int(indexdata.shape[1]-indexdata.shape[1]/2)):
            if indexdata[i,j] == classesl[j]:
                lrpdata_good.append(lrpdata[i,j,:])
            else:
                lrpdata_bad.append(lrpdata[i,j,:])
    lrpdata_goodmap = np.asarray(lrpdata_good).reshape(len(lrpdata_good),lat1.shape[0],lon1.shape[0])
    lrpdata_badmap = np.asarray(lrpdata_bad).reshape(len(lrpdata_bad),lat1.shape[0],lon1.shape[0])
    meangood_data = np.nanmean(lrpdata_goodmap,axis=0)
    meanbad_data = np.nanmean(lrpdata_badmap,axis=0)

    sizegood = len(lrpdata_goodmap)
    sizebad = len(lrpdata_badmap)

    return meangood_data,meanbad_data,sizegood,sizebad

def lrpType2(indexdata,lrpdata,classesl,lat1,lon1):
    lrpdata_good = []
    lrpdata_bad = []
    for i in range(indexdata.shape[0]):
        for j in range(int(indexdata.shape[1]-indexdata.shape[1]/2)):
            if indexdata[i,j] == classesl[int(indexdata.shape[1]/2)+j]:
                lrpdata_good.append(lrpdata[i,int(indexdata.shape[1]/2)+j,:])
            else:
                lrpdata_bad.append(lrpdata[i,j,:])
    lrpdata_goodmap = np.asarray(lrpdata_good).reshape(len(lrpdata_good),lat1.shape[0],lon1.shape[0])
    lrpdata_badmap = np.asarray(lrpdata_bad).reshape(len(lrpdata_bad),lat1.shape[0],lon1.shape[0])
    meangood_data = np.nanmean(lrpdata_goodmap,axis=0)
    meanbad_data = np.nanmean(lrpdata_badmap,axis=0)

    sizegood = len(lrpdata_goodmap)
    sizebad = len(lrpdata_badmap)

    return meangood_data,meanbad_data,sizegood,sizebad

goodtrain1,badtrain1,sizetraingood1,sizetrainbad1 = lrpType1(indextrain,lrptrain,classesl,lat1,lon1)
goodtest1,badtest1,sizetestgood1,sizetestbad1  = lrpType1(indextest,lrptest,classesl,lat1,lon1)
difftrain_1 = goodtrain1 - badtrain1
difftest_1 = goodtest1 - badtest1

goodtrain2,badtrain2,sizetraingood2,sizetrainbad2 = lrpType2(indextrain,lrptrain,classesl,lat1,lon1)
goodtest2,badtest2,sizetestgood2,sizetestbad2  = lrpType2(indextest,lrptest,classesl,lat1,lon1)
difftrain_2 = goodtrain2 - badtrain2
difftest_2 = goodtest2 - badtest2

trainplot = [goodtrain1,badtrain1,difftrain_1,
              goodtrain2,badtrain2,difftrain_2]
testplot = [goodtest1,badtest1,difftest_1,
              goodtest2,badtest2,difftest_2]

#######################################################################
#######################################################################
#######################################################################
### Plot subplot of LRP means
plt.rc('text',usetex=True)
plt.rc('font',**{'family':'sans-serif','sans-serif':['<NAME>']})

limit = np.arange(0,0.5001,0.005)
barlim = np.round(np.arange(0,0.6,0.1),2)
limitdiff = np.arange(-0.3,0.301,0.0025)
barlimdiff = np.round(np.arange(-0.3,0.31,0.15),2)
cmap = [cm.cubehelix2_16.mpl_colormap,cm.cubehelix2_16.mpl_colormap,dddd.Berlin_12.mpl_colormap,
         cm.cubehelix2_16.mpl_colormap,cm.cubehelix2_16.mpl_colormap,dddd.Berlin_12.mpl_colormap]
limits = [limit,limit,limitdiff,limit,limit,limitdiff]
barlimits = [barlim,barlimdiff]
label = r'\textbf{RELEVANCE}'
labeldiff = r'\textbf{LRP DIFFERENCE}'

fig = plt.figure(figsize=(9,4))
for r in range(len(trainplot)):
    var = trainplot[r]

    ax1 = plt.subplot(2,3,r+1)
    m = Basemap(projection='moll',lon_0=0,resolution='l',area_thresh=10000)
    circle = m.drawmapboundary(fill_color='dimgrey')
    circle.set_clip_on(False)
    if any([r==0,r==1,r==3,r==4]):
        m.drawcoastlines(color='darkgrey',linewidth=0.3)
    elif any([r==2,r==5]):
        m.drawcoastlines(color='darkgrey',linewidth=0.3)

    var, lons_cyclic = addcyclic(var, lon1)
    var, lons_cyclic = shiftgrid(180., var, lons_cyclic, start=False)
    lon2d, lat2d = np.meshgrid(lons_cyclic, lat1)
    x, y = m(lon2d, lat2d)

    circle = m.drawmapboundary(fill_color='dimgrey',color='dimgray',
                      linewidth=0.7)
    circle.set_clip_on(False)

    if any([r==0,r==1,r==3,r==4]):
        cs1 = m.contourf(x,y,var,limits[r],extend='max')
        cs1.set_cmap(cmap[r])
    elif any([r==2,r==5]):
        cs2 = m.contourf(x,y,var,limits[r],extend='both')
        cs2.set_cmap(cmap[r])

    if any([r==0,r==3]):
        ax1.annotate(r'\textbf{%s}' % labely[r],xy=(0,0),xytext=(-0.07,0.5),
                      textcoords='axes fraction',color='k',fontsize=10,
                      rotation=90,ha='center',va='center')
    if any([r==0,r==1,r==2]):
        ax1.annotate(r'\textbf{%s}' % headers[r],xy=(0,0),xytext=(0.5,1.10),
                      textcoords='axes fraction',color='dimgrey',fontsize=20,
                      rotation=0,ha='center',va='center')
    ax1.annotate(r'\textbf{[%s]}' % letters[r],xy=(0,0),xytext=(0.86,0.97),
                  textcoords='axes fraction',color='k',fontsize=6,
                  rotation=330,ha='center',va='center')

###############################################################################
cbar_ax1 = fig.add_axes([0.2,0.09,0.3,0.03])
cbar1 = fig.colorbar(cs1,cax=cbar_ax1,orientation='horizontal',
                    extend='max',extendfrac=0.07,drawedges=False)
cbar1.set_label(label,fontsize=9,color='dimgrey',labelpad=1.4)
cbar1.set_ticks(barlimits[0])
cbar1.set_ticklabels(list(map(str,barlimits[0])))
cbar1.ax.tick_params(axis='x', size=.01,labelsize=5)
cbar1.outline.set_edgecolor('dimgrey')

###############################################################################
cbar_ax2 = fig.add_axes([0.675,0.09,0.3,0.03])
cbar2 = fig.colorbar(cs2,cax=cbar_ax2,orientation='horizontal',
                    extend='both',extendfrac=0.07,drawedges=False)
cbar2.set_label(labeldiff,fontsize=9,color='dimgrey',labelpad=1.4)
cbar2.set_ticks(barlimits[1])
cbar2.set_ticklabels(list(map(str,barlimits[1])))
cbar2.ax.tick_params(axis='x', size=.01,labelsize=5)
cbar2.outline.set_edgecolor('dimgrey')

plt.tight_layout()
plt.subplots_adjust(top=0.85,wspace=0.02,hspace=0.02,bottom=0.14)

plt.savefig(directoryfigure + 'training_differenceLRP_composites_LENS.png',dpi=300)

#######################################################################
#######################################################################
#######################################################################
### Plot subplot of LRP means testing
plt.rc('text',usetex=True)
plt.rc('font',**{'family':'sans-serif','sans-serif':['Avant Garde']})

limit = np.arange(0,0.5001,0.005)
barlim = np.round(np.arange(0,0.6,0.1),2)
limitdiff = np.arange(-0.3,0.301,0.0025)
barlimdiff = np.round(np.arange(-0.3,0.31,0.15),2)
cmap = [cm.cubehelix2_16.mpl_colormap,cm.cubehelix2_16.mpl_colormap,dddd.Berlin_12.mpl_colormap,
          cm.cubehelix2_16.mpl_colormap,cm.cubehelix2_16.mpl_colormap,dddd.Berlin_12.mpl_colormap]
limits = [limit,limit,limitdiff,limit,limit,limitdiff]
barlimits = [barlim,barlimdiff]
label = r'\textbf{RELEVANCE}'
labeldiff = r'\textbf{LRP DIFFERENCE}'

fig = plt.figure(figsize=(9,4))
for r in range(len(testplot)):
    var = testplot[r]

    ax1 = plt.subplot(2,3,r+1)
    m = Basemap(projection='moll',lon_0=0,resolution='l',area_thresh=10000)
    circle = m.drawmapboundary(fill_color='dimgrey')
    circle.set_clip_on(False)
    if any([r==0,r==1,r==3,r==4]):
        m.drawcoastlines(color='darkgrey',linewidth=0.27)
    elif any([r==2,r==5]):
        m.drawcoastlines(color='darkgrey',linewidth=0.27)

    var, lons_cyclic = addcyclic(var, lon1)
    var, lons_cyclic = shiftgrid(180., var, lons_cyclic, start=False)
    lon2d, lat2d = np.meshgrid(lons_cyclic, lat1)
    x, y = m(lon2d, lat2d)

    circle = m.drawmapboundary(fill_color='dimgrey',color='dimgray',
                      linewidth=0.7)
    circle.set_clip_on(False)

    if any([r==0,r==1,r==3,r==4]):
        cs1 = m.contourf(x,y,var,limits[r],extend='max')
        cs1.set_cmap(cmap[r])
    elif any([r==2,r==5]):
        cs2 = m.contourf(x,y,var,limits[r],extend='both')
        cs2.set_cmap(cmap[r])

    if any([r==0,r==3]):
        ax1.annotate(r'\textbf{%s}' % labely[r],xy=(0,0),xytext=(-0.07,0.5),
                      textcoords='axes fraction',color='k',fontsize=10,
                      rotation=90,ha='center',va='center')
    if any([r==0,r==1,r==2]):
        ax1.annotate(r'\textbf{%s}' % headers[r],xy=(0,0),xytext=(0.5,1.10),
                      textcoords='axes fraction',color='dimgrey',fontsize=20,
                      rotation=0,ha='center',va='center')
    ax1.annotate(r'\textbf{[%s]}' % letters[r],xy=(0,0),xytext=(0.86,0.97),
                  textcoords='axes fraction',color='k',fontsize=6,
                  rotation=330,ha='center',va='center')

###############################################################################
cbar_ax1 = fig.add_axes([0.2,0.09,0.3,0.03])
cbar1 = fig.colorbar(cs1,cax=cbar_ax1,orientation='horizontal',
                    extend='max',extendfrac=0.07,drawedges=False)
cbar1.set_label(label,fontsize=9,color='dimgrey',labelpad=1.4)
cbar1.set_ticks(barlimits[0])
cbar1.set_ticklabels(list(map(str,barlimits[0])))
cbar1.ax.tick_params(axis='x', size=.01,labelsize=5)
cbar1.outline.set_edgecolor('dimgrey')

###############################################################################
cbar_ax2 = fig.add_axes([0.675,0.09,0.3,0.03])
cbar2 = fig.colorbar(cs2,cax=cbar_ax2,orientation='horizontal',
                    extend='both',extendfrac=0.07,drawedges=False)
cbar2.set_label(labeldiff,fontsize=9,color='dimgrey',labelpad=1.4)
cbar2.set_ticks(barlimits[1])
cbar2.set_ticklabels(list(map(str,barlimits[1])))
cbar2.ax.tick_params(axis='x', size=.01,labelsize=5)
cbar2.outline.set_edgecolor('dimgrey')

plt.tight_layout()
plt.subplots_adjust(top=0.85,wspace=0.02,hspace=0.02,bottom=0.14)

plt.savefig(directoryfigure + 'testing_differenceLRP_composites_LENS.png',dpi=300)


<reponame>ptro-development/trader
import pydot
import numpy

from scipy.stats.stats import pearsonr

from trader.libs.samples import get_leaf_atrribute_number, \
    get_key_sample_data

class Leaf(object):

    def __init__(self, key, data, previous_leaf):
        self.key = key
        self.data = data
        self.previous_leaf = previous_leaf
        self.next_leafs = []

    def __str__(self):
        return "<id:" + str(id(self)) + " key:" + str(self.key) + " data:" + str(self.data) + \
            " next_leafs count:" + str(len(self.next_leafs)) + ">"

    def is_last_leaf(self):
        return len(self.get_next_leafs()) == 0

    def get_key(self):
        return self.key

    def get_data(self):
        return self.data

    def set_data(self, data):
        self.data = data

    def get_next_leafs(self):
        return self.next_leafs

    def get_previous_leaf(self):
        return self.previous_leaf

    def add_leaf(self, key, data):
        new_leaf = Leaf(key, data, self)
        self.next_leafs.append(new_leaf)
        return new_leaf

    def get_leafs(self, keys):
        return filter(
            lambda x: x.get_key() in keys,
            self.next_leafs
        )

class TradeSequenceTree(object):

    def __init__(
            self, found_sample_keys,
            not_found_sample_keys,
            common_leafs_atributes,
            common_leafs_atributes_key_index_map):
        self.found_sample_keys = found_sample_keys
        self.not_found_sample_keys = not_found_sample_keys
        self.common_leafs_atributes = common_leafs_atributes
        self.common_leafs_atributes_key_index_map = common_leafs_atributes_key_index_map  # noqa
        self.first_leaf = Leaf(None, None, None)

    def get_root(self):
        return self.first_leaf

    def get_found_sample_keys(self):
        return self.found_sample_keys

    def get_not_found_sample_keys(self):
        return self.not_found_sample_keys

    def get_all_sample_keys(self):
        return self.found_sample_keys + self.not_found_sample_key

    def get_common_leafs_atributes(self):
        return self.common_leafs_atributes

    def get_common_leafs_atributes_key_index_map(self):
        return self.common_leafs_atributes_key_index_map

def get_next_not_found_leafs_extended(last_leaf, not_found_keys):
    assert isinstance(last_leaf, Leaf) and last_leaf is not None
    assert isinstance(not_found_keys, list) and len(not_found_keys) != 0
    # [[leaf, length, probability], ...]
    data = []
    overall_count = 0
    for n_leaf in last_leaf.get_leafs(not_found_keys):
        count = n_leaf.get_data()["count"]
        if count > 0:
            overall_count += count
            #data.append([n_leaf, int(n_leaf.get_key()[1:]), count])
            for volatility_set in n_leaf.get_data()["volatilities"]:
                data_to_add = [n_leaf, int(n_leaf.get_key()[1:]), count]
                data_to_add.extend(volatility_set)
                data.append(data_to_add)
            # data.append([n_leaf, int(n_leaf.get_key()[1:])])
    # compute probability
    for index, bit in enumerate(data):
        data[index][2] = float(bit[2]) / overall_count
        # pass
    # print "not_found data", data
    return data

def get_next_found_leafs_extended(
        key, last_leaf, found_keys, common_leaf_attributes, samples):
    """
    To compute array of leafs characteristics against incoming key.

    [[leaf, attribute_number, probability, correlation, volatilities_1, volatilities_2, ...], ...]

    last_leaf - represents previous leaf
    found_keys - are used to do selection on next leafs
    common_leaf_attributes -
    samples -
    """
    assert isinstance(last_leaf, Leaf) and last_leaf is not None
    assert isinstance(found_keys, list) and len(found_keys) != 0
    assert isinstance(common_leaf_attributes, dict) and len(common_leaf_attributes) != 0  # noqa
    assert isinstance(samples, list) and len(found_keys) != 0
    data = []
    overall_count = 0
    key_sample_data = get_key_sample_data(key, samples)
    #print last_leaf.get_leafs(found_keys)
    #import sys; sys.exit();
    for n_leaf in last_leaf.get_leafs(found_keys):
        volatilities = []
        count = n_leaf.get_data()["count"]
        # print "n_leaf", n_leaf, count
        if count > 0:
            overall_count += count
            n_attribute = get_leaf_atrribute_number(
                n_leaf.get_key(), common_leaf_attributes)
            n_key_sample_data_b = get_key_sample_data(
                n_leaf.get_key(), samples)
            # print "key_sample_data", key_sample_data, "n_key_sample_data_b", n_key_sample_data_b  # noqa
            correlation, other = pearsonr(key_sample_data, n_key_sample_data_b)
            # TODO: compute volatilities
            # only volatilities after re-sampling should be used, so from tree
            # print "correlation", correlation
            for volatility_set in n_leaf.get_data()["volatilities"]:
                data_to_add = [n_leaf, n_attribute, count, correlation]
                data_to_add.extend(volatility_set)
                data.append(data_to_add)
            # data.append([n_leaf, n_attribute])
    # compute probability
    for index, bit in enumerate(data):
        data[index][2] = float(bit[2]) / overall_count
        # pass
    return data

def get_leafs_chain(tree, leafs_key_sequence):
    success = True
    chain = []
    chain.append(tree.get_root())
    for index, key in enumerate(leafs_key_sequence):
        tmp_leaf = chain[-1].get_leafs([key])
        if tmp_leaf:
            chain.append(tmp_leaf[0])
        else:
            success = False
            break
    return (success, chain)

def get_last_leaf_in_leafs_chain(tree, leafs_key_sequence):
    success = True
    last_leaf = tree.get_root()
    for key in leafs_key_sequence:
        tmp_leaf = last_leaf.get_leafs([key])
        if tmp_leaf:
            last_leaf = tmp_leaf
        else:
            success = False
            break
    if success:
        return last_leaf
    else:
        return None

def get_leaf_samples_count(leaf, keys):
    return sum(
        leaf.get_data()["count"] for leaf in leaf.get_leafs(keys))

def get_leafs_chain_probability_report(tree, leafs_key_sequence):
    success, chain = get_leafs_chain(
        tree, leafs_key_sequence)
    if success:
        print "Whole sample probability chain:"
    else:
        print "Not whole sample probability chain was found."
    data = []
    for leaf in chain:
        found_probability, not_found_probability = get_leaf_next_probabilities(
            leaf, tree)
        data.append(
            leaf.get_key(),
            found_probability, not_found_probability)
    print data

def get_last_leafs(leaf, depth):
    if depth != 0:
        depth -= 1
        for leaf in leaf.get_next_leafs():
            if leaf.is_last_leaf():
                yield leaf
            else:
                for n_leaf in get_last_leafs(leaf, depth):
                    yield n_leaf
    else:
        yield leaf

def get_leafs_samples_chain(leaf, root_leaf):
    yield leaf
    if not leaf.get_previous_leaf() == root_leaf:
        for n_leaf in get_leafs_samples_chain(
                leaf.get_previous_leaf(), root_leaf):
            yield n_leaf

def reset_tree_counters(root_leaf):
    for leaf in root_leaf.get_next_leafs():
        data = leaf.get_data()
        data["count"] = 0
        leaf.set_data(data)
        if not leaf.is_last_leaf():
            reset_tree_counters(leaf)

def reset_tree_volatilities(root_leaf):
    for leaf in root_leaf.get_next_leafs():
        data = leaf.get_data()
        data["volatilities"] = []
        leaf.set_data(data)
        if not leaf.is_last_leaf():
            reset_tree_counters(leaf)

def reset_tree_values(root_leaf):
    reset_tree_counters(root_leaf)
    reset_tree_volatilities(root_leaf)

def normalise_tree_volatilities(root_leaf, linear_normalisation):
    for leaf in root_leaf.get_next_leafs():
        data = leaf.get_data()
        if "volatilities" in data:
            normalised_volatilities = []
            for volatilities in data["volatilities"]:
                normalised_volatilities.append(
                    linear_normalisation.normalise_array(volatilities))
            data["volatilities"] = normalised_volatilities
            leaf.set_data(data)
            if not leaf.is_last_leaf():
                normalise_tree_volatilities(leaf, linear_normalisation)

def get_leafs_layers(leafs, depth=0):
    next_leafs = []
    for leaf in leafs:
        next_leafs.extend(
            leaf.get_next_leafs())
    if next_leafs:
        yield next_leafs
        if depth > 0:
            for layer in get_leafs_layers(next_leafs, depth-1):
                yield layer

def get_next_leafs_extended(leaf):
    known = []
    unknown = []
    for n_leaf in leaf.get_next_leafs():
        if isinstance(n_leaf.get_key(), basestring):
            unknown.append(n_leaf)
        else:
            known.append(n_leaf)
    return known, unknown

def get_leafs_layers_extended(leafs, depth=0):
    known_next = []
    unknown_next = []
    for leaf in leafs:
        known, unknown = get_next_leafs_extended(leaf)
        known_next.extend(known)
        unknown_next.extend(unknown)
    if known_next or unknown_next:
        yield known_next, unknown_next
        if depth > 0:
            for known, unknown in get_leafs_layers_extended(
                    known_next + unknown_next, depth-1):
                yield known, unknown

def get_leafs_probabilities(leafs, tree):
    counts = numpy.array((9 * [0.0]))
    key_count_map = tree.get_common_leafs_atributes_key_index_map()
    atributes = tree.get_common_leafs_atributes()
    all_count = 0.0
    for leaf in leafs:
        count = leaf.get_data()["count"]
        counts[
            key_count_map[
                atributes[leaf.get_key()]]] += count
        all_count += count
    if all_count > 0:
        counts /= all_count
    return counts

def get_in_layer_match_probability(
        tree, layer_position, starting_layers_counts):
    current_layers_counts = get_tree_layers_counts(tree, layer_position)
    probability = 0.0
    if layer_position > -1 and layer_position < len(current_layers_counts):
        current_count = float(current_layers_counts[layer_position])
        if current_count > 0:
            probability += current_count / starting_layers_counts[
                layer_position]
    return probability

def get_in_layer_match_probability_extended(
        tree, layer_position,
        starting_layers_counts_known, starting_layers_counts_unknown):
    assert isinstance(starting_layers_counts_known, list) and \
        len(starting_layers_counts_known) != 0
    assert isinstance(starting_layers_counts_unknown, list) and \
        len(starting_layers_counts_unknown) != 0
    assert isinstance(layer_position, int)
    assert isinstance(tree, TradeSequenceTree)
    current_layers_counts_known, current_layers_counts_unknown = get_tree_layers_counts_extended(  # noqa
        tree, layer_position)
    # [know, unknown]
    probability = [0.0, 0.0]
    starting_layers = (
        starting_layers_counts_known, starting_layers_counts_unknown)
    for index, current_layers in enumerate((
            current_layers_counts_known, current_layers_counts_unknown)):
        if layer_position > -1 and layer_position < len(current_layers):
            current_count = float(current_layers[layer_position])
            if current_count > 0:
                probability[index] += current_count / starting_layers[index][
                    layer_position]
    return probability

def get_leafs_count(leafs):
    count = 0.0
    for leaf in leafs:
        count += leaf.get_data()["count"]
    return count

def get_leaf_next_probabilities(leaf, tree):
    # counts array
    # [found_count, not_found_count]
    counts = [0, 0]
    counts[0] = get_leaf_samples_count(
        leaf, tree.get_found_sample_keys())
    counts[1] = get_leaf_samples_count(
        leaf, tree.get_not_found_sample_keys())
    all_count = sum(counts)
    probabilities = 2 * [-1.0]
    for index, count in enumerate(counts):
        if count > 0:
            probabilities[index] = count / float(all_count)
    return probabilities

def get_extended_leaf_next_probabilities(leaf, tree, depth=0):
    counts = numpy.array((9 * [0.0]))
    key_count_map = tree.get_common_leafs_atributes_key_index_map()
    atributes = tree.get_common_leafs_atributes()
    all_count = 0.0
    for next_leaf in leaf.get_next_leafs():
        count = next_leaf.get_data()["count"]
        counts[
            key_count_map[
                atributes[next_leaf.get_key()]]] += count
        all_count += count
    if all_count > 0:
        counts /= all_count
    print "depth A", depth
    yield counts
    if depth > 0:
        print "depth B", depth
        counts_next = numpy.array((9 * [0.0]))
        for next_leaf in leaf.get_next_leafs():
            if not next_leaf.is_last_leaf():
                for prob in get_extended_leaf_next_probabilities(
                        next_leaf, tree, depth-1):
                    counts_next += prob
                all_counts = sum(counts_next)
                if all_counts > 0:
                    counts_next /= float(all_counts)
        yield counts_next

def report_extended_leaf_next_probabilities_old(leaf, tree, depth=0):
    key_count_map = tree.get_common_leafs_atributes_key_index_map()
    prob = [p for p in get_extended_leaf_next_probabilities(
        leaf, tree, depth)]
    new_chain = reversed(
        list(get_leafs_samples_chain(leaf, tree.get_root())))
    common_leaf_atributes = tree.get_common_leafs_atributes()
    print "Chain: " + "->".join(
        [str(l.get_key()) + "(%s)" % common_leaf_atributes[l.get_key()] for l in new_chain])  # noqa
    for index, p in enumerate(prob):
        data = zip(
            sorted(key_count_map.items(), key=lambda x: x[1]), p)
        report = [str(d[0][0]) + ":" + "%.3f" % d[1] for d in data]
        print "Level " + str(index) + " " + " ".join(report)

def report_extended_leaf_next_probabilities(leaf, tree, depth=0):
    key_count_map = tree.get_common_leafs_atributes_key_index_map()
    new_chain = reversed(
        list(get_leafs_samples_chain(leaf, tree.get_root())))
    common_leaf_atributes = tree.get_common_leafs_atributes()
    print "Chain: " + "->".join(
        [str(l.get_key()) + "(%s)" % common_leaf_atributes[l.get_key()] for l in new_chain])  # noqa
    for index, layer in enumerate(get_leafs_layers([leaf], depth)):
        prob = get_leafs_probabilities(layer, tree)
        data = zip(
            sorted(key_count_map.items(), key=lambda x: x[1]), prob)
        report = [str(d[0][0]) + ":" + "%.3f" % d[1] for d in data]
        print "Level " + str(index) + " " + " ".join(report)

def get_next_leafs_chains_report(tree, leafs_key_sequence, depth):
    success, chain = get_leafs_chain(
        tree, leafs_key_sequence)
    if success:
        found_chains = []
        not_found_chains = []
        for leaf in get_last_leafs(chain[-1], depth):
            new_chain = reversed(
                list(get_leafs_samples_chain(leaf, tree.get_root())))
            if leaf.get_key() in tree.get_found_sample_keys():
                found_chains.append(new_chain)
            else:
                not_found_chains.append(new_chain)
        probabilities = get_leaf_next_probabilities(
            chain[-1], tree)
        print "For sequence:%s depth:%s is found_chains:%s not_found_chains:%s found_probability:%.3f not_found_probability:%.3f" % (  # noqa
            leafs_key_sequence, depth,
            len(found_chains), len(not_found_chains),
            probabilities[0], probabilities[1])

def update_top_leaf_counts(
        tree, future_trails, top_leaf_key,
        next_leaf_key, positive_or_negative,
        report_probabilities=False, depth=0):
    tmp = []
    while(future_trails[top_leaf_key]):
        leaf = future_trails[top_leaf_key].pop()
        found_leaf = leaf.get_leafs([next_leaf_key])
        if found_leaf:
            leaf = found_leaf[0]
            data = leaf.get_data()
            if positive_or_negative:
                data["count"] += 1
            else:
                data["count"] -= 1
            leaf.set_data(data)
            if report_probabilities:
                report_extended_leaf_next_probabilities(
                    leaf, tree, depth)
            if not leaf.is_last_leaf():
                tmp.append(leaf)
            else:
                pass
                # add finished trail into results
                # for accounting
        else:
            pass
            """ This scenario should not be covered
            as there is supposed to be always the first
            match and next is not guarantied.
            """
    if tmp:
        future_trails[top_leaf_key] = tmp

def update_top_leafs_counts(
        future_trails, incoming_sequence_key,
        tree, top_leafs_keys, positive_or_negative=True,
        report_probabilities=False, depth=0):
    if incoming_sequence_key in top_leafs_keys:
        root_leaf = tree.get_root()
        future_trails[incoming_sequence_key].append(root_leaf)
    # do processing of new key for existing future trail
    for top_leaf_key in top_leafs_keys:
        update_top_leaf_counts(
            tree, future_trails, top_leaf_key,
            incoming_sequence_key, positive_or_negative,
            report_probabilities, depth)

def update_top_leafs_counts_2(
        tree_trails, start_end_positions,
        incoming_sequence_key, positive_or_negative,
        depth=10):
    tree_trails.append(incoming_sequence_key)
    tree_trails.update(
        incoming_sequence_key, start_end_positions,
        positive_or_negative, depth)

def generate_tree_edge(leaf, graph, depth=0, color="black", common_leafs=True):
    if leaf.get_previous_leaf():
        key = str(leaf.get_key()) + "_" + str(depth)
        if not common_leafs:
            key += "_" + str(leaf.get_previous_leaf().get_key())
    else:
        key = str(leaf.get_key()) + "_" + str(depth)
    depth += 1
    for n_leaf in leaf.get_next_leafs():
        key_next = str(n_leaf.get_key()) + "_" + str(depth)
        if not common_leafs:
            key_next += "_" + str(leaf.get_key())
        edge = pydot.Edge(
            key, key_next,
            label=str(n_leaf.get_data()["count"]),
            labelfontcolor="#009933",
            fontsize="10.0",
            color=color)
        graph.add_edge(edge)
        generate_tree_edge(n_leaf, graph, depth, color)

def generate_tree_chart(file_path, tree, size=None):
    if size:
        graph = pydot.Dot(graph_type="graph", size=size)
    else:
        graph = pydot.Dot(graph_type="graph", overlap="False")
    generate_tree_edge(tree.get_root(), graph)
    graph.write_png(file_path, prog="dot")

def get_tree_layers_counts(tree, max_depth=50):
    layer_counts = []
    for layer in get_leafs_layers([tree.get_root()], max_depth):
        layer_counts.append(get_leafs_count(layer))
    return layer_counts

def get_tree_layers_counts_extended(tree, max_depth=50):
    known_layer_counts = []
    unknown_layer_counts = []
    for known, unknown in get_leafs_layers_extended(
            [tree.get_root()], max_depth):
        known_layer_counts.append(get_leafs_count(known))
        unknown_layer_counts.append(get_leafs_count(unknown))
    return known_layer_counts, unknown_layer_counts

def increment_leaf_key_count(leaf, data_key="count"):
    data = leaf.get_data()
    if data_key in leaf.data:
        data[data_key] += 1
        leaf.set_data(data)
    else:
        data[data_key] = 1
        leaf.set_data(data)

def decrement_leaf_key_count(leaf, data_key="count"):
    if data_key in leaf.data:
        data = leaf.get_data()
        data[data_key] -= 1
        leaf.set_data(data)

def delete_leaf_volatilities(leaf, delete_index, data_key="volatilities"):
    if data_key in leaf.data:
        data = leaf.get_data()
        del data[data_key][delete_index]
        leaf.set_data(data)

def append_leaf_key_volatilities(leaf, volatilities, data_key="volatilities"):
    data = leaf.get_data()
    # print leaf.key, "volatilities", volatilities
    if data_key in leaf.data:
        data[data_key].append(volatilities)
        leaf.set_data(data)
    else:
        data[data_key] = [volatilities]
        leaf.set_data(data)
    # print "data", data

def get_leaf_key_volatilities(leaf, data_key="volatilities"):
    return leaf.get_data()["volatilities"]


from logging import Logger

import torch
from scipy.stats import spearmanr
from torch import nn
from torch.optim.adamw import AdamW
from torch.utils.data.dataloader import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from transformers import get_linear_schedule_with_warmup

from tasks.korsts.config import TrainConfig
from tasks.korsts.model import KorSTSModel

class Trainer:
    def __init__(
        self,
        config: TrainConfig,
        model: KorSTSModel,
        train_data_loader: DataLoader,
        dev_data_loader: DataLoader,
        test_data_loader: DataLoader,
        logger: Logger,
        summary_writer: SummaryWriter,
    ):
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.model = model
        self.model.to(self.device)

        self.train_data_loader = train_data_loader
        self.dev_data_loader = dev_data_loader
        self.test_data_loader = test_data_loader
        self.logger = logger
        self.summary_writer = summary_writer

        self.criterion = nn.MSELoss()
        self.optimizer = AdamW(model.parameters(), lr=config.learning_rate)

        # total step 계산
        self.steps_per_epoch = len(train_data_loader)
        self.total_steps = self.steps_per_epoch * config.num_epochs
        self.warmup_steps = config.warmup_step_ratio * self.total_steps

        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer, num_warmup_steps=self.warmup_steps, num_training_steps=self.total_steps
        )
        self.global_step = 0

    def train(self):
        # train
        self.logger.info("========== train ==========")
        self.logger.info(f"device                : {self.device}")
        self.logger.info(f"dataset length/ train : {len(self.train_data_loader.dataset)}")
        self.logger.info(f"dataset length/ dev   : {len(self.dev_data_loader.dataset)}")
        self.logger.info(f"dataset length/ test  : {len(self.test_data_loader.dataset)}")
        self.logger.info(f"batch size            : {self.config.batch_size}")
        self.logger.info(f"learning rate         : {self.config.learning_rate}")
        self.logger.info(f"dropout prob          : {self.config.dropout_prob}")
        self.logger.info(f"total epoch           : {self.config.num_epochs}")
        self.logger.info(f"steps per epoch       : {self.steps_per_epoch}")
        self.logger.info(f"total steps           : {self.total_steps}")
        self.logger.info(f"warmup steps          : {self.warmup_steps}\n")

        for epoch in range(self.config.num_epochs):
            running_loss = 0.0
            train_targets = []
            train_predictions = []

            for step, data in enumerate(tqdm(self.train_data_loader)):
                self.model.train()

                self.global_step += 1

                input_token_ids = data[0].to(self.device)
                attention_mask = data[1].to(self.device)
                token_type_ids = data[2].to(self.device)
                labels = data[3].to(self.device)

                loss, outputs = self._train_step(input_token_ids, attention_mask, token_type_ids, labels)

                running_loss += loss
                train_targets.extend(labels.tolist())
                train_predictions.extend(outputs.tolist())

                if (step + 1) % self.config.logging_interval == 0:
                    train_loss = running_loss / self.config.logging_interval
                    train_corr = spearmanr(train_targets, train_predictions)[0]
                    self.logger.info(
                        f"Epoch {epoch}, Step {step + 1}\t| Loss {train_loss:.4f}  "
                        f"Spearman Correlation {train_corr:.4f}"
                    )

                    self.summary_writer.add_scalar("korsts/train/loss", train_loss, self.global_step)
                    self.summary_writer.add_scalar("korsts/train/spearman", train_corr, self.global_step)

                    running_loss = 0.0
                    train_targets = []
                    train_predictions = []

            # dev every epoch
            dev_loss, dev_targets, dev_predictions = self._validation(self.dev_data_loader)
            dev_corr = spearmanr(dev_targets, dev_predictions)[0]
            self.logger.info(f"######### DEV REPORT #EP{epoch} #########")
            self.logger.info(f"Loss {dev_loss:.4f}")
            self.logger.info(f"Spearman Correlation {dev_corr:.4f}\n")

            self.summary_writer.add_scalar("korsts/dev/loss", dev_loss, self.global_step)
            self.summary_writer.add_scalar("korsts/dev/spearman", dev_corr, self.global_step)

            # test every epoch
            test_loss, test_targets, test_predictions = self._validation(self.test_data_loader)
            test_corr = spearmanr(test_targets, test_predictions)[0]
            self.logger.info(f"######### TEST REPORT #EP{epoch} #########")
            self.logger.info(f"Loss {test_loss:.4f}")
            self.logger.info(f"Spearman Correlation {test_corr:.4f}\n")

            self.summary_writer.add_scalar("korsts/test/loss", test_loss, self.global_step)
            self.summary_writer.add_scalar("korsts/test/spearman", test_corr, self.global_step)

            # output_path = os.path.join(self.config.checkpoint_dir, f"model-epoch-{epoch}.pth")
            # torch.save(self.model.state_dict(), output_path)
            # self.logger.info(f"MODEL IS SAVED AT {output_path}\n")

    def _train_step(self, input_token_ids, attention_mask, token_type_ids, labels):
        self.optimizer.zero_grad()

        outputs = self.model(input_token_ids, attention_mask, token_type_ids)
        outputs = outputs.view(-1)

        loss = self.criterion(outputs, labels)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)

        self.optimizer.step()
        self.scheduler.step()

        return loss.item(), outputs

    def _validation(self, data_loader):
        self.model.eval()

        running_loss = 0.0
        targets = []
        predictions = []

        with torch.no_grad():
            for data in data_loader:
                input_token_ids = data[0].to(self.device)
                attention_mask = data[1].to(self.device)
                token_type_ids = data[2].to(self.device)
                labels = data[3].to(self.device)

                outputs = self.model(input_token_ids, attention_mask, token_type_ids)
                outputs = outputs.view(-1)

                loss = self.criterion(outputs, labels)

                running_loss += loss.item()
                targets.extend(labels.tolist())
                predictions.extend(outputs.tolist())

        assert len(targets) == len(predictions)

        mean_loss = running_loss / len(data_loader)

        return mean_loss, targets, predictions


from __future__ import division

import logging

import jax.numpy as np
import jax
import numpy as raw_np
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.preprocessing import label_binarize

import scipy
import scipy.optimize
import scipy.linalg

from ..utils import clip_jax

from jax.config import config
config.update("jax_enable_x64", True)

class MultinomialRegression(BaseEstimator, RegressorMixin):
    def __init__(self, weights_0=None, method=None, initializer='identity',
                 reg_format=None,
                 reg_lambda=0.0, reg_mu=None, reg_norm=False, ref_row=True):
        if method not in [None, 'Full', 'Diag', 'FixDiag']:
            raise(ValueError('method {} not avaliable'.format(method)))

        self.weights_ = weights_0
        self.weights_0_ = weights_0
        self.method_ = method
        self.initializer = initializer
        self.reg_format = reg_format
        self.reg_lambda = reg_lambda
        self.reg_mu = reg_mu  # If number, then ODIR is applied
        self.reg_norm = reg_norm
        self.ref_row = ref_row
        self.classes = None

    @property
    def coef_(self):
        return self.weights_[:, :-1]

    @property
    def intercept_(self):
        return self.weights_[:, -1]

    def predict_proba(self, S):

        S_ = np.hstack((S, np.ones((len(S), 1))))

        return np.asarray(_calculate_outputs(self.weights_, S_))

    # FIXME Should we change predict for the argmax?
    def predict(self, S):

        return np.asarray(self.predict_proba(S))

    def fit(self, X, y, *args, **kwargs):

        X_ = np.hstack((X, np.ones((len(X), 1))))

        self.classes = raw_np.unique(y)

        k = len(self.classes)

        if self.reg_norm:
            if self.reg_mu is None:
                self.reg_lambda = self.reg_lambda / (k * (k + 1))
            else:
                self.reg_lambda = self.reg_lambda / (k * (k - 1))
                self.reg_mu = self.reg_mu / k

        target = label_binarize(y, self.classes)

        if k == 2:
            target = np.hstack([1-target, target])

        n, m = X_.shape

        XXT = (X_.repeat(m, axis=1) * np.hstack([X_]*m)).reshape((n, m, m))

        logging.debug(self.method_)

        self.weights_0_ = self._get_initial_weights(self.initializer)

        if k <= 36:
            weights = _newton_update(self.weights_0_, X_, XXT, target, k,
                                     self.method_, reg_lambda=self.reg_lambda,
                                     reg_mu=self.reg_mu, ref_row=self.ref_row,
                                     initializer=self.initializer,
                                     reg_format=self.reg_format)
        else:
            res = scipy.optimize.fmin_l_bfgs_b(func=_objective, fprime=_gradient,
                                               x0=self.weights_0_,
                                               args=(X_, XXT, target, k,
                                                     self.method_,
                                                     self.reg_lambda,
                                                     self.reg_mu, self.ref_row,
                                                     self.initializer,
                                                     self.reg_format),
                                               maxls=128,
                                               factr=1.0)
            weights = res[0]

        self.weights_ = _get_weights(weights, k, self.ref_row, self.method_)

        return self

    def _get_initial_weights(self, ref_row, initializer='identity'):
        ''' Returns an array containing only the weights of the full weight
        matrix.

        '''

        if initializer not in ['identity', None]:
            raise ValueError

        k = len(self.classes)

        if self.weights_0_ is None:

            if self.method_ == 'Full':
                if initializer == 'identity':
                    weights_0 = _get_identity_weights(k, ref_row, self.method_)
                else:
                    weights_0 = np.zeros(k * (k + 1))

            if self.method_ == 'Diag':
                if initializer == 'identity':
                    weights_0 = _get_identity_weights(k, ref_row, self.method_)
                else:
                    weights_0 = np.zeros(2*k)

            elif self.method_ == 'FixDiag':
                if initializer == 'identity':
                    weights_0 = _get_identity_weights(k, ref_row, self.method_)
                else:
                    weights_0 = np.zeros(1)

            if self.method_ is None:
                if initializer == 'identity':
                    weights_0 = _get_identity_weights(k, ref_row, self.method_)
                else:
                    weights_0 = np.zeros(k * (k + 1))
        else:
            weights_0 = self.weights_0_

        return weights_0

def _objective(params, *args):
    (X, _, y, k, method, reg_lambda, reg_mu, ref_row, _, reg_format) = args
    weights = _get_weights(params, k, ref_row, method)
    outputs = clip_jax(_calculate_outputs(weights, X))
    loss = np.mean(-np.log(np.sum(y * outputs, axis=1)))

    if reg_mu is None:
        if reg_format == 'identity':
            reg = np.hstack([np.eye(k), np.zeros((k, 1))])
        else:
            reg = np.zeros((k, k+1))
        loss = loss + reg_lambda * np.sum((weights - reg)**2)
    else:
        weights_hat = weights - np.hstack([weights[:, :-1] * np.eye(k),
                                           np.zeros((k, 1))])
        loss = loss + reg_lambda * np.sum(weights_hat[:, :-1] ** 2) + \
            reg_mu * np.sum(weights_hat[:, -1] ** 2)

    return loss

_gradient = jax.grad(_objective, argnums=0)

_hessian = jax.hessian(_objective, argnums=0)

def _get_weights(params, k, ref_row, method):
        ''' Reshapes the given params (weights) into the full matrix including 0
        '''

        if method in ['Full', None]:
            raw_weights = params.reshape(-1, k+1)
            # weights = np.zeros([k, k+1])
            # weights[:-1, :] = params.reshape(-1, k + 1)

        elif method == 'Diag':
            raw_weights = np.hstack([np.diag(params[:k]),
                                     params[k:].reshape(-1, 1)])
            # weights[:, :-1][np.diag_indices(k)] = params[:]

        elif method == 'FixDiag':
            raw_weights = np.hstack([np.eye(k) * params[0], np.zeros((k, 1))])
            # weights[np.dgag_indices(k - 1)] = params[0]
            # weights[np.diag_indices(k)] = params[0]
        else:
            raise(ValueError("Unknown calibration method {}".format(method)))

        if ref_row:
            weights = raw_weights - np.repeat(
                raw_weights[-1, :].reshape(1, -1), k, axis=0)
        else:
            weights = raw_weights

        return weights

def _get_identity_weights(n_classes, ref_row, method):

    raw_weights = None

    if (method is None) or (method == 'Full'):
        raw_weights = np.zeros((n_classes, n_classes + 1)) + \
                      np.hstack([np.eye(n_classes), np.zeros((n_classes, 1))])
        raw_weights = raw_weights.ravel()

    elif method == 'Diag':
        raw_weights = np.hstack([np.ones(n_classes), np.zeros(n_classes)])

    elif method == 'FixDiag':
        raw_weights = np.ones(1)

    return raw_weights.ravel()

def _calculate_outputs(weights, X):
    mul = np.dot(X, weights.transpose())
    return _softmax(mul)

def _softmax(X):
    """Compute the softmax of matrix X in a numerically stable way."""
    shiftx = X - np.max(X, axis=1).reshape(-1, 1)
    exps = np.exp(shiftx)
    return exps / np.sum(exps, axis=1).reshape(-1, 1)

def _newton_update(weights_0, X, XX_T, target, k, method_, maxiter=int(1024),
                   ftol=1e-12, gtol=1e-8, reg_lambda=0.0, reg_mu=None,
                   ref_row=True, initializer=None, reg_format=None):

    L_list = [raw_np.float(_objective(weights_0, X, XX_T, target, k, method_,
                                      reg_lambda, reg_mu, ref_row, initializer,
                                      reg_format))]

    weights = weights_0.copy()

    # TODO move this to the initialization
    if method_ is None:
        weights = np.zeros_like(weights)

    for i in range(0, maxiter):

        gradient = _gradient(weights, X, XX_T, target, k, method_, reg_lambda,
                             reg_mu, ref_row, initializer, reg_format)

        if np.abs(gradient).sum() < gtol:
            break

        # FIXME hessian is ocasionally NaN
        hessian = _hessian(weights, X, XX_T, target, k, method_, reg_lambda,
                           reg_mu, ref_row, initializer, reg_format)

        if method_ == 'FixDiag':
            updates = gradient / hessian
        else:
            try:
                inverse = scipy.linalg.pinv2(hessian)
                updates = np.matmul(inverse, gradient)
            except (raw_np.linalg.LinAlgError, ValueError) as err:
                logging.error(err)
                updates = gradient

        for step_size in np.hstack((np.linspace(1, 0.1, 10),
                                    np.logspace(-2, -32, 31))):

            tmp_w = weights - (updates * step_size).ravel()

            if np.any(np.isnan(tmp_w)):
                logging.debug("{}: There are NaNs in tmp_w".format(method_))

            L = _objective(tmp_w, X, XX_T, target, k, method_, reg_lambda,
                           reg_mu, ref_row, initializer, reg_format)

            if (L - L_list[-1]) < 0:
                break

        L_list.append(raw_np.float(L))

        logging.debug("{}: after {} iterations log-loss = {:.7e}, sum_grad = {:.7e}".format(
            method_, i, L, np.abs(gradient).sum()))

        if np.isnan(L):
            logging.error("{}: log-loss is NaN".format(method_))
            break

        if i >= 5:
            if (raw_np.float(raw_np.min(raw_np.diff(L_list[-5:]))) > -ftol) & \
               (raw_np.float(raw_np.sum(raw_np.diff(L_list[-5:])) > 0) == 0):
                weights = tmp_w.copy()
                logging.debug('{}: Terminate as there is not enough changes on loss.'.format(
                    method_))
                break

        if (L_list[-1] - L_list[-2]) > 0:
            logging.debug('{}: Terminate as the loss increased {}.'.format(
                method_, np.diff(L_list[-2:])))
            break
        else:
            weights = tmp_w.copy()

    L = _objective(weights, X, XX_T, target, k, method_,
                   reg_lambda, reg_mu, ref_row, initializer, reg_format)

    logging.debug("{}: after {} iterations final log-loss = {:.7e}, sum_grad = {:.7e}".format(
        method_, i, L, np.abs(gradient).sum()))

    return weights


<reponame>feynmanliang/travelling-cluster-chain<gh_stars>0
#!/usr/bin/env python3
#vim:set et sw=4 ts=8:

import scipy.io
import matplotlib.pyplot as plt
import numpy as np

if __name__ == '__main__':
    N_docs = 1740
    N_mb = 5
    plt.grid()
    num_workers = len(set(
        x[:16] for x in glob('output/samples-*-*.mm')))
    legend = []
    # plt.semilogy()
    for i in range(1, num_workers+1):
        # Perplexities here are recorded once per minibatch
        perplexities = scipy.io.mmread('./perplexities-{}.mm'.format(i))
        # plt.plot(perplexities)

        # average over entire batch
        pp = []
        for batch in range(int(perplexities.shape[0] / (N_docs / N_mb))):
            pp.append(perplexities[int(N_docs/N_mb*batch):int(N_docs/N_mb*(batch+1))].mean())
        plt.plot(pp)
        legend.append("Worker {}".format(i))

    plt.legend(legend)
    # plt.xlabel('Trajectory number')
    plt.xlabel('Iteration over dataset')
    plt.ylabel('Perplexity (nats)')
    plt.savefig('fig-perplexities.png'.format(i))