<commit_before>import argparse
import sys

from coalib.misc import Constants
from coalib.collecting.Collectors import get_all_bears_names

try:
    from argcomplete.completers import ChoicesCompleter
except ImportError:
    class ChoicesCompleter:

        def __init__(self, *args, **kwargs):
            pass

class CustomFormatter(argparse.RawDescriptionHelpFormatter):
    """
    A Custom Formatter that will keep the metavars in the usage but remove them
    in the more detailed arguments section.
    """

    def _format_action_invocation(self, action):
        if not action.option_strings:
            # For arguemnts that don't have options strings
            metavar, = self._metavar_formatter(action, action.dest)(1)
            return metavar
        else:
            # Option string arguments (like "-f, --files")
            parts = action.option_strings
            return ', '.join(parts)

def default_arg_parser(formatter_class=None):
    """
    This function creates an ArgParser to parse command line arguments.

    :param formatter_class: Formatting the arg_parser output into a specific
                            form. For example: In the manpage format.
    """
    formatter_class = formatter_class or CustomFormatter

    entry_point = sys.argv[0]
    for entry in ['coala-ci', 'coala-dbus', 'coala-format', 'coala-json',
                  'coala-delete-orig']:
        if entry_point.endswith(entry):
            parser_type = entry
            break
    else:
        parser_type = 'coala'

    arg_parser = argparse.ArgumentParser(
        formatter_class=formatter_class,
        prog="coala",
        description="coala is a simple COde AnaLysis Application. Its goal "
                    "is to make static code analysis easy and convenient "
                    "for all languages. coala uses bears, which are analysis "
                    "routines that can be combined arbitrarily.")

    arg_parser.add_argument('TARGETS',
                            nargs='*',
                            help="Sections to be executed exclusively.")
    arg_parser.add_argument('-c',
                            '--config',
                            nargs=1,
                            metavar='FILE',
                            help='Configuration file to be used, defaults to '
                                 + repr(Constants.default_coafile))
    FIND_CONFIG_HELP = ('Attempt to find config file by checking parent '
                        'directories of the current working directory. It is '
                        'assumed that the config file is named '
                        + repr(Constants.default_coafile) + '. This arg is '
                        'ignored if --config is also given')
    arg_parser.add_argument('-F',
                            '--find-config',
                            nargs='?',
                            const=True,
                            metavar='BOOL',
                            help=FIND_CONFIG_HELP)
    arg_parser.add_argument('-I',
                            '--no-config',
                            nargs='?',
                            const=True,
                            metavar='BOOL',
                            help="Run without using any config file")
    arg_parser.add_argument('-f',
                            '--files',
                            nargs='+',
                            metavar='FILE',
                            help='Files that should be checked')
    arg_parser.add_argument('-i',
                            '--ignore',
                            nargs='+',
                            metavar='FILE',
                            help='Files that should be ignored')
    arg_parser.add_argument('--limit-files',
                            nargs='+',
                            metavar='FILE',
                            help='Files that will be analyzed will be '
                                 'restricted to those in the globs listed '
                                 'in this argument as well the files setting')
    arg_parser.add_argument('-b',
                            '--bears',
                            nargs='+',
                            metavar='NAME',
                            help='Names of bears to use').completer =\
        ChoicesCompleter(get_all_bears_names())
    BEAR_DIRS_HELP = 'Additional directories where bears may lie'
    arg_parser.add_argument('-d',
                            '--bear-dirs',
                            nargs='+',
                            metavar='DIR',
                            help=BEAR_DIRS_HELP)
    LOG_LEVEL_HELP = ("Enum('ERROR','INFO','WARNING','DEBUG') to set level of "
                      "log output")
    arg_parser.add_argument('-L',
                            '--log-level',
                            nargs=1,
                            choices=['ERROR', 'INFO', 'WARNING', 'DEBUG'],
                            metavar='ENUM',
                            help=LOG_LEVEL_HELP)
    MIN_SEVERITY_HELP = ("Enum('INFO', 'NORMAL', 'MAJOR') to set the minimal "
                         "result severity.")
    arg_parser.add_argument('-m',
                            '--min-severity',
                            nargs=1,
                            choices=('INFO', 'NORMAL', 'MAJOR'),
                            metavar='ENUM',
                            help=MIN_SEVERITY_HELP)
    SETTINGS_HELP = 'Arbitrary settings in the form of section.key=value'
    arg_parser.add_argument('-S',
                            '--settings',
                            nargs='+',
                            metavar='SETTING',
                            help=SETTINGS_HELP)
    if parser_type == 'coala-json':
        arg_parser.add_argument('--text-logs',
                                nargs='?',
                                const=True,
                                metavar='BOOL',
                                help='Don\'t display logs as json, display '
                                     'them as we normally do in the console.')
        arg_parser.add_argument('-o',
                                '--output',
                                nargs=1,
                                metavar='FILE',
                                help='Write the logs as json to a file '
                                'where filename is specified as argument.')
    if parser_type == 'coala':
        SHOW_BEARS_HELP = ("Display bears and its metadata with the sections "
                           "that they belong to")
        arg_parser.add_argument('-B',
                                '--show-bears',
                                nargs='?',
                                const=True,
                                metavar='BOOL',
                                help=SHOW_BEARS_HELP)
        arg_parser.add_argument('-A',
                                '--show-all-bears',
                                nargs='?',
                                const=True,
                                metavar='BOOL',
                                help="Display all bears.")
        arg_parser.add_argument('-l',
                                '--show-language-bears',
                                nargs='+',
                                metavar='LANG',
                                help="Display all bears for the given "
                                "languages.")
    SAVE_HELP = ('Save the used arguments to a config file at the given path, '
                 "or at the value of -c, which is '.coafile' by default.")
    arg_parser.add_argument('-s',
                            '--save',
                            nargs='?',
                            const=True,
                            metavar='FILE',
                            help=SAVE_HELP)
    TAG_HELP = ('Tag results with a specific name. You can access the results'
                ' later with that tag.')
    arg_parser.add_argument('-t',
                            '--tag',
                            nargs='?',
                            const=True,
                            metavar='STRING',
                            help=TAG_HELP)

    DELETE_TAG_HELP = 'Delete pre-tagged results with tag name.'
    arg_parser.add_argument('-g',
                            '--dtag',
                            nargs='?',
                            const=True,
                            metavar='STRING',
                            help=DELETE_TAG_HELP)

    arg_parser.add_argument("-j",
                            "--jobs",
                            type=int,
                            help="Number of jobs to use in parallel.")

    arg_parser.add_argument('-v',
                            '--version',
                            action='version',
                            version=Constants.VERSION)

    arg_parser.add_argument('-n',
                            '--no-orig',
                            nargs='?',
                            const=True,
                            help="Deactivate creation of .orig files,"
                                 ".orig backup files before applying patches")
    arg_parser.add_argument('-r',
                            '--relpath',
                            nargs='?',
                            const=True,
                            help="return relative paths for files")

    try:  # pragma: no cover
        # Auto completion should be optional, because of somewhat complicated
        # setup.
        import argcomplete
        argcomplete.autocomplete(arg_parser)
    except ImportError:
        pass
    return arg_parser
<commit_msg>DefaultArgParser: Add logical argument groups<commit_after>import argparse
import sys

from coalib.misc import Constants
from coalib.collecting.Collectors import get_all_bears_names

try:
    from argcomplete.completers import ChoicesCompleter
except ImportError:
    class ChoicesCompleter:

        def __init__(self, *args, **kwargs):
            pass

class CustomFormatter(argparse.RawDescriptionHelpFormatter):
    """
    A Custom Formatter that will keep the metavars in the usage but remove them
    in the more detailed arguments section.
    """

    def _format_action_invocation(self, action):
        if not action.option_strings:
            # For arguemnts that don't have options strings
            metavar, = self._metavar_formatter(action, action.dest)(1)
            return metavar
        else:
            # Option string arguments (like "-f, --files")
            parts = action.option_strings
            return ', '.join(parts)

def default_arg_parser(formatter_class=None):
    """
    This function creates an ArgParser to parse command line arguments.

    :param formatter_class: Formatting the arg_parser output into a specific
                            form. For example: In the manpage format.
    """
    formatter_class = formatter_class or CustomFormatter

    entry_point = sys.argv[0]
    for entry in ['coala-ci', 'coala-dbus', 'coala-format', 'coala-json',
                  'coala-delete-orig']:
        if entry_point.endswith(entry):
            parser_type = entry
            break
    else:
        parser_type = 'coala'

    arg_parser = argparse.ArgumentParser(
        formatter_class=formatter_class,
        prog="coala",
        description="coala is a simple COde AnaLysis Application. Its goal "
                    "is to make static code analysis easy and convenient "
                    "for all languages. coala uses bears, which are analysis "
                    "routines that can be combined arbitrarily.",
        # Use our own help so that we can put it in the group we want
        add_help=False)

    arg_parser.add_argument('TARGETS',
                            nargs='*',
                            help="Sections to be executed exclusively.")

    info_group = arg_parser.add_argument_group('Info')

    info_group.add_argument('-h',
                            '--help',
                            action='help',
                            help='show this help message and exit')

    info_group.add_argument('-v',
                            '--version',
                            action='version',
                            version=Constants.VERSION)

    config_group = arg_parser.add_argument_group('Configuration')

    CONFIG_HELP = ('Configuration file to be used, defaults to '
                   + repr(Constants.default_coafile))
    config_group.add_argument('-c',
                              '--config',
                              nargs=1,
                              metavar='FILE',
                              help=CONFIG_HELP)

    FIND_CONFIG_HELP = ('Attempt to find config file by checking parent '
                        'directories of the current working directory. It is '
                        'assumed that the config file is named '
                        + repr(Constants.default_coafile) + '. This arg is '
                        'ignored if --config is also given')
    config_group.add_argument('-F',
                              '--find-config',
                              nargs='?',
                              const=True,
                              metavar='BOOL',
                              help=FIND_CONFIG_HELP)

    config_group.add_argument('-I',
                              '--no-config',
                              nargs='?',
                              const=True,
                              metavar='BOOL',
                              help="Run without using any config file")

    SAVE_HELP = ('Save the used arguments to a config file at the given path, '
                 "or at the value of -c, which is '.coafile' by default.")
    config_group.add_argument('-s',
                              '--save',
                              nargs='?',
                              const=True,
                              metavar='FILE',
                              help=SAVE_HELP)

    inputs_group = arg_parser.add_argument_group('Inputs')

    inputs_group.add_argument('-b',
                              '--bears',
                              nargs='+',
                              metavar='NAME',
                              help='Names of bears to use').completer =\
        ChoicesCompleter(get_all_bears_names())

    inputs_group.add_argument('-f',
                              '--files',
                              nargs='+',
                              metavar='FILE',
                              help='Files that should be checked')

    inputs_group.add_argument('-i',
                              '--ignore',
                              nargs='+',
                              metavar='FILE',
                              help='Files that should be ignored')

    LIMIT_FILES_HELP = ('Files that will be analyzed will be restricted to '
                        'those in the globs listed in this argument as well '
                        'the files setting')
    inputs_group.add_argument('--limit-files',
                              nargs='+',
                              metavar='FILE',
                              help=LIMIT_FILES_HELP)

    BEAR_DIRS_HELP = 'Additional directories where bears may lie'
    inputs_group.add_argument('-d',
                              '--bear-dirs',
                              nargs='+',
                              metavar='DIR',
                              help=BEAR_DIRS_HELP)

    outputs_group = arg_parser.add_argument_group('Outputs')

    LOG_LEVEL_HELP = ("Enum('ERROR','INFO','WARNING','DEBUG') to set level of "
                      "log output")
    outputs_group.add_argument('-L',
                               '--log-level',
                               nargs=1,
                               choices=['ERROR', 'INFO', 'WARNING', 'DEBUG'],
                               metavar='ENUM',
                               help=LOG_LEVEL_HELP)

    MIN_SEVERITY_HELP = ("Enum('INFO', 'NORMAL', 'MAJOR') to set the minimal "
                         "result severity.")
    outputs_group.add_argument('-m',
                               '--min-severity',
                               nargs=1,
                               choices=('INFO', 'NORMAL', 'MAJOR'),
                               metavar='ENUM',
                               help=MIN_SEVERITY_HELP)

    # The following are "coala" specific arguments
    if parser_type == 'coala':
        SHOW_BEARS_HELP = ("Display bears and its metadata with the sections "
                           "that they belong to")
        outputs_group.add_argument('-B',
                                   '--show-bears',
                                   nargs='?',
                                   const=True,
                                   metavar='BOOL',
                                   help=SHOW_BEARS_HELP)

        outputs_group.add_argument('-A',
                                   '--show-all-bears',
                                   nargs='?',
                                   const=True,
                                   metavar='BOOL',
                                   help="Display all bears.")

        SHOW_LANGUAGE_BEARS = ("Display all bears for the given languages.")
        outputs_group.add_argument('-l',
                                   '--show-language-bears',
                                   nargs='+',
                                   metavar='LANG',
                                   help=SHOW_LANGUAGE_BEARS)

    # The following are "coala-json" specific arguments
    if parser_type == 'coala-json':
        TEXT_LOGS_HELP = ("Don't display logs as json, display them as we "
                          "normally do in the console.")
        outputs_group.add_argument('--text-logs',
                                   nargs='?',
                                   const=True,
                                   metavar='BOOL',
                                   help=TEXT_LOGS_HELP)

        OUTPUT_HELP = ('Write the logs as json to a file where filename is '
                       'specified as argument.')
        outputs_group.add_argument('-o',
                                   '--output',
                                   nargs=1,
                                   metavar='FILE',
                                   help=OUTPUT_HELP)

        outputs_group.add_argument('-r',
                                   '--relpath',
                                   nargs='?',
                                   const=True,
                                   help="Return relative paths for files")

    misc_group = arg_parser.add_argument_group('Miscellaneous')

    SETTINGS_HELP = 'Arbitrary settings in the form of section.key=value'
    misc_group.add_argument('-S',
                            '--settings',
                            nargs='+',
                            metavar='SETTING',
                            help=SETTINGS_HELP)

    misc_group.add_argument("-j",
                            "--jobs",
                            type=int,
                            help="Number of jobs to use in parallel.")

    NO_ORIG_HELP = ('Deactivate creation of .orig files and .orig backup '
                    'files before applying patches')
    misc_group.add_argument('-n',
                            '--no-orig',
                            nargs='?',
                            const=True,
                            help=NO_ORIG_HELP)

    deprecated_group = arg_parser.add_argument_group('Deprecated')

    TAG_HELP = ('Tag results with a specific name. You can access the results '
                'later with that tag.')
    deprecated_group.add_argument('-t',
                                  '--tag',
                                  nargs='?',
                                  const=True,
                                  metavar='STRING',
                                  help=TAG_HELP)

    DELETE_TAG_HELP = 'Delete pre-tagged results with tag name.'
    deprecated_group.add_argument('-g',
                                  '--dtag',
                                  nargs='?',
                                  const=True,
                                  metavar='STRING',
                                  help=DELETE_TAG_HELP)

    try:  # pragma: no cover
        # Auto completion should be optional, because of somewhat complicated
        # setup.
        import argcomplete
        argcomplete.autocomplete(arg_parser)
    except ImportError:
        pass
    return arg_parser
<|endoftext|>


<commit_before>"""
Heatmapping algorithm
"""
from __future__ import division

import os
import pandas as pd

import cea.inputlocator
import cea.globalvar
import cea.config
from cea.interfaces.arcgis.modules import arcpy

__author__ = "Jimeno A. Fonseca"
__copyright__ = "Copyright 2015, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Jimeno A. Fonseca", "Daren Thomas"]
__license__ = "MIT"
__version__ = "0.1"
__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

def heatmaps(locator, analysis_fields, file_to_analyze):
    """
    algorithm to calculate heat maps out of n variables of interest

    :param file_to_analyze: path to the file with the variables to be mapped. This is used as the ``in_rows`` parameter
        to ``arcpy.CopyRows_management``, so anything supported by that tool could, in theory, work, but
        we use the csv files. Either the ``Total_Demand.csv`` file produced by the demand script or one of
        the files in the emissions results folder.
    :param analysis_fields:  when the path variables is selected, an array of n variables 'string'
        will be elaborated based on the selection of n input fields. For this a
        form like the one used in the Arcgis function ``merge/FieldMap`` could
        be used.

    :returns: - heat map of variable n: .tif
              - heat map file per variable of interest n.
    """
    arcpy.env.overwriteOutput = True
    arcpy.CheckOutExtension("Spatial")

    # figure out folder for the results based on file to analyze
    if locator.are_equal(file_to_analyze, locator.get_total_demand()):
        path_results = locator._ensure_folder(locator.get_plots_folder('heatmaps'))
    elif locator.are_equal(os.path.dirname(file_to_analyze), locator.get_lca_emissions_results_folder()):
        path_results = locator._ensure_folder(locator.get_plots_folder('heatmaps'))
    else:
        raise ValueError(
            'file_to_analyze must be either the demand totals file or a file in the emissions results folder: %s'
            % file_to_analyze)

    # local variables
    # create dbf file
    tempfile_name = "tempfile"
    tempfile = locator.get_temporary_file('tempfile.shp')
    tempfile_db_name = "data"
    tempfile_db = locator.get_temporary_file('data.dbf')

    gis_field_lookup = {}  # map csv_field -> gis_field

    # create a csv file in temp folder with gis_field names and only the analysis-fields (+ 'Name')
    export_fields = ['Name'] + analysis_fields
    file_to_analyse_df = pd.read_csv(file_to_analyze)[export_fields]
    file_to_analyse_df.columns = [get_gis_field(c, gis_field_lookup) for c in file_to_analyse_df.columns]
    tempfile_csv = locator.get_temporary_file('data.csv')
    file_to_analyse_df.to_csv(tempfile_csv)

    print('Copying rows from tempfile_csv={tempfile_csv} to out_table={tempfile_db}'.format(
        tempfile_csv=tempfile_csv, tempfile_db=tempfile_db))
    if os.path.exists(tempfile_db):
        os.remove(tempfile_db)
    arcpy.CopyRows_management(tempfile_csv, out_table=tempfile_db)

    arcpy.FeatureToPoint_management(locator.get_zone_geometry(), tempfile, "CENTROID")
    arcpy.MakeFeatureLayer_management(tempfile, "lyr", "#", "#")

    for csv_field in analysis_fields:
        gis_field = get_gis_field(csv_field, gis_field_lookup)
        arcpy.AddField_management("lyr", gis_field, "DOUBLE", "#", "#", "#", "#", "NULLABLE", "NON_REQUIRED", "#")

    # vector.append([])
    arcpy.AddJoin_management("lyr", "Name", tempfile_db, "Name", "KEEP_ALL")
    for csv_field in analysis_fields:
        gis_field = get_gis_field(csv_field, gis_field_lookup)
        arcpy.CalculateField_management(in_table="lyr", field="%(tempfile_name)s.%(gis_field)s" % locals(),
                                        expression="calc_non_null(!%(tempfile_db_name)s.%(gis_field)s!)" % locals(),
                                        expression_type="PYTHON_9.3",
                                        code_block="def calc_non_null(x):\n     if x is None:\n         return 0\n     elif x == '':\n         return 0\n     else:\n         return x\n")

    # calculate heatmaps
    for csv_field in analysis_fields:
        gis_field = get_gis_field(csv_field, gis_field_lookup)
        arcpy.gp.Idw_sa(tempfile, gis_field, os.path.join(path_results, gis_field), "1", "2", "VARIABLE 12")

def get_gis_field(csv_field, gis_field_lookup):
    """return a (max) 10 character representation of csv_field that is unique to the list of analysis fields"""
    import string
    if csv_field in gis_field_lookup:
        return gis_field_lookup[csv_field]
    gis_field_set = set(gis_field_lookup.values())
    gis_field = csv_field[:10]
    chars = iter(string.ascii_uppercase)
    while gis_field in gis_field_set:
        letters = list(gis_field)
        try:
            letters[-1] = next(chars)
        except StopIteration:
            raise Exception('Too many fields for analysis')
        gis_field = ''.join(letters)
    return gis_field

def main(config):
    assert os.path.exists(config.scenario), 'Scenario not found: %s' % config.scenario
    locator = cea.inputlocator.InputLocator(scenario=config.scenario)

    print('Running heatmaps with scenario = %s' % config.scenario)
    # print('Running heatmaps with file-to-analyze = %s' % config.heatmaps.file_to_analyze)
    analysis_fields = ['QH_sys_MWhyr', 'QC_sys_MWhyr', 'E_sys_MWhyr']
    # print('Running heatmaps with analysis-fields = %s' % config.heatmaps.analysis_fields)
    file_to_analyze = locator.get_total_demand()

    heatmaps(locator=locator, analysis_fields=analysis_fields, file_to_analyze=file_to_analyze)

if __name__ == '__main__':
    main(cea.config.Configuration())
<commit_msg>Add LCA to heatmaps<commit_after>"""
Heatmapping algorithm
"""
from __future__ import division

import os
import pandas as pd

import cea.inputlocator
import cea.globalvar
import cea.config
from cea.interfaces.arcgis.modules import arcpy

__author__ = "Jimeno A. Fonseca"
__copyright__ = "Copyright 2015, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Jimeno A. Fonseca", "Daren Thomas"]
__license__ = "MIT"
__version__ = "0.1"
__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

def heatmaps(locator, analysis_fields, file_to_analyze):
    """
    algorithm to calculate heat maps out of n variables of interest

    :param file_to_analyze: path to the file with the variables to be mapped. This is used as the ``in_rows`` parameter
        to ``arcpy.CopyRows_management``, so anything supported by that tool could, in theory, work, but
        we use the csv files. Either the ``Total_Demand.csv`` file produced by the demand script or one of
        the files in the emissions results folder.
    :param analysis_fields:  when the path variables is selected, an array of n variables 'string'
        will be elaborated based on the selection of n input fields. For this a
        form like the one used in the Arcgis function ``merge/FieldMap`` could
        be used.

    :returns: - heat map of variable n: .tif
              - heat map file per variable of interest n.
    """
    arcpy.env.overwriteOutput = True
    arcpy.CheckOutExtension("Spatial")

    # figure out folder for the results based on file to analyze
    if locator.are_equal(file_to_analyze, locator.get_total_demand()):
        path_results = locator._ensure_folder(locator.get_plots_folder('heatmaps'))
    elif locator.are_equal(os.path.dirname(file_to_analyze), locator.get_lca_emissions_results_folder()):
        path_results = locator._ensure_folder(locator.get_plots_folder('heatmaps'))
    else:
        raise ValueError(
            'file_to_analyze must be either the demand totals file or a file in the emissions results folder: %s'
            % file_to_analyze)

    # local variables
    # create dbf file
    tempfile_name = "tempfile"
    tempfile = locator.get_temporary_file('tempfile.shp')
    tempfile_db_name = "data"
    tempfile_db = locator.get_temporary_file('data.dbf')

    gis_field_lookup = {}  # map csv_field -> gis_field

    # create a csv file in temp folder with gis_field names and only the analysis-fields (+ 'Name')
    export_fields = ['Name'] + analysis_fields
    file_to_analyse_df = pd.read_csv(file_to_analyze)[export_fields]
    file_to_analyse_df.columns = [get_gis_field(c, gis_field_lookup) for c in file_to_analyse_df.columns]
    tempfile_csv = locator.get_temporary_file('data.csv')
    file_to_analyse_df.to_csv(tempfile_csv)

    print('Copying rows from tempfile_csv={tempfile_csv} to out_table={tempfile_db}'.format(
        tempfile_csv=tempfile_csv, tempfile_db=tempfile_db))
    if os.path.exists(tempfile_db):
        os.remove(tempfile_db)
    arcpy.CopyRows_management(tempfile_csv, out_table=tempfile_db)

    arcpy.FeatureToPoint_management(locator.get_zone_geometry(), tempfile, "CENTROID")
    arcpy.MakeFeatureLayer_management(tempfile, "lyr", "#", "#")

    for csv_field in analysis_fields:
        gis_field = get_gis_field(csv_field, gis_field_lookup)
        arcpy.AddField_management("lyr", gis_field, "DOUBLE", "#", "#", "#", "#", "NULLABLE", "NON_REQUIRED", "#")

    # vector.append([])
    arcpy.AddJoin_management("lyr", "Name", tempfile_db, "Name", "KEEP_ALL")
    for csv_field in analysis_fields:
        gis_field = get_gis_field(csv_field, gis_field_lookup)
        arcpy.CalculateField_management(in_table="lyr", field="%(tempfile_name)s.%(gis_field)s" % locals(),
                                        expression="calc_non_null(!%(tempfile_db_name)s.%(gis_field)s!)" % locals(),
                                        expression_type="PYTHON_9.3",
                                        code_block="def calc_non_null(x):\n     if x is None:\n         return 0\n     elif x == '':\n         return 0\n     else:\n         return x\n")

    # calculate heatmaps
    for csv_field in analysis_fields:
        gis_field = get_gis_field(csv_field, gis_field_lookup)
        arcpy.gp.Idw_sa(tempfile, gis_field, os.path.join(path_results, gis_field), "1", "2", "VARIABLE 12")

def get_gis_field(csv_field, gis_field_lookup):
    """return a (max) 10 character representation of csv_field that is unique to the list of analysis fields"""
    import string
    if csv_field in gis_field_lookup:
        return gis_field_lookup[csv_field]
    gis_field_set = set(gis_field_lookup.values())
    gis_field = csv_field[:10]
    chars = iter(string.ascii_uppercase)
    while gis_field in gis_field_set:
        letters = list(gis_field)
        try:
            letters[-1] = next(chars)
        except StopIteration:
            raise Exception('Too many fields for analysis')
        gis_field = ''.join(letters)
    return gis_field

def main(config):
    assert os.path.exists(config.scenario), 'Scenario not found: %s' % config.scenario
    locator = cea.inputlocator.InputLocator(scenario=config.scenario)

    print('Running heatmaps with scenario = %s' % config.scenario)
    # print('Running heatmaps with file-to-analyze = %s' % config.heatmaps.file_to_analyze)
    analysis_fields = ['QH_sys_MWhyr', 'QC_sys_MWhyr', 'E_sys_MWhyr']
    # print('Running heatmaps with analysis-fields = %s' % config.heatmaps.analysis_fields)
    file_to_analyze = locator.get_total_demand()

    heatmaps(locator=locator, analysis_fields=analysis_fields, file_to_analyze=file_to_analyze)

    analysis_fields = ['O_ghg_ton', 'O_ghg_kgm2', 'O_nre_pen_GJ', 'O_nre_pen_MJm2']
    file_to_analyze = locator.get_lca_operation()

    heatmaps(locator=locator, analysis_fields=analysis_fields, file_to_analyze=file_to_analyze)

if __name__ == '__main__':
    main(cea.config.Configuration())
<|endoftext|>


<commit_before>"""
This contains functions and classes used to evaluate if images are acceptable (do not show improper content, etc), and
to send them to S3.
"""

from PIL import Image
import urlparse
import requests
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from django.conf import settings
import pickle
import logging

log = logging.getLogger(__name__)

#Domains where any image linked to can be trusted to have acceptable content.
TRUSTED_IMAGE_DOMAINS = [
    'wikipedia.com',
    'wikipedia.net',
    'wikipedia.org',
    'edxuploads.s3.amazonaws.com'
]

#Suffixes that are allowed in image urls
ALLOWABLE_IMAGE_SUFFIXES = [
    'jpg',
    'png',
    'gif'
]

#Maximum allowed dimensions (x and y) for an uploaded image
MAX_ALLOWED_IMAGE_DIM = 1500

#Dimensions to which image is resized before it is evaluated for color count, etc
MAX_IMAGE_DIM = 150

#Maximum number of colors that should be counted in ImageProperties
MAX_COLORS_TO_COUNT = 16

#Maximum number of colors allowed in an uploaded image
MAX_COLORS = 20

class ImageProperties(object):
    """
    Class to check properties of an image and to validate if they are allowed.
    """
    def __init__(self, image):
        """
        Initializes class variables
        @param image: Image object (from PIL)
        @return: None
        """
        self.image = image
        image_size = self.image.size
        self.image_too_large = False
        if image_size[0] > MAX_ALLOWED_IMAGE_DIM or image_size[1] > MAX_ALLOWED_IMAGE_DIM:
            self.image_too_large = True
        if image_size[0] > MAX_IMAGE_DIM or image_size[1] > MAX_IMAGE_DIM:
            self.image = self.image.resize((MAX_IMAGE_DIM, MAX_IMAGE_DIM))
            self.image_size = self.image.size

    def count_colors(self):
        """
        Counts the number of colors in an image, and matches them to the max allowed
        @return: boolean true if color count is acceptable, false otherwise
        """
        colors = self.image.getcolors(MAX_COLORS_TO_COUNT)
        if colors is None:
            color_count = MAX_COLORS_TO_COUNT
        else:
            color_count = len(colors)

        too_many_colors = (color_count <= MAX_COLORS)
        return too_many_colors

    def get_skin_ratio(self):
        """
        Gets the ratio of skin tone colors in an image
        @return: True if the ratio is low enough to be acceptable, false otherwise
        """
        colors = self.image.getcolors(MAX_COLORS_TO_COUNT)
        is_okay = True
        if colors is not None:
            skin = sum([count for count, rgb in colors if
                        rgb[0] > 60 and rgb[1] < (rgb[0] * 0.85) and rgb[2] < (rgb[0] * 0.7) and rgb[1] > (rgb[0] * 0.4) and
                        rgb[2] > (rgb[0] * 0.2)])
            bad_color_val = float(skin) / len(colors)
            if bad_color_val > .4:
                is_okay = False

        return is_okay

    def run_tests(self):
        """
        Does all available checks on an image to ensure that it is okay (size, skin ratio, colors)
        @return: Boolean indicating whether or not image passes all checks
        """
        image_is_okay = False
        try:
            image_is_okay = self.count_colors() and self.get_skin_ratio() and not self.image_too_large
        except:
            pass

        return image_is_okay

class URLProperties(object):
    """
    Checks to see if a URL points to acceptable content.  Added to check if students are submitting reasonable
    links to the peer grading image functionality of the external grading service.
    """
    def __init__(self, url_string):
        self.url_string = url_string

    def check_if_parses(self):
        """
        Check to see if a URL parses properly
        @return: success (True if parses, false if not)
        """
        success = False
        try:
            self.parsed_url = urlparse.urlparse(url_string)
            success = True
        except:
            pass

        return success

    def check_suffix(self):
        """
        Checks the suffix of a url to make sure that it is allowed
        @return: True if suffix is okay, false if not
        """
        good_suffix = False
        for suffix in ALLOWABLE_IMAGE_SUFFIXES:
            if self.url_string.endswith(suffix):
                good_suffix = True
                break
        return good_suffix

    def run_tests(self):
        """
        Runs all available url tests
        @return: True if URL passes tests, false if not.
        """
        url_is_okay = self.check_suffix() and self.check_if_parses()
        return url_is_okay

def run_url_tests(url_string):
    """
    Creates a URLProperties object and runs all tests
    @param url_string: A URL in string format
    @return: Boolean indicating whether or not URL has passed all tests
    """
    url_properties = URLProperties(url_string)
    return url_properties.run_tests()

def run_image_tests(image):
    """
    Runs all available image tests
    @param image: PIL Image object
    @return: Boolean indicating whether or not all tests have been passed
    """
    image_properties = ImageProperties(image)
    return image_properties.run_tests()

def upload_to_s3(file_to_upload, keyname):
    '''
    Upload file to S3 using provided keyname.

    Returns:
        public_url: URL to access uploaded file
    '''
    #im = Image.open(file_to_upload)
    #out_im = cStringIO.StringIO()
    #im.save(out_im, 'PNG')

    try:
        conn = S3Connection(settings.AWS_ACCESS_KEY_ID, settings.AWS_SECRET_ACCESS_KEY)
        bucketname = str(settings.AWS_STORAGE_BUCKET_NAME)
        bucket = conn.create_bucket(bucketname.lower())

        k = Key(bucket)
        k.key = keyname
        k.set_metadata('filename', file_to_upload.name)
        k.set_contents_from_file(file_to_upload)
        #k.set_contents_from_string(out_im.getvalue())
        #k.set_metadata("Content-Type", 'images/png')

        k.set_acl("public-read")
        public_url = k.generate_url(60 * 60 * 24 * 365) # URL timeout in seconds.

        return True, public_url
    except:
        return False, "Could not connect to S3."

def get_from_s3(s3_public_url):
    """
    Gets an image from a given S3 url
    @param s3_public_url: The URL where an image is located
    @return: The image data
    """
    r = requests.get(s3_public_url, timeout=2)
    data = r.text
    return data

<commit_msg>Fix skin tone check<commit_after>"""
This contains functions and classes used to evaluate if images are acceptable (do not show improper content, etc), and
to send them to S3.
"""

from PIL import Image
import urlparse
import requests
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from django.conf import settings
import pickle
import logging

log = logging.getLogger(__name__)

#Domains where any image linked to can be trusted to have acceptable content.
TRUSTED_IMAGE_DOMAINS = [
    'wikipedia.com',
    'wikipedia.net',
    'wikipedia.org',
    'edxuploads.s3.amazonaws.com'
]

#Suffixes that are allowed in image urls
ALLOWABLE_IMAGE_SUFFIXES = [
    'jpg',
    'png',
    'gif'
]

#Maximum allowed dimensions (x and y) for an uploaded image
MAX_ALLOWED_IMAGE_DIM = 1500

#Dimensions to which image is resized before it is evaluated for color count, etc
MAX_IMAGE_DIM = 150

#Maximum number of colors that should be counted in ImageProperties
MAX_COLORS_TO_COUNT = 16

#Maximum number of colors allowed in an uploaded image
MAX_COLORS = 20

class ImageProperties(object):
    """
    Class to check properties of an image and to validate if they are allowed.
    """
    def __init__(self, image):
        """
        Initializes class variables
        @param image: Image object (from PIL)
        @return: None
        """
        self.image = image
        image_size = self.image.size
        self.image_too_large = False
        if image_size[0] > MAX_ALLOWED_IMAGE_DIM or image_size[1] > MAX_ALLOWED_IMAGE_DIM:
            self.image_too_large = True
        if image_size[0] > MAX_IMAGE_DIM or image_size[1] > MAX_IMAGE_DIM:
            self.image = self.image.resize((MAX_IMAGE_DIM, MAX_IMAGE_DIM))
            self.image_size = self.image.size

    def count_colors(self):
        """
        Counts the number of colors in an image, and matches them to the max allowed
        @return: boolean true if color count is acceptable, false otherwise
        """
        colors = self.image.getcolors(MAX_COLORS_TO_COUNT)
        if colors is None:
            color_count = MAX_COLORS_TO_COUNT
        else:
            color_count = len(colors)

        too_many_colors = (color_count <= MAX_COLORS)
        return too_many_colors

    def check_if_rgb_is_skin(self, rgb):
        """
        Checks if a given input rgb tuple/list is a skin tone
        @param rgb: RGB tuple
        @return: Boolean true false
        """
        r,g,b = rgb
        check_r = (r > 60)
        check_b =  (r * 0.4) < b < (r * 0.85)
        check_g = (r * 0.2) < g < (r * 0.7)

        return check_r and check_b and check_g

    def get_skin_ratio(self):
        """
        Gets the ratio of skin tone colors in an image
        @return: True if the ratio is low enough to be acceptable, false otherwise
        """
        colors = self.image.getcolors(MAX_COLORS_TO_COUNT)
        is_okay = True
        if colors is not None:
            skin = sum([count for count, rgb in colors if self.check_if_rgb_is_skin(rgb)])
            total_colored_pixels = sum([count for count, rgb in colors])
            bad_color_val = float(skin) / total_colored_pixels
            if bad_color_val > .4:
                is_okay = False

        return is_okay

    def run_tests(self):
        """
        Does all available checks on an image to ensure that it is okay (size, skin ratio, colors)
        @return: Boolean indicating whether or not image passes all checks
        """
        image_is_okay = False
        try:
            image_is_okay = self.count_colors() and self.get_skin_ratio() and not self.image_too_large
        except:
            pass

        return image_is_okay

class URLProperties(object):
    """
    Checks to see if a URL points to acceptable content.  Added to check if students are submitting reasonable
    links to the peer grading image functionality of the external grading service.
    """
    def __init__(self, url_string):
        self.url_string = url_string

    def check_if_parses(self):
        """
        Check to see if a URL parses properly
        @return: success (True if parses, false if not)
        """
        success = False
        try:
            self.parsed_url = urlparse.urlparse(url_string)
            success = True
        except:
            pass

        return success

    def check_suffix(self):
        """
        Checks the suffix of a url to make sure that it is allowed
        @return: True if suffix is okay, false if not
        """
        good_suffix = False
        for suffix in ALLOWABLE_IMAGE_SUFFIXES:
            if self.url_string.endswith(suffix):
                good_suffix = True
                break
        return good_suffix

    def run_tests(self):
        """
        Runs all available url tests
        @return: True if URL passes tests, false if not.
        """
        url_is_okay = self.check_suffix() and self.check_if_parses()
        return url_is_okay

def run_url_tests(url_string):
    """
    Creates a URLProperties object and runs all tests
    @param url_string: A URL in string format
    @return: Boolean indicating whether or not URL has passed all tests
    """
    url_properties = URLProperties(url_string)
    return url_properties.run_tests()

def run_image_tests(image):
    """
    Runs all available image tests
    @param image: PIL Image object
    @return: Boolean indicating whether or not all tests have been passed
    """
    image_properties = ImageProperties(image)
    return image_properties.run_tests()

def upload_to_s3(file_to_upload, keyname):
    '''
    Upload file to S3 using provided keyname.

    Returns:
        public_url: URL to access uploaded file
    '''
    #im = Image.open(file_to_upload)
    #out_im = cStringIO.StringIO()
    #im.save(out_im, 'PNG')

    try:
        conn = S3Connection(settings.AWS_ACCESS_KEY_ID, settings.AWS_SECRET_ACCESS_KEY)
        bucketname = str(settings.AWS_STORAGE_BUCKET_NAME)
        bucket = conn.create_bucket(bucketname.lower())

        k = Key(bucket)
        k.key = keyname
        k.set_metadata('filename', file_to_upload.name)
        k.set_contents_from_file(file_to_upload)
        #k.set_contents_from_string(out_im.getvalue())
        #k.set_metadata("Content-Type", 'images/png')

        k.set_acl("public-read")
        public_url = k.generate_url(60 * 60 * 24 * 365) # URL timeout in seconds.

        return True, public_url
    except:
        return False, "Could not connect to S3."

def get_from_s3(s3_public_url):
    """
    Gets an image from a given S3 url
    @param s3_public_url: The URL where an image is located
    @return: The image data
    """
    r = requests.get(s3_public_url, timeout=2)
    data = r.text
    return data

<|endoftext|>


<commit_before>
class COL_BASECLASS(object):
    def __init__(self, value):
        """
        The original string version (as it appeared in the .vcf file)
        of 'value' is preserved, with type 'string', as 'self.chr_value'.
        The "working" version of 'value' is stored as 'self.value'. This allows
        the type to change depending on the data type whilst maintaining the
        original formatting.
        """
        self.chr_value, self.value = value, value

    def __str__(self):
        return self.chr_value

#######################

class INFO(COL_BASECLASS):
    __slots__ = ['chr_value', 'value']

    def __getitem__(self, index):
        """
        The actual parsing of the INFO line only occurs when
        __getitem__ is called. This is to save on unnecesary
        processing of the VCF file that occurs when not all
        columns are used. The same also applies to each item
        within INFO.

        Each INFO entry is delimited by ';'. Each entry consist of the descriptor and
        value separated by '='. E.g:
        INFO
        DP=1;AF=0;AC1=0;DP4=1,0,0,0;MQ=25;FQ=-24.3
        """
        if isinstance(self.value, str):
            self.value = dict([tuple(i.split('=')) for i in chr_value.split(';') if '=' in i])

        INFO_parse = {'DP': int, 'MQ': int}

        value = self.value[index]

        if index in INFO_parse:
            value = INFO_parse[index](value)

        return value

###### ######

class FORMAT(COL_BASECLASS):

    def __init__(self):
        self.format_dict = {}

    def __getitem__(self, value):

        if value not in self.format_dict:
            new_value = dict([(j, i) for (i, j) in enumerate(value.split(":"))])
            self.format_dict.update({value: new_value})

        return self.format_dict[value]

###### ######

class SAMPLE(COL_BASECLASS):
    __slots__ = ['chr_value', 'value', 'Format', 'name']

    def __init__(self, value, Format, name = None):
        self.chr_value = value
        self.value = value
        self.Format = Format
        self.name = name

    def __getitem__(self, key):
        """

        """
        if self.chr_value == './.':
            item = None
        else:
            if isinstance(self.value, str):
                self.value = self.chr_value.split(':')
                # problems may occur if the number of items in the
                # FORMAT column != number items in the sample
                self.SAMPLE_parse = {'DP': int}
                self.SAMPLE_default = {'DP': 0, 'GT': None}

            nformat = len(self.Format.value.split(':'))
            findex = self.Format[self.Format.value][key]
            if findex >= nformat:
                item = None
            else:
                try:
                    item = self.value[self.Format[self.Format.value][key]]
                except KeyError:
                    item = self.SAMPLE_default[key]

            if key in self.SAMPLE_parse:
                try:
                    item = self.SAMPLE_parse[key](item)
                except ValueError:
                    item = self.SAMPLE_default[key]

        return item

    def __str__(self):
        return self.chr_value

    def is_het(self):

        gt = self['GT']

        if gt is None:
            het = None
        else:
            if gt == '0/1' or gt == '1/0':
                het = True
            else:
                het = False

        return het

    def str_fetch(self, key):

        if self.chr_value == "./.":
            item = 'NA'
        else:
            #if self.split_value is None:
            self.split_value = self.chr_value.split(":")

            item = self.split_value[self.Format[self.Format.value][key]]

        return item

    def genotype_str(self, ref, alt):
        possibleAlleles = [ref] + alt.split(',')
        GT = self['GT'].split('/')

        one, two = possibleAlleles[int(GT[0])], possibleAlleles[int(GT[1])]

        return one + two
<commit_msg>exception<commit_after>
class COL_BASECLASS(object):
    def __init__(self, value):
        """
        The original string version (as it appeared in the .vcf file)
        of 'value' is preserved, with type 'string', as 'self.chr_value'.
        The "working" version of 'value' is stored as 'self.value'. This allows
        the type to change depending on the data type whilst maintaining the
        original formatting.
        """
        self.chr_value, self.value = value, value

    def __str__(self):
        return self.chr_value

#######################

class INFO(COL_BASECLASS):
    __slots__ = ['chr_value', 'value']

    def __getitem__(self, index):
        """
        The actual parsing of the INFO line only occurs when
        __getitem__ is called. This is to save on unnecesary
        processing of the VCF file that occurs when not all
        columns are used. The same also applies to each item
        within INFO.

        Each INFO entry is delimited by ';'. Each entry consist of the descriptor and
        value separated by '='. E.g:
        INFO
        DP=1;AF=0;AC1=0;DP4=1,0,0,0;MQ=25;FQ=-24.3
        """
        if isinstance(self.value, str):
            self.value = dict([tuple(i.split('=')) for i in chr_value.split(';') if '=' in i])

        INFO_parse = {'DP': int, 'MQ': int}

        value = self.value[index]

        if index in INFO_parse:
            value = INFO_parse[index](value)

        return value

###### ######

class FORMAT(COL_BASECLASS):

    def __init__(self):
        self.format_dict = {}

    def __getitem__(self, value):

        if value not in self.format_dict:
            new_value = dict([(j, i) for (i, j) in enumerate(value.split(":"))])
            self.format_dict.update({value: new_value})

        return self.format_dict[value]

###### ######

class SAMPLE(COL_BASECLASS):
    __slots__ = ['chr_value', 'value', 'Format', 'name']

    def __init__(self, value, Format, name = None):
        self.chr_value = value
        self.value = value
        self.Format = Format
        self.name = name

    def __getitem__(self, key):
        """

        """
        if self.chr_value == './.':
            item = None
        else:
            if isinstance(self.value, str):
                self.value = self.chr_value.split(':')
                # problems may occur if the number of items in the
                # FORMAT column != number items in the sample
                self.SAMPLE_parse = {'DP': int}
                self.SAMPLE_default = {'DP': 0, 'GT': None}

            nformat = len(self.Format.value.split(':'))
            findex = self.Format[self.Format.value][key]
            if findex >= nformat:
                item = None
            else:
                try:
                    item = self.value[findex]
                except KeyError:
                    item = self.SAMPLE_default[key]

            if key in self.SAMPLE_parse:
                try:
                    item = self.SAMPLE_parse[key](item)
                except ValueError:
                    item = self.SAMPLE_default[key]

        return item

    def __str__(self):
        return self.chr_value

    def is_het(self):

        gt = self['GT']

        if gt is None:
            het = None
        else:
            if gt == '0/1' or gt == '1/0':
                het = True
            else:
                het = False

        return het

    def str_fetch(self, key):

        if self.chr_value == "./.":
            item = 'NA'
        else:
            #if self.split_value is None:
            self.split_value = self.chr_value.split(":")

            item = self.split_value[self.Format[self.Format.value][key]]

        return item

    def genotype_str(self, ref, alt):
        possibleAlleles = [ref] + alt.split(',')
        GT = self['GT'].split('/')

        one, two = possibleAlleles[int(GT[0])], possibleAlleles[int(GT[1])]

        return one + two
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
PURPOSE:

To return the theoretical Limb Darkening coefficients and/or model
for observations taken with the MOST Satellite given the input
parameters.

This code makes use of the Claret, Dragomir, and Matthews (2014) A&A 567, A3
paper, which should be cited if you use this code:

http://adsabs.harvard.edu/abs/2014A%26A...567A...3C

Created on 2015-05-08T10:26:38
"""

from __future__ import division, print_function
import sys
import argparse

try:
    import numpy as np
except ImportError:
    print('You need numpy installed')
    sys.exit(1)

import pandas as pd
import limbDarkening
ld_dir = '/'.join(limbDarkening.__file__.split('/')[:-1])+'/'

__author__ = "Matt Giguere (github: @mattgiguere)"
__license__ = "MIT"
__version__ = '0.0.1'
__maintainer__ = "Matt Giguere"
__email__ = "matthew.giguere@yale.edu"
__status__ = " Development NOT(Prototype or Production)"

class LimbDarkening:
    """
    The main class that creates a limb darkening object.

    :param teff:
        The effective temperature in Kelvin

    :param logg:
        The log of the surface gravity
    """

    def __init__(self, teff, logg, monh, turbvel):
        self.fn = ld_dir+'data/Claret2014MOST_ATLAS.tsv'

        self.teff = teff
        self.monh = monh
        self.logg = logg
        self.turbvel = turbvel
        self.round_values()

    def round_values(self):
        """
        PURPOSE:
        To round the input parameters for lookup in the Claret table.
        """
        self.rteff = np.round(self.teff/250) * 250
        self.rmonh = np.round(self.monh/0.1) * 0.1
        self.rlogg = np.round(self.logg/0.5) * 0.5

    def get_data(self):
        """
        PURPOSE:
        Restore the data from file and return them as a pandas DataFrame.
        """
        return pd.read_csv(self.fn, sep='\s*\s*', engine='python', header=70)[2:]

    def get_ldcs(self):
        """
        PURPOSE:
        This function will extract and return the limb darkening coefficients
        from the Claret table.
        """
        ldcs = np.zeros(4)
        coeffNames = ['a1', 'a2', 'a3', 'a4']
        idx = 0
        lddf = self.get_data()
        for cNm in coeffNames:
            ldcs[idx] = lddf[((lddf['Teff'] == '{:n}'.format(self.rteff)) &
                             (lddf['logg'] == '{:.2f}'.format(self.rlogg)) &
                             (lddf['logZ'] == '{:.1f}'.format(self.rmonh)))][cNm].values[0]
            idx += 1
        return ldcs

    def claret_model(self, ldcs):
        def limb_darkening_function(theta):
            mu = np.cos(theta)
            limbdark = 1. - ldcs[0] * (1. - mu**0.5)
            -  ldcs[1] * (1. - mu**1.0)
            -  ldcs[2] * (1. - mu**1.5)
            -  ldcs[3] * (1. - mu**2.0)
            return limbdark
        return limb_darkening_function

def claret14_limb_darkening_coeffs(teff, monh, logg, turbvel):
    """
    PURPOSE:
    To round the input values in order to extract the limb darkening
    coefficients from the Claret et al. (2014) table.

    :param teff:
    The stellar effective temperature in Kelvin

    :param monh:
    The stellar metalicity

    :param logg:
    The log of the stellar surface gravity

    :param turbvel:
    The stellar turbulent velocity

    See the Claret et al. (2014) paper for more details.
    """

    assert teff >= 3000, 'teff should be greater than 3000: {}'.format(teff)
    assert teff <= 50000, 'teff should be less than 50000: {}'.format(teff)
    assert monh >= -5.0, 'monh should be greater than -5.0: {}'.format(monh)
    assert monh <= 1.0, 'monh should be greater than 1.0: {}'.format(monh)
    assert logg >= 0.0, 'logg should be greater than 0.0: {}'.format(logg)
    assert logg <= 5.0, 'logg should be less than 5.0: {}'.format(logg)
    assert turbvel <= 8.0, 'turbvel should be less than 8.0: {}'.format(turbvel)
    assert turbvel >= 0.0, 'turbvel should be greater than 0.0: {}'.format(turbvel)

    print('input teff: {}'.format(teff))
    print('input monh: {}'.format(monh))
    print('input logg: {}'.format(logg))
    print('input turbvel: {}'.format(turbvel))
    ldo = LimbDarkening(teff, monh, logg, turbvel)
    print(ldo.get_ldcs())

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='argparse object.')
    parser.add_argument(
        'teff',
        help='The stellar effective temperature.')
    parser.add_argument(
        'monh',
        help='The stellar metalicity')
    parser.add_argument(
        'logg',
        help='The log of the stellar surface gravity')
    parser.add_argument(
        'turbvel',
        help='The stellar turbulent velocity')

    args = parser.parse_args()

    claret14_limb_darkening_coeffs(int(args.teff), float(args.monh),
                                   float(args.logg), float(args.turbvel))
 <commit_msg>remove debug prints<commit_after>#!/usr/bin/env python

"""
PURPOSE:

To return the theoretical Limb Darkening coefficients and/or model
for observations taken with the MOST Satellite given the input
parameters.

This code makes use of the Claret, Dragomir, and Matthews (2014) A&A 567, A3
paper, which should be cited if you use this code:

http://adsabs.harvard.edu/abs/2014A%26A...567A...3C

Created on 2015-05-08T10:26:38
"""

from __future__ import division, print_function
import sys
import argparse

try:
    import numpy as np
except ImportError:
    print('You need numpy installed')
    sys.exit(1)

import pandas as pd
import limbDarkening
ld_dir = '/'.join(limbDarkening.__file__.split('/')[:-1])+'/'

__author__ = "Matt Giguere (github: @mattgiguere)"
__license__ = "MIT"
__version__ = '0.0.1'
__maintainer__ = "Matt Giguere"
__email__ = "matthew.giguere@yale.edu"
__status__ = " Development NOT(Prototype or Production)"

class LimbDarkening:
    """
    The main class that creates a limb darkening object.

    :param teff:
        The effective temperature in Kelvin

    :param logg:
        The log of the surface gravity
    """

    def __init__(self, teff, logg, monh, turbvel):
        self.fn = ld_dir+'data/Claret2014MOST_ATLAS.tsv'

        self.teff = teff
        self.monh = monh
        self.logg = logg
        self.turbvel = turbvel
        self.round_values()

    def round_values(self):
        """
        PURPOSE:
        To round the input parameters for lookup in the Claret table.
        """
        self.rteff = np.round(self.teff/250) * 250
        self.rmonh = np.round(self.monh/0.1) * 0.1
        self.rlogg = np.round(self.logg/0.5) * 0.5

    def get_data(self):
        """
        PURPOSE:
        Restore the data from file and return them as a pandas DataFrame.
        """
        return pd.read_csv(self.fn, sep='\s*\s*', engine='python', header=70)[2:]

    def get_ldcs(self):
        """
        PURPOSE:
        This function will extract and return the limb darkening coefficients
        from the Claret table.
        """
        ldcs = np.zeros(4)
        coeffNames = ['a1', 'a2', 'a3', 'a4']
        idx = 0
        lddf = self.get_data()
        for cNm in coeffNames:
            ldcs[idx] = lddf[((lddf['Teff'] == '{:n}'.format(self.rteff)) &
                             (lddf['logg'] == '{:.2f}'.format(self.rlogg)) &
                             (lddf['logZ'] == '{:.1f}'.format(self.rmonh)))][cNm].values[0]
            idx += 1
        return ldcs

    def claret_model(self, ldcs):
        def limb_darkening_function(theta):
            mu = np.cos(theta)
            limbdark = 1. - ldcs[0] * (1. - mu**0.5)
            -  ldcs[1] * (1. - mu**1.0)
            -  ldcs[2] * (1. - mu**1.5)
            -  ldcs[3] * (1. - mu**2.0)
            return limbdark
        return limb_darkening_function

def claret14_limb_darkening_coeffs(teff, monh, logg, turbvel):
    """
    PURPOSE:
    To round the input values in order to extract the limb darkening
    coefficients from the Claret et al. (2014) table.

    :param teff:
    The stellar effective temperature in Kelvin

    :param monh:
    The stellar metalicity

    :param logg:
    The log of the stellar surface gravity

    :param turbvel:
    The stellar turbulent velocity

    See the Claret et al. (2014) paper for more details.
    """

    assert teff >= 3000, 'teff should be greater than 3000: {}'.format(teff)
    assert teff <= 50000, 'teff should be less than 50000: {}'.format(teff)
    assert monh >= -5.0, 'monh should be greater than -5.0: {}'.format(monh)
    assert monh <= 1.0, 'monh should be greater than 1.0: {}'.format(monh)
    assert logg >= 0.0, 'logg should be greater than 0.0: {}'.format(logg)
    assert logg <= 5.0, 'logg should be less than 5.0: {}'.format(logg)
    assert turbvel <= 8.0, 'turbvel should be less than 8.0: {}'.format(turbvel)
    assert turbvel >= 0.0, 'turbvel should be greater than 0.0: {}'.format(turbvel)

    ldo = LimbDarkening(teff, monh, logg, turbvel)
    print(ldo.get_ldcs())

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='argparse object.')
    parser.add_argument(
        'teff',
        help='The stellar effective temperature.')
    parser.add_argument(
        'monh',
        help='The stellar metalicity')
    parser.add_argument(
        'logg',
        help='The log of the stellar surface gravity')
    parser.add_argument(
        'turbvel',
        help='The stellar turbulent velocity')

    args = parser.parse_args()

    claret14_limb_darkening_coeffs(int(args.teff), float(args.monh),
                                   float(args.logg), float(args.turbvel))
 <|endoftext|>


<commit_before>"""Repository rules for automagically configuring the host toolchain."""

_BUILD_TEMPLATE = """
package(default_visibility=["//visibility:public"])

load("@//tools/build_rules/external_tools:external_tools_toolchain.bzl", "external_tools_toolchain")

external_tools_toolchain(
  name = "host_toolchain_impl",
  asciidoc = "{asciidoc}",
  path = "{path}"
)

toolchain(
    name = "host_toolchain",
    exec_compatible_with = [
      # These expect constraint_setting targets, not platforms.
      # But the desired constraints are visibility restricted.
      #"@bazel_tools//platforms:host_platform",
    ],
    target_compatible_with = [
      #"@bazel_tools//platforms:host_platform",
    ],
    toolchain = ":host_toolchain_impl",
    toolchain_type = "@//tools/build_rules/external_tools:external_tools_toolchain_type",
)
"""

def _external_toolchain_autoconf_impl(repository_ctx):
    asciidoc = repository_ctx.which("asciidoc")
    if asciidoc == None:
        fail("Unable to find 'asciidoc' executable on path.")

    # These are the tools that the doc/schema generation need beyond the
    # explicit call to asciidoc.
    tools = [
        "awk",
        "cat",
        "cut",
        "dot",
        "grep",
        "mkdir",
        "mktemp",
        "python",
        "readlink",
        "rm",
        "sed",
        "source-highlight",
        "tee",
        "touch",
        "zip",
    ]
    for tool in tools:
        symlink_command(repository_ctx, tool)

    repository_ctx.file("BUILD", _BUILD_TEMPLATE.format(
        asciidoc = asciidoc,
        path = repository_ctx.path(""),
    ))

def symlink_command(repository_ctx, command):
    binary = repository_ctx.which(command)
    if binary == None:
        fail("Unable to find '%s' executable on path." % command)
    repository_ctx.symlink(binary, command)

external_toolchain_autoconf = repository_rule(
    implementation = _external_toolchain_autoconf_impl,
    local = True,
    environ = ["PATH"],
)

def external_tools_configure():
    external_toolchain_autoconf(name = "local_config_tools")
    native.register_toolchains("@local_config_tools//:host_toolchain")
<commit_msg>fix(build): use @io_kythe// instead of @// in bazel build file template (#4711)<commit_after>"""Repository rules for automagically configuring the host toolchain."""

_BUILD_TEMPLATE = """
package(default_visibility=["//visibility:public"])

load("@io_kythe//tools/build_rules/external_tools:external_tools_toolchain.bzl", "external_tools_toolchain")

external_tools_toolchain(
  name = "host_toolchain_impl",
  asciidoc = "{asciidoc}",
  path = "{path}"
)

toolchain(
    name = "host_toolchain",
    exec_compatible_with = [
      # These expect constraint_setting targets, not platforms.
      # But the desired constraints are visibility restricted.
      #"@bazel_tools//platforms:host_platform",
    ],
    target_compatible_with = [
      #"@bazel_tools//platforms:host_platform",
    ],
    toolchain = ":host_toolchain_impl",
    toolchain_type = "@io_kythe//tools/build_rules/external_tools:external_tools_toolchain_type",
)
"""

def _external_toolchain_autoconf_impl(repository_ctx):
    asciidoc = repository_ctx.which("asciidoc")
    if asciidoc == None:
        fail("Unable to find 'asciidoc' executable on path.")

    # These are the tools that the doc/schema generation need beyond the
    # explicit call to asciidoc.
    tools = [
        "awk",
        "cat",
        "cut",
        "dot",
        "grep",
        "mkdir",
        "mktemp",
        "python",
        "readlink",
        "rm",
        "sed",
        "source-highlight",
        "tee",
        "touch",
        "zip",
    ]
    for tool in tools:
        symlink_command(repository_ctx, tool)

    repository_ctx.file("BUILD", _BUILD_TEMPLATE.format(
        asciidoc = asciidoc,
        path = repository_ctx.path(""),
    ))

def symlink_command(repository_ctx, command):
    binary = repository_ctx.which(command)
    if binary == None:
        fail("Unable to find '%s' executable on path." % command)
    repository_ctx.symlink(binary, command)

external_toolchain_autoconf = repository_rule(
    implementation = _external_toolchain_autoconf_impl,
    local = True,
    environ = ["PATH"],
)

def external_tools_configure():
    external_toolchain_autoconf(name = "local_config_tools")
    native.register_toolchains("@local_config_tools//:host_toolchain")
<|endoftext|>


<commit_before># Copyright 2021 The Kubric Authors.
#
# 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.

# Copyright 2021 The Kubric Authors
#
# 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
#
#    https://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.

import json
import numpy as np
import os
import pybullet as pb
import re
import shutil
import ssl
import tarfile
import trimesh as tm
import trimesh.exchange.obj as tri_obj
import urllib.request
import zipfile
import argparse
import logging
import itertools

from pathlib import Path

URDF_TEMPLATE = """
<robot name="{id}">
    <link name="base">
        <contact>
            <lateral_friction value="{friction}" />
        </contact>
        <inertial>
            <origin xyz="{center_mass[0]} {center_mass[1]} {center_mass[2]}" />
            <mass value="{mass}" />
            <inertia ixx="{inertia[0][0]}" ixy="{inertia[0][1]}"
                     ixz="{inertia[0][2]}" iyy="{inertia[1][1]}"
                     iyz="{inertia[1][2]}" izz="{inertia[2][2]}" />
        </inertial>
        <visual>
            <origin xyz="0 0 0" />
            <geometry>
                <mesh filename="visual_geometry.obj" />
            </geometry>
        </visual>
        <collision>
            <origin xyz="0 0 0" />
            <geometry>
                <mesh filename="collision_geometry.obj" />
            </geometry>
        </collision>
    </link>
</robot>
"""

def bypass_ssl_certification():
  if getattr(ssl, "_create_unverified_context", None):
    ssl._create_default_https_context = ssl._create_unverified_context

def collect_list_of_available_assets():
  bypass_ssl_certification()
  # Get a list of available assets
  list_of_assets = []
  for page in itertools.count(start=1):
    url = f"https://fuel.ignitionrobotics.org/1.0/GoogleResearch/models?page={page}"
    try:
      with urllib.request.urlopen(url) as f:
        list_of_assets.extend(json.load(f))
    except urllib.error.HTTPError:
      return list_of_assets

def download_gso_assets(download_dir: Path):
  download_dir.mkdir(parents=True, exist_ok=True)
  list_of_assets = collect_list_of_available_assets()
  # download
  for i, asset in enumerate(list_of_assets):
    name = asset["name"]
    url = "https://fuel.ignitionrobotics.org/1.0/GoogleResearch/models/{name}/1/{name}.zip".format(name=urllib.request.quote(name))
    filename = download_dir / f"{name}.zip"
    if not filename.exists():
      try:
        logging.info(f"{i+1:4d}/{len(list_of_assets)}: Fetching '{name}' from {url} into {filename}")
        urllib.request.urlretrieve(url, filename)
      except (urllib.error.HTTPError, urllib.error.URLError) as e:
        logging.warning(f"FAILED! skipping '{name}'", e)

def get_object_properties(tmesh, name, density=None):
  if density is None:
    tmesh.density = 1000.0

  rounda = lambda x: np.round(x, decimals=6).tolist()
  roundf = lambda x: float(np.round(x, decimals=6))

  return {
      "id": name,
      "density": roundf(tmesh.density),
      "friction": 0.5,
      "nr_vertices": len(tmesh.vertices),
      "nr_faces": len(tmesh.faces),
      "bounds": rounda(tmesh.bounds),
      "area": roundf(tmesh.area),
      "volume": roundf(tmesh.volume),
      "mass": roundf(tmesh.mass),
      "center_mass": rounda(tmesh.center_mass),
      "inertia": rounda(tmesh.moment_inertia),
      "is_convex": tmesh.is_convex,
      "euler_number": tmesh.euler_number,  # used for topological analysis (see: http://max-limper.de/publications/Euler/index.html),
  }

def convert_gso_asset(path, target_dir):
  name = path.stem
  asset_dir = path.parent / name
  with zipfile.ZipFile(path, "r") as zip_ref:
    zip_ref.extractall(asset_dir)
  texture_path = asset_dir / "materials" / "textures" / "texture.png"
  obj_path = asset_dir / "meshes" / "model.obj"
  mat_path = asset_dir / "meshes" / "model.mtl"

  target_asset_dir = target_dir / name
  if target_asset_dir.exists():
    shutil.rmtree(target_asset_dir)
  target_asset_dir.mkdir(parents=True, exist_ok=True)

  vis_path = target_asset_dir / "visual_geometry.obj"
  coll_path = target_asset_dir / "collision_geometry.obj"
  urdf_path = target_asset_dir / "object.urdf"
  json_path = target_asset_dir / "data.json"
  tex_path = target_asset_dir / "texture.png"
  tar_path = target_dir / (name + ".tar.gz")

  shutil.copy(texture_path, asset_dir / "meshes" / "texture.png")
  # import mesh into trimesh
  tmesh = tm.load_mesh(obj_path, file_type="obj")
  # center it around center of mass
  tmesh.apply_translation(-tmesh.center_mass)
  # export to obj again
  obj_content = tri_obj.export_obj(tmesh)
  obj_content = re.sub("mtllib material0.mtl\nusemtl material0\n",
                       "mtllib visual_geometry.mtl\nusemtl material_0\n", obj_content)
  with open(vis_path, "w") as f:
    f.write(obj_content)

  # compute a collision mesh using pybullets VHACD
  pb.vhacd(str(vis_path),
           str(coll_path),
           str(target_asset_dir / "pybullet_logs.txt"))

  # move material and texture
  shutil.move(mat_path, target_asset_dir / "visual_geometry.mtl")
  shutil.move(texture_path, tex_path)

  properties = get_object_properties(tmesh, name)

  with open(urdf_path, "w") as f:
    f.write(URDF_TEMPLATE.format(**properties))

  properties["paths"] = {
      "visual_geometry": ["visual_geometry.obj"],
      "collision_geometry": ["collision_geometry.obj"],
      "urdf": ["object.urdf"],
      "texture": ["texture.png"],
  }

  with open(json_path, "w") as f:
    json.dump(properties, f, indent=4, sort_keys=True)

  logging.info("  saving as", tar_path)
  with tarfile.open(tar_path, "w:gz") as tar:
    tar.add(target_asset_dir, arcname=name)

  shutil.rmtree(asset_dir)
  shutil.rmtree(target_asset_dir)
  return properties

def convert_all_gso_assets_from_dir(download_dir: Path, target_dir: Path):
  details_list = []
  list_of_asset_files = sorted(download_dir.glob("*.zip"))
  for i, path in enumerate(list_of_asset_files):
    logging.info(f"{i:4d}/{len(list_of_asset_files)}: Converting {path}...")
    properties = convert_gso_asset(path, target_dir)
    details_list.append(properties)

  with open(target_dir / "manifest.json", "w") as f:
    json.dump(details_list, f, indent=4, sort_keys=True)

def main(download_dir: Path = Path("../Assets/GSO_raw"), target_dir: Path = Path("../Assets/GSO"),
         keep_raw_assets=False):
  download_gso_assets(download_dir)
  target_dir.mkdir(parents=True, exist_ok=True)
  convert_all_gso_assets_from_dir(download_dir, target_dir)
  logging.info("Done!")
  if not keep_raw_assets:
    logging.info("Deleting the raw (unconverted) assets...")
    shutil.rmtree(download_dir)

if __name__ == '__main__':
  parser = argparse.Parser()
  parser.add_argument("--download_dir", type=str, default="Assets/GSO_raw")
  parser.add_argument("--target_dir", type=str, default="Assets/GSO")
  parser.add_argument("--keep_raw_assets", type=bool, default=False)
  FLAGS, unused = parser.parse_known_args()
  main(download_dir=FLAGS.download_dir, target_dir=FLAGS.target_dir,
       keep_raw_assets=FLAGS.keep_raw_assets)
<commit_msg>refactored download_gso script<commit_after># Copyright 2021 The Kubric Authors.
#
# 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.

# Copyright 2021 The Kubric Authors
#
# 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
#
#    https://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.
import argparse
import functools
import itertools
import json
import logging
import multiprocessing
import pathlib
import re
import shutil
import ssl
import sys
import tarfile
import urllib
import urllib.request
import urllib.error
import zipfile

import pybullet as pb
import tqdm
import trimesh as tm
import trimesh.exchange.obj as tri_obj

from kubric import file_io
from kubric import redirect_io
from kubric.kubric_typing import PathLike

URDF_TEMPLATE = """
<robot name="{id}">
    <link name="base">
        <inertial>
            <origin xyz="{center_mass[0]} {center_mass[1]} {center_mass[2]}" />
            <mass value="{mass}" />
            <inertia ixx="{inertia[0][0]}" ixy="{inertia[0][1]}"
                     ixz="{inertia[0][2]}" iyy="{inertia[1][1]}"
                     iyz="{inertia[1][2]}" izz="{inertia[2][2]}" />
        </inertial>
        <visual>
            <origin xyz="0 0 0" />
            <geometry>
                <mesh filename="visual_geometry.obj" />
            </geometry>
        </visual>
        <collision>
            <origin xyz="0 0 0" />
            <geometry>
                <mesh filename="collision_geometry.obj" />
            </geometry>
        </collision>
    </link>
</robot>
"""

def get_object_properties(tmesh, density=1.0):
  properties = {
      "nr_vertices": len(tmesh.vertices),
      "nr_faces": len(tmesh.faces),
      "bounds": tmesh.bounds.tolist(),
      "surface_area": tmesh.area,
      "volume": tmesh.volume,
      "center_mass": tmesh.center_mass.tolist(),
      "inertia": tmesh.moment_inertia.tolist(),
      "mass": tmesh.volume * density,
  }
  return properties

def bypass_ssl_certification():
  if getattr(ssl, "_create_unverified_context", None):
    ssl._create_default_https_context = ssl._create_unverified_context

def get_metadata(a):
  name = a["name"]
  return {
      "name": name,
      "category": a.get("categories", ["None"])[0],
      "description": a["description"],
      "url": "https://fuel.ignitionrobotics.org/1.0/GoogleResearch/models/{name}/1/{name}.zip".format(name=urllib.request.quote(name)),
  }

def collect_list_of_available_assets(catalogue_path="GSO_catalogue.json"):
  catalogue_path = file_io.as_path(catalogue_path)
  if catalogue_path.exists():
    return file_io.read_json(catalogue_path)

  bypass_ssl_certification()
  # Get a list of available assets
  list_of_assets = []
  for page in itertools.count(start=1):
    url = f"https://fuel.ignitionrobotics.org/1.0/GoogleResearch/models?page={page}"
    try:
      with urllib.request.urlopen(url) as f:
        list_of_assets.extend(json.load(f))
    except urllib.error.HTTPError:
      assets = [get_metadata(a) for a in list_of_assets]
      file_io.write_json(assets, catalogue_path)
      return assets

def download_asset(a, download_dir):
  filename = download_dir / f"{a['name']}.zip"
  if not filename.exists():
    try:
      urllib.request.urlretrieve(a["url"], filename)
    except (urllib.error.HTTPError, urllib.error.URLError) as e:
      logging.warning(f"FAILED! skipping '{a['name']}'", e)

def download_all(
    assets_list,
    num_processes: int = 16,
    download_dir: PathLike = 'GSO_raw'
):
  download_dir = pathlib.Path(download_dir)
  download_dir.mkdir(parents=True, exist_ok=True)
  with tqdm.tqdm(total=len(assets_list)) as pbar:
    with multiprocessing.Pool(num_processes, maxtasksperchild=1) as pool:
      promise = pool.imap_unordered(functools.partial(download_asset, download_dir=download_dir), assets_list)
      for _ in promise:
        pbar.update(1)

def kubricify(asset, source_dir, target_dir):
  source_dir = pathlib.Path(source_dir)
  target_dir = pathlib.Path(target_dir)
  target_dir.mkdir(parents=True, exist_ok=True)
  name = asset["name"]

  # --- source paths
  asset_raw_path = source_dir / (name + ".zip")
  asset_tmp_dir = source_dir / name

  texture_source_path = asset_tmp_dir / "materials" / "textures" / "texture.png"
  obj_source_path = asset_tmp_dir / "meshes" / "model.obj"
  mat_source_path = asset_tmp_dir / "meshes" / "model.mtl"

  # --- target paths
  target_asset_dir = target_dir / name
  vis_path = target_asset_dir / "visual_geometry.obj"
  mat_path = target_asset_dir / "visual_geometry.mtl"
  coll_path = target_asset_dir / "collision_geometry.obj"
  urdf_path = target_asset_dir / "object.urdf"
  json_path = target_asset_dir / "data.json"
  tex_path = target_asset_dir / "texture.png"
  tar_path = target_dir / (name + ".tar.gz")

  if tar_path.exists():
    with tarfile.open(tar_path, "r:gz") as tar:
      tar.extract("data.json", target_asset_dir)
    asset_entry = file_io.read_json(json_path)
    if asset_entry and "id" in asset_entry:
      shutil.rmtree(target_asset_dir)
      return asset_entry["id"], asset_entry

  # extract source to asset_tmp_dir
  with zipfile.ZipFile(asset_raw_path, "r") as zip_ref:
    zip_ref.extractall(asset_tmp_dir)

  # clear target path
  if target_asset_dir.exists():
    shutil.rmtree(target_asset_dir)
  target_asset_dir.mkdir(parents=True, exist_ok=True)

  # --- convert
  # copy the texture to a place that trimesh will find
  shutil.copy(texture_source_path, asset_tmp_dir / "meshes" / "texture.png")

  # import mesh into trimesh
  tmesh = tm.load_mesh(obj_source_path, file_type="obj")
  # center it around center of mass
  tmesh.apply_translation(-tmesh.center_mass)
  # export to obj again
  obj_content = tri_obj.export_obj(tmesh)
  obj_content = re.sub("mtllib material_0.mtl\nusemtl material_0\n",
                       f"mtllib {mat_path.name}\nusemtl material_0\n", obj_content)
  with open(vis_path, "w") as f:
    f.write(obj_content)

  shutil.move(mat_source_path, mat_path)
  shutil.move(texture_source_path, tex_path)

  with redirect_io.RedirectStream(stream=sys.stdout):
    pb.vhacd(str(vis_path), str(coll_path), str(asset_tmp_dir / "pybullet.log"))

  properties = get_object_properties(tmesh)

  with open(urdf_path, "w") as f:
    f.write(URDF_TEMPLATE.format(id=asset["name"], **properties))

  asset_entry = {
      "id": asset["name"],
      "asset_type": "FileBasedObject",
      "kwargs": {
          "bounds": properties["bounds"],
          "mass": properties["mass"],
          "render_filename": vis_path.name,
          "simulation_filename": urdf_path.name,
      },
      "license": "CC BY-SA 4.0",
      "metadata": {
          "nr_faces": properties["nr_faces"],
          "nr_vertices": properties["nr_vertices"],
          "surface_area": properties["surface_area"],
          "volume": properties["volume"],
          "category": asset["category"],
          "description": asset["description"],
      },
  }

  with open(json_path, "w") as f:
    json.dump(asset_entry, f, indent=4, sort_keys=True)

  with tarfile.open(tar_path, "w:gz") as tar:
    for p in target_asset_dir.glob("*"):
      tar.add(p, arcname=p.name)

  shutil.rmtree(asset_tmp_dir)
  shutil.rmtree(target_asset_dir)

  return name, asset_entry

def main(
    download_dir: PathLike = "GSO_raw",
    target_dir: PathLike = "GSO",
    keep_raw_assets=False
):
  download_dir = file_io.as_path(download_dir)
  target_dir = file_io.as_path(target_dir)
  assets = collect_list_of_available_assets()
  download_all(assets, download_dir=download_dir)
  assets_list = []
  with tqdm.tqdm(total=len(assets)) as pbar:
    with multiprocessing.Pool(32, maxtasksperchild=1) as pool:
      promise = pool.imap_unordered(functools.partial(kubricify,
                                                      source_dir=download_dir,
                                                      target_dir=target_dir),
                                    assets)
      for result in promise:
        assets_list.append(result)
        pbar.update(1)

  manifest_path = "GSO.json"
  manifest = {
      "name": "GSO",
      "base_path": target_dir,
      "version": "1.0",
      "assets": {k: v for k, v in assets_list}
  }
  file_io.write_json(manifest, manifest_path)
  if not keep_raw_assets:
    logging.info("Deleting the raw (unconverted) assets...")
    shutil.rmtree(download_dir)

if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument("--download_dir", type=str, default="GSO_raw")
  parser.add_argument("--target_dir", type=str, default="GSO")
  parser.add_argument("--keep_raw_assets", type=bool, default=False)
  FLAGS, unused = parser.parse_known_args()
  main(download_dir=FLAGS.download_dir, target_dir=FLAGS.target_dir,
       keep_raw_assets=FLAGS.keep_raw_assets)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
#
#  Copyright (c) 2021, The OpenThread Authors.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#  3. Neither the name of the copyright holder nor the
#     names of its contributors may be used to endorse or promote products
#     derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 'AS IS'
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
#
import logging
import unittest

import ipaddress
import config
import thread_cert

# Test description:
#   This test verifies that when external routes are changed in Network Data, a
#   border router will make according updates to routes at the linux kernel.
#
# Topology:
#    ----------------(eth)----------------------
#           |                  |          |
#          BR1 (Leader) ----- BR2       HOST
#           |                  |
#         ROUTER1             ROUTER2

BR1 = 1
BR2 = 2
ROUTER1 = 3
ROUTER2 = 4
HOST = 5

ROUTE1 = '2402:1234:1234:1234::/64'

class MultiBorderRouters(thread_cert.TestCase):
    USE_MESSAGE_FACTORY = False

    TOPOLOGY = {
        BR1: {
            'name': 'BR_1',
            'allowlist': [BR2, ROUTER1],
            'is_otbr': True,
            'version': '1.2',
        },
        BR2: {
            'name': 'BR_2',
            'allowlist': [BR1, ROUTER2],
            'is_otbr': True,
            'version': '1.2',
        },
        ROUTER1: {
            'name': 'Router_1',
            'allowlist': [BR1],
            'version': '1.2',
        },
        ROUTER2: {
            'name': 'Router_2',
            'allowlist': [BR2],
            'version': '1.2',
        },
        HOST: {
            'name': 'Host',
            'is_host': True,
        }
    }

    def test(self):
        br1 = self.nodes[BR1]
        br2 = self.nodes[BR2]
        router1 = self.nodes[ROUTER1]
        router2 = self.nodes[ROUTER2]
        host = self.nodes[HOST]

        br1.start()
        self.simulator.go(5)
        self.assertEqual('leader', br1.get_state())

        br2.start()
        self.simulator.go(5)
        self.assertEqual('router', br2.get_state())

        router1.start()
        router2.start()
        host.start()
        self.simulator.go(5)
        self.assertEqual('router', router1.get_state())
        self.assertEqual('router', router2.get_state())

        # Manually add ROUTE1 at BR1
        br1.add_route(ROUTE1)
        br1.register_netdata()

        # BRs has installed all external routes in their kernels respectively
        netdata = br1.get_netdata()
        br1_kernel_routes = self.get_routes_from_kernel(br1)
        for route in self.get_routes_from_netdata(netdata, br1.get_addr16()):
            self.assertIn(route, br1_kernel_routes)
        br2_kernel_routes = self.get_routes_from_kernel(br2)
        for route in self.get_routes_from_netdata(netdata, br2.get_addr16()):
            self.assertIn(route, br2_kernel_routes)

        # ROUTE1 has been installed in BR2 but not BR1
        self.assertNotIn(ipaddress.IPv6Network(ROUTE1), br1_kernel_routes)
        self.assertIn(ipaddress.IPv6Network(ROUTE1), br2_kernel_routes)

        # Remove ROUTE1
        br1.remove_route(ROUTE1)
        br1.register_netdata()

        # Verify that external routes are still in kernel routes
        netdata = br1.get_netdata()
        br1_kernel_routes = self.get_routes_from_kernel(br1)
        for route in self.get_routes_from_netdata(netdata, br1.get_addr16()):
            self.assertIn(route, br1_kernel_routes)
        br2_kernel_routes = self.get_routes_from_kernel(br2)
        for route in self.get_routes_from_netdata(netdata, br2.get_addr16()):
            self.assertIn(route, br2_kernel_routes)

        # ROUTE1 has been removed from kernel routes
        self.assertNotIn(ipaddress.IPv6Network(ROUTE1), br1_kernel_routes)
        self.assertNotIn(ipaddress.IPv6Network(ROUTE1), br2_kernel_routes)

        br2.bash('ip link set eth0 down')
        self.simulator.go(10)

        # HOST pings BR2's OMR, the ping should succeed
        self.assertTrue(host.ping_ether(br2.get_ip6_address(address_type=config.ADDRESS_TYPE.OMR)[0]))

    def get_routes_from_netdata(self, netdata, exclude_rloc16):
        routes = []
        for entry in netdata['Routes']:
            items = entry.split()
            prefix = items[0]
            rloc16 = int(items[-1], 16)
            if rloc16 != exclude_rloc16:
                routes.append(ipaddress.IPv6Network(prefix))
        return routes

    def get_routes_from_kernel(self, br):
        routes = []
        for entry in br.bash('ip -6 -d route list dev wpan0'):
            routes.append(ipaddress.IPv6Network(entry.split()[1]))
        return routes

if __name__ == '__main__':
    unittest.main()
<commit_msg>[tests] fix test_external_route test case name (#6828)<commit_after>#!/usr/bin/env python3
#
#  Copyright (c) 2021, The OpenThread Authors.
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#  3. Neither the name of the copyright holder nor the
#     names of its contributors may be used to endorse or promote products
#     derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 'AS IS'
#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
#  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
#
import logging
import unittest

import ipaddress
import config
import thread_cert

# Test description:
#   This test verifies that when external routes are changed in Network Data, a
#   border router will make according updates to routes at the linux kernel.
#
# Topology:
#    ----------------(eth)----------------------
#           |                  |          |
#          BR1 (Leader) ----- BR2       HOST
#           |                  |
#         ROUTER1             ROUTER2

BR1 = 1
BR2 = 2
ROUTER1 = 3
ROUTER2 = 4
HOST = 5

ROUTE1 = '2402:1234:1234:1234::/64'

class ExternalRoutes(thread_cert.TestCase):
    USE_MESSAGE_FACTORY = False

    TOPOLOGY = {
        BR1: {
            'name': 'BR_1',
            'allowlist': [BR2, ROUTER1],
            'is_otbr': True,
            'version': '1.2',
        },
        BR2: {
            'name': 'BR_2',
            'allowlist': [BR1, ROUTER2],
            'is_otbr': True,
            'version': '1.2',
        },
        ROUTER1: {
            'name': 'Router_1',
            'allowlist': [BR1],
            'version': '1.2',
        },
        ROUTER2: {
            'name': 'Router_2',
            'allowlist': [BR2],
            'version': '1.2',
        },
        HOST: {
            'name': 'Host',
            'is_host': True,
        }
    }

    def test(self):
        br1 = self.nodes[BR1]
        br2 = self.nodes[BR2]
        router1 = self.nodes[ROUTER1]
        router2 = self.nodes[ROUTER2]
        host = self.nodes[HOST]

        br1.start()
        self.simulator.go(5)
        self.assertEqual('leader', br1.get_state())

        br2.start()
        self.simulator.go(5)
        self.assertEqual('router', br2.get_state())

        router1.start()
        router2.start()
        host.start()
        self.simulator.go(5)
        self.assertEqual('router', router1.get_state())
        self.assertEqual('router', router2.get_state())

        # Manually add ROUTE1 at BR1
        br1.add_route(ROUTE1)
        br1.register_netdata()

        # BRs has installed all external routes in their kernels respectively
        netdata = br1.get_netdata()
        br1_kernel_routes = self.get_routes_from_kernel(br1)
        for route in self.get_routes_from_netdata(netdata, br1.get_addr16()):
            self.assertIn(route, br1_kernel_routes)
        br2_kernel_routes = self.get_routes_from_kernel(br2)
        for route in self.get_routes_from_netdata(netdata, br2.get_addr16()):
            self.assertIn(route, br2_kernel_routes)

        # ROUTE1 has been installed in BR2 but not BR1
        self.assertNotIn(ipaddress.IPv6Network(ROUTE1), br1_kernel_routes)
        self.assertIn(ipaddress.IPv6Network(ROUTE1), br2_kernel_routes)

        # Remove ROUTE1
        br1.remove_route(ROUTE1)
        br1.register_netdata()

        # Verify that external routes are still in kernel routes
        netdata = br1.get_netdata()
        br1_kernel_routes = self.get_routes_from_kernel(br1)
        for route in self.get_routes_from_netdata(netdata, br1.get_addr16()):
            self.assertIn(route, br1_kernel_routes)
        br2_kernel_routes = self.get_routes_from_kernel(br2)
        for route in self.get_routes_from_netdata(netdata, br2.get_addr16()):
            self.assertIn(route, br2_kernel_routes)

        # ROUTE1 has been removed from kernel routes
        self.assertNotIn(ipaddress.IPv6Network(ROUTE1), br1_kernel_routes)
        self.assertNotIn(ipaddress.IPv6Network(ROUTE1), br2_kernel_routes)

        br2.bash('ip link set eth0 down')
        self.simulator.go(10)

        # HOST pings BR2's OMR, the ping should succeed
        self.assertTrue(host.ping_ether(br2.get_ip6_address(address_type=config.ADDRESS_TYPE.OMR)[0]))

    def get_routes_from_netdata(self, netdata, exclude_rloc16):
        routes = []
        for entry in netdata['Routes']:
            items = entry.split()
            prefix = items[0]
            rloc16 = int(items[-1], 16)
            if rloc16 != exclude_rloc16:
                routes.append(ipaddress.IPv6Network(prefix))
        return routes

    def get_routes_from_kernel(self, br):
        routes = []
        for entry in br.bash('ip -6 -d route list dev wpan0'):
            routes.append(ipaddress.IPv6Network(entry.split()[1]))
        return routes

if __name__ == '__main__':
    unittest.main()
<|endoftext|>