<commit_before># Copyright 2018 Google LLC
#
# 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.

__all__ = [
    'load_component',
    'load_component_from_text',
    'load_component_from_url',
    'load_component_from_file',
]

import sys
from collections import OrderedDict
from ._naming import _sanitize_file_name, _sanitize_python_function_name, generate_unique_name_conversion_table
from ._yaml_utils import load_yaml
from ._structures import ComponentSpec
from ._structures import *
from ._data_passing import serialize_value, type_name_to_type
from kfp.dsl import PipelineParam
from kfp.dsl.types import verify_type_compatibility

_default_component_name = 'Component'

def load_component(filename=None, url=None, text=None):
    '''
    Loads component from text, file or URL and creates a task factory function

    Only one argument should be specified.

    Args:
        filename: Path of local file containing the component definition.
        url: The URL of the component file data
        text: A string containing the component file data.

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    #This function should be called load_task_factory since it returns a factory function.
    #The real load_component function should produce an object with component properties (e.g. name, description, inputs/outputs).
    #TODO: Change this function to return component spec object but it should be callable to construct tasks.
    non_null_args_count = len([name for name, value in locals().items() if value != None])
    if non_null_args_count != 1:
        raise ValueError('Need to specify exactly one source')
    if filename:
        return load_component_from_file(filename)
    elif url:
        return load_component_from_url(url)
    elif text:
        return load_component_from_text(text)
    else:
        raise ValueError('Need to specify a source')

def load_component_from_url(url):
    '''
    Loads component from URL and creates a task factory function

    Args:
        url: The URL of the component file data

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    if url is None:
        raise TypeError

    #Handling Google Cloud Storage URIs
    if url.startswith('gs://'):
        #Replacing the gs:// URI with https:// URI (works for public objects)
        url = 'https://storage.googleapis.com/' + url[len('gs://'):]

    import requests
    resp = requests.get(url)
    resp.raise_for_status()
    component_ref = ComponentReference(url=url)
    return _load_component_from_yaml_or_zip_bytes(resp.content, url, component_ref)

def load_component_from_file(filename):
    '''
    Loads component from file and creates a task factory function

    Args:
        filename: Path of local file containing the component definition.

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    if filename is None:
        raise TypeError
    with open(filename, 'rb') as component_stream:
        return _load_component_from_yaml_or_zip_stream(component_stream, filename)

def load_component_from_text(text):
    '''
    Loads component from text and creates a task factory function

    Args:
        text: A string containing the component file data.

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    if text is None:
        raise TypeError
    return _create_task_factory_from_component_text(text, None)

_COMPONENT_FILE_NAME_IN_ARCHIVE = 'component.yaml'

def _load_component_from_yaml_or_zip_bytes(bytes, component_filename=None, component_ref: ComponentReference = None):
    import io
    component_stream = io.BytesIO(bytes)
    return _load_component_from_yaml_or_zip_stream(component_stream, component_filename, component_ref)

def _load_component_from_yaml_or_zip_stream(stream, component_filename=None, component_ref: ComponentReference = None):
    '''Loads component from a stream and creates a task factory function.
    The stream can be YAML or a zip file with a component.yaml file inside.
    '''
    import zipfile
    stream.seek(0)
    if zipfile.is_zipfile(stream):
        stream.seek(0)
        with zipfile.ZipFile(stream) as zip_obj:
            with zip_obj.open(_COMPONENT_FILE_NAME_IN_ARCHIVE) as component_stream:
                return _create_task_factory_from_component_text(component_stream, component_filename, component_ref)
    else:
        stream.seek(0)
        return _create_task_factory_from_component_text(stream, component_filename, component_ref)

def _create_task_factory_from_component_text(text_or_file, component_filename=None, component_ref: ComponentReference = None):
    component_dict = load_yaml(text_or_file)
    return _create_task_factory_from_component_dict(component_dict, component_filename, component_ref)

def _create_task_factory_from_component_dict(component_dict, component_filename=None, component_ref: ComponentReference = None):
    component_spec = ComponentSpec.from_dict(component_dict)
    return _create_task_factory_from_component_spec(component_spec, component_filename, component_ref)

_inputs_dir = '/tmp/inputs'
_outputs_dir = '/tmp/outputs'
_single_io_file_name = 'data'

def _generate_input_file_name(port_name):
    return _inputs_dir + '/' + _sanitize_file_name(port_name) + '/' + _single_io_file_name

def _generate_output_file_name(port_name):
    return _outputs_dir + '/' + _sanitize_file_name(port_name) + '/' + _single_io_file_name

#Holds the transformation functions that are called each time TaskSpec instance is created from a component. If there are multiple handlers, the last one is used.
_created_task_transformation_handler = []

#TODO: Move to the dsl.Pipeline context class
from . import _dsl_bridge
_created_task_transformation_handler.append(_dsl_bridge.create_container_op_from_task)

class _DefaultValue:
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return repr(self.value)

#TODO: Refactor the function to make it shorter
def _create_task_factory_from_component_spec(component_spec:ComponentSpec, component_filename=None, component_ref: ComponentReference = None):
    name = component_spec.name or _default_component_name

    func_docstring_lines = []
    if component_spec.name:
        func_docstring_lines.append(component_spec.name)
    if component_spec.description:
        func_docstring_lines.append(component_spec.description)

    inputs_list = component_spec.inputs or [] #List[InputSpec]
    input_names = [input.name for input in inputs_list]

    #Creating the name translation tables : Original <-> Pythonic
    input_name_to_pythonic = generate_unique_name_conversion_table(input_names, _sanitize_python_function_name)
    pythonic_name_to_input_name = {v: k for k, v in input_name_to_pythonic.items()}

    if component_ref is None:
        component_ref = ComponentReference(spec=component_spec, url=component_filename)
    else:
        component_ref.spec = component_spec

    def create_task_from_component_and_arguments(pythonic_arguments):
        arguments = {}
        # Not checking for missing or extra arguments since the dynamic factory function checks that
        for argument_name, argument_value in pythonic_arguments.items():
            if isinstance(argument_value, _DefaultValue): # Skipping passing arguments for optional values that have not been overridden.
                continue
            input_name = pythonic_name_to_input_name[argument_name]
            input_type = component_spec._inputs_dict[input_name].type

            if isinstance(argument_value, (GraphInputArgument, TaskOutputArgument, PipelineParam)):
                # argument_value is a reference

                if isinstance(argument_value, PipelineParam):
                    reference_type = argument_value.param_type
                    argument_value = str(argument_value)
                elif isinstance(argument_value, TaskOutputArgument):
                    reference_type = argument_value.task_output.type
                else:
                    reference_type = None

                verify_type_compatibility(reference_type, input_type, 'Incompatible argument passed to the input "{}" of component "{}": '.format(input_name, component_spec.name))

                arguments[input_name] = argument_value
            else:
                # argument_value is a constant value
                serialized_argument_value = serialize_value(argument_value, input_type)
                arguments[input_name] = serialized_argument_value

        task = TaskSpec(
            component_ref=component_ref,
            arguments=arguments,
        )
        task._init_outputs()

        if isinstance(component_spec.implementation, GraphImplementation):
            return _resolve_graph_task(task, component_spec)

        if _created_task_transformation_handler:
            task = _created_task_transformation_handler[-1](task)
        return task

    import inspect
    from . import _dynamic

    #Reordering the inputs since in Python optional parameters must come after required parameters
    reordered_input_list = [input for input in inputs_list if input.default is None and not input.optional] + [input for input in inputs_list if not (input.default is None and not input.optional)]

    def component_default_to_func_default(component_default: str, is_optional: bool):
        if is_optional:
            return _DefaultValue(component_default)
        if component_default is not None:
            return component_default
        return inspect.Parameter.empty

    input_parameters = [
        _dynamic.KwParameter(
            input_name_to_pythonic[port.name],
            annotation=(type_name_to_type.get(str(port.type), str(port.type)) if port.type else inspect.Parameter.empty),
            default=component_default_to_func_default(port.default, port.optional),
        )
        for port in reordered_input_list
    ]
    factory_function_parameters = input_parameters #Outputs are no longer part of the task factory function signature. The paths are always generated by the system.

    task_factory = _dynamic.create_function_from_parameters(
        create_task_from_component_and_arguments,
        factory_function_parameters,
        documentation='\n'.join(func_docstring_lines),
        func_name=name,
        func_filename=component_filename if (component_filename and (component_filename.endswith('.yaml') or component_filename.endswith('.yml'))) else None,
    )
    task_factory.component_spec = component_spec
    return task_factory

def _resolve_graph_task(graph_task: TaskSpec, graph_component_spec: ComponentSpec) -> TaskSpec:
    from ..components import ComponentStore
    component_store = ComponentStore.default_store

    graph = graph_component_spec.implementation.graph

    outputs_of_tasks = {}
    def resolve_argument(argument):
        if isinstance(argument, (str, int, float, bool)):
            return argument
        elif isinstance(argument, GraphInputArgument):
            return graph_task.arguments[argument.input_name]
        elif isinstance(argument, TaskOutputArgument):
            upstream_task_output_ref = argument.task_output
            upstream_task_outputs = outputs_of_tasks[upstream_task_output_ref.task_id]
            upstream_task_output = upstream_task_outputs[upstream_task_output_ref.output_name]
            return upstream_task_output
        else:
            raise TypeError('Argument for input has unexpected type "{}".'.format(type(argument)))

    for task_id, task_spec in graph._toposorted_tasks.items(): # Cannot use graph.tasks here since they might be listed not in dependency order. Especially on python <3.6 where the dicts do not preserve ordering
        task_factory = component_store._load_component_from_ref(task_spec.component_ref)
        task_arguments = {input_name: resolve_argument(argument) for input_name, argument in task_spec.arguments.items()}
        task_component_spec = task_spec.component_ref.spec

        input_name_to_pythonic = generate_unique_name_conversion_table([input.name for input in task_component_spec.inputs], _sanitize_python_function_name)
        output_name_to_pythonic = generate_unique_name_conversion_table([output.name for output in task_component_spec.outputs], _sanitize_python_function_name)
        pythonic_output_name_to_original = {pythonic_name: original_name for original_name, pythonic_name in output_name_to_pythonic.items()}
        pythonic_task_arguments = {input_name_to_pythonic[input_name]: argument for input_name, argument in task_arguments.items()}

        task_obj = task_factory(**pythonic_task_arguments)
        task_outputs_with_pythonic_names = task_obj.outputs
        task_outputs_with_original_names = {pythonic_output_name_to_original[pythonic_output_name]: output_value for pythonic_output_name, output_value in task_outputs_with_pythonic_names.items()}
        outputs_of_tasks[task_id] = task_outputs_with_original_names

    resolved_graph_outputs = OrderedDict([(output_name, resolve_argument(argument)) for output_name, argument in graph.output_values.items()])

    # For resolved graph component tasks task.outputs point to the actual tasks that originally produced the output that is later returned from the graph
    graph_task.output_references = graph_task.outputs
    graph_task.outputs = resolved_graph_outputs

    return graph_task
<commit_msg>SDK - Components - Fixed small bugs in graph component resolving (#2269)<commit_after># Copyright 2018 Google LLC
#
# 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.

__all__ = [
    'load_component',
    'load_component_from_text',
    'load_component_from_url',
    'load_component_from_file',
]

import sys
from collections import OrderedDict
from ._naming import _sanitize_file_name, _sanitize_python_function_name, generate_unique_name_conversion_table
from ._yaml_utils import load_yaml
from ._structures import ComponentSpec
from ._structures import *
from ._data_passing import serialize_value, type_name_to_type
from kfp.dsl import PipelineParam
from kfp.dsl.types import verify_type_compatibility

_default_component_name = 'Component'

def load_component(filename=None, url=None, text=None):
    '''
    Loads component from text, file or URL and creates a task factory function

    Only one argument should be specified.

    Args:
        filename: Path of local file containing the component definition.
        url: The URL of the component file data
        text: A string containing the component file data.

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    #This function should be called load_task_factory since it returns a factory function.
    #The real load_component function should produce an object with component properties (e.g. name, description, inputs/outputs).
    #TODO: Change this function to return component spec object but it should be callable to construct tasks.
    non_null_args_count = len([name for name, value in locals().items() if value != None])
    if non_null_args_count != 1:
        raise ValueError('Need to specify exactly one source')
    if filename:
        return load_component_from_file(filename)
    elif url:
        return load_component_from_url(url)
    elif text:
        return load_component_from_text(text)
    else:
        raise ValueError('Need to specify a source')

def load_component_from_url(url):
    '''
    Loads component from URL and creates a task factory function

    Args:
        url: The URL of the component file data

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    if url is None:
        raise TypeError

    #Handling Google Cloud Storage URIs
    if url.startswith('gs://'):
        #Replacing the gs:// URI with https:// URI (works for public objects)
        url = 'https://storage.googleapis.com/' + url[len('gs://'):]

    import requests
    resp = requests.get(url)
    resp.raise_for_status()
    component_ref = ComponentReference(url=url)
    return _load_component_from_yaml_or_zip_bytes(resp.content, url, component_ref)

def load_component_from_file(filename):
    '''
    Loads component from file and creates a task factory function

    Args:
        filename: Path of local file containing the component definition.

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    if filename is None:
        raise TypeError
    with open(filename, 'rb') as component_stream:
        return _load_component_from_yaml_or_zip_stream(component_stream, filename)

def load_component_from_text(text):
    '''
    Loads component from text and creates a task factory function

    Args:
        text: A string containing the component file data.

    Returns:
        A factory function with a strongly-typed signature.
        Once called with the required arguments, the factory constructs a pipeline task instance (ContainerOp).
    '''
    if text is None:
        raise TypeError
    return _create_task_factory_from_component_text(text, None)

_COMPONENT_FILE_NAME_IN_ARCHIVE = 'component.yaml'

def _load_component_from_yaml_or_zip_bytes(bytes, component_filename=None, component_ref: ComponentReference = None):
    import io
    component_stream = io.BytesIO(bytes)
    return _load_component_from_yaml_or_zip_stream(component_stream, component_filename, component_ref)

def _load_component_from_yaml_or_zip_stream(stream, component_filename=None, component_ref: ComponentReference = None):
    '''Loads component from a stream and creates a task factory function.
    The stream can be YAML or a zip file with a component.yaml file inside.
    '''
    import zipfile
    stream.seek(0)
    if zipfile.is_zipfile(stream):
        stream.seek(0)
        with zipfile.ZipFile(stream) as zip_obj:
            with zip_obj.open(_COMPONENT_FILE_NAME_IN_ARCHIVE) as component_stream:
                return _create_task_factory_from_component_text(component_stream, component_filename, component_ref)
    else:
        stream.seek(0)
        return _create_task_factory_from_component_text(stream, component_filename, component_ref)

def _create_task_factory_from_component_text(text_or_file, component_filename=None, component_ref: ComponentReference = None):
    component_dict = load_yaml(text_or_file)
    return _create_task_factory_from_component_dict(component_dict, component_filename, component_ref)

def _create_task_factory_from_component_dict(component_dict, component_filename=None, component_ref: ComponentReference = None):
    component_spec = ComponentSpec.from_dict(component_dict)
    return _create_task_factory_from_component_spec(component_spec, component_filename, component_ref)

_inputs_dir = '/tmp/inputs'
_outputs_dir = '/tmp/outputs'
_single_io_file_name = 'data'

def _generate_input_file_name(port_name):
    return _inputs_dir + '/' + _sanitize_file_name(port_name) + '/' + _single_io_file_name

def _generate_output_file_name(port_name):
    return _outputs_dir + '/' + _sanitize_file_name(port_name) + '/' + _single_io_file_name

#Holds the transformation functions that are called each time TaskSpec instance is created from a component. If there are multiple handlers, the last one is used.
_created_task_transformation_handler = []

#TODO: Move to the dsl.Pipeline context class
from . import _dsl_bridge
_created_task_transformation_handler.append(_dsl_bridge.create_container_op_from_task)

class _DefaultValue:
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return repr(self.value)

#TODO: Refactor the function to make it shorter
def _create_task_factory_from_component_spec(component_spec:ComponentSpec, component_filename=None, component_ref: ComponentReference = None):
    name = component_spec.name or _default_component_name

    func_docstring_lines = []
    if component_spec.name:
        func_docstring_lines.append(component_spec.name)
    if component_spec.description:
        func_docstring_lines.append(component_spec.description)

    inputs_list = component_spec.inputs or [] #List[InputSpec]
    input_names = [input.name for input in inputs_list]

    #Creating the name translation tables : Original <-> Pythonic
    input_name_to_pythonic = generate_unique_name_conversion_table(input_names, _sanitize_python_function_name)
    pythonic_name_to_input_name = {v: k for k, v in input_name_to_pythonic.items()}

    if component_ref is None:
        component_ref = ComponentReference(spec=component_spec, url=component_filename)
    else:
        component_ref.spec = component_spec

    def create_task_from_component_and_arguments(pythonic_arguments):
        arguments = {}
        # Not checking for missing or extra arguments since the dynamic factory function checks that
        for argument_name, argument_value in pythonic_arguments.items():
            if isinstance(argument_value, _DefaultValue): # Skipping passing arguments for optional values that have not been overridden.
                continue
            input_name = pythonic_name_to_input_name[argument_name]
            input_type = component_spec._inputs_dict[input_name].type

            if isinstance(argument_value, (GraphInputArgument, TaskOutputArgument, PipelineParam)):
                # argument_value is a reference

                if isinstance(argument_value, PipelineParam):
                    reference_type = argument_value.param_type
                    argument_value = str(argument_value)
                elif isinstance(argument_value, TaskOutputArgument):
                    reference_type = argument_value.task_output.type
                else:
                    reference_type = None

                verify_type_compatibility(reference_type, input_type, 'Incompatible argument passed to the input "{}" of component "{}": '.format(input_name, component_spec.name))

                arguments[input_name] = argument_value
            else:
                # argument_value is a constant value
                serialized_argument_value = serialize_value(argument_value, input_type)
                arguments[input_name] = serialized_argument_value

        task = TaskSpec(
            component_ref=component_ref,
            arguments=arguments,
        )
        task._init_outputs()

        if isinstance(component_spec.implementation, GraphImplementation):
            return _resolve_graph_task(task, component_spec)

        if _created_task_transformation_handler:
            task = _created_task_transformation_handler[-1](task)
        return task

    import inspect
    from . import _dynamic

    #Reordering the inputs since in Python optional parameters must come after required parameters
    reordered_input_list = [input for input in inputs_list if input.default is None and not input.optional] + [input for input in inputs_list if not (input.default is None and not input.optional)]

    def component_default_to_func_default(component_default: str, is_optional: bool):
        if is_optional:
            return _DefaultValue(component_default)
        if component_default is not None:
            return component_default
        return inspect.Parameter.empty

    input_parameters = [
        _dynamic.KwParameter(
            input_name_to_pythonic[port.name],
            annotation=(type_name_to_type.get(str(port.type), str(port.type)) if port.type else inspect.Parameter.empty),
            default=component_default_to_func_default(port.default, port.optional),
        )
        for port in reordered_input_list
    ]
    factory_function_parameters = input_parameters #Outputs are no longer part of the task factory function signature. The paths are always generated by the system.

    task_factory = _dynamic.create_function_from_parameters(
        create_task_from_component_and_arguments,
        factory_function_parameters,
        documentation='\n'.join(func_docstring_lines),
        func_name=name,
        func_filename=component_filename if (component_filename and (component_filename.endswith('.yaml') or component_filename.endswith('.yml'))) else None,
    )
    task_factory.component_spec = component_spec
    return task_factory

def _resolve_graph_task(graph_task: TaskSpec, graph_component_spec: ComponentSpec) -> TaskSpec:
    from ..components import ComponentStore
    component_store = ComponentStore.default_store

    graph = graph_component_spec.implementation.graph

    graph_input_arguments = {input.name: input.default for input in graph_component_spec.inputs if input.default is not None}
    graph_input_arguments.update(graph_task.arguments)

    outputs_of_tasks = {}
    def resolve_argument(argument):
        if isinstance(argument, (str, int, float, bool)):
            return argument
        elif isinstance(argument, GraphInputArgument):
            return graph_input_arguments[argument.input_name]
        elif isinstance(argument, TaskOutputArgument):
            upstream_task_output_ref = argument.task_output
            upstream_task_outputs = outputs_of_tasks[upstream_task_output_ref.task_id]
            upstream_task_output = upstream_task_outputs[upstream_task_output_ref.output_name]
            return upstream_task_output
        else:
            raise TypeError('Argument for input has unexpected type "{}".'.format(type(argument)))

    for task_id, task_spec in graph._toposorted_tasks.items(): # Cannot use graph.tasks here since they might be listed not in dependency order. Especially on python <3.6 where the dicts do not preserve ordering
        task_factory = component_store._load_component_from_ref(task_spec.component_ref)
        # TODO: Handle the case when optional graph component input is passed to optional task component input
        task_arguments = {input_name: resolve_argument(argument) for input_name, argument in task_spec.arguments.items()}
        task_component_spec = task_factory.component_spec

        input_name_to_pythonic = generate_unique_name_conversion_table([input.name for input in task_component_spec.inputs or []], _sanitize_python_function_name)
        output_name_to_pythonic = generate_unique_name_conversion_table([output.name for output in task_component_spec.outputs or []], _sanitize_python_function_name)
        pythonic_output_name_to_original = {pythonic_name: original_name for original_name, pythonic_name in output_name_to_pythonic.items()}
        pythonic_task_arguments = {input_name_to_pythonic[input_name]: argument for input_name, argument in task_arguments.items()}

        task_obj = task_factory(**pythonic_task_arguments)
        task_outputs_with_pythonic_names = task_obj.outputs
        task_outputs_with_original_names = {pythonic_output_name_to_original[pythonic_output_name]: output_value for pythonic_output_name, output_value in task_outputs_with_pythonic_names.items()}
        outputs_of_tasks[task_id] = task_outputs_with_original_names

    resolved_graph_outputs = OrderedDict([(output_name, resolve_argument(argument)) for output_name, argument in graph.output_values.items()])

    # For resolved graph component tasks task.outputs point to the actual tasks that originally produced the output that is later returned from the graph
    graph_task.output_references = graph_task.outputs
    graph_task.outputs = resolved_graph_outputs

    return graph_task
<|endoftext|>


<commit_before>#!/usr/bin/env python
################################################################################
#
# pysearch.py
#
# A file search utility implemented in python (2.x)
#
################################################################################
import sys

from searcher import Searcher
from searchoptions import SearchOptions

def log(message):
    print(message)

def main():
    searchoptions = SearchOptions()

    settings = None
    try:
        settings = searchoptions.search_settings_from_args(sys.argv[1:])
    except Exception, e:
        log('\nException: {0!s}\n'.format(e))
        searchoptions.usage()

    if settings.printusage:
        searchoptions.usage()

    if settings.printversion:
        log('Version: 0.1')
        sys.exit(1)

    if settings.debug:
        log('settings: {0!s}'.format(settings))

    try:
        searcher = Searcher(settings)
        searcher.search()

        # print the results
        if settings.printresults:
            log('')
            searcher.print_results()

        if settings.listdirs:
            dir_list = searcher.get_matching_dirs()
            if dir_list:
                log('\nDirectories with matches (%d):' % len(dir_list))
                for d in dir_list:
                    log(d)

        if settings.listfiles:
            file_list = searcher.get_matching_files()
            if file_list:
                log('\nFiles with matches (%d):' % len(file_list))
                for f in file_list:
                    log(f)

        if settings.listlines:
            line_list = searcher.get_matching_lines()
            if line_list:
                log('\nLines with matches (%d):' % len(line_list))
                for line in line_list:
                    log(line)

    except AssertionError as e:
        log('ERROR: {0!s}\n'.format(e))
        searchoptions.usage()
    except KeyboardInterrupt:
        log('')
        sys.exit(0)

if __name__ == '__main__':
    main()
<commit_msg>Added docstring to log function<commit_after>#!/usr/bin/env python
################################################################################
#
# pysearch.py
#
# A file search utility implemented in python (2.x)
#
################################################################################
import sys

from searcher import Searcher
from searchoptions import SearchOptions

def log(message):
    """log a message (for now just print to stdout)"""
    print message

def main():
    searchoptions = SearchOptions()

    settings = None
    try:
        settings = searchoptions.search_settings_from_args(sys.argv[1:])
    except Exception, e:
        log('\nException: {0!s}\n'.format(e))
        searchoptions.usage()

    if settings.printusage:
        searchoptions.usage()

    if settings.printversion:
        log('Version: 0.1')
        sys.exit(1)

    if settings.debug:
        log('settings: {0!s}'.format(settings))

    try:
        searcher = Searcher(settings)
        searcher.search()

        # print the results
        if settings.printresults:
            log('')
            searcher.print_results()

        if settings.listdirs:
            dir_list = searcher.get_matching_dirs()
            if dir_list:
                log('\nDirectories with matches (%d):' % len(dir_list))
                for d in dir_list:
                    log(d)

        if settings.listfiles:
            file_list = searcher.get_matching_files()
            if file_list:
                log('\nFiles with matches (%d):' % len(file_list))
                for f in file_list:
                    log(f)

        if settings.listlines:
            line_list = searcher.get_matching_lines()
            if line_list:
                log('\nLines with matches (%d):' % len(line_list))
                for line in line_list:
                    log(line)

    except AssertionError as e:
        log('ERROR: {0!s}\n'.format(e))
        searchoptions.usage()
    except KeyboardInterrupt:
        log('')
        sys.exit(0)

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


<commit_before>import logging
from types import FunctionType

import ray
import ray.cloudpickle as pickle
from ray.experimental.internal_kv import _internal_kv_initialized, \
    _internal_kv_get, _internal_kv_put
from ray.tune.error import TuneError

TRAINABLE_CLASS = "trainable_class"
ENV_CREATOR = "env_creator"
RLLIB_MODEL = "rllib_model"
RLLIB_PREPROCESSOR = "rllib_preprocessor"
RLLIB_ACTION_DIST = "rllib_action_dist"
TEST = "__test__"
KNOWN_CATEGORIES = [
    TRAINABLE_CLASS, ENV_CREATOR, RLLIB_MODEL, RLLIB_PREPROCESSOR,
    RLLIB_ACTION_DIST, TEST
]

logger = logging.getLogger(__name__)

def has_trainable(trainable_name):
    return _global_registry.contains(TRAINABLE_CLASS, trainable_name)

def get_trainable_cls(trainable_name):
    validate_trainable(trainable_name)
    return _global_registry.get(TRAINABLE_CLASS, trainable_name)

def validate_trainable(trainable_name):
    if not has_trainable(trainable_name):
        # Make sure everything rllib-related is registered.
        from ray.rllib import _register_all
        _register_all()
        if not has_trainable(trainable_name):
            raise TuneError("Unknown trainable: " + trainable_name)

def register_trainable(name, trainable, warn=True):
    """Register a trainable function or class.

    This enables a class or function to be accessed on every Ray process
    in the cluster.

    Args:
        name (str): Name to register.
        trainable (obj): Function or tune.Trainable class. Functions must
            take (config, status_reporter) as arguments and will be
            automatically converted into a class during registration.
    """

    from ray.tune.trainable import Trainable
    from ray.tune.function_runner import wrap_function

    if isinstance(trainable, type):
        logger.debug("Detected class for trainable.")
    elif isinstance(trainable, FunctionType):
        logger.debug("Detected function for trainable.")
        trainable = wrap_function(trainable, warn=warn)
    elif callable(trainable):
        logger.info(
            "Detected unknown callable for trainable. Converting to class.")
        trainable = wrap_function(trainable, warn=warn)

    if not issubclass(trainable, Trainable):
        raise TypeError("Second argument must be convertable to Trainable",
                        trainable)
    _global_registry.register(TRAINABLE_CLASS, name, trainable)

def register_env(name, env_creator):
    """Register a custom environment for use with RLlib.

    This enables the environment to be accessed on every Ray process
    in the cluster.

    Args:
        name (str): Name to register.
        env_creator (obj): Function that creates an env.
    """

    if not isinstance(env_creator, FunctionType):
        raise TypeError("Second argument must be a function.", env_creator)
    _global_registry.register(ENV_CREATOR, name, env_creator)

def check_serializability(key, value):
    _global_registry.register(TEST, key, value)

def _make_key(category, key):
    """Generate a binary key for the given category and key.

    Args:
        category (str): The category of the item
        key (str): The unique identifier for the item

    Returns:
        The key to use for storing a the value.
    """
    return (b"TuneRegistry:" + category.encode("ascii") + b"/" +
            key.encode("ascii"))

class _Registry:
    def __init__(self):
        self._to_flush = {}

    def register(self, category, key, value):
        """Registers the value with the global registry.

        Raises:
            PicklingError if unable to pickle to provided file.
        """
        if category not in KNOWN_CATEGORIES:
            from ray.tune import TuneError
            raise TuneError("Unknown category {} not among {}".format(
                category, KNOWN_CATEGORIES))
        self._to_flush[(category, key)] = pickle.dumps_debug(value)
        if _internal_kv_initialized():
            self.flush_values()

    def contains(self, category, key):
        if _internal_kv_initialized():
            value = _internal_kv_get(_make_key(category, key))
            return value is not None
        else:
            return (category, key) in self._to_flush

    def get(self, category, key):
        if _internal_kv_initialized():
            value = _internal_kv_get(_make_key(category, key))
            if value is None:
                raise ValueError(
                    "Registry value for {}/{} doesn't exist.".format(
                        category, key))
            return pickle.loads(value)
        else:
            return pickle.loads(self._to_flush[(category, key)])

    def flush_values(self):
        for (category, key), value in self._to_flush.items():
            _internal_kv_put(_make_key(category, key), value, overwrite=True)
        self._to_flush.clear()

_global_registry = _Registry()
ray.worker._post_init_hooks.append(_global_registry.flush_values)

class _ParameterRegistry:
    def __init__(self):
        self.to_flush = {}
        self.references = {}

    def put(self, k, v):
        self.to_flush[k] = v
        if ray.is_initialized():
            self.flush()

    def get(self, k):
        if not ray.is_initialized():
            return self.to_flush[k]
        return ray.get(self.references[k])

    def flush(self):
        for k, v in self.to_flush.items():
            self.references[k] = ray.put(v)
        self.to_flush.clear()

parameter_registry = _ParameterRegistry()
ray.worker._post_init_hooks.append(parameter_registry.flush)
<commit_msg>[tune] Callable accepted for register_env (#12618)<commit_after>import logging
from types import FunctionType

import ray
import ray.cloudpickle as pickle
from ray.experimental.internal_kv import _internal_kv_initialized, \
    _internal_kv_get, _internal_kv_put
from ray.tune.error import TuneError

TRAINABLE_CLASS = "trainable_class"
ENV_CREATOR = "env_creator"
RLLIB_MODEL = "rllib_model"
RLLIB_PREPROCESSOR = "rllib_preprocessor"
RLLIB_ACTION_DIST = "rllib_action_dist"
TEST = "__test__"
KNOWN_CATEGORIES = [
    TRAINABLE_CLASS, ENV_CREATOR, RLLIB_MODEL, RLLIB_PREPROCESSOR,
    RLLIB_ACTION_DIST, TEST
]

logger = logging.getLogger(__name__)

def has_trainable(trainable_name):
    return _global_registry.contains(TRAINABLE_CLASS, trainable_name)

def get_trainable_cls(trainable_name):
    validate_trainable(trainable_name)
    return _global_registry.get(TRAINABLE_CLASS, trainable_name)

def validate_trainable(trainable_name):
    if not has_trainable(trainable_name):
        # Make sure everything rllib-related is registered.
        from ray.rllib import _register_all
        _register_all()
        if not has_trainable(trainable_name):
            raise TuneError("Unknown trainable: " + trainable_name)

def register_trainable(name, trainable, warn=True):
    """Register a trainable function or class.

    This enables a class or function to be accessed on every Ray process
    in the cluster.

    Args:
        name (str): Name to register.
        trainable (obj): Function or tune.Trainable class. Functions must
            take (config, status_reporter) as arguments and will be
            automatically converted into a class during registration.
    """

    from ray.tune.trainable import Trainable
    from ray.tune.function_runner import wrap_function

    if isinstance(trainable, type):
        logger.debug("Detected class for trainable.")
    elif isinstance(trainable, FunctionType):
        logger.debug("Detected function for trainable.")
        trainable = wrap_function(trainable, warn=warn)
    elif callable(trainable):
        logger.info(
            "Detected unknown callable for trainable. Converting to class.")
        trainable = wrap_function(trainable, warn=warn)

    if not issubclass(trainable, Trainable):
        raise TypeError("Second argument must be convertable to Trainable",
                        trainable)
    _global_registry.register(TRAINABLE_CLASS, name, trainable)

def register_env(name, env_creator):
    """Register a custom environment for use with RLlib.

    This enables the environment to be accessed on every Ray process
    in the cluster.

    Args:
        name (str): Name to register.
        env_creator (obj): Callable that creates an env.
    """

    if not callable(env_creator):
        raise TypeError("Second argument must be callable.", env_creator)
    _global_registry.register(ENV_CREATOR, name, env_creator)

def check_serializability(key, value):
    _global_registry.register(TEST, key, value)

def _make_key(category, key):
    """Generate a binary key for the given category and key.

    Args:
        category (str): The category of the item
        key (str): The unique identifier for the item

    Returns:
        The key to use for storing a the value.
    """
    return (b"TuneRegistry:" + category.encode("ascii") + b"/" +
            key.encode("ascii"))

class _Registry:
    def __init__(self):
        self._to_flush = {}

    def register(self, category, key, value):
        """Registers the value with the global registry.

        Raises:
            PicklingError if unable to pickle to provided file.
        """
        if category not in KNOWN_CATEGORIES:
            from ray.tune import TuneError
            raise TuneError("Unknown category {} not among {}".format(
                category, KNOWN_CATEGORIES))
        self._to_flush[(category, key)] = pickle.dumps_debug(value)
        if _internal_kv_initialized():
            self.flush_values()

    def contains(self, category, key):
        if _internal_kv_initialized():
            value = _internal_kv_get(_make_key(category, key))
            return value is not None
        else:
            return (category, key) in self._to_flush

    def get(self, category, key):
        if _internal_kv_initialized():
            value = _internal_kv_get(_make_key(category, key))
            if value is None:
                raise ValueError(
                    "Registry value for {}/{} doesn't exist.".format(
                        category, key))
            return pickle.loads(value)
        else:
            return pickle.loads(self._to_flush[(category, key)])

    def flush_values(self):
        for (category, key), value in self._to_flush.items():
            _internal_kv_put(_make_key(category, key), value, overwrite=True)
        self._to_flush.clear()

_global_registry = _Registry()
ray.worker._post_init_hooks.append(_global_registry.flush_values)

class _ParameterRegistry:
    def __init__(self):
        self.to_flush = {}
        self.references = {}

    def put(self, k, v):
        self.to_flush[k] = v
        if ray.is_initialized():
            self.flush()

    def get(self, k):
        if not ray.is_initialized():
            return self.to_flush[k]
        return ray.get(self.references[k])

    def flush(self):
        for k, v in self.to_flush.items():
            self.references[k] = ray.put(v)
        self.to_flush.clear()

parameter_registry = _ParameterRegistry()
ray.worker._post_init_hooks.append(parameter_registry.flush)
<|endoftext|>


<commit_before>import xformer
import xformer.data

import os
import torch
import time
import itertools
import argparse
import wandb
from contextlib import nullcontext

from torch.profiler import profile, ProfilerAction

class Dataset:
    def __init__(self, path):
        self.data = torch.load(path)

    def __len__(self):
        return len(next(iter(self.data.values())))

    def __getitem__(self, i):
        return {k: v[i] for (k, v) in self.data.items()}

def fwd_and_loss(model, xent, record):
    batch = record["positions"].to(device=model.device, dtype=torch.long)
    logits = model(batch[:, :-1])
    targets = batch[:, 1:]
    loss = (xent(logits.permute(0, 2, 1), targets) * record["mask"][:, :-1]).mean()
    return loss

def main():
    parser = argparse.ArgumentParser(description="Train a transformer")
    parser.add_argument("--layers", type=int, default=2, help="Number of layers")
    parser.add_argument("--d_model", type=int, default=None, help="embedding dimension")
    parser.add_argument("--d_head", type=int, default=32, help="head dimension")
    parser.add_argument(
        "--n_ctx", type=int, default=1024, help="maximum context length"
    )
    parser.add_argument(
        "--pe", type=str, default=None, help="positional encoding (sin, learned, none)"
    )

    parser.add_argument("--data", type=str, default="data/positions", help="datasource")

    parser.add_argument("--batch", type=int, default=64, help="batch size")
    parser.add_argument("--minibatch", type=int, default=4, help="minibatch")

    parser.add_argument(
        "--test-batches", type=int, default=16, help="number of test minibatches"
    )
    parser.add_argument(
        "--test-freq", type=int, default=1000, help="measure test loss every N steps"
    )

    parser.add_argument(
        "--device", type=str, choices=("cpu", "cuda"), default="cuda", help="device"
    )

    parser.add_argument("--wandb", action="store_true", default=False)
    parser.add_argument("--no-wandb", action="store_false", dest="wandb")

    parser.add_argument("--lr", type=float, default=0.001, help="learning rate")
    parser.add_argument("--steps", type=int, default=None)
    parser.add_argument("--profile-steps", type=str, default=None)
    parser.add_argument("--positions", type=int, default=None)

    args = parser.parse_args()

    cfg = xformer.Config(
        n_layer=args.layers,
        d_model=args.d_model or 128 * args.layers,
        d_head=args.d_head,
        n_ctx=args.n_ctx,
        n_vocab=256,
        autoregressive_mask=True,
    )
    if args.pe is not None:
        cfg.positional_encoding = args.pe

    train_ds = Dataset(args.data + "-train.pt")
    loader = torch.utils.data.DataLoader(
        train_ds,
        batch_size=args.minibatch,
        pin_memory=True,
        num_workers=1,
        shuffle=True,
    )

    test_ds = Dataset(args.data + "-test.pt")
    test_loader = torch.utils.data.DataLoader(
        test_ds,
        batch_size=args.minibatch,
        pin_memory=True,
        num_workers=1,
    )
    test_batches = list(itertools.islice(test_loader, args.test_batches))

    model = xformer.Transformer(cfg, dtype=torch.float32, device=args.device)

    xent = torch.nn.CrossEntropyLoss(reduction="none")
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr)

    assert args.batch % args.minibatch == 0, "minibatch must divide batch"
    steps_per_batch = args.batch // args.minibatch

    data = iter(loader)

    if args.wandb:
        run = wandb.init()  # noqa
        wandb.watch(model, log_freq=100, log="gradients")
        wandb.config.update(args)
        wandb.config.update({"n_parameters": cfg.n_parameters})

    model.init_weights()
    param_bytes = 0
    for p in model.parameters():
        param_bytes += p.numel() * p.element_size()

    print(
        f"Training a {cfg.n_layer}L model with {cfg.n_parameters:,} non-embedding parameters..."
    )
    print(f" Model params use {param_bytes/1024**3:.2f}GiB on device")

    start = time.time()
    positions = 0

    steps = range(args.steps) if args.steps is not None else itertools.count()

    profile_steps = set()
    if args.profile_steps is not None:
        profile_steps = set(int(s) for s in args.profile_steps.split(","))

    def schedule(step):
        if step in profile_steps:
            print(f"Profiling step {step}...")
            return ProfilerAction.RECORD_AND_SAVE
        if step + 1 in profile_steps:
            return ProfilerAction.WARMUP
        return ProfilerAction.NONE

    def save_profile(prof):
        os.makedirs("profile", 0o755, True)
        prof.export_chrome_trace(f"profile/step_{step_i}.pt.trace.json")

    profiler = profile(schedule=schedule, with_stack=True, on_trace_ready=save_profile)

    with profiler:
        for step_i in steps:
            step_start = time.time()

            avg_loss = torch.zeros((), device=args.device)
            opt.zero_grad(set_to_none=True)
            for _ in range(steps_per_batch):
                batch = next(data)
                loss = fwd_and_loss(model, xent, batch)
                avg_loss += loss
                positions += batch["positions"].size(0)
                (loss / steps_per_batch).backward()
            opt.step()
            profiler.step()

            now = time.time()
            avg_loss = (avg_loss / steps_per_batch).item()
            print(
                f"[step={step_i:06d} t={now-start:.1f}s positions={positions:08d}] loss={avg_loss:2.2f} ms_per_step={1000*(now-step_start):.0f}"
            )
            if args.wandb:
                wandb.log(
                    {
                        "positions": positions,
                        "elapsed_time": now - start,
                        "train_loss": avg_loss,
                    },
                    step=step_i,
                )
            if args.positions is not None and positions >= args.positions:
                break

            if step_i % args.test_freq == 0:
                test_loss = torch.tensor(
                    [fwd_and_loss(model, xent, b).item() for b in test_batches]
                ).mean()
                print(f"[step={step_i:06d}] test_loss={test_loss:4.2f}")

if __name__ == "__main__":
    main()
<commit_msg>move the mask to the device<commit_after>import xformer
import xformer.data

import os
import torch
import time
import itertools
import argparse
import wandb
from contextlib import nullcontext

from torch.profiler import profile, ProfilerAction

class Dataset:
    def __init__(self, path):
        self.data = torch.load(path)

    def __len__(self):
        return len(next(iter(self.data.values())))

    def __getitem__(self, i):
        return {k: v[i] for (k, v) in self.data.items()}

def fwd_and_loss(model, xent, record):
    batch = record["positions"].to(device=model.device, dtype=torch.long)
    logits = model(batch[:, :-1])
    targets = batch[:, 1:]
    loss = (
        xent(logits.permute(0, 2, 1), targets)
        * record["mask"].to(device=model.device)[:, :-1]
    ).mean()
    return loss

def main():
    parser = argparse.ArgumentParser(description="Train a transformer")
    parser.add_argument("--layers", type=int, default=2, help="Number of layers")
    parser.add_argument("--d_model", type=int, default=None, help="embedding dimension")
    parser.add_argument("--d_head", type=int, default=32, help="head dimension")
    parser.add_argument(
        "--n_ctx", type=int, default=1024, help="maximum context length"
    )
    parser.add_argument(
        "--pe", type=str, default=None, help="positional encoding (sin, learned, none)"
    )

    parser.add_argument("--data", type=str, default="data/positions", help="datasource")

    parser.add_argument("--batch", type=int, default=64, help="batch size")
    parser.add_argument("--minibatch", type=int, default=4, help="minibatch")

    parser.add_argument(
        "--test-batches", type=int, default=16, help="number of test minibatches"
    )
    parser.add_argument(
        "--test-freq", type=int, default=1000, help="measure test loss every N steps"
    )

    parser.add_argument(
        "--device", type=str, choices=("cpu", "cuda"), default="cuda", help="device"
    )

    parser.add_argument("--wandb", action="store_true", default=False)
    parser.add_argument("--no-wandb", action="store_false", dest="wandb")

    parser.add_argument("--lr", type=float, default=0.001, help="learning rate")
    parser.add_argument("--steps", type=int, default=None)
    parser.add_argument("--profile-steps", type=str, default=None)
    parser.add_argument("--positions", type=int, default=None)

    args = parser.parse_args()

    cfg = xformer.Config(
        n_layer=args.layers,
        d_model=args.d_model or 128 * args.layers,
        d_head=args.d_head,
        n_ctx=args.n_ctx,
        n_vocab=256,
        autoregressive_mask=True,
    )
    if args.pe is not None:
        cfg.positional_encoding = args.pe

    train_ds = Dataset(args.data + "-train.pt")
    loader = torch.utils.data.DataLoader(
        train_ds,
        batch_size=args.minibatch,
        pin_memory=True,
        num_workers=1,
        shuffle=True,
    )

    test_ds = Dataset(args.data + "-test.pt")
    test_loader = torch.utils.data.DataLoader(
        test_ds,
        batch_size=args.minibatch,
        pin_memory=True,
        num_workers=1,
    )
    test_batches = list(itertools.islice(test_loader, args.test_batches))

    model = xformer.Transformer(cfg, dtype=torch.float32, device=args.device)

    xent = torch.nn.CrossEntropyLoss(reduction="none")
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr)

    assert args.batch % args.minibatch == 0, "minibatch must divide batch"
    steps_per_batch = args.batch // args.minibatch

    data = iter(loader)

    if args.wandb:
        run = wandb.init()  # noqa
        wandb.watch(model, log_freq=100, log="gradients")
        wandb.config.update(args)
        wandb.config.update({"n_parameters": cfg.n_parameters})

    model.init_weights()
    param_bytes = 0
    for p in model.parameters():
        param_bytes += p.numel() * p.element_size()

    print(
        f"Training a {cfg.n_layer}L model with {cfg.n_parameters:,} non-embedding parameters..."
    )
    print(f" Model params use {param_bytes/1024**3:.2f}GiB on device")

    start = time.time()
    positions = 0

    steps = range(args.steps) if args.steps is not None else itertools.count()

    profile_steps = set()
    if args.profile_steps is not None:
        profile_steps = set(int(s) for s in args.profile_steps.split(","))

    def schedule(step):
        if step in profile_steps:
            print(f"Profiling step {step}...")
            return ProfilerAction.RECORD_AND_SAVE
        if step + 1 in profile_steps:
            return ProfilerAction.WARMUP
        return ProfilerAction.NONE

    def save_profile(prof):
        os.makedirs("profile", 0o755, True)
        prof.export_chrome_trace(f"profile/step_{step_i}.pt.trace.json")

    profiler = profile(schedule=schedule, with_stack=True, on_trace_ready=save_profile)

    with profiler:
        for step_i in steps:
            step_start = time.time()

            avg_loss = torch.zeros((), device=args.device)
            opt.zero_grad(set_to_none=True)
            for _ in range(steps_per_batch):
                batch = next(data)
                loss = fwd_and_loss(model, xent, batch)
                avg_loss += loss
                positions += batch["positions"].size(0)
                (loss / steps_per_batch).backward()
            opt.step()
            profiler.step()

            now = time.time()
            avg_loss = (avg_loss / steps_per_batch).item()
            print(
                f"[step={step_i:06d} t={now-start:.1f}s positions={positions:08d}] loss={avg_loss:2.2f} ms_per_step={1000*(now-step_start):.0f}"
            )
            if args.wandb:
                wandb.log(
                    {
                        "positions": positions,
                        "elapsed_time": now - start,
                        "train_loss": avg_loss,
                    },
                    step=step_i,
                )
            if args.positions is not None and positions >= args.positions:
                break

            if step_i % args.test_freq == 0:
                test_loss = torch.tensor(
                    [fwd_and_loss(model, xent, b).item() for b in test_batches]
                ).mean()
                print(f"[step={step_i:06d}] test_loss={test_loss:4.2f}")

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>import itertools
import unittest
from parameterized import parameterized
import numpy as np
from numpy.testing import assert_almost_equal
import pandas as pd
import smurff
from pprint import pprint

verbose = 0

class TestPredictSession(unittest.TestCase):
    # Python 2.7 @unittest.skip fix
    __name__ = "TestPredictSession"

    def run_train_session(self, nmodes, density):
        shape = range(11, nmodes+11) # 11, 12, 13, ...
        Y, X = smurff.generate.gen_tensor(shape, 3, density)
        self.Ytrain, self.Ytest = smurff.make_train_test(Y, 0.1)
        priors = ['normal'] * nmodes

        trainSession = smurff.TrainSession(priors = priors, num_latent=4,
                burnin=10, nsamples=25, verbose=verbose,
                save_freq = 1, save_name = smurff.helper.temp_savename())

        trainSession.addTrainAndTest(self.Ytrain, self.Ytest)
        for i, x in enumerate(X):
            trainSession.addSideInfo(i, x)

        trainSession.init()
        while trainSession.step():
            pass

        return trainSession, Y, X

    def run_predict_session(self, nmodes, density):
        train_session, Y, X = self.run_train_session(nmodes, density)
        predict_session = train_session.makePredictSession()
        self.run_predict_some_all_one(train_session, predict_session)
        self.run_predict_predict(predict_session, X)

    def run_predict_some_all_one(self, train_session, predict_session):
        p1 = sorted(train_session.getTestPredictions())
        p2 = sorted(predict_session.predict_some(self.Ytest))
        p3 = predict_session.predict_one(p1[0].coords, p1[0].val)
        p4 = predict_session.predict_all()

        self.assertEqual(len(p1), len(p2))

        # check train_session vs predict_session for Ytest
        self.assertEqual(p1[0].coords, p2[0].coords)
        assert_almost_equal(p1[0].val, p2[0].val, decimal = 2)
        assert_almost_equal(p1[0].pred_1sample, p2[0].pred_1sample, decimal = 2)
        assert_almost_equal(p1[0].pred_avg, p2[0].pred_avg, decimal = 2)

        # check predict_session.predict_some vs predict_session.predict_one
        self.assertEqual(p1[0].coords, p3.coords)
        assert_almost_equal(p1[0].val, p3.val, decimal = 2)
        assert_almost_equal(p1[0].pred_1sample, p3.pred_1sample, decimal = 2)
        assert_almost_equal(p1[0].pred_avg, p3.pred_avg, decimal = 2)

        # check predict_session.predict_some vs predict_session.predict_all
        for s in p2:
            ecoords = (Ellipsis,) + s.coords
            for p in zip(s.pred_all, p4[ecoords]):
                self.assertAlmostEqual(*p, places=2)

        p1_rmse_avg = smurff.calc_rmse(p1)
        p2_rmse_avg = smurff.calc_rmse(p2)

        self.assertAlmostEqual(train_session.getRmseAvg(), p2_rmse_avg, places = 2)
        self.assertAlmostEqual(train_session.getRmseAvg(), p1_rmse_avg, places = 2)

    def run_predict_predict(self, predict_session, X):
        """ Test the PredictSession.predict function """

        def run_n_samples(samples, expected_nsamples):
            operand_and_sizes = [
                [
                    ( slice(10)  , 10         ),
                    ( range(10)  , 10         ),
                    ( 10         , 1          ),
                    ( x          , x.shape[0] ),
                    ( x[10]      , 1          ),
                    ( x[8:11]    , 3          ),
                    ( x[2:11:2]  , 5          ),
                ]
            for x in X ]

            for o in itertools.product(*operand_and_sizes):
                operands, expected_sizes = zip(*o) # unzip
                expected_shape = (expected_nsamples,) + expected_sizes
                shape = predict_session.predict(operands, samples).shape
                self.assertEqual(shape, expected_shape)

        run_n_samples(None, len(predict_session.samples))
        run_n_samples(slice(10), 10)
        run_n_samples(slice(10, 20), 10)
        run_n_samples(slice(10, 20, 2), 5)

    @parameterized.expand(map(lambda x: ("dims%d" % x, x), range(2,5))) # 2, 3, ..., 6
    def test_predict_dense(self, name, nmodes):
        self.run_predict_session(nmodes, 1.0)

    @parameterized.expand(map(lambda x: ("dims%d" % x, x), range(2,5))) # 2, 3, ..., 6
    def test_predict_sparse(self, name, nmodes):
        self.run_predict_session(nmodes, 0.5)

if __name__ == '__main__':
    unittest.main()
<commit_msg>ENH: limit test_predict<commit_after>import itertools
import unittest
from parameterized import parameterized
import numpy as np
from numpy.testing import assert_almost_equal
import pandas as pd
import smurff
from pprint import pprint

verbose = 0

class TestPredictSession(unittest.TestCase):
    # Python 2.7 @unittest.skip fix
    __name__ = "TestPredictSession"

    def run_train_session(self, nmodes, density):
        shape = range(5, nmodes+5) # 5, 6, 7, ...
        Y, X = smurff.generate.gen_tensor(shape, 3, density)
        self.Ytrain, self.Ytest = smurff.make_train_test(Y, 0.1)
        priors = ['normal'] * nmodes

        trainSession = smurff.TrainSession(priors = priors, num_latent=4,
                burnin=10, nsamples=25, verbose=verbose,
                save_freq = 1, save_name = smurff.helper.temp_savename())

        trainSession.addTrainAndTest(self.Ytrain, self.Ytest)
        for i, x in enumerate(X):
            trainSession.addSideInfo(i, x)

        trainSession.init()
        while trainSession.step():
            pass

        return trainSession, Y, X

    def run_predict_session(self, nmodes, density):
        train_session, Y, X = self.run_train_session(nmodes, density)
        predict_session = train_session.makePredictSession()
        self.run_predict_some_all_one(train_session, predict_session)
        self.run_predict_predict(predict_session, X)

    def run_predict_some_all_one(self, train_session, predict_session):
        p1 = sorted(train_session.getTestPredictions())
        p2 = sorted(predict_session.predict_some(self.Ytest))
        p3 = predict_session.predict_one(p1[0].coords, p1[0].val)
        p4 = predict_session.predict_all()

        self.assertEqual(len(p1), len(p2))

        # check train_session vs predict_session for Ytest
        self.assertEqual(p1[0].coords, p2[0].coords)
        assert_almost_equal(p1[0].val, p2[0].val, decimal = 2)
        assert_almost_equal(p1[0].pred_1sample, p2[0].pred_1sample, decimal = 2)
        assert_almost_equal(p1[0].pred_avg, p2[0].pred_avg, decimal = 2)

        # check predict_session.predict_some vs predict_session.predict_one
        self.assertEqual(p1[0].coords, p3.coords)
        assert_almost_equal(p1[0].val, p3.val, decimal = 2)
        assert_almost_equal(p1[0].pred_1sample, p3.pred_1sample, decimal = 2)
        assert_almost_equal(p1[0].pred_avg, p3.pred_avg, decimal = 2)

        # check predict_session.predict_some vs predict_session.predict_all
        for s in p2:
            ecoords = (Ellipsis,) + s.coords
            for p in zip(s.pred_all, p4[ecoords]):
                self.assertAlmostEqual(*p, places=2)

        p1_rmse_avg = smurff.calc_rmse(p1)
        p2_rmse_avg = smurff.calc_rmse(p2)

        self.assertAlmostEqual(train_session.getRmseAvg(), p2_rmse_avg, places = 2)
        self.assertAlmostEqual(train_session.getRmseAvg(), p1_rmse_avg, places = 2)

    def run_predict_predict(self, predict_session, X):
        """ Test the PredictSession.predict function """

        def run_n_samples(samples, expected_nsamples):
            operand_and_sizes = [
                [
                    ( Ellipsis   , x.shape[0] ),
                    ( slice(3)   , 3          ),
                    ( range(3)   , 3          ),
                    ( 3          , 1          ),
                    ( x          , x.shape[0] ),
                    ( x[3]       , 1          ),
                    ( x[1:4]     , 3          ),
                    ( x[0:4:2]   , 2          ),
                ]
            for x in X ]

            MAX = 50
            for c, o in enumerate(itertools.product(*operand_and_sizes)):
                operands, expected_sizes = zip(*o) # unzip
                expected_shape = (expected_nsamples,) + expected_sizes
                shape = predict_session.predict(operands, samples).shape
                self.assertEqual(shape, expected_shape)

                if c > MAX: break

        run_n_samples(None, len(predict_session.samples))
        run_n_samples(slice(10), 10)
        run_n_samples(slice(10, 20), 10)
        run_n_samples(slice(10, 20, 2), 5)

    @parameterized.expand(map(lambda x: ("dims%d" % x, x), range(2,5))) # 2, 3, 4
    def test_predict_dense(self, name, nmodes):
        self.run_predict_session(nmodes, 1.0)

    @parameterized.expand(map(lambda x: ("dims%d" % x, x), range(2,5))) # 2, 3, 4
    def test_predict_sparse(self, name, nmodes):
        self.run_predict_session(nmodes, 0.5)

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


<commit_before>#!/usr/bin/env python3
# TODO: why are the keras models saved with python 2?
from __future__ import print_function

import tensorflow as tf  # pylint: disable=import-error
import os
import sys
import numpy as np
from tensorflow.keras.models import load_model  # pylint: disable=import-error

def read(sz):
  dd = []
  gt = 0
  while gt < sz * 4:
    st = os.read(0, sz * 4 - gt)
    assert(len(st) > 0)
    dd.append(st)
    gt += len(st)
  return np.frombuffer(b''.join(dd), dtype=np.float32)

def write(d):
  os.write(1, d.tobytes())

def run_loop(m):
  ishapes = [[1]+ii.shape[1:] for ii in m.inputs]
  print("ready to run keras model", ishapes, file=sys.stderr)
  while 1:
    inputs = []
    for shp in ishapes:
      ts = np.product(shp)
      #print("reshaping %s with offset %d" % (str(shp), offset), file=sys.stderr)
      inputs.append(read(ts).reshape(shp))
    ret = m.predict_on_batch(inputs)
    #print(ret, file=sys.stderr)
    for r in ret:
      write(r)

if __name__ == "__main__":
  print(tf.__version__, file=sys.stderr)
  # limit gram alloc
  gpus = tf.config.experimental.list_physical_devices('GPU')
  if len(gpus) > 0:
    tf.config.experimental.set_virtual_device_configuration(gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])

  m = load_model(sys.argv[1])
  print(m, file=sys.stderr)

  run_loop(m)

<commit_msg>1024 MB should be enough for anyone<commit_after>#!/usr/bin/env python3
# TODO: why are the keras models saved with python 2?
from __future__ import print_function

import tensorflow as tf  # pylint: disable=import-error
import os
import sys
import numpy as np
from tensorflow.keras.models import load_model  # pylint: disable=import-error

def read(sz):
  dd = []
  gt = 0
  while gt < sz * 4:
    st = os.read(0, sz * 4 - gt)
    assert(len(st) > 0)
    dd.append(st)
    gt += len(st)
  return np.frombuffer(b''.join(dd), dtype=np.float32)

def write(d):
  os.write(1, d.tobytes())

def run_loop(m):
  ishapes = [[1]+ii.shape[1:] for ii in m.inputs]
  print("ready to run keras model", ishapes, file=sys.stderr)
  while 1:
    inputs = []
    for shp in ishapes:
      ts = np.product(shp)
      #print("reshaping %s with offset %d" % (str(shp), offset), file=sys.stderr)
      inputs.append(read(ts).reshape(shp))
    ret = m.predict_on_batch(inputs)
    #print(ret, file=sys.stderr)
    for r in ret:
      write(r)

if __name__ == "__main__":
  print(tf.__version__, file=sys.stderr)
  # limit gram alloc
  gpus = tf.config.experimental.list_physical_devices('GPU')
  if len(gpus) > 0:
    tf.config.experimental.set_virtual_device_configuration(gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

  m = load_model(sys.argv[1])
  print(m, file=sys.stderr)

  run_loop(m)
<|endoftext|>


<commit_before>"""parse_wikidump

Usage:
    parse_wikidump [options] <dump> <model-filename>
    parse_wikidump --download=<wikiname> <model-filename>

Options:
    --download=wikiname     Download dump from dumps.wikimedia.org first
    --ngram=<order>, -N <order>
                            Maximum order of ngrams, set to None to disable
                            [default: 7]
    --help, -h              This help
"""
from __future__ import print_function

import logging
import re
import sqlite3
import sys

from six.moves.urllib.error import HTTPError
from six.moves.urllib.request import urlretrieve

from docopt import docopt

from . import parse_dump
from .._semanticizer import createtables_path

logger = logging.getLogger('semanticizest')
logger.addHandler(logging.StreamHandler(sys.stderr))
logger.setLevel('INFO')

class Progress(object):
    def __init__(self):
        self.threshold = .0

    def __call__(self, n_blocks, blocksize, totalsize):
        done = n_blocks * blocksize
        if done >= self.threshold * totalsize:
            logger.info("%3d%% done", int(100 * self.threshold))
            self.threshold += .05

DUMP_TEMPLATE = (
    "https://dumps.wikimedia.org/{0}/latest/{0}-latest-pages-articles.xml.bz2")

def die(msg):
    print("semanticizest.parse_wikidump: %s" % msg, file=sys.stderr)
    sys.exit(1)

def main(args):
    args = docopt(__doc__, args)

    if args["--download"]:
        wikidump = args["--download"] + ".xml.bz2"
    else:
        wikidump = args['<dump>']

    model_fname = args['<model-filename>']
    ngram = args['--ngram']
    if ngram == "None":
        ngram = None
    else:
        ngram = int(ngram)

    logger.info("Creating database at %r" % model_fname)
    try:
        db = sqlite3.connect(model_fname)
    except sqlite3.OperationalError as e:
        if 'unable to open' in str(e):
            # This exception doesn't store the path.
            die("%s: %r" % (e, model_fname))
        else:
            raise
    with open(createtables_path()) as f:
        create = f.read()

    c = db.cursor()
    try:
        c.executescript(create)
    except sqlite3.OperationalError as e:
        if re.search(r'table .* already exists', str(e)):
            die("database %r already populated" % model_fname)
        else:
            raise

    if args["--download"]:
        url = DUMP_TEMPLATE.format(args["--download"])
        logger.info("Saving wikidump to %r", wikidump)
        try:
            urlretrieve(url, wikidump, Progress())
        except HTTPError as e:
            die("Cannot download {0!r}: {1}".format(url, e))

    parse_dump(wikidump, db, N=ngram)
    db.close()

if __name__ == '__main__':
    main(sys.argv[1:])
<commit_msg>download first, then create database<commit_after>"""parse_wikidump

Usage:
    parse_wikidump [options] <dump> <model-filename>
    parse_wikidump --download=<wikiname> <model-filename>

Options:
    --download=wikiname     Download dump from dumps.wikimedia.org first
    --ngram=<order>, -N <order>
                            Maximum order of ngrams, set to None to disable
                            [default: 7]
    --help, -h              This help
"""
from __future__ import print_function

import logging
import re
import sqlite3
import sys

from six.moves.urllib.error import HTTPError
from six.moves.urllib.request import urlretrieve

from docopt import docopt

from . import parse_dump
from .._semanticizer import createtables_path

logger = logging.getLogger('semanticizest')
logger.addHandler(logging.StreamHandler(sys.stderr))
logger.setLevel('INFO')

class Progress(object):
    def __init__(self):
        self.threshold = .0

    def __call__(self, n_blocks, blocksize, totalsize):
        done = n_blocks * blocksize
        if done >= self.threshold * totalsize:
            logger.info("%3d%% done", int(100 * self.threshold))
            self.threshold += .05

DUMP_TEMPLATE = (
    "https://dumps.wikimedia.org/{0}/latest/{0}-latest-pages-articles.xml.bz2")

def die(msg):
    print("semanticizest.parse_wikidump: %s" % msg, file=sys.stderr)
    sys.exit(1)

def main(args):
    args = docopt(__doc__, args)

    if args["--download"]:
        wikidump = args["--download"] + ".xml.bz2"
    else:
        wikidump = args['<dump>']

    model_fname = args['<model-filename>']
    ngram = args['--ngram']
    if ngram == "None":
        ngram = None
    else:
        ngram = int(ngram)

    if args["--download"]:
        url = DUMP_TEMPLATE.format(args["--download"])
        logger.info("Saving wikidump to %r", wikidump)
        try:
            urlretrieve(url, wikidump, Progress())
        except HTTPError as e:
            die("Cannot download {0!r}: {1}".format(url, e))

    logger.info("Creating database at %r" % model_fname)
    try:
        db = sqlite3.connect(model_fname)
    except sqlite3.OperationalError as e:
        if 'unable to open' in str(e):
            # This exception doesn't store the path.
            die("%s: %r" % (e, model_fname))
        else:
            raise
    with open(createtables_path()) as f:
        create = f.read()

    c = db.cursor()
    try:
        c.executescript(create)
    except sqlite3.OperationalError as e:
        if re.search(r'table .* already exists', str(e)):
            die("database %r already populated" % model_fname)
        else:
            raise

    parse_dump(wikidump, db, N=ngram)
    db.close()

if __name__ == '__main__':
    main(sys.argv[1:])
<|endoftext|>


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

from runner.koan import *

class AboutSets(Koan):
    def test_sets_make_keep_lists_unique(self):
        highlanders = ['MacLeod', 'Ramirez', 'MacLeod', 'Matunas', 'MacLeod', 'Malcolm', 'MacLeod']

        there_can_only_be_only_one = set(highlanders)

        self.assertEqual(__, there_can_only_be_only_one)

    def test_empty_sets_have_different_syntax_to_populated_sets(self):
        self.assertEqual(__, {1, 2, 3})
        self.assertEqual(__, set())

    def test_dictionaries_and_sets_use_same_curly_braces(self):
        # Note: Literal sets using braces were introduced in python 3.
        #       They were also backported to python 2.7.

        self.assertEqual(__, {1, 2, 3}.__class__)
        self.assertEqual(__, {'one': 1, 'two': 2}.__class__)

        self.assertEqual(__, {}.__class__)

    def test_creating_sets_using_strings(self):
        self.assertEqual(__, {'12345'})
        self.assertEqual(__, set('12345'))

    def test_convert_the_set_into_a_list_to_sort_it(self):
        self.assertEqual(__, sorted(set('12345')))

    # ------------------------------------------------------------------

    def test_set_have_arithmetic_operators(self):
        scotsmen = {'MacLeod', 'Wallace', 'Willie'}
        warriors = {'MacLeod', 'Wallace', 'Leonidas'}

        self.assertEqual(__, scotsmen - warriors)
        self.assertEqual(__, scotsmen | warriors)
        self.assertEqual(__, scotsmen & warriors)
        self.assertEqual(__, scotsmen ^ warriors)

    # ------------------------------------------------------------------

    def test_we_can_query_set_membership(self):
        self.assertEqual(__, 127 in {127, 0, 0, 1} )
        self.assertEqual(__, 'cow' not in set('apocalypse now') )

    def test_we_can_compare_subsets(self):
        self.assertEqual(__, set('cake') <= set('cherry cake'))
        self.assertEqual(__, set('cake').issubset(set('cherry cake')) )

        self.assertEqual(__, set('cake') > set('pie'))
<commit_msg>Completed:12. Sets<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

#https://docs.python.org/2/library/sets.html

from runner.koan import *

class AboutSets(Koan):

    # The sets module provides classes for constructing and manipulating unordered collections of unique elements.
    # Common uses include membership testing, removing duplicates from a sequence,
    # and computing standard math operations on sets such as intersection, union, difference, and symmetric difference.

    def test_sets_make_keep_lists_unique(self):
        highlanders = ['MacLeod', 'Ramirez', 'MacLeod', 'Matunas', 'MacLeod', 'Malcolm', 'MacLeod']

        there_can_only_be_only_one = set(highlanders)

        self.assertEqual({'MacLeod', 'Ramirez', 'Matunas', 'Malcolm'}, there_can_only_be_only_one)

    def test_empty_sets_have_different_syntax_to_populated_sets(self):
        self.assertEqual(set([1,2,3]), {1, 2, 3})
        self.assertEqual(set([]), set())

    def test_dictionaries_and_sets_use_same_curly_braces(self):
        # Note: Literal sets using braces were introduced in python 3.
        #       They were also backported to python 2.7.

        self.assertEqual(set, {1, 2, 3}.__class__)
        self.assertEqual(dict, {'one': 1, 'two': 2}.__class__)

        self.assertEqual(dict, {}.__class__)

    def test_creating_sets_using_strings(self):
        self.assertEqual(set(['12345']), {'12345'})
        self.assertEqual({'1', '2', '3', '4', '5'}, set('12345'))

    def test_convert_the_set_into_a_list_to_sort_it(self):
        self.assertEqual(['1', '2', '3', '4', '5'], sorted(set('12345')))

    # ------------------------------------------------------------------

    def test_set_have_arithmetic_operators(self):
        scotsmen = {'MacLeod', 'Wallace', 'Willie'}
        warriors = {'MacLeod', 'Wallace', 'Leonidas'}

        self.assertEqual({'Willie'}, scotsmen - warriors) 	# (s - t) new set with elements in s but not in t
        self.assertEqual({'MacLeod', 'Wallace', 'Willie','Leonidas'}, scotsmen | warriors) # (s | t) new set with elements from both s and t
        self.assertEqual({'MacLeod', 'Wallace'}, scotsmen & warriors) # (s & t) new set with elements common to s and t
        self.assertEqual({'Willie','Leonidas'}, scotsmen ^ warriors) # (s ^ t) 	new set with elements in either s or t but not both

    # ------------------------------------------------------------------

    def test_we_can_query_set_membership(self):
        self.assertEqual(True, 127 in {127, 0, 0, 1} )
        self.assertEqual(True, 'cow' not in set('apocalypse now') )

    def test_we_can_compare_subsets(self):
        self.assertEqual(True, set('cake') <= set('cherry cake')) # (s <= t ) test whether every element in s is in t
        self.assertEqual(True, set('cake').issubset(set('cherry cake')) ) # (s <= t ) test whether every element in s is in t

        self.assertEqual(False, set('cake') > set('pie'))
<|endoftext|>


<commit_before>

import numpy as np
import pandas as pd
import pytest

from simulation import simulation
from random_init import generate_random_dict
from print_init import print_dict
from import_process import import_process

class TestClass():
    def test1(self):
        '''Testing relations in simulated dataset'''

        is_deterministic, dict_ = generate_random_dict()
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)
        print(dict_)
        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        # endogeneous variables

        assert Y_1.all() == (
            np.dot(dict_['TREATED']['all'], X.T) + U[0:, 1]).all()
        assert Y_0.all() == (
            np.dot(dict_['UNTREATED']['all'], X.T) + U[0:, 0]).all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

    def test2(self):
        '''Testing if relationships hold if process is deterministic'''

        is_deterministic, dict_= generate_random_dict(deterministic = 1.0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)
        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()
    def test3(self):
        '''Testing if relationships hold if coefficients are zero in different setups'''

        # All Coefficients
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        for key_ in ['TREATED', 'UNTREATED', 'COST']:
            dict_[key_]['all'] = np.array([0]*len(dict_[key_]['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Treated and Untreated
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        for key_ in ['TREATED', 'UNTREATED']:
            dict_[key_]['all'] = np.array([0]*len(dict_[key_]['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Only Treated
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        dict_['TREATED']['all'] = np.array([0]*len(dict_['TREATED']['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == (
            np.dot(dict_['UNTREATED']['all'], X.T) + U[0:, 0]).all()

        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Only Untreated
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        dict_['UNTREATED']['all'] = np.array([0]*len(dict_['UNTREATED']['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == (
            np.dot(dict_['TREATED']['all'], X.T) + U[0:, 1]).all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Only Cost
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        dict_['COST']['all'] = np.array([0]*len(dict_['COST']['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == (
            np.dot(dict_['TREATED']['all'], X.T) + U[0:, 1]).all()
        assert Y_0.all() == (
            np.dot(dict_['UNTREATED']['all'], X.T) + U[0:, 0]).all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

<commit_msg>Update test_unit.py<commit_after>

import numpy as np
import pandas as pd
import pytest

from simulation import simulation
from random_init import generate_random_dict
from print_init import print_dict
from import_process import import_process

class TestClass():
    def test1(self):
        '''Testing relations in simulated dataset'''

        is_deterministic, dict_ = generate_random_dict()
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)
        print(dict_)
        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        # endogeneous variables

        assert Y_1.all() == (
            np.dot(dict_['TREATED']['all'], X.T) + U[0:, 1]).all()
        assert Y_0.all() == (
            np.dot(dict_['UNTREATED']['all'], X.T) + U[0:, 0]).all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

    def test2(self):
        '''Testing whether relationships hold if process is deterministic'''

        is_deterministic, dict_= generate_random_dict(deterministic = 1.0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)
        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()
    def test3(self):
        '''Testing whether relationships hold if coefficients are zero in different setups'''

        # All Coefficients
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        for key_ in ['TREATED', 'UNTREATED', 'COST']:
            dict_[key_]['all'] = np.array([0]*len(dict_[key_]['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Treated and Untreated
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        for key_ in ['TREATED', 'UNTREATED']:
            dict_[key_]['all'] = np.array([0]*len(dict_[key_]['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Only Treated
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        dict_['TREATED']['all'] = np.array([0]*len(dict_['TREATED']['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == U[0:, 1].all()
        assert Y_0.all() == (
            np.dot(dict_['UNTREATED']['all'], X.T) + U[0:, 0]).all()

        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Only Untreated
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        dict_['UNTREATED']['all'] = np.array([0]*len(dict_['UNTREATED']['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == (
            np.dot(dict_['TREATED']['all'], X.T) + U[0:, 1]).all()
        assert Y_0.all() == U[0:, 0].all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

        # Only Cost
        is_deterministic, dict_= generate_random_dict(deterministic = .0)
        init_file_name = dict_['SIMULATION']['source']
        print_dict(dict_, init_file_name)
        dict_ = import_process(init_file_name + '.grmpy.ini', is_deterministic)

        dict_['COST']['all'] = np.array([0]*len(dict_['COST']['all']))

        df, Y, Y_1, Y_0, D, X, Z, U, V = simulation(dict_)

        assert Y_1.all() == (
            np.dot(dict_['TREATED']['all'], X.T) + U[0:, 1]).all()
        assert Y_0.all() == (
            np.dot(dict_['UNTREATED']['all'], X.T) + U[0:, 0]).all()
        assert Y[D == 1].all() == Y_1.all()
        assert Y[D == 0].all() == Y_0.all()
        assert V.all() == (U[0:, 2] - U[0:, 1] + U[0:, 0]).all()

<|endoftext|>