<commit_before># -*- coding: utf-8 -*-

# OpenFisca -- A versatile microsimulation software
# By: OpenFisca Team <contact@openfisca.fr>
#
# Copyright (C) 2011, 2012, 2013, 2014, 2015 OpenFisca Team
# https://github.com/openfisca
#
# This file is part of OpenFisca.
#
# OpenFisca is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# OpenFisca is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from openfisca_core.columns import AgeCol, DateCol, FloatCol, IntCol, EnumCol
from openfisca_core.enumerations import Enum
from openfisca_core.formulas import (
    make_reference_formula_decorator, dated_function, DatedFormulaColumn, reference_input_variable, SimpleFormulaColumn
    )
from openfisca_france_data.statshelpers import mark_weighted_percentiles

from ..entities import entity_class_by_symbol, Individus, Menages

__all__ = [
    'AgeCol',
    'DateCol',
    'DatedFormulaColumn',
    'dated_function',
    'Enum',
    'EnumCol',
    'FloatCol',
    'Individus',
    'IntCol',
    'mark_weighted_percentiles',
    'Menages',
    'reference_formula',
    'reference_input_variable',
    'SimpleFormulaColumn',
    ]

# Functions and decorators

reference_formula = make_reference_formula_decorator(entity_class_by_symbol = entity_class_by_symbol)
<commit_msg>Move statshelper to openfisca_survey_manager<commit_after># -*- coding: utf-8 -*-

# OpenFisca -- A versatile microsimulation software
# By: OpenFisca Team <contact@openfisca.fr>
#
# Copyright (C) 2011, 2012, 2013, 2014, 2015 OpenFisca Team
# https://github.com/openfisca
#
# This file is part of OpenFisca.
#
# OpenFisca is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# OpenFisca is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from openfisca_core.columns import AgeCol, DateCol, FloatCol, IntCol, EnumCol
from openfisca_core.enumerations import Enum
from openfisca_core.formulas import (
    make_reference_formula_decorator, dated_function, DatedFormulaColumn, reference_input_variable, SimpleFormulaColumn
    )
from openfisca_survey_manager.statshelpers import mark_weighted_percentiles

from ..entities import entity_class_by_symbol, Individus, Menages

__all__ = [
    'AgeCol',
    'DateCol',
    'DatedFormulaColumn',
    'dated_function',
    'Enum',
    'EnumCol',
    'FloatCol',
    'Individus',
    'IntCol',
    'mark_weighted_percentiles',
    'Menages',
    'reference_formula',
    'reference_input_variable',
    'SimpleFormulaColumn',
    ]

# Functions and decorators

reference_formula = make_reference_formula_decorator(entity_class_by_symbol = entity_class_by_symbol)
<|endoftext|>


<commit_before>from rest_framework.relations import RelatedField
from rest_framework_gis import serializers as gis_serializers
from models import *
from rest_framework import serializers

class SourceSerializer(serializers.ModelSerializer):
    class Meta:
        model = Source

# Cultural Trait Variables
class VariableCodeDescriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = VariableCodeDescription
        fields = ('id', 'code', 'description', 'short_description', 'variable')

class VariableDescriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = VariableDescription
        fields = ('id', 'label', 'name', 'codebook_info', 'data_type', 'source', 'index_categories')

class VariableCategorySerializer(serializers.ModelSerializer):
    class Meta:
        model = VariableCategory
        fields = ('id', 'name',)

class VariableDescriptionDetailSerializer(serializers.ModelSerializer):
    index_categories = VariableCategorySerializer(many=True)
    niche_categories = VariableCategorySerializer(many=True)
    class Meta:
        model = VariableDescription
        fields = ('id', 'label', 'name', 'index_categories', 'niche_categories')

class VariableCategoryDetailSerializer(serializers.ModelSerializer):
    # Use a Primary key related field or just get the variable
    index_variables = VariableDescriptionSerializer(many=True)
    niche_variables = VariableDescriptionSerializer(many=True)
    class Meta:
        model = VariableCategory
        fields = ('id', 'name', 'index_variables','niche_variables')

class VariableCodedValueSerializer(serializers.ModelSerializer):
    code_description = VariableCodeDescriptionSerializer(source='code')
    references = SourceSerializer(many=True)
    class Meta:
        model = VariableCodedValue
        fields = ('id', 'variable', 'society', 'coded_value', 'code_description', 'source', 'references', 'subcase', 'focal_year', 'comment')

# ISO Codes
class ISOCodeSerializer(gis_serializers.GeoModelSerializer):
    class Meta:
        model = ISOCode

# Glotto Codes
class GlottoCodeSerializer(serializers.ModelSerializer):
    class Meta:
        model = GlottoCode

# Environmental Data
class EnvironmentalCategorySerializer(serializers.ModelSerializer):
    class Meta:
        model = EnvironmentalCategory

class EnvironmentalVariableSerializer(serializers.ModelSerializer):
    class Meta:
        model = EnvironmentalVariable

class EnvironmentalValueSerializer(serializers.ModelSerializer):
    class Meta:
        model = EnvironmentalValue

class EnvironmentalSerializer(gis_serializers.GeoModelSerializer):
    class Meta:
        model = Environmental

# Languages
class LanguageClassSerializer(serializers.ModelSerializer):
    class Meta:
        model = LanguageClass

class LanguageSerializer(serializers.ModelSerializer):
    language_family = LanguageClassSerializer(source='languageclassification_set.first.class_family')
    glotto_code = serializers.CharField(source='glotto_code.glotto_code')
    iso_code = serializers.CharField(source='iso_code.iso_code')
    class Meta:
        model = Language

class LanguageClassificationSerializer(serializers.ModelSerializer):
    language = LanguageSerializer()
    class Meta:
        model = LanguageClassification

# Societies
class SocietySerializer(gis_serializers.GeoModelSerializer):
    language = LanguageSerializer(source='language')
    source = SourceSerializer(source='source')
    class Meta:
        model = Society
        fields = ('id', 'ext_id', 'xd_id', 'name', 'location','language', 'focal_year', 'source')

# Geographic Regions
class GeographicRegionSerializer(gis_serializers.GeoModelSerializer):
    class Meta:
        model = GeographicRegion
        fields = ('id','level_2_re','count','region_nam','continent','tdwg_code')

class LanguageTreeSerializer(serializers.ModelSerializer):
   languages = LanguageSerializer(source='languages', many=True)
   class Meta:
        model = LanguageTree
        fields = ('id','name','languages', 'newick_string')

SEARCH_LANGUAGE = 'l'
SEARCH_ENVIRONMENTAL = 'e'
SEARCH_VARIABLES = 'v'
SEARCH_GEOGRAPHIC = 'g'

class SocietyResult(object):
    '''
    Encapsulates societies and their related search values for serializing
    '''
    def __init__(self, society):
        self.society = society
        self.variable_coded_values = set()
        self.environmental_values = set()
        self.languages = set()
        self.geographic_regions = set()
    def add_variable_coded_value(self,variable_coded_value):
        self.variable_coded_values.add(variable_coded_value)
    def add_environmental_value(self,environmental_value):
        self.environmental_values.add(environmental_value)
    def add_language(self,language):
        self.languages.add(language)
    def add_geographic_region(self,geographic_region):
        self.geographic_regions.add(geographic_region)
    def includes_criteria(self,criteria=[]):
        result = True
        if SEARCH_VARIABLES in criteria and len(self.variable_coded_values) == 0:
            result = False
        if SEARCH_ENVIRONMENTAL in criteria and len(self.environmental_values) == 0:
            result = False
        if SEARCH_LANGUAGE in criteria and len(self.languages) == 0:
            result = False
        if SEARCH_GEOGRAPHIC in criteria and len(self.geographic_regions) == 0:
            result = False
        return result
class VariableCode(object):
    '''
    Encapsulates societies and their related search values for serializing
    '''
    def __init__(self, codes, variable):
        self.codes = codes
        self.variable = variable

class SocietyResultSet(object):
    '''
    Provides a mapping of Societies to SocietyResult objects
    Used in building the search response
    '''

    def __init__(self):
        # Use a dictionary to map society_id -> SocietyResult
        self._society_results = dict()
        self._codes = dict()
        self.societies = None # not valid until finalize() is called
        # These are the column headers in the search results
        self.variable_descriptions = None
        self.environmental_variables = set()
        self.languages = set()
        self.geographic_regions = set()
        self.language_trees = set()
        #self.var_codes = None

    def _get_society_result(self,society):
        if society.id not in self._society_results.keys():
            self._society_results[society.id] = SocietyResult(society)
        return self._society_results[society.id]

    def add_code(self, codes, variable):
        if variable.id not in self._codes.keys():
            self._codes[variable.id] = VariableCode(codes, variable)
     #   return self._codes[variable.id]

    def add_cultural(self,society,variable_description,codes,variable_coded_value):
        self._get_society_result(society).add_variable_coded_value(variable_coded_value)
        self.add_code(codes, variable_description)

    def add_environmental(self,society,environmental_variable,environmental_value):
        self.environmental_variables.add(environmental_variable)
        self._get_society_result(society).add_environmental_value(environmental_value)

    def add_language(self,society,language):
        self.languages.add(language)
        self._get_society_result(society).add_language(language)

    def add_geographic_region(self,society,geographic_region):
        self.geographic_regions.add(geographic_region)
        self._get_society_result(society).add_geographic_region(geographic_region)

    def add_language_tree(self,language_tree):
        self.language_trees.add(language_tree)

    def finalize(self,criteria):
        self.societies = [x for x in self._society_results.values() if x.includes_criteria(criteria)]
        self.variable_descriptions = [x for x in self._codes.values()]

class VariableCodeSerializer(serializers.Serializer):
    codes = VariableCodeDescriptionSerializer(many=True)
    variable = VariableDescriptionSerializer()

class SocietyResultSerializer(serializers.Serializer):
    '''
    Serializer for the SocietyResult object
    '''
    society = SocietySerializer()
    variable_coded_values = VariableCodedValueSerializer(many=True)
    environmental_values = EnvironmentalValueSerializer(many=True)
    languages = LanguageSerializer(many=True)
    geographic_regions = GeographicRegionSerializer(many=True)

class SocietyResultSetSerializer(serializers.Serializer):
    '''
    Serialize a set of society results and the search criteria
    '''
    # Results contains a society and variable values that matched
    societies = SocietyResultSerializer(many=True)
    # These contain the search parameters
    # variable descriptions -> variable codes
    variable_descriptions = VariableCodeSerializer(many=True)
    # environmental variables -> environmental values
    environmental_variables = EnvironmentalVariableSerializer(many=True)
    # languages - Does not map to a more specific value
    languages = LanguageSerializer(many=True)
    # Geographic Regions - does not map to a more specific value
    geographic_regions = GeographicRegionSerializer(many=True)
    #language trees for this society result set
    language_trees = LanguageTreeSerializer(many=True)
    #var_codes = VariableCodeSerializer(many=True)

class Legend(object):
    def __init__(self, name, svg):
        self.name = name
        self.svg = svg

class LegendSerializer(serializers.Serializer):
    name = serializers.CharField()
    svg = serializers.CharField()

class ZipResultSet(object):
    def __init__(self):
        self.tree = None
        self.name = None
        self.legends = []

    def add_tree(self,tree):
        self.tree = tree

class ZipResultSetSerializer (serializers.Serializer):
    tree = serializers.CharField()
    name = serializers.CharField()
    legends = LegendSerializer(many=True)

<commit_msg>fixed serializers<commit_after>from rest_framework.relations import RelatedField
from rest_framework_gis import serializers as gis_serializers
from models import *
from rest_framework import serializers

class SourceSerializer(serializers.ModelSerializer):
    class Meta:
        model = Source

# Cultural Trait Variables
class VariableCodeDescriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = VariableCodeDescription
        fields = ('id', 'code', 'description', 'short_description', 'variable')

class VariableDescriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = VariableDescription
        fields = ('id', 'label', 'name', 'codebook_info', 'data_type', 'source', 'index_categories')

class VariableCategorySerializer(serializers.ModelSerializer):
    class Meta:
        model = VariableCategory
        fields = ('id', 'name',)

class VariableDescriptionDetailSerializer(serializers.ModelSerializer):
    index_categories = VariableCategorySerializer(many=True)
    niche_categories = VariableCategorySerializer(many=True)
    class Meta:
        model = VariableDescription
        fields = ('id', 'label', 'name', 'index_categories', 'niche_categories')

class VariableCategoryDetailSerializer(serializers.ModelSerializer):
    # Use a Primary key related field or just get the variable
    index_variables = VariableDescriptionSerializer(many=True)
    niche_variables = VariableDescriptionSerializer(many=True)
    class Meta:
        model = VariableCategory
        fields = ('id', 'name', 'index_variables','niche_variables')

class VariableCodedValueSerializer(serializers.ModelSerializer):
    code_description = VariableCodeDescriptionSerializer(source='code')
    references = SourceSerializer(many=True)
    class Meta:
        model = VariableCodedValue
        fields = ('id', 'variable', 'society', 'coded_value', 'code_description', 'source', 'references', 'subcase', 'focal_year', 'comment')

# ISO Codes
class ISOCodeSerializer(gis_serializers.GeoModelSerializer):
    class Meta:
        model = ISOCode

# Glotto Codes
class GlottoCodeSerializer(serializers.ModelSerializer):
    class Meta:
        model = GlottoCode

# Environmental Data
class EnvironmentalCategorySerializer(serializers.ModelSerializer):
    class Meta:
        model = EnvironmentalCategory

class EnvironmentalVariableSerializer(serializers.ModelSerializer):
    class Meta:
        model = EnvironmentalVariable

class EnvironmentalValueSerializer(serializers.ModelSerializer):
    class Meta:
        model = EnvironmentalValue

class EnvironmentalSerializer(gis_serializers.GeoModelSerializer):
    class Meta:
        model = Environmental

# Languages
class LanguageClassSerializer(serializers.ModelSerializer):
    class Meta:
        model = LanguageClass

class LanguageSerializer(serializers.ModelSerializer):
    language_family = LanguageClassSerializer(source='languageclassification_set.first.class_family')
    glotto_code = serializers.CharField(source='glotto_code.glotto_code')
    iso_code = serializers.CharField(source='iso_code.iso_code')
    class Meta:
        model = Language

class LanguageClassificationSerializer(serializers.ModelSerializer):
    language = LanguageSerializer()
    class Meta:
        model = LanguageClassification

# Societies
class SocietySerializer(gis_serializers.GeoModelSerializer):
    language = LanguageSerializer()
    source = SourceSerializer()
    class Meta:
        model = Society
        fields = ('id', 'ext_id', 'xd_id', 'name', 'location','language', 'focal_year', 'source')

# Geographic Regions
class GeographicRegionSerializer(gis_serializers.GeoModelSerializer):
    class Meta:
        model = GeographicRegion
        fields = ('id','level_2_re','count','region_nam','continent','tdwg_code')

class LanguageTreeSerializer(serializers.ModelSerializer):
   languages = LanguageSerializer(source='languages', many=True)
   class Meta:
        model = LanguageTree
        fields = ('id','name','languages', 'newick_string')

SEARCH_LANGUAGE = 'l'
SEARCH_ENVIRONMENTAL = 'e'
SEARCH_VARIABLES = 'v'
SEARCH_GEOGRAPHIC = 'g'

class SocietyResult(object):
    '''
    Encapsulates societies and their related search values for serializing
    '''
    def __init__(self, society):
        self.society = society
        self.variable_coded_values = set()
        self.environmental_values = set()
        self.languages = set()
        self.geographic_regions = set()
    def add_variable_coded_value(self,variable_coded_value):
        self.variable_coded_values.add(variable_coded_value)
    def add_environmental_value(self,environmental_value):
        self.environmental_values.add(environmental_value)
    def add_language(self,language):
        self.languages.add(language)
    def add_geographic_region(self,geographic_region):
        self.geographic_regions.add(geographic_region)
    def includes_criteria(self,criteria=[]):
        result = True
        if SEARCH_VARIABLES in criteria and len(self.variable_coded_values) == 0:
            result = False
        if SEARCH_ENVIRONMENTAL in criteria and len(self.environmental_values) == 0:
            result = False
        if SEARCH_LANGUAGE in criteria and len(self.languages) == 0:
            result = False
        if SEARCH_GEOGRAPHIC in criteria and len(self.geographic_regions) == 0:
            result = False
        return result
class VariableCode(object):
    '''
    Encapsulates societies and their related search values for serializing
    '''
    def __init__(self, codes, variable):
        self.codes = codes
        self.variable = variable

class SocietyResultSet(object):
    '''
    Provides a mapping of Societies to SocietyResult objects
    Used in building the search response
    '''

    def __init__(self):
        # Use a dictionary to map society_id -> SocietyResult
        self._society_results = dict()
        self._codes = dict()
        self.societies = None # not valid until finalize() is called
        # These are the column headers in the search results
        self.variable_descriptions = None
        self.environmental_variables = set()
        self.languages = set()
        self.geographic_regions = set()
        self.language_trees = set()
        #self.var_codes = None

    def _get_society_result(self,society):
        if society.id not in self._society_results.keys():
            self._society_results[society.id] = SocietyResult(society)
        return self._society_results[society.id]

    def add_code(self, codes, variable):
        if variable.id not in self._codes.keys():
            self._codes[variable.id] = VariableCode(codes, variable)
     #   return self._codes[variable.id]

    def add_cultural(self,society,variable_description,codes,variable_coded_value):
        self._get_society_result(society).add_variable_coded_value(variable_coded_value)
        self.add_code(codes, variable_description)

    def add_environmental(self,society,environmental_variable,environmental_value):
        self.environmental_variables.add(environmental_variable)
        self._get_society_result(society).add_environmental_value(environmental_value)

    def add_language(self,society,language):
        self.languages.add(language)
        self._get_society_result(society).add_language(language)

    def add_geographic_region(self,society,geographic_region):
        self.geographic_regions.add(geographic_region)
        self._get_society_result(society).add_geographic_region(geographic_region)

    def add_language_tree(self,language_tree):
        self.language_trees.add(language_tree)

    def finalize(self,criteria):
        self.societies = [x for x in self._society_results.values() if x.includes_criteria(criteria)]
        self.variable_descriptions = [x for x in self._codes.values()]

class VariableCodeSerializer(serializers.Serializer):
    codes = VariableCodeDescriptionSerializer(many=True)
    variable = VariableDescriptionSerializer()

class SocietyResultSerializer(serializers.Serializer):
    '''
    Serializer for the SocietyResult object
    '''
    society = SocietySerializer()
    variable_coded_values = VariableCodedValueSerializer(many=True)
    environmental_values = EnvironmentalValueSerializer(many=True)
    languages = LanguageSerializer(many=True)
    geographic_regions = GeographicRegionSerializer(many=True)

class SocietyResultSetSerializer(serializers.Serializer):
    '''
    Serialize a set of society results and the search criteria
    '''
    # Results contains a society and variable values that matched
    societies = SocietyResultSerializer(many=True)
    # These contain the search parameters
    # variable descriptions -> variable codes
    variable_descriptions = VariableCodeSerializer(many=True)
    # environmental variables -> environmental values
    environmental_variables = EnvironmentalVariableSerializer(many=True)
    # languages - Does not map to a more specific value
    languages = LanguageSerializer(many=True)
    # Geographic Regions - does not map to a more specific value
    geographic_regions = GeographicRegionSerializer(many=True)
    #language trees for this society result set
    language_trees = LanguageTreeSerializer(many=True)
    #var_codes = VariableCodeSerializer(many=True)

class Legend(object):
    def __init__(self, name, svg):
        self.name = name
        self.svg = svg

class LegendSerializer(serializers.Serializer):
    name = serializers.CharField()
    svg = serializers.CharField()

class ZipResultSet(object):
    def __init__(self):
        self.tree = None
        self.name = None
        self.legends = []

    def add_tree(self,tree):
        self.tree = tree

class ZipResultSetSerializer (serializers.Serializer):
    tree = serializers.CharField()
    name = serializers.CharField()
    legends = LegendSerializer(many=True)

<|endoftext|>


<commit_before>#!/usr/bin/python

import board
import game_state
import alpha_beta
import game

from length_lookup_table import *
from utility_stats import *
from utility_calculator import *

class ABState():
    """ Bridge for state, for use by alpha_beta code """
    def __init__(self, parent=None):
        if parent == None:
            self.utility_stats = UtilityStats()
            self.utility_calculator = UtilityCalculator()
        else:
            self.utility_stats = UtilityStats(parent.utility_stats)
            self.utility_calculator = parent.utility_calculator

    def reset_state(self):
        self.utility_stats = UtilityStats()

    def get_black_line_counts(self):
        return self.utility_stats.lines[BLACK]

    def get_white_line_counts(self):
        return self.utility_stats.lines[WHITE]

    def get_takes(self):
        return self.utility_stats.takes

    def get_threats(self):
        return self.utility_stats.threats

    def get_iter(self, to_move):
        return self.utility_stats.search_filter

    def set_state(self, s):
        self.state = s
        s.add_observer(self)

    def to_move_colour(self):
        return self.state.to_move_colour()

    def get_move_number(self):
        return self.state.get_move_number()

    def to_move(self):
        """ This is only to keep the AB code unchanged; the value is unused. """
        return None

    def search_player_colour(self):
        """ The AI player who is performing the search """
        game = self.game()
        return game.to_move_colour()

    # TODO?
    def __repr__(self):
        return ""

    def game(self):
        return self.state.game

    # All the utility calculations belong in the UtilityCalculator
    # TODO: Cache stuff somehow?
    def utility(self):
        # The search_colour is the colour of the
        # AI player doing the search.
        # The turn_colour is the colour of the player to
        # move at the leaf state of the search.
        # eval_colour is the colour that we are assessing
        # the utility value for -
        # we do both for each call to utility
        turn_colour = self.to_move_colour()
        search_colour = self.search_player_colour()
        other_colour = opposite_colour(turn_colour)

        # Check for immediate wins first, then forceable wins
        for win_eval_func in (self.zero_turn_win, self.one_turn_win):
            for eval_colour in (turn_colour, other_colour):
                eval_captured = self.state.get_all_captured()[eval_colour]
                eval_lines = self.utility_stats.lines[eval_colour]

                util, won = win_eval_func(
                    eval_lines, eval_captured, eval_colour, turn_colour)

                if won:
                    self.set_win_for(eval_colour)
                    if search_colour == eval_colour:
                        return util
                    else:
                        return -util

        # No forceable win has been found, so fudge up a score
        util_scores = [None, None, None]

        for eval_colour in (turn_colour, other_colour):
            eval_captured = self.state.get_all_captured()[eval_colour]
            eval_lines = self.utility_stats.lines[eval_colour]
            util = self.utility_score(
                    eval_lines,
                    eval_captured,
                    eval_colour, turn_colour)
            util_scores[eval_colour] = util

        # It is a very significant advantage to have the move
        util_scores[turn_colour] *= 100

        if search_colour == turn_colour:
            return util_scores[turn_colour] - util_scores[other_colour]
        else:
            return util_scores[other_colour] - util_scores[turn_colour]

    def zero_turn_win(self, lines, captured, eval_colour, turn_colour):
        """ Detect a win in this position """
        if lines[4] > 0:
            # This position has been won already, mark it as such so
            # that the search is not continued from this node.
            self.state.set_won_by(eval_colour)
            return self.winning_score(), True

        rules = self.get_rules()
        sfcw = rules.stones_for_capture_win
        ccp = rules.can_capture_pairs

        if sfcw > 0 and ccp:
            captured = self.state.get_all_captured()
            if captured[eval_colour] >= sfcw:
                # This position has been won already, mark it as such so
                # that the search is not continued from this node.
                self.state.set_won_by(eval_colour)
                return self.set_win_for(eval_colour), True

        return 0, False

    def one_turn_win(self, lines, captured, eval_colour, turn_colour):
        """ Detect a forceable win after one turn each """
        rules = self.get_rules()
        sfcw = rules.stones_for_capture_win
        ccp = rules.can_capture_pairs

        if lines[3] > 0:
            if eval_colour == turn_colour:
                # An unanswered line of four out of five will win
                return self.winning_score() / 100, True

            if lines[3] > 1:
                # Two or more lines of four, with no danger of being
                # captured is a win.
                if ccp:
                    if self.get_takes()[opposite_colour(eval_colour)] == 0:
                        return self.winning_score() / 100, True

        if ccp and sfcw > 0:
            # Can win by captures
            if eval_colour == turn_colour:
                if (sfcw - captured) <= 2 and \
                        self.get_takes()[eval_colour] > 0:
                    # eval_colour can take the last pair for a win
                    return self.winning_score() / 100, True
            else:
                if (sfcw - captured) <= 2 and \
                        self.get_takes()[eval_colour] > 2:
                    # eval_colour can take the last pair for a win
                    return self.winning_score() / 100, True

        return 0, False

    def utility_score(self, lines, captured, eval_colour, turn_colour):
        """ Calculate a score for eval_colour for this state. """
        rules = self.get_rules()
        sfcw = rules.stones_for_capture_win
        ccp = rules.can_capture_pairs

        score = 0

        for i in range(len(lines)):
            score *= 120
            rev = 4 - i
            score += lines[rev]

        if ccp:
            if sfcw > 0:
                cc = self.utility_calculator.captured_contrib(captured)
                score += cc
            # else: Captured stones are not worth anything

            # Give takes and threats some value for their ability to help
            # get 5 in a row.
            tc = self.utility_calculator.take_contrib(self.get_takes()[eval_colour])
            score += tc

            tc = self.utility_calculator.threat_contrib(self.get_threats()[eval_colour])
            score += tc

        return score

    def set_win_for(self, colour):
        return self.winning_score()

    # Scale these INFINITIES down to discourage sadistic
    # won game lengthening.
    def winning_score(self):
        # TODO: Sadistic mode for Rich - just multiply by move number ;)
        return INFINITY / self.state.get_move_number()

    # TODO: Rename to get_board
    def board(self):
        return self.state.board

    def get_rules(self):
        return self.game().rules

    def before_set_occ(self, pos, colour):
        self.utility_stats.set_or_reset_occs(self.board(), self.get_rules(), pos, -1)

    def after_set_occ(self, pos, colour):
        self.utility_stats.set_or_reset_occs(self.board(), self.get_rules(), pos, 1)

    def create_state(self, move_pos):
        ab_child = ABState(self)

        # clone the base level state object
        base_child = game_state.GameState(self.state.game, self.state)

        # connect the two (including move hook)
        ab_child.set_state(base_child)

        # make the move for the base (which updates ab_child)
        base_child.make_move(move_pos)

        return ab_child

    def terminal(self):
        return self.state.get_won_by() != game.EMPTY

<commit_msg>Mostly renaming<commit_after>#!/usr/bin/python

import board
import game_state
import alpha_beta
import game

from length_lookup_table import *
from utility_stats import *
from utility_calculator import *

class ABState():
    """ Bridge for state, for use by alpha_beta code """
    def __init__(self, parent=None):
        if parent == None:
            self.utility_stats = UtilityStats()
            self.utility_calculator = UtilityCalculator()
        else:
            self.utility_stats = UtilityStats(parent.utility_stats)
            self.utility_calculator = parent.utility_calculator

    def reset_state(self):
        self.utility_stats = UtilityStats()

    def get_black_line_counts(self):
        return self.utility_stats.lines[BLACK]

    def get_white_line_counts(self):
        return self.utility_stats.lines[WHITE]

    def get_takes(self):
        return self.utility_stats.takes

    def get_threats(self):
        return self.utility_stats.threats

    def get_iter(self, to_move):
        return self.utility_stats.search_filter

    def set_state(self, s):
        self.state = s
        s.add_observer(self)

    def to_move_colour(self):
        return self.state.to_move_colour()

    def get_move_number(self):
        return self.state.get_move_number()

    def to_move(self):
        """ This is only to keep the AB code unchanged; the value is unused. """
        return None

    def search_player_colour(self):
        """ The AI player who is performing the search """
        game = self.game()
        return game.to_move_colour()

    # TODO?
    def __repr__(self):
        return ""

    def game(self):
        return self.state.game

    # All the utility calculations belong in the UtilityCalculator
    # TODO: Cache stuff somehow?
    def utility(self):
        # The search_colour is the colour of the
        # AI player doing the search.
        # The turn_colour is the colour of the player to
        # move at the leaf state of the search.
        # eval_colour is the colour that we are assessing
        # the utility value for -
        # we do both for each call to utility
        turn_colour = self.to_move_colour()
        search_colour = self.search_player_colour()
        other_colour = opposite_colour(turn_colour)

        # Check for immediate wins first, then forceable wins
        for win_eval_func in (self.zero_turn_win, self.one_turn_win):
            for eval_colour in (turn_colour, other_colour):
                eval_captured = self.state.get_all_captured()[eval_colour]
                eval_lines = self.utility_stats.lines[eval_colour]

                util, won = win_eval_func(
                    eval_lines, eval_captured, eval_colour, turn_colour)

                if won:
                    self.set_win_for(eval_colour)
                    if search_colour == eval_colour:
                        return util
                    else:
                        return -util

        # No forceable win has been found, so fudge up a score
        util_scores = [None, None, None]

        for eval_colour in (turn_colour, other_colour):
            eval_captured = self.state.get_all_captured()[eval_colour]
            eval_lines = self.utility_stats.lines[eval_colour]
            util = self.utility_score(
                    eval_lines,
                    eval_captured,
                    eval_colour, turn_colour)
            util_scores[eval_colour] = util

        # It is a very significant advantage to have the move
        util_scores[turn_colour] *= 100

        if search_colour == turn_colour:
            return util_scores[turn_colour] - util_scores[other_colour]
        else:
            return util_scores[other_colour] - util_scores[turn_colour]

    def zero_turn_win(self, eval_lines, eval_captured, eval_colour, turn_colour):
        """ Detect a win in this position """
        if eval_lines[4] > 0:
            # This position has been won already, mark it as such so
            # that the search is not continued from this node.
            self.state.set_won_by(eval_colour)
            return self.winning_score(), True

        rules = self.get_rules()
        sfcw = rules.stones_for_capture_win
        ccp = rules.can_capture_pairs

        if sfcw > 0 and ccp:
            if eval_captured >= sfcw:
                # This position has been won already, mark it as such so
                # that the search is not continued from this node.
                self.state.set_won_by(eval_colour)
                return self.set_win_for(eval_colour), True

        return 0, False

    def one_turn_win(self, eval_lines, eval_captured, eval_colour, turn_colour):
        """ Detect a forceable win after one turn each """
        rules = self.get_rules()
        sfcw = rules.stones_for_capture_win
        ccp = rules.can_capture_pairs

        if eval_lines[3] > 0:
            if eval_colour == turn_colour:
                # An unanswered line of four out of five will win
                return self.winning_score() / 100, True

            if eval_lines[3] > 1:
                # Two or more lines of four, with no danger of being
                # captured is a win.
                if ccp:
                    if self.get_takes()[opposite_colour(eval_colour)] == 0:
                        return self.winning_score() / 100, True

        if ccp and sfcw > 0:
            # Can win by captures
            if eval_colour == turn_colour:
                if (sfcw - eval_captured) <= 2 and \
                        self.get_takes()[eval_colour] > 0:
                    # eval_colour can take the last pair for a win
                    return self.winning_score() / 100, True
            else:
                if (sfcw - eval_captured) <= 2 and \
                        self.get_takes()[eval_colour] > 2:
                    # eval_colour can take the last pair for a win
                    return self.winning_score() / 100, True

        return 0, False

    def utility_score(self, eval_lines, eval_captured, eval_colour, turn_colour):
        """ Calculate a score for eval_colour for this state. """
        rules = self.get_rules()
        sfcw = rules.stones_for_capture_win
        ccp = rules.can_capture_pairs

        score = 0

        for i in range(len(eval_lines)):
            score *= 120
            rev = 4 - i
            score += eval_lines[rev]

        if ccp:
            if sfcw > 0:
                cc = self.utility_calculator.captured_contrib(eval_captured)
                score += cc
            # else: Captured stones are not worth anything

            # Give takes and threats some value for their ability to help
            # get 5 in a row.
            tc = self.utility_calculator.take_contrib(self.get_takes()[eval_colour])
            score += tc

            tc = self.utility_calculator.threat_contrib(self.get_threats()[eval_colour])
            score += tc

        return score

    def set_win_for(self, colour):
        return self.winning_score()

    # Scale these INFINITIES down to discourage sadistic
    # won game lengthening.
    def winning_score(self):
        # TODO: Sadistic mode for Rich - just multiply by move number ;)
        return INFINITY / self.state.get_move_number()

    # TODO: Rename to get_board
    def board(self):
        return self.state.board

    def get_rules(self):
        return self.game().rules

    # TODO: Could these two be moved to utility_stats too?
    def before_set_occ(self, pos, colour):
        self.utility_stats.set_or_reset_occs( \
                self.board(), self.get_rules(), pos, -1)

    def after_set_occ(self, pos, colour):
        self.utility_stats.set_or_reset_occs( \
                self.board(), self.get_rules(), pos, 1)

    def create_state(self, move_pos):
        ab_child = ABState(self)

        # clone the base level state object
        base_child = game_state.GameState(self.state.game, self.state)

        # connect the two (including move hook)
        ab_child.set_state(base_child)

        # make the move for the base (which updates ab_child)
        base_child.make_move(move_pos)

        return ab_child

    def terminal(self):
        return self.state.get_won_by() != game.EMPTY

<|endoftext|>


<commit_before>from django.shortcuts import render, redirect
# from django import forms
from django.http import HttpResponse, HttpResponseNotFound

from standbase.models import *

from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger
from django.core.serializers.json import DjangoJSONEncoder

from django.db.models import Count

from django.views.decorators.csrf import csrf_exempt
from django.views.decorators.http import require_POST
from django.views.decorators.cache import cache_page

from django.utils import timezone
from django.utils.text import slugify
from django.utils.cache import add_never_cache_headers

import datetime
import json
import string

import logging
logger = logging.getLogger('testlogger')

def expire_view_cache(view_name, args=[], namespace=None, key_prefix=None):
    """
    This function allows you to invalidate any view-level cache.
        view_name: view function you wish to invalidate or it's named url pattern
        args: any arguments passed to the view function
        namepace: optioal, if an application namespace is needed
        key prefix: for the @cache_page decorator for the function (if any)

    from: http://stackoverflow.com/questions/2268417/expire-a-view-cache-in-django
    """
    from django.core.urlresolvers import reverse
    from django.http import HttpRequest
    from django.utils.cache import get_cache_key
    from django.core.cache import cache
    # create a fake request object
    request = HttpRequest()
    # Loookup the request path:
    if namespace:
        view_name = namespace + ":" + view_name
    request.path = reverse(view_name, args=args)
    # get cache key, expire if the cached item exists:
    key = get_cache_key(request, key_prefix=key_prefix)
    if key:
        if cache.get(key):
            # Delete the cache entry.
            #
            # Note that there is a possible race condition here, as another
            # process / thread may have refreshed the cache between
            # the call to cache.get() above, and the cache.set(key, None)
            # below.  This may lead to unexpected performance problems under
            # severe load.
            cache.set(key, None, 0)
        return True
    return False

from django.db.models.signals import post_save

def invalidate_api_state(sender, **kwargs):
    expire_view_cache('api_state')

post_save.connect(invalidate_api_state, sender=StandSession)
post_save.connect(invalidate_api_state, sender=Topic)

def get_active_sessions():
    return StandSession.objects.filter(datefinished=None).filter(datelive__gt=timezone.now()-datetime.timedelta(seconds=300)).order_by('-datecreated')

def get_completed_sessions():
    return StandSession.public_objects.exclude(datefinished=None).order_by('-datefinished')[:10]

def index(request):
    return render(request, 'standbase/index.html', {
        'active_sessions': get_active_sessions(),
        'completed_sessions': get_completed_sessions(),
        'trending_topics': Topic.public_objects.exclude(slug='something').annotate(Count('standsession')).order_by('-standsession__count')[:5],
        'total_time': formatted_duration(StandSession.objects.all().aggregate(Sum('duration'))['duration__sum'])
    })

# This is cached for four minutes max
# Cache is invalidated on write to any object, so it should be always up to date
@cache_page(60 * 4)
def api_state(request):
    # .values('id', 'datecreated', 'lat', 'lon', 'datelive', 'datefinished', 'topic__name', 'parsed_geocode')
    response = {
        'active_sessions': [{
            'id': s.id,
            'datecreated': s.datecreated,
            'lat': s.lat,
            'lon': s.lon,
            'datelive': s.datelive,
            'datefinished': s.datefinished,
            'topic__name': s.topic.public and s.topic.name or 'something',
            'parsed_geocode': s.parsed_geocode()
        } for s in get_active_sessions()],
        'completed_sessions': [{
            'id': s.id,
            'datecreated': s.datecreated,
            'lat': s.lat,
            'lon': s.lon,
            'datelive': s.datelive,
            'datefinished': s.datefinished,
            'topic__name': s.topic.name,
            'parsed_geocode': s.parsed_geocode(),
            'get_absolute_url': s.get_absolute_url(),
        } for s in get_completed_sessions()]
    }

    # However we don't want the client to cache this page ever
    response = HttpResponse(json.dumps(response, cls=DjangoJSONEncoder), content_type='application/json')

    add_never_cache_headers(response)

    return response

def session(request, sessionid):
    try:
        s = StandSession.objects.get(id=sessionid)

        if s.datefinished:
            return render(request, 'standbase/stand.html', {
                's': s,
                'total_time': formatted_duration(StandSession.objects.all().aggregate(Sum('duration'))['duration__sum'])
            })
        else:
            return HttpResponseNotFound("Couldn't find such a session.")
    except StandSession.DoesNotExist:
        return HttpResponseNotFound("Couldn't find such a session.")

def topic(request, topic_slug):
    try:
        t = Topic.objects.get(slug=topic_slug, public=True)

        page = request.GET.get('page')

        session_list = StandSession.public_objects.filter(topic=t).exclude(datefinished=None).order_by('-datefinished')
        paginator = Paginator(session_list, 25)

        try:
            sessions = paginator.page(page)
        except PageNotAnInteger:
            # If page is not an integer, deliver first page.
            sessions = paginator.page(1)
        except EmptyPage:
            # If page is out of range (e.g. 9999), deliver last page of results.
            sessions = paginator.page(paginator.num_pages)

        return render(request, 'standbase/topic.html', {
            't': t,
            'sessions': sessions,
            'total_time': formatted_duration(StandSession.objects.all().aggregate(Sum('duration'))['duration__sum'])
        })
    except Topic.DoesNotExist:
        return HttpResponseNotFound("Couldn't find such a topic.")

@require_POST
@csrf_exempt
def catch(request):
    lat = float(request.POST.get('lat', ''))
    lon = float(request.POST.get('lon', ''))

    vendorid = request.POST.get('vendorid', '')

    message = request.POST.get('message', 'something')
    slug = slugify(message)
    try:
        topic = Topic.objects.get(slug=slug)
    except Topic.DoesNotExist:
        topic = Topic.objects.create(name=message, slug=slug)

    s = StandSession.objects.create(lat=lat, lon=lon, vendorid=vendorid, datelive=timezone.now(), topic=topic)

    import django_rq
    django_rq.enqueue(s.retrieve_reverse_geocode)

    response = {
        'sessionid': s.id,
        'secret': s.secret
    }
    return HttpResponse(json.dumps(response), content_type='application/json')

@csrf_exempt
@require_POST
def live(request):
    secret = request.POST.get('secret', '')
    sessionid = request.POST.get('sessionid', '')

    try:
        s = StandSession.objects.get(secret=secret, id=sessionid)

        s.datelive = timezone.now()
        s.save()

        code = 1
    except StandSession.DoesNotExist:
        code = 0

    response = {
        'status': code
    }

    return HttpResponse(json.dumps(response), content_type='application/json')

@csrf_exempt
@require_POST
def done(request):
    secret = request.POST.get('secret', '')
    sessionid = request.POST.get('sessionid', '')

    try:
        s = StandSession.objects.get(secret=secret, id=sessionid)

        if 'duration' in request.POST:
            duration = int(request.POST.get('duration', '0'))

            s.duration = duration
            s.datefinished = s.datecreated + datetime.timedelta(seconds=duration)

        s.save()

        code = 1
    except StandSession.DoesNotExist:
        code = 0

    response = {
        'status': code
    }

    return HttpResponse(json.dumps(response), content_type='application/json')

<commit_msg>whitespace<commit_after>from django.shortcuts import render, redirect
# from django import forms
from django.http import HttpResponse, HttpResponseNotFound

from standbase.models import *

from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger
from django.core.serializers.json import DjangoJSONEncoder

from django.db.models import Count

from django.views.decorators.csrf import csrf_exempt
from django.views.decorators.http import require_POST
from django.views.decorators.cache import cache_page

from django.utils import timezone
from django.utils.text import slugify
from django.utils.cache import add_never_cache_headers

import datetime
import json
import string

import logging
logger = logging.getLogger('testlogger')

def expire_view_cache(view_name, args=[], namespace=None, key_prefix=None):
    """
    This function allows you to invalidate any view-level cache.
        view_name: view function you wish to invalidate or it's named url pattern
        args: any arguments passed to the view function
        namepace: optioal, if an application namespace is needed
        key prefix: for the @cache_page decorator for the function (if any)

    from: http://stackoverflow.com/questions/2268417/expire-a-view-cache-in-django
    """
    from django.core.urlresolvers import reverse
    from django.http import HttpRequest
    from django.utils.cache import get_cache_key
    from django.core.cache import cache
    # create a fake request object
    request = HttpRequest()
    # Loookup the request path:
    if namespace:
        view_name = namespace + ":" + view_name
    request.path = reverse(view_name, args=args)
    # get cache key, expire if the cached item exists:
    key = get_cache_key(request, key_prefix=key_prefix)
    if key:
        if cache.get(key):
            # Delete the cache entry.
            #
            # Note that there is a possible race condition here, as another
            # process / thread may have refreshed the cache between
            # the call to cache.get() above, and the cache.set(key, None)
            # below.  This may lead to unexpected performance problems under
            # severe load.
            cache.set(key, None, 0)
        return True
    return False

from django.db.models.signals import post_save

def invalidate_api_state(sender, **kwargs):
    expire_view_cache('api_state')

post_save.connect(invalidate_api_state, sender=StandSession)
post_save.connect(invalidate_api_state, sender=Topic)

def get_active_sessions():
    return StandSession.objects.filter(datefinished=None).filter(datelive__gt=timezone.now()-datetime.timedelta(seconds=300)).order_by('-datecreated')

def get_completed_sessions():
    return StandSession.public_objects.exclude(datefinished=None).order_by('-datefinished')[:10]

def index(request):
    return render(request, 'standbase/index.html', {
        'active_sessions': get_active_sessions(),
        'completed_sessions': get_completed_sessions(),
        'trending_topics': Topic.public_objects.exclude(slug='something').annotate(Count('standsession')).order_by('-standsession__count')[:5],
        'total_time': formatted_duration(StandSession.objects.all().aggregate(Sum('duration'))['duration__sum'])
    })

# This is cached for four minutes max
# Cache is invalidated on write to any object, so it should be always up to date
@cache_page(60 * 4)
def api_state(request):
    # .values('id', 'datecreated', 'lat', 'lon', 'datelive', 'datefinished', 'topic__name', 'parsed_geocode')
    response = {
        'active_sessions': [{
            'id': s.id,
            'datecreated': s.datecreated,
            'lat': s.lat,
            'lon': s.lon,
            'datelive': s.datelive,
            'datefinished': s.datefinished,
            'topic__name': s.topic.public and s.topic.name or 'something',
            'parsed_geocode': s.parsed_geocode()
        } for s in get_active_sessions()],
        'completed_sessions': [{
            'id': s.id,
            'datecreated': s.datecreated,
            'lat': s.lat,
            'lon': s.lon,
            'datelive': s.datelive,
            'datefinished': s.datefinished,
            'topic__name': s.topic.name,
            'parsed_geocode': s.parsed_geocode(),
            'get_absolute_url': s.get_absolute_url(),
        } for s in get_completed_sessions()]
    }

    # However we don't want the client to cache this page ever
    response = HttpResponse(json.dumps(response, cls=DjangoJSONEncoder), content_type='application/json')

    add_never_cache_headers(response)

    return response

def session(request, sessionid):
    try:
        s = StandSession.objects.get(id=sessionid)

        if s.datefinished:
            return render(request, 'standbase/stand.html', {
                's': s,
                'total_time': formatted_duration(StandSession.objects.all().aggregate(Sum('duration'))['duration__sum'])
            })
        else:
            return HttpResponseNotFound("Couldn't find such a session.")
    except StandSession.DoesNotExist:
        return HttpResponseNotFound("Couldn't find such a session.")

def topic(request, topic_slug):
    try:
        t = Topic.objects.get(slug=topic_slug, public=True)

        page = request.GET.get('page')

        session_list = StandSession.public_objects.filter(topic=t).exclude(datefinished=None).order_by('-datefinished')
        paginator = Paginator(session_list, 25)

        try:
            sessions = paginator.page(page)
        except PageNotAnInteger:
            # If page is not an integer, deliver first page.
            sessions = paginator.page(1)
        except EmptyPage:
            # If page is out of range (e.g. 9999), deliver last page of results.
            sessions = paginator.page(paginator.num_pages)

        return render(request, 'standbase/topic.html', {
            't': t,
            'sessions': sessions,
            'total_time': formatted_duration(StandSession.objects.all().aggregate(Sum('duration'))['duration__sum'])
        })
    except Topic.DoesNotExist:
        return HttpResponseNotFound("Couldn't find such a topic.")

@require_POST
@csrf_exempt
def catch(request):
    lat = float(request.POST.get('lat', ''))
    lon = float(request.POST.get('lon', ''))

    vendorid = request.POST.get('vendorid', '')

    message = request.POST.get('message', 'something')
    slug = slugify(message)
    try:
        topic = Topic.objects.get(slug=slug)
    except Topic.DoesNotExist:
        topic = Topic.objects.create(name=message, slug=slug)

    s = StandSession.objects.create(lat=lat, lon=lon, vendorid=vendorid, datelive=timezone.now(), topic=topic)

    import django_rq
    django_rq.enqueue(s.retrieve_reverse_geocode)

    response = {
        'sessionid': s.id,
        'secret': s.secret
    }
    return HttpResponse(json.dumps(response), content_type='application/json')

@csrf_exempt
@require_POST
def live(request):
    secret = request.POST.get('secret', '')
    sessionid = request.POST.get('sessionid', '')

    try:
        s = StandSession.objects.get(secret=secret, id=sessionid)

        s.datelive = timezone.now()
        s.save()

        code = 1
    except StandSession.DoesNotExist:
        code = 0

    response = {
        'status': code
    }

    return HttpResponse(json.dumps(response), content_type='application/json')

@csrf_exempt
@require_POST
def done(request):
    secret = request.POST.get('secret', '')
    sessionid = request.POST.get('sessionid', '')

    try:
        s = StandSession.objects.get(secret=secret, id=sessionid)

        if 'duration' in request.POST:
            duration = int(request.POST.get('duration', '0'))

            s.duration = duration
            s.datefinished = s.datecreated + datetime.timedelta(seconds=duration)

        s.save()

        code = 1
    except StandSession.DoesNotExist:
        code = 0

    response = {
        'status': code
    }

    return HttpResponse(json.dumps(response), content_type='application/json')

<|endoftext|>


<commit_before># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core functionality for the classical PSHA risk calculator.
"""

from openquake.calculators.risk import general
from openquake.calculators.risk.classical import core as classical
from openquake.utils import stats
from openquake.utils import tasks
from openquake import logs

@tasks.oqtask
@general.with_assets
@stats.count_progress('r')
def classical_bcr(job_id, assets, hazard_getter, hazard_id,
                  bcr_distribution_id,
                  lrem_steps_per_interval,
                  asset_life_expectancy,
                  interest_rate):
    logs.LOG.debug('Implement me')

class ClassicalRiskCalculatorWithBCR(classical.ClassicalRiskCalculator):
    celery_task = classical_bcr

<commit_msg>added core calculator class<commit_after># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core functionality for the classical PSHA risk calculator.
"""

from risklib import api

from openquake.calculators.risk import general
from openquake.calculators.risk.classical import core as classical
from openquake.utils import stats
from openquake.utils import tasks
from openquake import logs
from openquake.db import models

@tasks.oqtask
@general.with_assets
@stats.count_progress('r')
def classical_bcr(job_id, assets, hazard_getter, hazard_id,
                  bcr_distribution_id,
                  lrem_steps_per_interval,
                  asset_life_expectancy,
                  interest_rate):
    """
    Celery task for the BCR risk calculator based on the classical
    calculator.

    Gets the vulnerability models, instantiates risklib calculators
    and stores results to db in a single transaction.

    :param int job_id:
      ID of the currently running job
    :param assets:
      list of Assets to take into account
    :param hazard_getter:
      Strategy used to get the hazard curves
    :param int hazard_id
      ID of the Hazard Output the risk calculation is based on
    :param bcr_distribution_id
      ID of the `openquake.db.models.BCRDistribution` output container used
      to store the computed bcr distribution
    :param int lrem_steps_per_interval
      Steps per interval used to compute the Loss Ratio Exceedance matrix
    :param float interest_rate
      The interest rate used in the Cost Benefit Analysis
    :param float asset_life_expectancy
      The life expectancy used for every asset
    """
    logs.LOG.debug('Implement me')

class ClassicalRiskCalculatorWithBCR(classical.ClassicalRiskCalculator):
    celery_task = classical_bcr

    @property
    def calculation_parameters(self):
        rc = self.job.risk_calculation

        return {
            'lrem_steps_per_interval': rc.lrem_steps_per_interval,
            'interest_rate': rc.interest_rate,
            'asset_life_expectancy': rc.asset_life_expectancy
            }

    @property
    def create_outputs(self):
        return dict(
            bcr_distribution_id=models.BCRDistribution.objects.create(
                output=models.Output.objects.create_output(
                    self.job, "BCR Distribution", "bcr_distribution")).pk)
<|endoftext|>


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

import sys

import csv
import argparse

ALEVEL_FACTOR = 2
EXAM_FACTOR = 4
POINTS_MULTIPLIER = 40
POINTS_DIVIDER = 48

def read_semester_grades(filename):
    """ Reads the semester grades from the given file """
    with open(filename[0], 'r') as grades:
        reader = csv.reader(grades, delimiter=';')
        points = 0
        lines = 0
        for semester in reader:
            lines += 1
            if len(semester) != 9:
                raise IndexError("One of the semesters doesn't have the "
                                 "expected nine grades.\nLength: {} => {}"
                                 .format(len(semester), semester))

            for index, grade in enumerate(semester):
                try:
                    grade = int(grade)
                except ValueError as error:
                    raise ValueError("One of the grades isn't a number. "
                                     "Please correct that.") from error

                if not 0 <= grade <= 15:
                    raise ValueError("Your grades should be between 0 and 15.")

                if index <= 2:
                    # First three are the a-level courses
                    points += grade * ALEVEL_FACTOR
                else:
                    points += grade

        if lines != 4:
            raise IndexError("The *.csv-file doesn't have the expected "
                             "four lines of semester grades.")

        return points

def read_exam_grades(grades):
    """ Reads the exam grades from the argument list """
    for grade in grades:
        if not 0 <= int(grade) <= 15:
            raise ValueError("Your grades should be between 0 and 15.")

    return sum(map(int, grades))

def calculate_average_grade(points):
    """ Calculates the average grade """
    return 17/3 - points/180

def main():
    parser = argparse.ArgumentParser(description='Calculates the average grade'
                                                 ' of your A-levels.')
    parser.add_argument('file', metavar='grades.csv',
                        type=str, nargs=1,
                        help='csv file with semester grades')
    parser.add_argument('grades', metavar='grade', type=int, nargs=5,
                        help='five exam grades', choices=range(0,16))
    if len(sys.argv) is 1:
        parser.print_help()
        sys.exit(0)
    args = parser.parse_args()

    points_block1 = read_semester_grades(args.file)
    points_block1 *= POINTS_MULTIPLIER
    points_block1 = round(points_block1 / POINTS_DIVIDER, 0)

    points_block2 = read_exam_grades(args.grades) * EXAM_FACTOR
    points_sum = points_block1 + points_block2
    average_grade = calculate_average_grade(points_sum)

    print(30 * '=')
    print('Points in Block 1: {}'.format(int(points_block1)))
    print('Points in Block 2: {}'.format(points_block2))
    print('Sum (Block1 + Block2): {}'.format(int(points_sum)))
    print(30 * '=')
    print('Average grade: {}'.format(round(average_grade, 2)))

if __name__ == "__main__":
    try:
        main()
    except (ValueError, IndexError) as error:
        print(error)
        sys.exit(1)<commit_msg>changed according to post https://www.python-forum.de/viewtopic.php?f=9&t=38597#p295918<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys

import csv
import argparse

ALEVEL_FACTOR = 2
EXAM_FACTOR = 4
POINTS_MULTIPLIER = 40
POINTS_DIVIDER = 48

def read_semester_grades(filename):
    """ Reads the semester grades from the given file """
    with open(filename[0], 'r') as grades:
        points = 0

        reader = csv.reader(grades, delimiter=';')
        for lines, semester in enumerate(reader):
            if len(semester) != 9:
                raise IndexError("One of the semesters doesn't have the "
                                 "expected nine grades.\nLength: {} => {}"
                                 .format(len(semester), semester))

            for index, grade in enumerate(semester):
                try:
                    grade = int(grade)
                except ValueError as error:
                    raise ValueError("One of the grades isn't a number. "
                                     "Please correct that.") from error

                if not 0 <= grade <= 15:
                    raise ValueError("Your grades should be between 0 and 15.")

                if index <= 2:
                    # First three are the a-level courses
                    points += grade * ALEVEL_FACTOR
                else:
                    points += grade

        if lines != 3:
            raise IndexError("The *.csv-file doesn't have the expected "
                             "four lines of semester grades.")

        return points

def read_exam_grades(grades):
    """ Reads the exam grades from the argument list """
    for grade in grades:
        if not 0 <= int(grade) <= 15:
            raise ValueError("Your grades should be between 0 and 15.")

    return sum(map(int, grades))

def calculate_average_grade(points):
    """ Calculates the average grade """
    return 17/3 - points/180

def main():
    parser = argparse.ArgumentParser(description='Calculates the average grade'
                                                 ' of your A-levels.')
    parser.add_argument('file', metavar='grades.csv',
                        type=str, nargs=1,
                        help='csv file with semester grades')
    parser.add_argument('grades', metavar='grade', type=int, nargs=5,
                        help='five exam grades', choices=range(0,16))
    if len(sys.argv) is 1:
        parser.print_help()
        sys.exit(0)
    args = parser.parse_args()

    points_block1 = read_semester_grades(args.file)
    points_block1 *= POINTS_MULTIPLIER
    points_block1 = round(points_block1 / POINTS_DIVIDER, 0)

    points_block2 = read_exam_grades(args.grades) * EXAM_FACTOR
    points_sum = points_block1 + points_block2
    average_grade = calculate_average_grade(points_sum)

    print(30 * '=')
    print('Points in Block 1: {}'.format(int(points_block1)))
    print('Points in Block 2: {}'.format(points_block2))
    print('Sum (Block1 + Block2): {}'.format(int(points_sum)))
    print(30 * '=')
    print('Average grade: {}'.format(round(average_grade, 2)))

if __name__ == "__main__":
    try:
        main()
    except (ValueError, IndexError) as error:
        print(error)
        sys.exit(1)<|endoftext|>


<commit_before># Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Client for interacting with Vertex AI.

https://cloud.google.com/vertex-ai/docs/reference/rest
"""
import asyncio
import logging
import math
import os
from typing import Any, Dict, List, Optional, Sequence, Tuple

import attr
from google.cloud import aiplatform
from google.cloud import aiplatform_v1 as aip_v1
from xmanager import xm
from xmanager.cloud import auth
from xmanager.xm import utils
from xmanager.xm_local import executables as local_executables
from xmanager.xm_local import execution as local_execution
from xmanager.xm_local import executors as local_executors
from xmanager.xm_local import status as local_status

_DEFAULT_LOCATION = 'us-central1'
# The only machines available on AI Platform are N1 machines and A2 machines.
# https://cloud.google.com/ai-platform-unified/docs/predictions/machine-types#machine_type_comparison
# TODO: Move lookup to a canonical place.
_MACHINE_TYPE_TO_CPU_RAM = {
    'n1-standard-4': (4, 15 * xm.GiB),
    'n1-standard-8': (8, 30 * xm.GiB),
    'n1-standard-16': (16, 60 * xm.GiB),
    'n1-standard-32': (32, 120 * xm.GiB),
    'n1-standard-64': (64, 240 * xm.GiB),
    'n1-standard-96': (96, 360 * xm.GiB),
    'n1-highmem-2': (2, 13 * xm.GiB),
    'n1-highmem-4': (4, 26 * xm.GiB),
    'n1-highmem-8': (8, 52 * xm.GiB),
    'n1-highmem-16': (16, 104 * xm.GiB),
    'n1-highmem-32': (32, 208 * xm.GiB),
    'n1-highmem-64': (64, 416 * xm.GiB),
    'n1-highmem-96': (96, 624 * xm.GiB),
    'n1-highcpu-16': (16, 14 * xm.GiB),
    'n1-highcpu-32': (32, 28 * xm.GiB),
    'n1-highcpu-64': (64, 57 * xm.GiB),
    'n1-highcpu-96': (96, 86 * xm.GiB),
}
_A100_GPUS_TO_MACHINE_TYPE = {
    1: 'a2-highgpu-1g',
    2: 'a2-highgpu-2g',
    4: 'a2-highgpu-4g',
    8: 'a2-highgpu-8g',
    16: 'a2-megagpu-16g',
}

_CLOUD_TPU_ACCELERATOR_TYPES = {
    xm.ResourceType.TPU_V2: 'TPU_V2',
    xm.ResourceType.TPU_V3: 'TPU_V3',
}

_STATE_TO_STATUS = {
    aip_v1.JobState.JOB_STATE_SUCCEEDED:
        local_status.LocalWorkUnitStatusEnum.COMPLETED,
    aip_v1.JobState.JOB_STATE_CANCELLED:
        local_status.LocalWorkUnitStatusEnum.CANCELLED,
    aip_v1.JobState.JOB_STATE_QUEUED:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_PENDING:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_RUNNING:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_CANCELLING:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_PAUSED:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_FAILED:
        local_status.LocalWorkUnitStatusEnum.FAILED,
}

# Hide noisy warning regarding:
# `file_cache is unavailable when using oauth2client >= 4.0.0 or google-auth`
logging.getLogger('googleapiclient.discovery_cache').setLevel(logging.ERROR)

_default_client = None

def _vertex_job_predicate(job: xm.Job) -> bool:
  return isinstance(job.executor, local_executors.Vertex)

class Client:
  """Client class for interacting with AI Platform (Unified)."""

  def __init__(self,
               project: Optional[str] = None,
               location: str = _DEFAULT_LOCATION) -> None:
    """Create a client.

    Args:
      project: GCP project ID.
      location: Location of the API endpoint. Defaults to region us-central1.
        https://cloud.google.com/vertex-ai/docs/reference/rest#service-endpoint
    """
    self.location = location
    self.project = project or auth.get_project_name()
    self.parent = f'projects/{self.project}/locations/{self.location}'
    aiplatform.init(project=self.project, location=self.location)

  def launch(self, name: str, jobs: Sequence[xm.Job]) -> str:
    """Launch jobs on AI Platform (Unified)."""
    pools = []
    tensorboard, output_dir = self.get_tensorboard_settings(jobs)
    for i, job in enumerate(jobs):
      executable = job.executable
      if not isinstance(executable,
                        local_executables.GoogleContainerRegistryImage):
        raise ValueError('Executable {} has type {}. Executable must be of '
                         'type GoogleContainerRegistryImage'.format(
                             executable, type(executable)))

      args = xm.merge_args(executable.args, job.args).to_list(utils.ARG_ESCAPER)
      env_vars = {**executable.env_vars, **job.env_vars}
      env = [{'name': k, 'value': v} for k, v in env_vars.items()]
      if i == 0 and job.executor.requirements.replicas > 1:
        raise ValueError('The first job in a JobGroup using the Vertex AI '
                         'executor cannot have requirements.replicas > 1.')
      pool = aip_v1.WorkerPoolSpec(
          machine_spec=get_machine_spec(job),
          container_spec=aip_v1.ContainerSpec(
              image_uri=executable.image_path,
              args=args,
              env=env,
          ),
          replica_count=job.executor.requirements.replicas,
      )
      pools.append(pool)

    # TOOD(chenandrew): Vertex Training only allows for 4 worker pools.
    # If Vertex does not implement more worker pools, another work-around
    # is to simply put all jobs into the same worker pool as replicas.
    # Each replica would contain the same image, but would execute different
    # modules based the replica index. A disadvantage of this implementation
    # is that every replica must have the same machine_spec.
    if len(pools) > 4:
      raise ValueError('Cloud Job for xm jobs {} contains {} worker types. '
                       'Only 4 worker types are supported'.format(
                           jobs, len(pools)))

    custom_job = aiplatform.CustomJob(
        project=self.project,
        location=self.location,
        display_name=name,
        worker_pool_specs=pools,
        base_output_dir=output_dir,
        staging_bucket=f'gs://{auth.get_bucket()}',
    )
    # TODO Vertex AI can't use TPUs with SA.
    # https://github.com/deepmind/xmanager/issues/11
    service_account = auth.get_service_account()
    for job in jobs:
      assert isinstance(job.executor, local_executors.Vertex)
      if job.executor.requirements.accelerator in xm.TpuType:
        service_account = None
        break
    custom_job.run(
        sync=False,
        service_account=service_account,
        tensorboard=tensorboard,
        enable_web_access=True,
    )
    custom_job.wait_for_resource_creation()
    print(f'Job launched at: {custom_job._dashboard_uri()}')  # pylint: disable=protected-access
    return custom_job.resource_name

  def get_tensorboard_settings(self, jobs: Sequence[xm.Job]) -> Tuple[str, str]:
    """Get the tensorboard settings for a sequence of Jobs."""
    executors = []
    for job in jobs:
      assert isinstance(job.executor, local_executors.Vertex)
      executors.append(job.executor)
    if all(not executor.tensorboard for executor in executors):
      return '', ''

    if not executors[0].tensorboard:
      raise ValueError(
          'Jobs in this job group must have the same tensorboard settings. ' +
          'jobs[0] has no tensorboard settings.')
    output_dir = executors[0].tensorboard.base_output_directory
    tensorboard = executors[0].tensorboard.name
    for i, executor in enumerate(executors):
      if not executor:
        raise ValueError(
            'Jobs in this job group must have the same tensorboard settings. ' +
            'jobs[i] has no tensorboard settings.')
      if (executor.tensorboard.name != tensorboard or
          executor.tensorboard.base_output_directory != output_dir):
        raise ValueError(
            'Jobs in this job group must have the same tensorboard settings. ' +
            f'jobs[0] has tensorboard = {tensorboard} and output_dir =' +
            f'{output_dir}. jobs[{i}] has tensorboard = ' +
            f'{executor.tensorboard.name} and output_dir = '
            f'{executor.tensorboard.base_output_directory}.')
    if output_dir and not output_dir.startswith('gs://'):
      output_dir = os.path.join('gs://', output_dir)
    return tensorboard, output_dir

  async def wait_for_job(self, job_name: str) -> None:
    job = aiplatform.CustomJob.get(job_name)
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, job._block_until_complete)  # pylint: disable=protected-access

  def cancel(self, job_name: str) -> None:
    aiplatform.CustomJob.get(job_name).cancel()

  async def get_or_create_tensorboard(self, name: str) -> str:
    """Gets or creates a Veretex Tensorboard instance."""
    tensorboard_client = aip_v1.TensorboardServiceAsyncClient(client_options={
        'api_endpoint': f'{self.location}-aiplatform.googleapis.com'
    })
    request = aip_v1.ListTensorboardsRequest(
        parent=self.parent, filter=f'displayName={name}')
    response = await tensorboard_client.list_tensorboards(request)
    async for page in response.pages:
      if page.tensorboards:
        return response.tensorboards[0].name
    return await self.create_tensorboard(name)

  async def create_tensorboard(self, name: str) -> str:
    """Creates a Vertex Tensorboard instance."""
    tensorboard_client = aip_v1.TensorboardServiceAsyncClient(client_options={
        'api_endpoint': f'{self.location}-aiplatform.googleapis.com'
    })
    tensorboard = aip_v1.Tensorboard(display_name=name)
    op = await tensorboard_client.create_tensorboard(
        aip_v1.CreateTensorboardRequest(
            parent=self.parent,
            tensorboard=tensorboard,
        ))
    return (await op.result()).name

  def get_state(self, job_name: str) -> aip_v1.JobState:
    return aiplatform.CustomJob.get(job_name).state

def set_default_client(client: Client) -> None:
  global _default_client
  _default_client = client

def get_default_client():
  global _default_client
  if _default_client is None:
    _default_client = Client()
  return _default_client

def get_machine_spec(job: xm.Job) -> Dict[str, Any]:
  """Get the GCP machine type that best matches the Job's requirements."""
  assert isinstance(job.executor, local_executors.Vertex)
  requirements = job.executor.requirements
  spec = {}
  for resource, value in requirements.task_requirements.items():
    accelerator_type = None
    if resource in xm.GpuType:
      accelerator_type = 'NVIDIA_TESLA_' + str(resource).upper()
    elif resource in xm.TpuType:
      accelerator_type = _CLOUD_TPU_ACCELERATOR_TYPES[resource]
    if accelerator_type:
      # TPU_V2 and TPU_V3 removed in
      # https://github.com/googleapis/python-aiplatform/commit/f3a3d03c8509dc49c24139155a572dacbe954f66
      # Hardcode their values until they are reintroduced.
      if accelerator_type == 'TPU_V2':
        spec['accelerator_type'] = 6
      elif accelerator_type == 'TPU_V3':
        spec['accelerator_type'] = 7
      else:
        spec['accelerator_type'] = aip_v1.AcceleratorType[accelerator_type]
      spec['accelerator_count'] = int(value)
  accelerator = spec.get('accelerator_type', None)
  if accelerator and accelerator == aip_v1.AcceleratorType.NVIDIA_TESLA_A100:
    for (gpus, machine_type) in sorted(_A100_GPUS_TO_MACHINE_TYPE.items()):
      if spec['accelerator_count'] <= gpus:
        spec['machine_type'] = machine_type
        break
    if not spec.get('machine_type', None):
      raise ValueError('a100={} does not fit in any valid machine type'.format(
          spec['accelerator_count']))
  elif accelerator == 6 or accelerator == 7:
    spec['machine_type'] = 'cloud-tpu'
  else:
    spec['machine_type'] = cpu_ram_to_machine_type(
        requirements.task_requirements.get(xm.ResourceType.CPU),
        requirements.task_requirements.get(xm.ResourceType.RAM))
  return spec

@attr.s(auto_attribs=True)
class VertexHandle(local_execution.ExecutionHandle):
  """A handle for referring to the launched container."""

  job_name: str

  async def wait(self) -> None:
    await get_default_client().wait_for_job(self.job_name)

  def stop(self) -> None:
    get_default_client().cancel(self.job_name)

  def get_status(self) -> local_status.LocalWorkUnitStatus:
    state = get_default_client().get_state(self.job_name)
    status = _STATE_TO_STATUS[state]
    return local_status.LocalWorkUnitStatus(status=status)

# Must act on all jobs with `local_executors.Vertex` executor.
def launch(experiment_title: str, work_unit_name: str,
           job_group: xm.JobGroup) -> List[VertexHandle]:
  """Launch Vertex jobs in the job_group and return a handler."""
  jobs = xm.job_operators.collect_jobs_by_filter(job_group,
                                                 _vertex_job_predicate)
  # As client creation may throw, do not initiate it if there are no jobs.
  if not jobs:
    return []

  job_name = get_default_client().launch(
      name=f'{experiment_title}_{work_unit_name}',
      jobs=jobs,
  )
  return [VertexHandle(job_name=job_name)]

def cpu_ram_to_machine_type(cpu: Optional[int], ram: Optional[int]) -> str:
  """Convert a cpu and memory spec into a machine type."""
  cpu = cpu or 0
  ram = ram or 0
  if cpu + ram == 0:
    return 'n1-standard-4'

  optimal_machine_type = ''
  optimal_excess_resources = math.inf
  for machine_type, (machine_cpu,
                     machine_ram) in _MACHINE_TYPE_TO_CPU_RAM.items():
    if machine_cpu >= cpu and machine_ram >= ram:
      excess = machine_cpu + machine_ram - cpu - ram
      if excess < optimal_excess_resources:
        optimal_machine_type = machine_type
        optimal_excess_resources = excess

  if optimal_machine_type:
    return optimal_machine_type
  raise ValueError(
      '(cpu={}, ram={}) does not fit in any valid machine type'.format(
          cpu, ram))
<commit_msg>Fix `get_or_create_tensorboard` docs typo<commit_after># Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Client for interacting with Vertex AI.

https://cloud.google.com/vertex-ai/docs/reference/rest
"""
import asyncio
import logging
import math
import os
from typing import Any, Dict, List, Optional, Sequence, Tuple

import attr
from google.cloud import aiplatform
from google.cloud import aiplatform_v1 as aip_v1
from xmanager import xm
from xmanager.cloud import auth
from xmanager.xm import utils
from xmanager.xm_local import executables as local_executables
from xmanager.xm_local import execution as local_execution
from xmanager.xm_local import executors as local_executors
from xmanager.xm_local import status as local_status

_DEFAULT_LOCATION = 'us-central1'
# The only machines available on AI Platform are N1 machines and A2 machines.
# https://cloud.google.com/ai-platform-unified/docs/predictions/machine-types#machine_type_comparison
# TODO: Move lookup to a canonical place.
_MACHINE_TYPE_TO_CPU_RAM = {
    'n1-standard-4': (4, 15 * xm.GiB),
    'n1-standard-8': (8, 30 * xm.GiB),
    'n1-standard-16': (16, 60 * xm.GiB),
    'n1-standard-32': (32, 120 * xm.GiB),
    'n1-standard-64': (64, 240 * xm.GiB),
    'n1-standard-96': (96, 360 * xm.GiB),
    'n1-highmem-2': (2, 13 * xm.GiB),
    'n1-highmem-4': (4, 26 * xm.GiB),
    'n1-highmem-8': (8, 52 * xm.GiB),
    'n1-highmem-16': (16, 104 * xm.GiB),
    'n1-highmem-32': (32, 208 * xm.GiB),
    'n1-highmem-64': (64, 416 * xm.GiB),
    'n1-highmem-96': (96, 624 * xm.GiB),
    'n1-highcpu-16': (16, 14 * xm.GiB),
    'n1-highcpu-32': (32, 28 * xm.GiB),
    'n1-highcpu-64': (64, 57 * xm.GiB),
    'n1-highcpu-96': (96, 86 * xm.GiB),
}
_A100_GPUS_TO_MACHINE_TYPE = {
    1: 'a2-highgpu-1g',
    2: 'a2-highgpu-2g',
    4: 'a2-highgpu-4g',
    8: 'a2-highgpu-8g',
    16: 'a2-megagpu-16g',
}

_CLOUD_TPU_ACCELERATOR_TYPES = {
    xm.ResourceType.TPU_V2: 'TPU_V2',
    xm.ResourceType.TPU_V3: 'TPU_V3',
}

_STATE_TO_STATUS = {
    aip_v1.JobState.JOB_STATE_SUCCEEDED:
        local_status.LocalWorkUnitStatusEnum.COMPLETED,
    aip_v1.JobState.JOB_STATE_CANCELLED:
        local_status.LocalWorkUnitStatusEnum.CANCELLED,
    aip_v1.JobState.JOB_STATE_QUEUED:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_PENDING:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_RUNNING:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_CANCELLING:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_PAUSED:
        local_status.LocalWorkUnitStatusEnum.RUNNING,
    aip_v1.JobState.JOB_STATE_FAILED:
        local_status.LocalWorkUnitStatusEnum.FAILED,
}

# Hide noisy warning regarding:
# `file_cache is unavailable when using oauth2client >= 4.0.0 or google-auth`
logging.getLogger('googleapiclient.discovery_cache').setLevel(logging.ERROR)

_default_client = None

def _vertex_job_predicate(job: xm.Job) -> bool:
  return isinstance(job.executor, local_executors.Vertex)

class Client:
  """Client class for interacting with AI Platform (Unified)."""

  def __init__(self,
               project: Optional[str] = None,
               location: str = _DEFAULT_LOCATION) -> None:
    """Create a client.

    Args:
      project: GCP project ID.
      location: Location of the API endpoint. Defaults to region us-central1.
        https://cloud.google.com/vertex-ai/docs/reference/rest#service-endpoint
    """
    self.location = location
    self.project = project or auth.get_project_name()
    self.parent = f'projects/{self.project}/locations/{self.location}'
    aiplatform.init(project=self.project, location=self.location)

  def launch(self, name: str, jobs: Sequence[xm.Job]) -> str:
    """Launch jobs on AI Platform (Unified)."""
    pools = []
    tensorboard, output_dir = self.get_tensorboard_settings(jobs)
    for i, job in enumerate(jobs):
      executable = job.executable
      if not isinstance(executable,
                        local_executables.GoogleContainerRegistryImage):
        raise ValueError('Executable {} has type {}. Executable must be of '
                         'type GoogleContainerRegistryImage'.format(
                             executable, type(executable)))

      args = xm.merge_args(executable.args, job.args).to_list(utils.ARG_ESCAPER)
      env_vars = {**executable.env_vars, **job.env_vars}
      env = [{'name': k, 'value': v} for k, v in env_vars.items()]
      if i == 0 and job.executor.requirements.replicas > 1:
        raise ValueError('The first job in a JobGroup using the Vertex AI '
                         'executor cannot have requirements.replicas > 1.')
      pool = aip_v1.WorkerPoolSpec(
          machine_spec=get_machine_spec(job),
          container_spec=aip_v1.ContainerSpec(
              image_uri=executable.image_path,
              args=args,
              env=env,
          ),
          replica_count=job.executor.requirements.replicas,
      )
      pools.append(pool)

    # TOOD(chenandrew): Vertex Training only allows for 4 worker pools.
    # If Vertex does not implement more worker pools, another work-around
    # is to simply put all jobs into the same worker pool as replicas.
    # Each replica would contain the same image, but would execute different
    # modules based the replica index. A disadvantage of this implementation
    # is that every replica must have the same machine_spec.
    if len(pools) > 4:
      raise ValueError('Cloud Job for xm jobs {} contains {} worker types. '
                       'Only 4 worker types are supported'.format(
                           jobs, len(pools)))

    custom_job = aiplatform.CustomJob(
        project=self.project,
        location=self.location,
        display_name=name,
        worker_pool_specs=pools,
        base_output_dir=output_dir,
        staging_bucket=f'gs://{auth.get_bucket()}',
    )
    # TODO Vertex AI can't use TPUs with SA.
    # https://github.com/deepmind/xmanager/issues/11
    service_account = auth.get_service_account()
    for job in jobs:
      assert isinstance(job.executor, local_executors.Vertex)
      if job.executor.requirements.accelerator in xm.TpuType:
        service_account = None
        break
    custom_job.run(
        sync=False,
        service_account=service_account,
        tensorboard=tensorboard,
        enable_web_access=True,
    )
    custom_job.wait_for_resource_creation()
    print(f'Job launched at: {custom_job._dashboard_uri()}')  # pylint: disable=protected-access
    return custom_job.resource_name

  def get_tensorboard_settings(self, jobs: Sequence[xm.Job]) -> Tuple[str, str]:
    """Get the tensorboard settings for a sequence of Jobs."""
    executors = []
    for job in jobs:
      assert isinstance(job.executor, local_executors.Vertex)
      executors.append(job.executor)
    if all(not executor.tensorboard for executor in executors):
      return '', ''

    if not executors[0].tensorboard:
      raise ValueError(
          'Jobs in this job group must have the same tensorboard settings. ' +
          'jobs[0] has no tensorboard settings.')
    output_dir = executors[0].tensorboard.base_output_directory
    tensorboard = executors[0].tensorboard.name
    for i, executor in enumerate(executors):
      if not executor:
        raise ValueError(
            'Jobs in this job group must have the same tensorboard settings. ' +
            'jobs[i] has no tensorboard settings.')
      if (executor.tensorboard.name != tensorboard or
          executor.tensorboard.base_output_directory != output_dir):
        raise ValueError(
            'Jobs in this job group must have the same tensorboard settings. ' +
            f'jobs[0] has tensorboard = {tensorboard} and output_dir =' +
            f'{output_dir}. jobs[{i}] has tensorboard = ' +
            f'{executor.tensorboard.name} and output_dir = '
            f'{executor.tensorboard.base_output_directory}.')
    if output_dir and not output_dir.startswith('gs://'):
      output_dir = os.path.join('gs://', output_dir)
    return tensorboard, output_dir

  async def wait_for_job(self, job_name: str) -> None:
    job = aiplatform.CustomJob.get(job_name)
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, job._block_until_complete)  # pylint: disable=protected-access

  def cancel(self, job_name: str) -> None:
    aiplatform.CustomJob.get(job_name).cancel()

  async def get_or_create_tensorboard(self, name: str) -> str:
    """Gets or creates a Vertex Tensorboard instance."""
    tensorboard_client = aip_v1.TensorboardServiceAsyncClient(client_options={
        'api_endpoint': f'{self.location}-aiplatform.googleapis.com'
    })
    request = aip_v1.ListTensorboardsRequest(
        parent=self.parent, filter=f'displayName={name}')
    response = await tensorboard_client.list_tensorboards(request)
    async for page in response.pages:
      if page.tensorboards:
        return response.tensorboards[0].name
    return await self.create_tensorboard(name)

  async def create_tensorboard(self, name: str) -> str:
    """Creates a Vertex Tensorboard instance."""
    tensorboard_client = aip_v1.TensorboardServiceAsyncClient(client_options={
        'api_endpoint': f'{self.location}-aiplatform.googleapis.com'
    })
    tensorboard = aip_v1.Tensorboard(display_name=name)
    op = await tensorboard_client.create_tensorboard(
        aip_v1.CreateTensorboardRequest(
            parent=self.parent,
            tensorboard=tensorboard,
        ))
    return (await op.result()).name

  def get_state(self, job_name: str) -> aip_v1.JobState:
    return aiplatform.CustomJob.get(job_name).state

def set_default_client(client: Client) -> None:
  global _default_client
  _default_client = client

def get_default_client():
  global _default_client
  if _default_client is None:
    _default_client = Client()
  return _default_client

def get_machine_spec(job: xm.Job) -> Dict[str, Any]:
  """Get the GCP machine type that best matches the Job's requirements."""
  assert isinstance(job.executor, local_executors.Vertex)
  requirements = job.executor.requirements
  spec = {}
  for resource, value in requirements.task_requirements.items():
    accelerator_type = None
    if resource in xm.GpuType:
      accelerator_type = 'NVIDIA_TESLA_' + str(resource).upper()
    elif resource in xm.TpuType:
      accelerator_type = _CLOUD_TPU_ACCELERATOR_TYPES[resource]
    if accelerator_type:
      # TPU_V2 and TPU_V3 removed in
      # https://github.com/googleapis/python-aiplatform/commit/f3a3d03c8509dc49c24139155a572dacbe954f66
      # Hardcode their values until they are reintroduced.
      if accelerator_type == 'TPU_V2':
        spec['accelerator_type'] = 6
      elif accelerator_type == 'TPU_V3':
        spec['accelerator_type'] = 7
      else:
        spec['accelerator_type'] = aip_v1.AcceleratorType[accelerator_type]
      spec['accelerator_count'] = int(value)
  accelerator = spec.get('accelerator_type', None)
  if accelerator and accelerator == aip_v1.AcceleratorType.NVIDIA_TESLA_A100:
    for (gpus, machine_type) in sorted(_A100_GPUS_TO_MACHINE_TYPE.items()):
      if spec['accelerator_count'] <= gpus:
        spec['machine_type'] = machine_type
        break
    if not spec.get('machine_type', None):
      raise ValueError('a100={} does not fit in any valid machine type'.format(
          spec['accelerator_count']))
  elif accelerator == 6 or accelerator == 7:
    spec['machine_type'] = 'cloud-tpu'
  else:
    spec['machine_type'] = cpu_ram_to_machine_type(
        requirements.task_requirements.get(xm.ResourceType.CPU),
        requirements.task_requirements.get(xm.ResourceType.RAM))
  return spec

@attr.s(auto_attribs=True)
class VertexHandle(local_execution.ExecutionHandle):
  """A handle for referring to the launched container."""

  job_name: str

  async def wait(self) -> None:
    await get_default_client().wait_for_job(self.job_name)

  def stop(self) -> None:
    get_default_client().cancel(self.job_name)

  def get_status(self) -> local_status.LocalWorkUnitStatus:
    state = get_default_client().get_state(self.job_name)
    status = _STATE_TO_STATUS[state]
    return local_status.LocalWorkUnitStatus(status=status)

# Must act on all jobs with `local_executors.Vertex` executor.
def launch(experiment_title: str, work_unit_name: str,
           job_group: xm.JobGroup) -> List[VertexHandle]:
  """Launch Vertex jobs in the job_group and return a handler."""
  jobs = xm.job_operators.collect_jobs_by_filter(job_group,
                                                 _vertex_job_predicate)
  # As client creation may throw, do not initiate it if there are no jobs.
  if not jobs:
    return []

  job_name = get_default_client().launch(
      name=f'{experiment_title}_{work_unit_name}',
      jobs=jobs,
  )
  return [VertexHandle(job_name=job_name)]

def cpu_ram_to_machine_type(cpu: Optional[int], ram: Optional[int]) -> str:
  """Convert a cpu and memory spec into a machine type."""
  cpu = cpu or 0
  ram = ram or 0
  if cpu + ram == 0:
    return 'n1-standard-4'

  optimal_machine_type = ''
  optimal_excess_resources = math.inf
  for machine_type, (machine_cpu,
                     machine_ram) in _MACHINE_TYPE_TO_CPU_RAM.items():
    if machine_cpu >= cpu and machine_ram >= ram:
      excess = machine_cpu + machine_ram - cpu - ram
      if excess < optimal_excess_resources:
        optimal_machine_type = machine_type
        optimal_excess_resources = excess

  if optimal_machine_type:
    return optimal_machine_type
  raise ValueError(
      '(cpu={}, ram={}) does not fit in any valid machine type'.format(
          cpu, ram))
<|endoftext|>