<commit_before>import logging

import numpy as np
from numpy.linalg.linalg import LinAlgError
import pandas as pd

from tardis.plasma.properties.base import ProcessingPlasmaProperty
from tardis.plasma.exceptions import PlasmaConfigError

logger = logging.getLogger(__name__)

__all__ = ['LevelBoltzmannFactorLTE', 'LevelBoltzmannFactorDiluteLTE',
           'LevelBoltzmannFactorNoNLTE', 'LevelBoltzmannFactorNLTE',
           'PartitionFunction']

class LevelBoltzmannFactorLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    general_level_boltzmann_factor : Pandas DataFrame, dtype float
                             Level population proportionality values. Indexed
                             by atomic number, ion number, level number.
                             Columns corresponding to zones. Does not consider
                             NLTE.
    """
    outputs = ('general_level_boltzmann_factor',)
    latex_name = ('bf_{i,j,k}',)
    latex_formula = ('g_{i,j,k}e^{\\dfrac{-\\epsilon_{i,j,k}}{k_{\
        \\textrm{B}}T_{\\textrm{rad}}}}',)

    @staticmethod
    def calculate(excitation_energy, g, beta_rad, levels):
        exponential = np.exp(np.outer(excitation_energy.values, -beta_rad))
        level_boltzmann_factor_array = (g.values[np.newaxis].T *
                                        exponential)
        level_boltzmann_factor = pd.DataFrame(level_boltzmann_factor_array,
                                              index=levels,
                                              columns=np.arange(len(beta_rad)),
                                              dtype=np.float64)
        return level_boltzmann_factor

class LevelBoltzmannFactorDiluteLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    general_level_boltzmann_factor : Pandas DataFrame, dtype float
                             Level population proportionality values. Indexed
                             by atomic number, ion number, level number.
                             Columns corresponding to zones. Dilute radiation
                             field means non-metastable level values are
                             multiplied by an additional factor W. Does not
                             consider NLTE.
    """
    outputs = ('general_level_boltzmann_factor',)
    latex_name = ('bf_{i,j,k}',)
    latex_formula = ('Wg_{i,j,k}e^{\\dfrac{-\\epsilon_{i,j,k}}{k_{\
        \\textrm{B}}T_{\\textrm{rad}}}}',)

    def calculate(
            self, levels, g, excitation_energy, beta_rad, w,
            metastability):
        level_boltzmann_factor = LevelBoltzmannFactorLTE.calculate(
            excitation_energy, g, beta_rad, levels)
        level_boltzmann_factor[~metastability] *= w
        return level_boltzmann_factor

class LevelBoltzmannFactorNoNLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    level_boltzmann_factor : Pandas DataFrame, dtype float
                             Returns general_level_boltzmann_factor as this
                             property is included if NLTE is not used.
    """
    outputs = ('level_boltzmann_factor',)

    @staticmethod
    def calculate(general_level_boltzmann_factor):
        return general_level_boltzmann_factor

class LevelBoltzmannFactorNLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    level_boltzmann_factor : Pandas DataFrame, dtype float
                             Returns general_level_boltzmann_factor but
                             updated for those species treated in NLTE.
    """
    outputs = ('level_boltzmann_factor',)

    def calculate(self):
        raise AttributeError(
                'This attribute is not defined on the parent class.'
                'Please use one of the subclasses.')

    @staticmethod
    def from_config(nlte_conf):
        if nlte_conf.classical_nebular and not nlte_conf.coronal_approximation:
            return LevelBoltzmannFactorNLTEClassic
        elif (
                nlte_conf.coronal_approximation and
                not nlte_conf.classical_nebular):
            return LevelBoltzmannFactorNLTECoronal
        elif nlte_conf.coronal_approximation and nlte_conf.classical_nebular:
            raise PlasmaConfigError('Both coronal approximation and '
                                    'classical nebular specified in the '
                                    'config.')
        else:
            return LevelBoltzmannFactorNLTEGeneral

    def __init__(self, plasma_parent):
        """
        Selects appropriate 'calculate' function based on NLTE config
        options selected.
        """
        super(LevelBoltzmannFactorNLTE, self).__init__(plasma_parent)

        self._update_inputs()

    def _main_nlte_calculation(
            self, atomic_data, nlte_data, t_electrons,
            j_blues, beta_sobolevs, general_level_boltzmann_factor,
            previous_electron_densities):
        """
        The core of the NLTE calculation, used with all possible config.
        options.
        """
        for species in nlte_data.nlte_species:
            logger.info('Calculating rates for species %s', species)
            number_of_levels = atomic_data.levels.energy.ix[species].count()
            lnl = nlte_data.lines_level_number_lower[species]
            lnu = nlte_data.lines_level_number_upper[species]
            lines_index, = nlte_data.lines_idx[species]

            try:
                j_blues_filtered = j_blues.iloc[lines_index]
            except AttributeError:
                j_blues_filtered = j_blues
            try:
                beta_sobolevs_filtered = beta_sobolevs.iloc[lines_index]
            except AttributeError:
                beta_sobolevs_filtered = beta_sobolevs
            A_uls = nlte_data.A_uls[species]
            B_uls = nlte_data.B_uls[species]
            B_lus = nlte_data.B_lus[species]
            r_lu_index = lnu * number_of_levels + lnl
            r_ul_index = lnl * number_of_levels + lnu
            r_ul_matrix = np.zeros(
                    (number_of_levels, number_of_levels, len(t_electrons)),
                    dtype=np.float64)
            r_ul_matrix_reshaped = r_ul_matrix.reshape(
                    (number_of_levels**2, len(t_electrons)))
            r_ul_matrix_reshaped[r_ul_index] = A_uls[np.newaxis].T + \
                B_uls[np.newaxis].T * j_blues_filtered
            r_ul_matrix_reshaped[r_ul_index] *= beta_sobolevs_filtered
            r_lu_matrix = np.zeros_like(r_ul_matrix)
            r_lu_matrix_reshaped = r_lu_matrix.reshape(
                    (number_of_levels**2, len(t_electrons)))
            r_lu_matrix_reshaped[r_lu_index] = B_lus[np.newaxis].T * \
                j_blues_filtered * beta_sobolevs_filtered
            if atomic_data.collision_data is None:
                collision_matrix = np.zeros_like(r_ul_matrix)
            else:
                if previous_electron_densities is None:
                    collision_matrix = np.zeros_like(r_ul_matrix)
                else:
                    collision_matrix = nlte_data.get_collision_matrix(
                            species, t_electrons
                            ) * previous_electron_densities.values
            rates_matrix = r_lu_matrix + r_ul_matrix + collision_matrix
            for i in xrange(number_of_levels):
                rates_matrix[i, i] = -rates_matrix[:, i].sum(axis=0)
            rates_matrix[0, :, :] = 1.0
            x = np.zeros(rates_matrix.shape[0])
            x[0] = 1.0
            for i in xrange(len(t_electrons)):
                try:
                    level_boltzmann_factor = np.linalg.solve(
                            rates_matrix[:, :, i], x)
                except LinAlgError as e:
                    if e.message == 'Singular matrix':
                        raise ValueError(
                                'SingularMatrixError during solving of the '
                                'rate matrix. Does the atomic data contain '
                                'collision data?')
                    else:
                        raise e
                general_level_boltzmann_factor[i].ix[species] = \
                    level_boltzmann_factor
        return general_level_boltzmann_factor

    def _calculate_classical_nebular(
            self, t_electrons, lines, atomic_data,
            nlte_data, general_level_boltzmann_factor, j_blues,
            previous_electron_densities):
        """
        Performs NLTE calculations using the classical nebular treatment.
        All beta sobolev values taken as 1.
        """
        beta_sobolevs = 1.0

        general_level_boltzmann_factor = self._main_nlte_calculation(
                atomic_data,
                nlte_data,
                t_electrons,
                j_blues,
                beta_sobolevs,
                general_level_boltzmann_factor,
                previous_electron_densities)
        return general_level_boltzmann_factor

    def _calculate_coronal_approximation(
            self, t_electrons, lines, atomic_data,
            nlte_data, general_level_boltzmann_factor,
            previous_electron_densities):
        """
        Performs NLTE calculations using the coronal approximation.
        All beta sobolev values taken as 1 and j_blues taken as 0.
        """
        beta_sobolevs = 1.0
        j_blues = 0.0
        general_level_boltzmann_factor = self._main_nlte_calculation(
            atomic_data, nlte_data, t_electrons, j_blues,
            beta_sobolevs, general_level_boltzmann_factor,
            previous_electron_densities)
        return general_level_boltzmann_factor

    def _calculate_general(
            self, t_electrons, lines, atomic_data, nlte_data,
            general_level_boltzmann_factor, j_blues,
            previous_beta_sobolev, previous_electron_densities):
        """
        Full NLTE calculation without approximations.
        """
        if previous_beta_sobolev is None:
            beta_sobolevs = 1.0
        else:
            beta_sobolevs = previous_beta_sobolev

        general_level_boltzmann_factor = self._main_nlte_calculation(
            atomic_data, nlte_data, t_electrons, j_blues,
            beta_sobolevs, general_level_boltzmann_factor,
            previous_electron_densities)
        return general_level_boltzmann_factor

class LevelBoltzmannFactorNLTECoronal(LevelBoltzmannFactorNLTE):
    calculate = LevelBoltzmannFactorNLTE._calculate_coronal_approximation

class LevelBoltzmannFactorNLTEClassic(LevelBoltzmannFactorNLTE):
    calculate = LevelBoltzmannFactorNLTE._calculate_classical_nebular

class LevelBoltzmannFactorNLTEGeneral(LevelBoltzmannFactorNLTE):
    calculate = LevelBoltzmannFactorNLTE._calculate_general

class PartitionFunction(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    partition_function : Pandas DataFrame, dtype float
                         Indexed by atomic number, ion number.
                         Columns are zones.
    """
    outputs = ('partition_function',)
    latex_name = ('Z_{i,j}',)
    latex_formula = ('\\sum_{k}bf_{i,j,k}',)

    def calculate(self, level_boltzmann_factor):
        return level_boltzmann_factor.groupby(
            level=['atomic_number', 'ion_number']).sum()
<commit_msg>Fix NLTE Normalization - issue #784<commit_after>import logging

import numpy as np
from numpy.linalg.linalg import LinAlgError
import pandas as pd

from tardis.plasma.properties.base import ProcessingPlasmaProperty
from tardis.plasma.exceptions import PlasmaConfigError

logger = logging.getLogger(__name__)

__all__ = ['LevelBoltzmannFactorLTE', 'LevelBoltzmannFactorDiluteLTE',
           'LevelBoltzmannFactorNoNLTE', 'LevelBoltzmannFactorNLTE',
           'PartitionFunction']

class LevelBoltzmannFactorLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    general_level_boltzmann_factor : Pandas DataFrame, dtype float
                             Level population proportionality values. Indexed
                             by atomic number, ion number, level number.
                             Columns corresponding to zones. Does not consider
                             NLTE.
    """
    outputs = ('general_level_boltzmann_factor',)
    latex_name = ('bf_{i,j,k}',)
    latex_formula = ('g_{i,j,k}e^{\\dfrac{-\\epsilon_{i,j,k}}{k_{\
        \\textrm{B}}T_{\\textrm{rad}}}}',)

    @staticmethod
    def calculate(excitation_energy, g, beta_rad, levels):
        exponential = np.exp(np.outer(excitation_energy.values, -beta_rad))
        level_boltzmann_factor_array = (g.values[np.newaxis].T *
                                        exponential)
        level_boltzmann_factor = pd.DataFrame(level_boltzmann_factor_array,
                                              index=levels,
                                              columns=np.arange(len(beta_rad)),
                                              dtype=np.float64)
        return level_boltzmann_factor

class LevelBoltzmannFactorDiluteLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    general_level_boltzmann_factor : Pandas DataFrame, dtype float
                             Level population proportionality values. Indexed
                             by atomic number, ion number, level number.
                             Columns corresponding to zones. Dilute radiation
                             field means non-metastable level values are
                             multiplied by an additional factor W. Does not
                             consider NLTE.
    """
    outputs = ('general_level_boltzmann_factor',)
    latex_name = ('bf_{i,j,k}',)
    latex_formula = ('Wg_{i,j,k}e^{\\dfrac{-\\epsilon_{i,j,k}}{k_{\
        \\textrm{B}}T_{\\textrm{rad}}}}',)

    def calculate(
            self, levels, g, excitation_energy, beta_rad, w,
            metastability):
        level_boltzmann_factor = LevelBoltzmannFactorLTE.calculate(
            excitation_energy, g, beta_rad, levels)
        level_boltzmann_factor[~metastability] *= w
        return level_boltzmann_factor

class LevelBoltzmannFactorNoNLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    level_boltzmann_factor : Pandas DataFrame, dtype float
                             Returns general_level_boltzmann_factor as this
                             property is included if NLTE is not used.
    """
    outputs = ('level_boltzmann_factor',)

    @staticmethod
    def calculate(general_level_boltzmann_factor):
        return general_level_boltzmann_factor

class LevelBoltzmannFactorNLTE(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    level_boltzmann_factor : Pandas DataFrame, dtype float
                             Returns general_level_boltzmann_factor but
                             updated for those species treated in NLTE.
    """
    outputs = ('level_boltzmann_factor',)

    def calculate(self):
        raise AttributeError(
                'This attribute is not defined on the parent class.'
                'Please use one of the subclasses.')

    @staticmethod
    def from_config(nlte_conf):
        if nlte_conf.classical_nebular and not nlte_conf.coronal_approximation:
            return LevelBoltzmannFactorNLTEClassic
        elif (
                nlte_conf.coronal_approximation and
                not nlte_conf.classical_nebular):
            return LevelBoltzmannFactorNLTECoronal
        elif nlte_conf.coronal_approximation and nlte_conf.classical_nebular:
            raise PlasmaConfigError('Both coronal approximation and '
                                    'classical nebular specified in the '
                                    'config.')
        else:
            return LevelBoltzmannFactorNLTEGeneral

    def __init__(self, plasma_parent):
        """
        Selects appropriate 'calculate' function based on NLTE config
        options selected.
        """
        super(LevelBoltzmannFactorNLTE, self).__init__(plasma_parent)

        self._update_inputs()

    def _main_nlte_calculation(
            self, atomic_data, nlte_data, t_electrons,
            j_blues, beta_sobolevs, general_level_boltzmann_factor,
            previous_electron_densities, g):
        """
        The core of the NLTE calculation, used with all possible config.
        options.
        """
        for species in nlte_data.nlte_species:
            logger.info('Calculating rates for species %s', species)
            number_of_levels = atomic_data.levels.energy.ix[species].count()
            lnl = nlte_data.lines_level_number_lower[species]
            lnu = nlte_data.lines_level_number_upper[species]
            lines_index, = nlte_data.lines_idx[species]

            try:
                j_blues_filtered = j_blues.iloc[lines_index]
            except AttributeError:
                j_blues_filtered = j_blues
            try:
                beta_sobolevs_filtered = beta_sobolevs.iloc[lines_index]
            except AttributeError:
                beta_sobolevs_filtered = beta_sobolevs
            A_uls = nlte_data.A_uls[species]
            B_uls = nlte_data.B_uls[species]
            B_lus = nlte_data.B_lus[species]
            r_lu_index = lnu * number_of_levels + lnl
            r_ul_index = lnl * number_of_levels + lnu
            r_ul_matrix = np.zeros(
                    (number_of_levels, number_of_levels, len(t_electrons)),
                    dtype=np.float64)
            r_ul_matrix_reshaped = r_ul_matrix.reshape(
                    (number_of_levels**2, len(t_electrons)))
            r_ul_matrix_reshaped[r_ul_index] = A_uls[np.newaxis].T + \
                B_uls[np.newaxis].T * j_blues_filtered
            r_ul_matrix_reshaped[r_ul_index] *= beta_sobolevs_filtered
            r_lu_matrix = np.zeros_like(r_ul_matrix)
            r_lu_matrix_reshaped = r_lu_matrix.reshape(
                    (number_of_levels**2, len(t_electrons)))
            r_lu_matrix_reshaped[r_lu_index] = B_lus[np.newaxis].T * \
                j_blues_filtered * beta_sobolevs_filtered
            if atomic_data.collision_data is None:
                collision_matrix = np.zeros_like(r_ul_matrix)
            else:
                if previous_electron_densities is None:
                    collision_matrix = np.zeros_like(r_ul_matrix)
                else:
                    collision_matrix = nlte_data.get_collision_matrix(
                            species, t_electrons
                            ) * previous_electron_densities.values
            rates_matrix = r_lu_matrix + r_ul_matrix + collision_matrix
            for i in xrange(number_of_levels):
                rates_matrix[i, i] = -rates_matrix[:, i].sum(axis=0)
            rates_matrix[0, :, :] = 1.0
            x = np.zeros(rates_matrix.shape[0])
            x[0] = 1.0
            for i in xrange(len(t_electrons)):
                try:
                    level_boltzmann_factor = np.linalg.solve(
                            rates_matrix[:, :, i], x)
                except LinAlgError as e:
                    if e.message == 'Singular matrix':
                        raise ValueError(
                                'SingularMatrixError during solving of the '
                                'rate matrix. Does the atomic data contain '
                                'collision data?')
                    else:
                        raise e
                general_level_boltzmann_factor[i].ix[species] = \
                    level_boltzmann_factor * g[species][0] / level_boltzmann_factor[0]
        return general_level_boltzmann_factor

    def _calculate_classical_nebular(
            self, t_electrons, lines, atomic_data,
            nlte_data, general_level_boltzmann_factor, j_blues,
            previous_electron_densities):
        """
        Performs NLTE calculations using the classical nebular treatment.
        All beta sobolev values taken as 1.
        """
        beta_sobolevs = 1.0

        general_level_boltzmann_factor = self._main_nlte_calculation(
                atomic_data,
                nlte_data,
                t_electrons,
                j_blues,
                beta_sobolevs,
                general_level_boltzmann_factor,
                previous_electron_densities)
        return general_level_boltzmann_factor

    def _calculate_coronal_approximation(
            self, t_electrons, lines, atomic_data,
            nlte_data, general_level_boltzmann_factor,
            previous_electron_densities):
        """
        Performs NLTE calculations using the coronal approximation.
        All beta sobolev values taken as 1 and j_blues taken as 0.
        """
        beta_sobolevs = 1.0
        j_blues = 0.0
        general_level_boltzmann_factor = self._main_nlte_calculation(
            atomic_data, nlte_data, t_electrons, j_blues,
            beta_sobolevs, general_level_boltzmann_factor,
            previous_electron_densities)
        return general_level_boltzmann_factor

    def _calculate_general(
            self, t_electrons, lines, atomic_data, nlte_data,
            general_level_boltzmann_factor, j_blues,
            previous_beta_sobolev, previous_electron_densities, g):
        """
        Full NLTE calculation without approximations.
        """
        if previous_beta_sobolev is None:
            beta_sobolevs = 1.0
        else:
            beta_sobolevs = previous_beta_sobolev

        general_level_boltzmann_factor = self._main_nlte_calculation(
            atomic_data, nlte_data, t_electrons, j_blues,
            beta_sobolevs, general_level_boltzmann_factor,
            previous_electron_densities, g)
        return general_level_boltzmann_factor

class LevelBoltzmannFactorNLTECoronal(LevelBoltzmannFactorNLTE):
    calculate = LevelBoltzmannFactorNLTE._calculate_coronal_approximation

class LevelBoltzmannFactorNLTEClassic(LevelBoltzmannFactorNLTE):
    calculate = LevelBoltzmannFactorNLTE._calculate_classical_nebular

class LevelBoltzmannFactorNLTEGeneral(LevelBoltzmannFactorNLTE):
    calculate = LevelBoltzmannFactorNLTE._calculate_general

class PartitionFunction(ProcessingPlasmaProperty):
    """
    Attributes
    ----------
    partition_function : Pandas DataFrame, dtype float
                         Indexed by atomic number, ion number.
                         Columns are zones.
    """
    outputs = ('partition_function',)
    latex_name = ('Z_{i,j}',)
    latex_formula = ('\\sum_{k}bf_{i,j,k}',)

    def calculate(self, level_boltzmann_factor):
        return level_boltzmann_factor.groupby(
            level=['atomic_number', 'ion_number']).sum()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    ``alignak_backend.template`` module

    This module manages the templates (host / services)
"""
from __future__ import print_function
from copy import deepcopy
from future.utils import iteritems
from flask import current_app, g, request, abort
from eve.methods.post import post_internal
from eve.methods.patch import patch_internal
from eve.methods.put import put_internal
from eve.methods.delete import deleteitem_internal
from bson.objectid import ObjectId
from alignak_backend.models.host import get_schema as host_schema
from alignak_backend.models.service import get_schema as service_schema

class Template(object):
    """
        Template class
    """

    @staticmethod
    def pre_post_host(user_request):
        """
        Called by EVE HOOK (app.on_pre_POST_host)

        If we use templates, we fill fields with template values

        :param user_request: request of the user
        :type user_request: object
        :return: None
        """
        if isinstance(user_request.json, dict):
            Template.fill_template_host(user_request.json)
        else:
            for i in user_request.json:
                Template.fill_template_host(i)

    @staticmethod
    def on_update_host(updates, original):
        """
        Called by EVE HOOK (app.on_update_host)

        On update host, if not template, remove in '_template_fields' fields in updates because
        we update these fields, so they are now not dependant of template

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        if g.get('ignore_hook_patch', False):
            return
        if not original['_is_template']:
            ignore_schema_fields = ['realm', '_template_fields', '_templates',
                                    '_is_template',
                                    '_templates_with_services']
            template_fields = original['_template_fields']
            do_put = False
            for (field_name, _) in iteritems(updates):
                if field_name not in ignore_schema_fields:
                    if field_name in template_fields:
                        del template_fields[field_name]
                        do_put = True
            if do_put:
                lookup = {"_id": original['_id']}
                putdata = deepcopy(original)
                putdata['_template_fields'] = template_fields
                del putdata['_etag']
                del putdata['_updated']
                del putdata['_created']
                response = put_internal('host', putdata, False, False, **lookup)
                updates['_etag'] = response[0]['_etag']
                original['_etag'] = response[0]['_etag']

    @staticmethod
    def on_updated_host(updates, original):
        """
        Called by EVE HOOK (app.on_updated_host)

        After host updated,
        if host is a template, report value of fields updated on host used this template
        if host is not template, add or remove services templates if _templates changed

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        # pylint: disable=too-many-locals
        if g.get('ignore_hook_patch', False):
            g.ignore_hook_patch = False
            return
        if original['_is_template']:
            # We must update all host use this template
            host_db = current_app.data.driver.db['host']
            hosts = host_db.find({'_templates': original['_id']})
            for host in hosts:
                Template.update_host_use_template(host, updates)
        else:
            if '_templates'in updates and updates['_templates'] != original['_templates']:
                if original['_templates_with_services']:
                    service_db = current_app.data.driver.db['service']
                    # Get all services of this host
                    myservices = service_db.find({'_is_template': False,
                                                  'host': original['_id']})
                    myservices_template_id = []
                    myservices_bis = {}
                    for myservice in myservices:
                        myservices_template_id.append(myservice['_templates'][0])
                        myservices_bis[myservice['_templates'][0]] = myservice

                    services = {}
                    service_template_id = []
                    # loop on host templates and add into services the service are templates
                    for hostid in updates['_templates']:
                        services_template = service_db.find({'_is_template': True,
                                                             'host': hostid})
                        for srv in services_template:
                            services[srv['name']] = Template.prepare_service_to_post(srv,
                                                                                original[
                                                                                '_id'])
                            service_template_id.append(services[srv['name']]['_templates'][0])
                    services_to_add = list(set(service_template_id) - set(myservices_template_id))
                    services_to_del = list(set(myservices_template_id) - set(service_template_id))
                    for (_, service) in iteritems(services):
                        if service['_templates'][0] in services_to_add:
                            post_internal('service', [service])
                    for template_id in services_to_del:
                        if template_id in myservices_bis:
                            lookup = {"_id": myservices_bis[template_id]['_id']}
                            deleteitem_internal('service', False, False, **lookup)

    @staticmethod
    def on_inserted_host(items):
        """
        Called by EVE HOOK (app.on_inserted_host)

        After host inserted, if it use a template (or templates) and the the host use template
        with services, we add templates services to this host

        :param items: list of hosts
        :type items: list
        :return: None
        """
        service_db = current_app.data.driver.db['service']
        for _, item in enumerate(items):
            if item['_templates'] != [] and item['_templates_with_services']:
                # add services
                services = {}
                # loop on host templates and add into services the service are templates
                for hostid in item['_templates']:
                    services_template = service_db.find({'_is_template': True,
                                                         'host': hostid})
                    for srv in services_template:
                        services[srv['name']] = Template.prepare_service_to_post(srv,
                                                                                item[
                                                                                '_id'])
                # when ok, add all services to this host
                post_internal('service', [services[k] for k in services])

    @staticmethod
    def on_inserted_service(items):
        """
        Called by EVE HOOK (app.on_inserted_service)

        After service inserted, if it is a template and the host linked is a template with services
        we add the service in all hosts have this host in template

        :param items: List of services
        :type items: list
        :return: None
        """
        host_db = current_app.data.driver.db['host']
        services = []
        for _, item in enumerate(items):
            if item['_templates_from_host_template'] and item['_is_template']:
                # case where this service is template host+service, so add this service on all
                # hosts
                # use the host template and have _templates_with_services=True
                hostid = item['host']
                hosts = host_db.find(
                    {'_templates': hostid, '_templates_with_services': True})
                for hs in hosts:
                    services.append(Template.prepare_service_to_post(item, hs['_id']))
        if services != []:
            post_internal('service', services)

    @staticmethod
    def on_deleted_item_service(item):
        """
        Called by EVE HOOK (app.on_deleted_item_service)

        After deleted a template service, we delete all services are linked to this template

        :param item: service dict
        :type item: dict
        :return: None
        """
        service_db = current_app.data.driver.db['service']
        if item['_is_template']:
            services = service_db.find({'_templates': item['_id']})
            for service in services:
                lookup = {"_id": service['_id']}
                deleteitem_internal('service', False, False, **lookup)

    @staticmethod
    def pre_post_service(user_request):
        """
        Called by EVE HOOK (app.on_pre_POST_service)

        If we use templates, we fill fields with template values

        :param user_request: request of the user
        :type user_request: object
        :return: None
        """
        if isinstance(user_request.json, dict):
            Template.fill_template_service(user_request.json)
        else:
            for i in user_request.json:
                Template.fill_template_service(i)

    @staticmethod
    def on_update_service(updates, original):
        """
        Called by EVE HOOK (app.on_update_service)

        On update service, if not template, remove in '_template_fields' fields in updates because
        we update these fields, so they are now not dependant of template

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        if g.get('ignore_hook_patch', False):
            return
        if not original['_is_template']:
            ignore_schema_fields = ['realm', '_template_fields', '_templates',
                                    '_is_template',
                                    '_templates_from_host_template']
            template_fields = original['_template_fields']
            do_put = False
            for (field_name, _) in iteritems(updates):
                if field_name not in ignore_schema_fields:
                    if field_name in template_fields:
                        del template_fields[field_name]
                        do_put = True
            if do_put:
                lookup = {"_id": original['_id']}
                putdata = deepcopy(original)
                putdata['_template_fields'] = template_fields
                del putdata['_etag']
                del putdata['_updated']
                del putdata['_created']
                response = put_internal('service', putdata, False, False, **lookup)
                updates['_etag'] = response[0]['_etag']
                original['_etag'] = response[0]['_etag']

    @staticmethod
    def on_updated_service(updates, original):
        """
        Called by EVE HOOK (app.on_updated_service)

        After service updated, if service is a template, report value of fields updated on
        service used this template

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        if g.get('ignore_hook_patch', False):
            g.ignore_hook_patch = False
            return
        if original['_is_template']:
            # We must update all service use this template
            service_db = current_app.data.driver.db['service']
            services = service_db.find({'_templates': original['_id']})
            for service in services:
                Template.update_service_use_template(service, updates)

    @staticmethod
    def fill_template_host(item):
        """
        Prepare fields of host with fields of host templates

        :param item: field name / values of the host
        :type item: dict
        :return: None
        """
        host = current_app.data.driver.db['host']
        ignore_fields = ['_id', '_etag', '_updated', '_created', '_template_fields', '_templates',
                         '_is_template', 'realm', '_templates_with_services']
        fields_not_update = []
        for (field_name, field_value) in iteritems(item):
            fields_not_update.append(field_name)
        item['_template_fields'] = {}
        if ('_is_template' not in item or not item['_is_template']) \
                and '_templates' in item and item['_templates'] != []:
            for host_template in item['_templates']:
                hosts = host.find_one({'_id': ObjectId(host_template)})
                if hosts is not None:
                    for (field_name, field_value) in iteritems(hosts):
                        if field_name not in fields_not_update \
                                and field_name not in ignore_fields:
                            item[field_name] = field_value
                            item['_template_fields'][field_name] = host_template
            schema = host_schema()
            ignore_schema_fields = ['realm', '_template_fields', '_templates', '_is_template',
                                    '_templates_with_services']
            for key in schema['schema']:
                if key not in ignore_schema_fields:
                    if key not in item:
                        item['_template_fields'][key] = 0

    @staticmethod
    def update_host_use_template(host, fields):
        """
        This update (patch) host with values of template

        :param host: fields / values of the host
        :type host: dict
        :param fields: fields updated in the template host
        :type fields: dict
        :return: None
        """
        host_db = current_app.data.driver.db['host']
        template_fields = {}
        for template_id in host['_templates']:
            temp = host_db.find_one({'_id': template_id})
            for (name, value) in iteritems(temp):
                template_fields[name] = value
        to_patch = {}
        for (name, value) in iteritems(fields):
            if name in host['_template_fields']:
                to_patch[name] = template_fields[name]
        if len(to_patch) > 0:
            g.ignore_hook_patch = True
            lookup = {"_id": host['_id']}
            patch_internal('host', to_patch, False, False, **lookup)

    @staticmethod
    def fill_template_service(item):
        """
        Prepare fields of service with fields of service templates

        :param item: field name / values of the service
        :type item: dict
        :return: None
        """
        service = current_app.data.driver.db['service']
        ignore_fields = ['_id', '_etag', '_updated', '_created', '_template_fields', '_templates',
                         '_is_template', '_realm', 'host', '_templates_from_host_template']
        fields_not_update = []
        for (field_name, field_value) in iteritems(item):
            fields_not_update.append(field_name)
        item['_template_fields'] = {}
        if ('_is_template' not in item or not item['_is_template']) \
                and '_templates' in item and item['_templates'] != []:
            for service_template in item['_templates']:
                services = service.find_one({'_id': ObjectId(service_template)})
                if services is not None:
                    for (field_name, field_value) in iteritems(services):
                        if field_name not in fields_not_update \
                                and field_name not in ignore_fields:
                            item[field_name] = field_value
                            item['_template_fields'][field_name] = service_template
            schema = service_schema()
            ignore_schema_fields = ['_realm', '_template_fields', '_templates', '_is_template',
                                    '_templates_from_host_template']
            for key in schema['schema']:
                if key not in ignore_schema_fields:
                    if key not in item:
                        item['_template_fields'][key] = 0

    @staticmethod
    def update_service_use_template(service, fields):
        """
        This update (patch) service with values of template

        :param service: fields / values of the service
        :type service: dict
        :param fields: fields updated in the template service
        :type fields: dict
        :return: None
        """
        service_db = current_app.data.driver.db['service']
        template_fields = {}
        for template_id in service['_templates']:
            temp = service_db.find_one({'_id': template_id})
            for (name, value) in iteritems(temp):
                template_fields[name] = value
        to_patch = {}
        for (name, value) in iteritems(fields):
            if name in service['_template_fields']:
                to_patch[name] = template_fields[name]
        if len(to_patch) > 0:
            g.ignore_hook_patch = True
            lookup = {"_id": service['_id']}
            patch_internal('service', to_patch, False, False, **lookup)

    @staticmethod
    def prepare_service_to_post(item, hostid):
        """
        Prepare a service linked with a service template in case the host is in template host +
        services

        :param item: the service template source
        :type item: dict
        :param hostid: the id of the host where put this new service
        :type hostid: str
        :return: The service with right fields
        :rtype: dict
        """
        ignore_schema_fields = ['_realm', '_template_fields', '_templates',
                                '_is_template', '_templates_from_host_template']
        schema = service_schema()
        item['host'] = hostid
        item['_templates'] = [item['_id']]
        del item['_etag']
        del item['_id']
        del item['_created']
        del item['_updated']
        if '_status' in item:
            del item['_status']
        if '_links' in item:
            del item['_links']
        item['_is_template'] = False
        item['_templates_from_host_template'] = True
        item['_template_fields'] = {}
        for key in schema['schema']:
            if item not in ignore_schema_fields:
                item['_template_fields'][key] = 0
        return item
<commit_msg>Fix #193: no services templates linked to hosts templates<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
    ``alignak_backend.template`` module

    This module manages the templates (host / services)
"""
from __future__ import print_function
from copy import deepcopy
from future.utils import iteritems
from flask import current_app, g, request, abort
from eve.methods.post import post_internal
from eve.methods.patch import patch_internal
from eve.methods.put import put_internal
from eve.methods.delete import deleteitem_internal
from bson.objectid import ObjectId
from alignak_backend.models.host import get_schema as host_schema
from alignak_backend.models.service import get_schema as service_schema

class Template(object):
    """
        Template class
    """

    @staticmethod
    def pre_post_host(user_request):
        """
        Called by EVE HOOK (app.on_pre_POST_host)

        If we use templates, we fill fields with template values

        :param user_request: request of the user
        :type user_request: object
        :return: None
        """
        if isinstance(user_request.json, dict):
            Template.fill_template_host(user_request.json)
        else:
            for i in user_request.json:
                Template.fill_template_host(i)

    @staticmethod
    def on_update_host(updates, original):
        """
        Called by EVE HOOK (app.on_update_host)

        On update host, if not template, remove in '_template_fields' fields in updates because
        we update these fields, so they are now not dependant of template

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        if g.get('ignore_hook_patch', False):
            return
        if not original['_is_template']:
            ignore_schema_fields = ['realm', '_template_fields', '_templates',
                                    '_is_template',
                                    '_templates_with_services']
            template_fields = original['_template_fields']
            do_put = False
            for (field_name, _) in iteritems(updates):
                if field_name not in ignore_schema_fields:
                    if field_name in template_fields:
                        del template_fields[field_name]
                        do_put = True
            if do_put:
                lookup = {"_id": original['_id']}
                putdata = deepcopy(original)
                putdata['_template_fields'] = template_fields
                del putdata['_etag']
                del putdata['_updated']
                del putdata['_created']
                response = put_internal('host', putdata, False, False, **lookup)
                updates['_etag'] = response[0]['_etag']
                original['_etag'] = response[0]['_etag']

    @staticmethod
    def on_updated_host(updates, original):
        """
        Called by EVE HOOK (app.on_updated_host)

        After host updated,
        if host is a template, report value of fields updated on host used this template
        if host is not template, add or remove services templates if _templates changed

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        # pylint: disable=too-many-locals
        if g.get('ignore_hook_patch', False):
            g.ignore_hook_patch = False
            return
        if original['_is_template']:
            # We must update all host use this template
            host_db = current_app.data.driver.db['host']
            hosts = host_db.find({'_templates': original['_id']})
            for host in hosts:
                Template.update_host_use_template(host, updates)
        else:
            if '_templates'in updates and updates['_templates'] != original['_templates']:
                if original['_templates_with_services']:
                    service_db = current_app.data.driver.db['service']
                    # Get all services of this host
                    myservices = service_db.find({'_is_template': False,
                                                  'host': original['_id']})
                    myservices_template_id = []
                    myservices_bis = {}
                    for myservice in myservices:
                        myservices_template_id.append(myservice['_templates'][0])
                        myservices_bis[myservice['_templates'][0]] = myservice

                    services = {}
                    service_template_id = []
                    # loop on host templates and add into services the service are templates
                    for hostid in updates['_templates']:
                        services_template = service_db.find({'_is_template': True,
                                                             'host': hostid})
                        for srv in services_template:
                            services[srv['name']] = Template.prepare_service_to_post(srv,
                                                                                original[
                                                                                '_id'])
                            service_template_id.append(services[srv['name']]['_templates'][0])
                    services_to_add = list(set(service_template_id) - set(myservices_template_id))
                    services_to_del = list(set(myservices_template_id) - set(service_template_id))
                    for (_, service) in iteritems(services):
                        if service['_templates'][0] in services_to_add:
                            post_internal('service', [service])
                    for template_id in services_to_del:
                        if template_id in myservices_bis:
                            lookup = {"_id": myservices_bis[template_id]['_id']}
                            deleteitem_internal('service', False, False, **lookup)

    @staticmethod
    def on_inserted_host(items):
        """
        Called by EVE HOOK (app.on_inserted_host)

        After host inserted, if it use a template (or templates) and the the host use template
        with services, we add templates services to this host

        :param items: list of hosts
        :type items: list
        :return: None
        """
        service_db = current_app.data.driver.db['service']
        for _, item in enumerate(items):
            if item['_templates'] != [] and item['_templates_with_services']:
                # Try to add services
                services = {}
                # loop on host templates and collect services that are templates
                for hostid in item['_templates']:
                    services_template = service_db.find({'_is_template': True, 'host': hostid})
                    for srv in services_template:
                        services[srv['name']] = Template.prepare_service_to_post(srv, item['_id'])

                # when ok, and if some exist, add all services to this host
                template_services = [services[k] for k in services]
                if template_services:
                    post_internal('service', template_services)

    @staticmethod
    def on_inserted_service(items):
        """
        Called by EVE HOOK (app.on_inserted_service)

        After service inserted, if it is a template and the host linked is a template with services
        we add the service in all hosts have this host in template

        :param items: List of services
        :type items: list
        :return: None
        """
        host_db = current_app.data.driver.db['host']
        services = []
        for _, item in enumerate(items):
            if item['_templates_from_host_template'] and item['_is_template']:
                # case where this service is template host+service, so add this service on all
                # hosts
                # use the host template and have _templates_with_services=True
                hostid = item['host']
                hosts = host_db.find(
                    {'_templates': hostid, '_templates_with_services': True})
                for hs in hosts:
                    services.append(Template.prepare_service_to_post(item, hs['_id']))
        if services != []:
            post_internal('service', services)

    @staticmethod
    def on_deleted_item_service(item):
        """
        Called by EVE HOOK (app.on_deleted_item_service)

        After deleted a template service, we delete all services are linked to this template

        :param item: service dict
        :type item: dict
        :return: None
        """
        service_db = current_app.data.driver.db['service']
        if item['_is_template']:
            services = service_db.find({'_templates': item['_id']})
            for service in services:
                lookup = {"_id": service['_id']}
                deleteitem_internal('service', False, False, **lookup)

    @staticmethod
    def pre_post_service(user_request):
        """
        Called by EVE HOOK (app.on_pre_POST_service)

        If we use templates, we fill fields with template values

        :param user_request: request of the user
        :type user_request: object
        :return: None
        """
        if isinstance(user_request.json, dict):
            Template.fill_template_service(user_request.json)
        else:
            for i in user_request.json:
                Template.fill_template_service(i)

    @staticmethod
    def on_update_service(updates, original):
        """
        Called by EVE HOOK (app.on_update_service)

        On update service, if not template, remove in '_template_fields' fields in updates because
        we update these fields, so they are now not dependant of template

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        if g.get('ignore_hook_patch', False):
            return
        if not original['_is_template']:
            ignore_schema_fields = ['realm', '_template_fields', '_templates',
                                    '_is_template',
                                    '_templates_from_host_template']
            template_fields = original['_template_fields']
            do_put = False
            for (field_name, _) in iteritems(updates):
                if field_name not in ignore_schema_fields:
                    if field_name in template_fields:
                        del template_fields[field_name]
                        do_put = True
            if do_put:
                lookup = {"_id": original['_id']}
                putdata = deepcopy(original)
                putdata['_template_fields'] = template_fields
                del putdata['_etag']
                del putdata['_updated']
                del putdata['_created']
                response = put_internal('service', putdata, False, False, **lookup)
                updates['_etag'] = response[0]['_etag']
                original['_etag'] = response[0]['_etag']

    @staticmethod
    def on_updated_service(updates, original):
        """
        Called by EVE HOOK (app.on_updated_service)

        After service updated, if service is a template, report value of fields updated on
        service used this template

        :param updates: modified fields
        :type updates: dict
        :param original: original fields
        :type original: dict
        :return: None
        """
        if g.get('ignore_hook_patch', False):
            g.ignore_hook_patch = False
            return
        if original['_is_template']:
            # We must update all service use this template
            service_db = current_app.data.driver.db['service']
            services = service_db.find({'_templates': original['_id']})
            for service in services:
                Template.update_service_use_template(service, updates)

    @staticmethod
    def fill_template_host(item):
        """
        Prepare fields of host with fields of host templates

        :param item: field name / values of the host
        :type item: dict
        :return: None
        """
        host = current_app.data.driver.db['host']
        ignore_fields = ['_id', '_etag', '_updated', '_created', '_template_fields', '_templates',
                         '_is_template', 'realm', '_templates_with_services']
        fields_not_update = []
        for (field_name, field_value) in iteritems(item):
            fields_not_update.append(field_name)
        item['_template_fields'] = {}
        if ('_is_template' not in item or not item['_is_template']) \
                and '_templates' in item and item['_templates'] != []:
            for host_template in item['_templates']:
                hosts = host.find_one({'_id': ObjectId(host_template)})
                if hosts is not None:
                    for (field_name, field_value) in iteritems(hosts):
                        if field_name not in fields_not_update \
                                and field_name not in ignore_fields:
                            item[field_name] = field_value
                            item['_template_fields'][field_name] = host_template
            schema = host_schema()
            ignore_schema_fields = ['realm', '_template_fields', '_templates', '_is_template',
                                    '_templates_with_services']
            for key in schema['schema']:
                if key not in ignore_schema_fields:
                    if key not in item:
                        item['_template_fields'][key] = 0

    @staticmethod
    def update_host_use_template(host, fields):
        """
        This update (patch) host with values of template

        :param host: fields / values of the host
        :type host: dict
        :param fields: fields updated in the template host
        :type fields: dict
        :return: None
        """
        host_db = current_app.data.driver.db['host']
        template_fields = {}
        for template_id in host['_templates']:
            temp = host_db.find_one({'_id': template_id})
            for (name, value) in iteritems(temp):
                template_fields[name] = value
        to_patch = {}
        for (name, value) in iteritems(fields):
            if name in host['_template_fields']:
                to_patch[name] = template_fields[name]
        if len(to_patch) > 0:
            g.ignore_hook_patch = True
            lookup = {"_id": host['_id']}
            patch_internal('host', to_patch, False, False, **lookup)

    @staticmethod
    def fill_template_service(item):
        """
        Prepare fields of service with fields of service templates

        :param item: field name / values of the service
        :type item: dict
        :return: None
        """
        service = current_app.data.driver.db['service']
        ignore_fields = ['_id', '_etag', '_updated', '_created', '_template_fields', '_templates',
                         '_is_template', '_realm', 'host', '_templates_from_host_template']
        fields_not_update = []
        for (field_name, field_value) in iteritems(item):
            fields_not_update.append(field_name)
        item['_template_fields'] = {}
        if ('_is_template' not in item or not item['_is_template']) \
                and '_templates' in item and item['_templates'] != []:
            for service_template in item['_templates']:
                services = service.find_one({'_id': ObjectId(service_template)})
                if services is not None:
                    for (field_name, field_value) in iteritems(services):
                        if field_name not in fields_not_update \
                                and field_name not in ignore_fields:
                            item[field_name] = field_value
                            item['_template_fields'][field_name] = service_template
            schema = service_schema()
            ignore_schema_fields = ['_realm', '_template_fields', '_templates', '_is_template',
                                    '_templates_from_host_template']
            for key in schema['schema']:
                if key not in ignore_schema_fields:
                    if key not in item:
                        item['_template_fields'][key] = 0

    @staticmethod
    def update_service_use_template(service, fields):
        """
        This update (patch) service with values of template

        :param service: fields / values of the service
        :type service: dict
        :param fields: fields updated in the template service
        :type fields: dict
        :return: None
        """
        service_db = current_app.data.driver.db['service']
        template_fields = {}
        for template_id in service['_templates']:
            temp = service_db.find_one({'_id': template_id})
            for (name, value) in iteritems(temp):
                template_fields[name] = value
        to_patch = {}
        for (name, value) in iteritems(fields):
            if name in service['_template_fields']:
                to_patch[name] = template_fields[name]
        if len(to_patch) > 0:
            g.ignore_hook_patch = True
            lookup = {"_id": service['_id']}
            patch_internal('service', to_patch, False, False, **lookup)

    @staticmethod
    def prepare_service_to_post(item, hostid):
        """
        Prepare a service linked with a service template in case the host is in template host +
        services

        :param item: the service template source
        :type item: dict
        :param hostid: the id of the host where put this new service
        :type hostid: str
        :return: The service with right fields
        :rtype: dict
        """
        ignore_schema_fields = ['_realm', '_template_fields', '_templates',
                                '_is_template', '_templates_from_host_template']
        schema = service_schema()
        item['host'] = hostid
        item['_templates'] = [item['_id']]
        del item['_etag']
        del item['_id']
        del item['_created']
        del item['_updated']
        if '_status' in item:
            del item['_status']
        if '_links' in item:
            del item['_links']
        item['_is_template'] = False
        item['_templates_from_host_template'] = True
        item['_template_fields'] = {}
        for key in schema['schema']:
            if item not in ignore_schema_fields:
                item['_template_fields'][key] = 0
        return item
<|endoftext|>


<commit_before>from tcga_encoder.utils.helpers import *
from tcga_encoder.data.data import *
from tcga_encoder.definitions.tcga import *
#from tcga_encoder.definitions.nn import *
from tcga_encoder.definitions.locations import *
#from tcga_encoder.algorithms import *
import seaborn as sns
from sklearn.manifold import TSNE, locally_linear_embedding

def main( data_location, results_location ):
  data_path    = os.path.join( HOME_DIR ,data_location ) #, "data.h5" )
  results_path = os.path.join( HOME_DIR, results_location )

  data_filename = os.path.join( data_path, "data.h5")
  fill_filename = os.path.join( results_path, "full_vae_fill.h5" )

  survival_dir = os.path.join( results_path, "survival" )
  check_and_mkdir(survival_dir)

  print "HOME_DIR: ", HOME_DIR
  print "data_filename: ", data_filename
  print "fill_filename: ", fill_filename

  print "LOADING stores"
  data_store = pd.HDFStore( data_filename, "r" )
  fill_store = pd.HDFStore( fill_filename, "r" )

  Z_train = fill_store["/Z/TRAIN/Z/mu"]
  Z_val = fill_store["/Z/VAL/Z/mu"]

  Z = np.vstack( (Z_train.values, Z_val.values) )
  n_z = Z.shape[1]
  #pdb.set_trace()
  z_names = ["z_%d"%z_idx for z_idx in range(Z.shape[1])]
  Z = pd.DataFrame( Z, index = np.hstack( (Z_train.index.values, Z_val.index.values)), columns = z_names )

  barcodes = np.union1d( Z_train.index.values, Z_val.index.values )

  Z=Z.loc[barcodes]

  tissues = data_store["/CLINICAL/TISSUE"].loc[barcodes]
  #Overall Survival (OS) The event call is derived from "vital status" parameter. The time_to_event is in days, equals to days_to_death if patient deceased; in the case of a patient is still living, the time variable is the maximum(days_to_last_known_alive, days_to_last_followup).  This pair of clinical parameters are called _EVENT and _TIME_TO_EVENT on the cancer browser.

  ALL_SURVIVAL = data_store["/CLINICAL/data"][["patient.days_to_last_followup","patient.days_to_death","patient.days_to_birth"]]
  tissue_barcodes = np.array( ALL_SURVIVAL.index.tolist(), dtype=str )
  surv_barcodes = np.array([ x+"_"+y for x,y in tissue_barcodes])
  NEW_SURVIVAL = pd.DataFrame( ALL_SURVIVAL.values, index =surv_barcodes, columns = ALL_SURVIVAL.columns )
  NEW_SURVIVAL = NEW_SURVIVAL.loc[barcodes]
  #clinical = data_store["/CLINICAL/data"].loc[barcodes]

  Age = NEW_SURVIVAL[ "patient.days_to_birth" ].values.astype(int)
  Times = NEW_SURVIVAL[ "patient.days_to_last_followup" ].fillna(0).values.astype(int)+NEW_SURVIVAL[ "patient.days_to_death" ].fillna(0).values.astype(int)
  Events = (1-np.isnan( NEW_SURVIVAL[ "patient.days_to_death" ].astype(float)) ).astype(int)

  ok_age_query = Age<-10
  ok_age = pp.find(ok_age_query )
  tissues = tissues[ ok_age_query ]
  #pdb.set_trace()
  Age=-Age[ok_age]
  Times = Times[ok_age]
  Events = Events[ok_age]
  barcodes = barcodes[ok_age]
  NEW_SURVIVAL = NEW_SURVIVAL.loc[barcodes]

  #ok_followup_query = NEW_SURVIVAL[ "patient.days_to_last_followup" ].fillna(0).values>=0
  #ok_followup = pp.find( ok_followup_query )

  bad_followup_query = NEW_SURVIVAL[ "patient.days_to_last_followup" ].fillna(0).values.astype(int)<0
  bad_followup = pp.find( bad_followup_query )

  ok_followup_query = 1-bad_followup_query
  ok_followup = pp.find( ok_followup_query )

  bad_death_query = NEW_SURVIVAL[ "patient.days_to_death" ].fillna(0).values.astype(int)<0
  bad_death = pp.find( bad_death_query )

  #pdb.set_trace()
  Age=Age[ok_followup]
  Times = Times[ok_followup]
  Events = Events[ok_followup]
  barcodes = barcodes[ok_followup]
  NEW_SURVIVAL = NEW_SURVIVAL.loc[barcodes]

  Z = Z.loc[barcodes]
  Z["E"] = Events
  Z["T"] = Times
  Z["Age"] = np.log(Age)

  tissues = data_store["/CLINICAL/TISSUE"].loc[barcodes]

  tissue_names = tissues.columns
  tissue_idx = np.argmax( tissues.values, 1 )

  Z["Tissue"] = tissue_idx

  from lifelines import CoxPHFitter
  from lifelines.datasets import load_regression_dataset
  from lifelines.utils import k_fold_cross_validation
  from lifelines import KaplanMeierFitter
  from lifelines.statistics import logrank_test
  # cf = CoxPHFitter()
  # cf.fit( Z, 'T', event_col='E', strata=['Tissue'])
  #
  # cf.print_summary()
  # S = np.hstack( [])
  # survival_data = pd.DataFrame( X)
  n_tissues = len(tissue_names)
  n_trials = 2*n_z
  trial_names = ["r_%d"%(trial_idx) for trial_idx in range(n_trials)]

  p_values_half  = np.ones( (n_tissues,n_z), dtype=float)
  p_values_third_random  = np.ones( (n_tissues,n_trials), dtype=float)
  p_values_third = np.ones( (n_tissues,n_z), dtype=float)
  p_values_fifth = np.ones( (n_tissues,n_z), dtype=float)
  p_values_tenth = np.ones( (n_tissues,n_z), dtype=float)

  for t_idx in range(n_tissues):
    t_ids = tissue_idx == t_idx
    tissue_name = tissue_names[t_idx]

    print "working %s"%(tissue_name)
    bcs = barcodes[t_ids]
    Z_tissue = Z.loc[ bcs ]

    events = Z_tissue["E"]
    times  = Z_tissue["T"]
    Z_values = Z_tissue[z_names].values

    n_tissue = len(bcs)
    if n_tissue < 20:
      continue

    half  = int(n_tissue/2.0)
    third = int(n_tissue/3.0)
    third = int(n_tissue/3.0)
    fifth = int(n_tissue/5.0)
    tenth = int(n_tissue/10.0)
    for z_idx in range(n_z):
      z = Z_values[:,z_idx]
      I = np.argsort(z)
      z1_half = I[:half]
      z2_half = I[half:]
      z1_third = I[:third]
      z2_third = I[-third:]
      z1_fifth = I[:fifth]
      z2_fifth = I[-fifth:]
      z1_tenth = I[:tenth]
      z2_tenth = I[-tenth:]

      #kmf = KaplanMeierFitter()
      #kmf.fit(times[z1_half], event_observed=z1_half[I] )

      results_half = logrank_test(times[z1_half], times[z2_half], events[z1_half], events[z2_half], alpha=.99 )
      p_values_half[t_idx,z_idx] = results_half.p_value
      results_third = logrank_test(times[z1_third], times[z2_third], events[z1_third], events[z2_third], alpha=.99 )
      p_values_third[t_idx,z_idx] = results_third.p_value

      results_fifth = logrank_test(times[z1_fifth], times[z2_fifth], events[z1_fifth], events[z2_fifth], alpha=.99 )
      p_values_fifth[t_idx,z_idx] = results_fifth.p_value
      results_tenth = logrank_test(times[z1_tenth], times[z2_tenth], events[z1_tenth], events[z2_tenth], alpha=.99 )
      p_values_tenth[t_idx,z_idx] = results_tenth.p_value
      #pdb.set_trace()

    # for trial_idx in range(n_trials):
    #   I = np.random.permutation(n_tissue)
    #   z1_third = I[:third]
    #   z2_third = I[-third:]
    #   results_third = logrank_test(times[z1_third], times[z2_third], events[z1_third], events[z2_third], alpha=.99 )
    #   p_values_third_random[t_idx,trial_idx] = results_third.p_value
    #events = Z["E"].loc[]

  #
  p_values_half  = pd.DataFrame( p_values_half, index = tissue_names, columns=z_names )
  p_values_third = pd.DataFrame( p_values_third, index = tissue_names, columns=z_names )
  p_values_fifth = pd.DataFrame( p_values_fifth, index = tissue_names, columns=z_names )
  p_values_tenth = pd.DataFrame( p_values_tenth, index = tissue_names, columns=z_names )
  #p_values_third_random = pd.DataFrame( p_values_third_random, index = tissue_names, columns=trial_names )
  #pdb.set_trace()
  #
  #
  # tissue_names = tissues.columns
  # tissue_idx = np.argmax( tissues.values, 1 )
  #
  # pdb.set_trace()
  # pp.savefig( tsne_dir + "/tsne_perplexity_%d.png"%(perplexity), format='png', dpi=300 )
  p_values_half.to_csv( survival_dir + "/p_values_half.csv" )
  p_values_third.to_csv( survival_dir + "/p_values_third.csv" )
  p_values_fifth.to_csv( survival_dir + "/p_values_fifth.csv" )
  p_values_tenth.to_csv( survival_dir + "/p_values_tenth.csv" )
  #p_values_third_random.to_csv( survival_dir + "/p_values_third_random.csv" )

  pv = p_values_third
  splits = "1/3"
  split_word="third"
  binses = [20,50,100]
  for pv,splits,split_word in zip( [p_values_half,p_values_third,p_values_fifth,p_values_tenth],["1/2","1/3","1/5","1/10"],["half","third","fifth","tenth"]):
    for bins in binses:
      pp.figure()
      pp.hist( pv.values.flatten(), bins, range=(0,1), normed=True, histtype="step", lw=3 )
      pp.plot( [0,1],[1,1], 'r-', lw=3)
      pp.legend( ["Z","random"])
      pp.xlabel("p-value logrank test")
      pp.ylabel("Pr(p_value)")
      pp.title("Comparison between p-values using latent space (%s splits)"%splits)
      pp.savefig( survival_dir + "/p_values_comparison_%s_%dbins.png"%(split_word,bins), format='png', dpi=300 )

  pp.close('all')
  #pdb.set_trace()

if __name__ == "__main__":

  data_location = sys.argv[1]
  results_location = sys.argv[2]

  main( data_location, results_location )<commit_msg>10 bins<commit_after>from tcga_encoder.utils.helpers import *
from tcga_encoder.data.data import *
from tcga_encoder.definitions.tcga import *
#from tcga_encoder.definitions.nn import *
from tcga_encoder.definitions.locations import *
#from tcga_encoder.algorithms import *
import seaborn as sns
from sklearn.manifold import TSNE, locally_linear_embedding

def main( data_location, results_location ):
  data_path    = os.path.join( HOME_DIR ,data_location ) #, "data.h5" )
  results_path = os.path.join( HOME_DIR, results_location )

  data_filename = os.path.join( data_path, "data.h5")
  fill_filename = os.path.join( results_path, "full_vae_fill.h5" )

  survival_dir = os.path.join( results_path, "survival" )
  check_and_mkdir(survival_dir)

  print "HOME_DIR: ", HOME_DIR
  print "data_filename: ", data_filename
  print "fill_filename: ", fill_filename

  print "LOADING stores"
  data_store = pd.HDFStore( data_filename, "r" )
  fill_store = pd.HDFStore( fill_filename, "r" )

  Z_train = fill_store["/Z/TRAIN/Z/mu"]
  Z_val = fill_store["/Z/VAL/Z/mu"]

  Z = np.vstack( (Z_train.values, Z_val.values) )
  n_z = Z.shape[1]
  #pdb.set_trace()
  z_names = ["z_%d"%z_idx for z_idx in range(Z.shape[1])]
  Z = pd.DataFrame( Z, index = np.hstack( (Z_train.index.values, Z_val.index.values)), columns = z_names )

  barcodes = np.union1d( Z_train.index.values, Z_val.index.values )

  Z=Z.loc[barcodes]

  tissues = data_store["/CLINICAL/TISSUE"].loc[barcodes]
  #Overall Survival (OS) The event call is derived from "vital status" parameter. The time_to_event is in days, equals to days_to_death if patient deceased; in the case of a patient is still living, the time variable is the maximum(days_to_last_known_alive, days_to_last_followup).  This pair of clinical parameters are called _EVENT and _TIME_TO_EVENT on the cancer browser.

  ALL_SURVIVAL = data_store["/CLINICAL/data"][["patient.days_to_last_followup","patient.days_to_death","patient.days_to_birth"]]
  tissue_barcodes = np.array( ALL_SURVIVAL.index.tolist(), dtype=str )
  surv_barcodes = np.array([ x+"_"+y for x,y in tissue_barcodes])
  NEW_SURVIVAL = pd.DataFrame( ALL_SURVIVAL.values, index =surv_barcodes, columns = ALL_SURVIVAL.columns )
  NEW_SURVIVAL = NEW_SURVIVAL.loc[barcodes]
  #clinical = data_store["/CLINICAL/data"].loc[barcodes]

  Age = NEW_SURVIVAL[ "patient.days_to_birth" ].values.astype(int)
  Times = NEW_SURVIVAL[ "patient.days_to_last_followup" ].fillna(0).values.astype(int)+NEW_SURVIVAL[ "patient.days_to_death" ].fillna(0).values.astype(int)
  Events = (1-np.isnan( NEW_SURVIVAL[ "patient.days_to_death" ].astype(float)) ).astype(int)

  ok_age_query = Age<-10
  ok_age = pp.find(ok_age_query )
  tissues = tissues[ ok_age_query ]
  #pdb.set_trace()
  Age=-Age[ok_age]
  Times = Times[ok_age]
  Events = Events[ok_age]
  barcodes = barcodes[ok_age]
  NEW_SURVIVAL = NEW_SURVIVAL.loc[barcodes]

  #ok_followup_query = NEW_SURVIVAL[ "patient.days_to_last_followup" ].fillna(0).values>=0
  #ok_followup = pp.find( ok_followup_query )

  bad_followup_query = NEW_SURVIVAL[ "patient.days_to_last_followup" ].fillna(0).values.astype(int)<0
  bad_followup = pp.find( bad_followup_query )

  ok_followup_query = 1-bad_followup_query
  ok_followup = pp.find( ok_followup_query )

  bad_death_query = NEW_SURVIVAL[ "patient.days_to_death" ].fillna(0).values.astype(int)<0
  bad_death = pp.find( bad_death_query )

  #pdb.set_trace()
  Age=Age[ok_followup]
  Times = Times[ok_followup]
  Events = Events[ok_followup]
  barcodes = barcodes[ok_followup]
  NEW_SURVIVAL = NEW_SURVIVAL.loc[barcodes]

  Z = Z.loc[barcodes]
  Z["E"] = Events
  Z["T"] = Times
  Z["Age"] = np.log(Age)

  tissues = data_store["/CLINICAL/TISSUE"].loc[barcodes]

  tissue_names = tissues.columns
  tissue_idx = np.argmax( tissues.values, 1 )

  Z["Tissue"] = tissue_idx

  from lifelines import CoxPHFitter
  from lifelines.datasets import load_regression_dataset
  from lifelines.utils import k_fold_cross_validation
  from lifelines import KaplanMeierFitter
  from lifelines.statistics import logrank_test
  # cf = CoxPHFitter()
  # cf.fit( Z, 'T', event_col='E', strata=['Tissue'])
  #
  # cf.print_summary()
  # S = np.hstack( [])
  # survival_data = pd.DataFrame( X)
  n_tissues = len(tissue_names)
  n_trials = 2*n_z
  trial_names = ["r_%d"%(trial_idx) for trial_idx in range(n_trials)]

  p_values_half  = np.ones( (n_tissues,n_z), dtype=float)
  p_values_third_random  = np.ones( (n_tissues,n_trials), dtype=float)
  p_values_third = np.ones( (n_tissues,n_z), dtype=float)
  p_values_fifth = np.ones( (n_tissues,n_z), dtype=float)
  p_values_tenth = np.ones( (n_tissues,n_z), dtype=float)

  for t_idx in range(n_tissues):
    t_ids = tissue_idx == t_idx
    tissue_name = tissue_names[t_idx]

    print "working %s"%(tissue_name)
    bcs = barcodes[t_ids]
    Z_tissue = Z.loc[ bcs ]

    events = Z_tissue["E"]
    times  = Z_tissue["T"]
    Z_values = Z_tissue[z_names].values

    n_tissue = len(bcs)
    if n_tissue < 10:
      continue

    half  = int(n_tissue/2.0)
    third = int(n_tissue/3.0)
    third = int(n_tissue/3.0)
    fifth = int(n_tissue/5.0)
    tenth = int(n_tissue/10.0)
    for z_idx in range(n_z):
      z = Z_values[:,z_idx]
      I = np.argsort(z)
      z1_half = I[:half]
      z2_half = I[half:]
      z1_third = I[:third]
      z2_third = I[-third:]
      z1_fifth = I[:fifth]
      z2_fifth = I[-fifth:]
      z1_tenth = I[:tenth]
      z2_tenth = I[-tenth:]

      #kmf = KaplanMeierFitter()
      #kmf.fit(times[z1_half], event_observed=z1_half[I] )

      results_half = logrank_test(times[z1_half], times[z2_half], events[z1_half], events[z2_half], alpha=.99 )
      p_values_half[t_idx,z_idx] = results_half.p_value
      results_third = logrank_test(times[z1_third], times[z2_third], events[z1_third], events[z2_third], alpha=.99 )
      p_values_third[t_idx,z_idx] = results_third.p_value

      results_fifth = logrank_test(times[z1_fifth], times[z2_fifth], events[z1_fifth], events[z2_fifth], alpha=.99 )
      p_values_fifth[t_idx,z_idx] = results_fifth.p_value
      results_tenth = logrank_test(times[z1_tenth], times[z2_tenth], events[z1_tenth], events[z2_tenth], alpha=.99 )
      p_values_tenth[t_idx,z_idx] = results_tenth.p_value
      #pdb.set_trace()

    # for trial_idx in range(n_trials):
    #   I = np.random.permutation(n_tissue)
    #   z1_third = I[:third]
    #   z2_third = I[-third:]
    #   results_third = logrank_test(times[z1_third], times[z2_third], events[z1_third], events[z2_third], alpha=.99 )
    #   p_values_third_random[t_idx,trial_idx] = results_third.p_value
    #events = Z["E"].loc[]

  #
  p_values_half  = pd.DataFrame( p_values_half, index = tissue_names, columns=z_names )
  p_values_third = pd.DataFrame( p_values_third, index = tissue_names, columns=z_names )
  p_values_fifth = pd.DataFrame( p_values_fifth, index = tissue_names, columns=z_names )
  p_values_tenth = pd.DataFrame( p_values_tenth, index = tissue_names, columns=z_names )
  #p_values_third_random = pd.DataFrame( p_values_third_random, index = tissue_names, columns=trial_names )
  #pdb.set_trace()
  #
  #
  # tissue_names = tissues.columns
  # tissue_idx = np.argmax( tissues.values, 1 )
  #
  # pdb.set_trace()
  # pp.savefig( tsne_dir + "/tsne_perplexity_%d.png"%(perplexity), format='png', dpi=300 )
  p_values_half.to_csv( survival_dir + "/p_values_half.csv" )
  p_values_third.to_csv( survival_dir + "/p_values_third.csv" )
  p_values_fifth.to_csv( survival_dir + "/p_values_fifth.csv" )
  p_values_tenth.to_csv( survival_dir + "/p_values_tenth.csv" )
  #p_values_third_random.to_csv( survival_dir + "/p_values_third_random.csv" )

  pv = p_values_third
  splits = "1/3"
  split_word="third"
  binses = [20,50,100]
  for pv,splits,split_word in zip( [p_values_half,p_values_third,p_values_fifth,p_values_tenth],["1/2","1/3","1/5","1/10"],["half","third","fifth","tenth"]):
    for bins in binses:
      pp.figure()
      pp.hist( pv.values.flatten(), bins, range=(0,1), normed=True, histtype="step", lw=3 )
      pp.plot( [0,1],[1,1], 'r-', lw=3)
      pp.legend( ["Z","random"])
      pp.xlabel("p-value logrank test")
      pp.ylabel("Pr(p_value)")
      pp.title("Comparison between p-values using latent space (%s splits)"%splits)
      pp.savefig( survival_dir + "/p_values_comparison_%s_%dbins.png"%(split_word,bins), format='png', dpi=300 )

  pp.close('all')
  #pdb.set_trace()

if __name__ == "__main__":

  data_location = sys.argv[1]
  results_location = sys.argv[2]

  main( data_location, results_location )<|endoftext|>


<commit_before>from repository import changeset
from repository import fsrepos
import filecontainer
import md5
import os
import sqlite
import tempfile
import xmlshims

class NetworkRepositoryServer(xmlshims.NetworkConvertors):

    def allTroveNames(self, authToken):
	if not self.auth.check(authToken, write = False):
	    raise InsufficientPermission

	return [ x for x in self.iterAllTroveNames(authToken) ]

    def createBranch(self, authToken, newBranch, kind, frozenLocation,
		     troveList):
	if not self.auth.check(authToken, write = True):
	    raise InsufficientPermission

	newBranch = self.toLabel(newBranch)
	if kind == 'v':
	    location = self.toVersion(frozenLocation)
	elif kind == 'l':
	    location = self.toLabel(frozenLocation)
	else:
	    return 0

	self.repos.createBranch(newBranch, location, troveList)
	return 1

    def hasPackage(self, authToken, pkgName):
	if not self.auth.check(authToken, write = False, trove = pkgName):
	    raise InsufficientPermission

	return self.repos.troveStore.hasTrove(pkgName)

    def hasTrove(self, authToken, pkgName, version, flavor):
	if not self.auth.check(authToken, write = False, trove = pkgName):
	    raise InsufficientPermission

	return self.repos.troveStore.hasTrove(pkgName, troveVersion = version,
					troveFlavor = flavor)

    def getTroveVersionList(self, authToken, troveNameList):
	d = {}
	for troveName in troveNameList:
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    d[troveName] = [ x for x in
			    self.repos.troveStore.iterTroveVersions(troveName) ]

	return d

    def getFilesInTrove(self, authToken, troveName, version, flavor,
                        sortByPath = False, withFiles = False):
	if not self.auth.check(authToken, write = False, trove = troveName,
			       label = troveName.branch().label()):
	    raise InsufficientPermission

        gen = self.repos.troveStore.iterFilesInTrove(troveName,
                                               self.toVersion(version),
                                               self.toFlavor(flavor),
                                               sortByPath,
                                               withFiles)
        if withFiles:
            return [ (x[0], x[1], self.fromVersion(x[2]), self.fromFile(x[3]))
                     for x in gen ]
        else:
            return [ (x[0], x[1], self.fromVersion(x[2])) for x in gen ]

    def getFileContents(self, authToken, sha1list):
	# XXX X this isn't properly checked!
	if not self.auth.check(authToken, write = False, trove = troveName,
			       label = troveName.branch().label()):
	    raise InsufficientPermission

	(fd, path) = tempfile.mkstemp(dir = self.tmpPath, suffix = '.cfc-out')
	f = os.fdopen(fd, "w")

	fc = filecontainer.FileContainer(f)
	del f
	d = self.repos.getFileContents(sha1list)

	for sha1 in sha1list:
	    fc.addFile(sha1, d[sha1], "", d[sha1].fullSize)
	fc.close()

	fileName = os.path.basename(path)
	return "%s/%s" % (self.urlBase, fileName[:-4])

    def getAllTroveLeafs(self, authToken, troveNames):
	d = {}
	for troveName in troveNames:
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    d[troveName] = [ x for x in
			    self.repos.troveStore.iterAllTroveLeafs(troveName) ]
	return d

    def getTroveLeavesByLabel(self, authToken, troveNameList, labelStr):
	d = {}
	for troveName in troveNameList:
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    d[troveName] = [ x for x in
			self.repos.troveStore.iterTroveLeafsByLabel(troveName,
								   labelStr) ]

	return d

    def getTroveVersionFlavors(self, authToken, troveDict):
	newD = {}
	for (troveName, versionList) in troveDict.iteritems():
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    innerD = {}
	    for versionStr in versionList:
		innerD[versionStr] = [ self.fromFlavor(x) for x in
		    self.repos.troveStore.iterTroveFlavors(troveName,
						 self.toVersion(versionStr)) ]
	    newD[troveName] = innerD

	return newD

    def getTroveLatestVersion(self, authToken, pkgName, branchStr):
	branch = self.toBranch(branchStr)

	if not self.auth.check(authToken, write = False, trove = pkgName,
			       label = branch.label()):
	    raise InsufficientPermission

        try:
            return self.fromVersion(self.repos.troveStore.troveLatestVersion(pkgName,
					     self.toBranch(branchStr)))
        except KeyError:
            return 0

    def getTroveFlavorsLatestVersion(self, authToken, troveName, branchStr):
	branch = self.toBranch(branchStr)

	if not self.auth.check(authToken, write = False, trove = pkgName,
			       label = branch.label()):
	    raise InsufficientPermission

	return self.repos.troveStore.iterTrovePerFlavorLeafs(troveName, branchStr)

    def getChangeSet(self, authToken, chgSetList, recurse, withFiles):
	l = []
	for (name, flavor, old, new, absolute) in chgSetList:
	    newVer = self.toVersion(new)

	    if not self.auth.check(authToken, write = False, trove = name,
				   label = newVer.branch().label()):
		raise InsufficientPermission

	    if old == 0:
		l.append((name, self.toFlavor(flavor), None,
			 self.toVersion(new), absolute))
	    else:
		l.append((name, self.toFlavor(flavor), self.toVersion(old),
			 self.toVersion(new), absolute))

	cs = self.repos.createChangeSet(l, recurse = recurse,
					withFiles = withFiles)
	(fd, path) = tempfile.mkstemp(dir = self.tmpPath, suffix = '.ccs-out')
	os.close(fd)
	cs.writeToFile(path)
	fileName = os.path.basename(path)
	return "%s/%s" % (self.urlBase, fileName[:-4])

    def iterAllTroveNames(self, authToken):
	if not self.auth.check(authToken, write = False):
	    raise InsufficientPermission

	return self.repos.iterAllTroveNames()

    def prepareChangeSet(self, authToken):
	# make sure they have a valid account and permission to commit to
	# *something*
	if not self.auth.check(authToken, write = True):
	    raise InsufficientPermission

	(fd, path) = tempfile.mkstemp(dir = self.tmpPath, suffix = '.ccs-in')
	os.close(fd)
	fileName = os.path.basename(path)
	return "%s/%s" % (self.urlBase, fileName[:-3])

    def commitChangeSet(self, authToken, url):
	assert(url.startswith(self.urlBase))
	fileName = url[len(self.urlBase):] + "-in"
	path = "%s/%s" % (self.tmpPath, fileName)

	try:
	    cs = changeset.ChangeSetFromFile(path)
	finally:
	    pass
	    os.unlink(path)

	# walk through all of the branches this change set commits to
	# and make sure the user has enough permissions for the operation
	items = {}
	for pkgCs in cs.iterNewPackageList():
	    items[(pkgCs.getName(), pkgCs.getNewVersion())] = True
	    if not self.auth.check(authToken, write = True,
			       label = pkgCs.getNewVersion().branch().label(),
			       trove = pkgCs.getName()):
		raise InsufficientPermission

	self.repos.commitChangeSet(cs)

	return True

    def getFileVersion(self, authToken, fileId, version, withContents = 0):
	# XXX needs to authentication against the trove the file is part of,
	# which is unfortunate
	f = self.repos.troveStore.getFile(fileId, self.toVersion(version))
	return self.fromFile(f)

    def checkVersion(self, authToken, clientVersion):
	if not self.auth.check(authToken, write = False):
	    raise InsufficientPermission

        if clientVersion < 0:
            raise RuntimeError, "client is too old"
        return 0

    def __init__(self, path, tmpPath, urlBase, authDbPath):
	self.repos = fsrepos.FilesystemRepository(path)
	self.tmpPath = tmpPath
	self.urlBase = urlBase
	self.auth = NetworkAuthorization(authDbPath)

class NetworkAuthorization:

    def check(self, authToken, write = False, label = None, trove = None):
	if not write and self.anonReads:
	    return True

	if not authToken[0]:
	    return False

	stmt = """
	    SELECT count(*) FROM
	       (SELECT userId as uuserId FROM Users WHERE user=%s AND
		    password=%s)
	    JOIN Permissions ON uuserId=Permissions.userId
	"""
	m = md5.new()
	m.update(authToken[1])
	params = [authToken[0], m.hexdigest()]

	where = []
	if label:
	    where.append(" labelId=(SELECT labelId FROM Labels WHERE " \
			    "label=%s) OR labelId is Null")
	    params.append(label.asString())

	if trove:
	    where.append(" troveNameId=(SELECT troveNameId FROM TroveNames "
			        "WHERE troveName=%s) OR troveNameId is Null" )
	    params.append(trove)

	if write:
	    where.append("write=1")

	if where:
	    stmt += "WHERE " + " AND ".join(where)

	cu = self.db.cursor()
	cu.execute(stmt, params)
	result = cu.fetchone()[0]

	return result != 0

    def __init__(self, dbpath, anonymousReads = False):
	self.db = sqlite.connect(dbpath)
	self.anonReads = False

class InsufficientPermission(Exception):

    pass
<commit_msg>s/pkgName/troveName/<commit_after>from repository import changeset
from repository import fsrepos
import filecontainer
import md5
import os
import sqlite
import tempfile
import xmlshims

class NetworkRepositoryServer(xmlshims.NetworkConvertors):

    def allTroveNames(self, authToken):
	if not self.auth.check(authToken, write = False):
	    raise InsufficientPermission

	return [ x for x in self.iterAllTroveNames(authToken) ]

    def createBranch(self, authToken, newBranch, kind, frozenLocation,
		     troveList):
	if not self.auth.check(authToken, write = True):
	    raise InsufficientPermission

	newBranch = self.toLabel(newBranch)
	if kind == 'v':
	    location = self.toVersion(frozenLocation)
	elif kind == 'l':
	    location = self.toLabel(frozenLocation)
	else:
	    return 0

	self.repos.createBranch(newBranch, location, troveList)
	return 1

    def hasPackage(self, authToken, pkgName):
	if not self.auth.check(authToken, write = False, trove = pkgName):
	    raise InsufficientPermission

	return self.repos.troveStore.hasTrove(pkgName)

    def hasTrove(self, authToken, pkgName, version, flavor):
	if not self.auth.check(authToken, write = False, trove = pkgName):
	    raise InsufficientPermission

	return self.repos.troveStore.hasTrove(pkgName, troveVersion = version,
					troveFlavor = flavor)

    def getTroveVersionList(self, authToken, troveNameList):
	d = {}
	for troveName in troveNameList:
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    d[troveName] = [ x for x in
			    self.repos.troveStore.iterTroveVersions(troveName) ]

	return d

    def getFilesInTrove(self, authToken, troveName, version, flavor,
                        sortByPath = False, withFiles = False):
	if not self.auth.check(authToken, write = False, trove = troveName,
			       label = troveName.branch().label()):
	    raise InsufficientPermission

        gen = self.repos.troveStore.iterFilesInTrove(troveName,
                                               self.toVersion(version),
                                               self.toFlavor(flavor),
                                               sortByPath,
                                               withFiles)
        if withFiles:
            return [ (x[0], x[1], self.fromVersion(x[2]), self.fromFile(x[3]))
                     for x in gen ]
        else:
            return [ (x[0], x[1], self.fromVersion(x[2])) for x in gen ]

    def getFileContents(self, authToken, sha1list):
	# XXX X this isn't properly checked!
	if not self.auth.check(authToken, write = False, trove = troveName,
			       label = troveName.branch().label()):
	    raise InsufficientPermission

	(fd, path) = tempfile.mkstemp(dir = self.tmpPath, suffix = '.cfc-out')
	f = os.fdopen(fd, "w")

	fc = filecontainer.FileContainer(f)
	del f
	d = self.repos.getFileContents(sha1list)

	for sha1 in sha1list:
	    fc.addFile(sha1, d[sha1], "", d[sha1].fullSize)
	fc.close()

	fileName = os.path.basename(path)
	return "%s/%s" % (self.urlBase, fileName[:-4])

    def getAllTroveLeafs(self, authToken, troveNames):
	d = {}
	for troveName in troveNames:
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    d[troveName] = [ x for x in
			    self.repos.troveStore.iterAllTroveLeafs(troveName) ]
	return d

    def getTroveLeavesByLabel(self, authToken, troveNameList, labelStr):
	d = {}
	for troveName in troveNameList:
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    d[troveName] = [ x for x in
			self.repos.troveStore.iterTroveLeafsByLabel(troveName,
								   labelStr) ]

	return d

    def getTroveVersionFlavors(self, authToken, troveDict):
	newD = {}
	for (troveName, versionList) in troveDict.iteritems():
	    if not self.auth.check(authToken, write = False, trove = troveName):
		raise InsufficientPermission

	    innerD = {}
	    for versionStr in versionList:
		innerD[versionStr] = [ self.fromFlavor(x) for x in
		    self.repos.troveStore.iterTroveFlavors(troveName,
						 self.toVersion(versionStr)) ]
	    newD[troveName] = innerD

	return newD

    def getTroveLatestVersion(self, authToken, pkgName, branchStr):
	branch = self.toBranch(branchStr)

	if not self.auth.check(authToken, write = False, trove = pkgName,
			       label = branch.label()):
	    raise InsufficientPermission

        try:
            return self.fromVersion(self.repos.troveStore.troveLatestVersion(pkgName,
					     self.toBranch(branchStr)))
        except KeyError:
            return 0

    def getTroveFlavorsLatestVersion(self, authToken, troveName, branchStr):
	branch = self.toBranch(branchStr)

	if not self.auth.check(authToken, write = False, trove = troveName,
			       label = branch.label()):
	    raise InsufficientPermission

	return self.repos.troveStore.iterTrovePerFlavorLeafs(troveName, branchStr)

    def getChangeSet(self, authToken, chgSetList, recurse, withFiles):
	l = []
	for (name, flavor, old, new, absolute) in chgSetList:
	    newVer = self.toVersion(new)

	    if not self.auth.check(authToken, write = False, trove = name,
				   label = newVer.branch().label()):
		raise InsufficientPermission

	    if old == 0:
		l.append((name, self.toFlavor(flavor), None,
			 self.toVersion(new), absolute))
	    else:
		l.append((name, self.toFlavor(flavor), self.toVersion(old),
			 self.toVersion(new), absolute))

	cs = self.repos.createChangeSet(l, recurse = recurse,
					withFiles = withFiles)
	(fd, path) = tempfile.mkstemp(dir = self.tmpPath, suffix = '.ccs-out')
	os.close(fd)
	cs.writeToFile(path)
	fileName = os.path.basename(path)
	return "%s/%s" % (self.urlBase, fileName[:-4])

    def iterAllTroveNames(self, authToken):
	if not self.auth.check(authToken, write = False):
	    raise InsufficientPermission

	return self.repos.iterAllTroveNames()

    def prepareChangeSet(self, authToken):
	# make sure they have a valid account and permission to commit to
	# *something*
	if not self.auth.check(authToken, write = True):
	    raise InsufficientPermission

	(fd, path) = tempfile.mkstemp(dir = self.tmpPath, suffix = '.ccs-in')
	os.close(fd)
	fileName = os.path.basename(path)
	return "%s/%s" % (self.urlBase, fileName[:-3])

    def commitChangeSet(self, authToken, url):
	assert(url.startswith(self.urlBase))
	fileName = url[len(self.urlBase):] + "-in"
	path = "%s/%s" % (self.tmpPath, fileName)

	try:
	    cs = changeset.ChangeSetFromFile(path)
	finally:
	    pass
	    os.unlink(path)

	# walk through all of the branches this change set commits to
	# and make sure the user has enough permissions for the operation
	items = {}
	for pkgCs in cs.iterNewPackageList():
	    items[(pkgCs.getName(), pkgCs.getNewVersion())] = True
	    if not self.auth.check(authToken, write = True,
			       label = pkgCs.getNewVersion().branch().label(),
			       trove = pkgCs.getName()):
		raise InsufficientPermission

	self.repos.commitChangeSet(cs)

	return True

    def getFileVersion(self, authToken, fileId, version, withContents = 0):
	# XXX needs to authentication against the trove the file is part of,
	# which is unfortunate
	f = self.repos.troveStore.getFile(fileId, self.toVersion(version))
	return self.fromFile(f)

    def checkVersion(self, authToken, clientVersion):
	if not self.auth.check(authToken, write = False):
	    raise InsufficientPermission

        if clientVersion < 0:
            raise RuntimeError, "client is too old"
        return 0

    def __init__(self, path, tmpPath, urlBase, authDbPath):
	self.repos = fsrepos.FilesystemRepository(path)
	self.tmpPath = tmpPath
	self.urlBase = urlBase
	self.auth = NetworkAuthorization(authDbPath)

class NetworkAuthorization:

    def check(self, authToken, write = False, label = None, trove = None):
	if not write and self.anonReads:
	    return True

	if not authToken[0]:
	    return False

	stmt = """
	    SELECT count(*) FROM
	       (SELECT userId as uuserId FROM Users WHERE user=%s AND
		    password=%s)
	    JOIN Permissions ON uuserId=Permissions.userId
	"""
	m = md5.new()
	m.update(authToken[1])
	params = [authToken[0], m.hexdigest()]

	where = []
	if label:
	    where.append(" labelId=(SELECT labelId FROM Labels WHERE " \
			    "label=%s) OR labelId is Null")
	    params.append(label.asString())

	if trove:
	    where.append(" troveNameId=(SELECT troveNameId FROM TroveNames "
			        "WHERE troveName=%s) OR troveNameId is Null" )
	    params.append(trove)

	if write:
	    where.append("write=1")

	if where:
	    stmt += "WHERE " + " AND ".join(where)

	cu = self.db.cursor()
	cu.execute(stmt, params)
	result = cu.fetchone()[0]

	return result != 0

    def __init__(self, dbpath, anonymousReads = False):
	self.db = sqlite.connect(dbpath)
	self.anonReads = False

class InsufficientPermission(Exception):

    pass
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""representation.py

Compute the representation of the different populations in the areal units
provided as an input.

See [Louf:2015]_ for the definition of the different quantities
"""

__all__ = ['representation']

#
# Helper functions
#
def compute_totals(distribution, classes):
    "Compute the number of individuals per class, per unit and in total"
    N_unit = {au:sum([dist_a[cl] for cl in classes]) for au in distribution}
    N_class = {cl:sum([dist_a[cl] for dist_a in distribution.values()] for cl in classes}
    N_tot = sum(N_class,values())
    return N_unit, N_class, N_tot

def single_representation(n, n_unit, N_class, N_tot):
    pass

def single_deviation(n, n_unit, N_class, N_tot):
    pass

#
# Callable functions
#
def representation(distribution, classes=None):

    # Compute the total numbers per class and per individual
    N_unit, N_class, N_tot = compute_totals(distribution)

    representation = {au:{cl:(single_representation(dist_au[cl],
                                                    N_unit[au],
                                                    N_class[cl],
                                                    N_tot),
                              single_deviation(dist_au[cl],
                                               N_unit[au],
                                               N_class[cl],
                                               N_tot) )
                          for cl in classes}
                      for au, dist_au in distribution.iteritems()}

    return representation
<commit_msg>Functions to compute single representation and standard deviation. Finished to code the whole function as well.<commit_after># -*- coding: utf-8 -*-
"""representation.py

Compute the representation of the different populations in the areal units
provided as an input.

See [Louf:2015]_ for the definition of the different quantities
"""
from __future__ import division
import math

__all__ = ['representation']

#
# Helper functions
#
def compute_totals(distribution, classes):
    "Compute the number of individuals per class, per unit and in total"
    N_unit = {au:sum([dist_a[cl] for cl in classes]) for au in distribution}
    N_class = {cl:sum([dist_a[cl] for dist_a in distribution.values()] for cl in classes}
    N_tot = sum(N_class,values())
    return N_unit, N_class, N_tot

def single_representation(n, n_unit, N_class, N_tot):
    return (n/n_unit) / (N_class/N_tot)

def single_deviation(n_unit, N_class, N_tot):
    return math.sqrt( (1/N_class)*((N_tot/n_unit) - 1) )

#
# Callable functions
#
def representation(distribution, classes=None):

    # Compute the total numbers per class and per individual
    N_unit, N_class, N_tot = compute_totals(distribution)

    representation = {au:{cl:(single_representation(dist_au[cl],
                                                    N_unit[au],
                                                    N_class[cl],
                                                    N_tot),
                              single_deviation(N_unit[au],
                                               N_class[cl],
                                               N_tot) )
                          for cl in classes}
                      for au, dist_au in distribution.iteritems()}

    return representation
<|endoftext|>