<commit_before># -*- encoding: utf-8 -*-
import collections
import copy
import hashlib
from supriya.tools.servertools.ServerObjectProxy import ServerObjectProxy

class SynthDef(ServerObjectProxy):
    r'''A synth definition.

    ::

        >>> from supriya.tools import synthdeftools
        >>> from supriya.tools import ugentools
        >>> builder = synthdeftools.SynthDefBuilder(frequency=440)
        >>> sin_osc = ugentools.SinOsc.ar(frequency=builder['frequency'])
        >>> out = ugentools.Out.ar(bus=0, source=sin_osc)
        >>> builder.add_ugen(out)
        >>> synthdef = builder.build()

    ::

        >>> from supriya import servertools
        >>> server = servertools.Server().boot()

    ::

        >>> synthdef.allocate(server=server)

    ::

        >>> synthdef in server
        True

    ::

        >>> synthdef.free()

    ::

        >>> synthdef in server
        False

    ::

        >>> server.quit()
        <Server: offline>

    '''

    ### CLASS VARIABLES ###

    __slots__ = (
        '_compiled_ugen_graph',
        '_constants',
        '_name',
        '_parameters',
        '_ugens',
        )

    ### INITIALIZER ###

    def __init__(
        self,
        ugens,
        name=None,
        ):
        from supriya.tools import synthdeftools
        ServerObjectProxy.__init__(self)
        compiler = synthdeftools.SynthDefCompiler
        self._name = name
        ugens = copy.deepcopy(ugens)
        ugens = self._flatten_ugens(ugens)
        ugens = self._optimize_ugen_graph(ugens)
        ugens, parameters = self._extract_parameters(ugens)
        control_ugens, control_mapping = self._collect_controls(
            parameters,
            )
        self._parameters = tuple(control_mapping.keys())
        self._remap_controls(ugens, control_mapping)
        ugens = control_ugens + ugens
        ugens = self._sort_ugens_topologically(ugens)
        self._ugens = tuple(ugens)
        self._constants = self._collect_constants(self._ugens)
        self._compiled_ugen_graph = compiler.compile_ugen_graph(self)

    ### SPECIAL METHODS ###

    def __eq__(self, expr):
        if type(expr) != type(self):
            return False
        if expr.name != self.name:
            return False
        if expr._compiled_ugen_graph != self._compiled_ugen_graph:
            return False
        return True

    def __hash__(self):
        hash_values = (
            type(self),
            self._name,
            self._compiled_ugen_graph,
            )
        return hash(hash_values)

    def __str__(self):
        r'''Gets string representation of synth definition.

        ::

            >>> from supriya.tools import synthdeftools
            >>> from supriya.tools import ugentools

        ::

            >>> builder = synthdeftools.SynthDefBuilder()
            >>> sin_one = ugentools.SinOsc.ar()
            >>> sin_two = ugentools.SinOsc.ar(frequency=443)
            >>> sum = sin_one + sin_two
            >>> out = ugentools.Out.ar(bus=0, source=sum)
            >>> builder.add_ugen(out)
            >>> synthdef = builder.build(name='test')

        ::

            >>> print(synthdef)
            SynthDef test {
                const_0:440.0 -> 0_SinOsc[0:frequency]
                const_1:0.0 -> 0_SinOsc[1:phase]
                const_2:443.0 -> 1_SinOsc[0:frequency]
                const_1:0.0 -> 1_SinOsc[1:phase]
                0_SinOsc[0] -> 2_BinaryOpUGen:ADDITION[0:left]
                1_SinOsc[0] -> 2_BinaryOpUGen:ADDITION[1:right]
                const_1:0.0 -> 3_Out[0]
                2_BinaryOpUGen:ADDITION[0] -> 3_Out[1]
            }

        Returns string.
        '''
        def get_ugen_name(ugen):
            ugen_index = self._ugens.index(ugen)
            ugen_class = type(ugen).__name__
            if isinstance(ugen, ugentools.BinaryOpUGen):
                ugen_op = synthdeftools.BinaryOperator.from_expr(
                    ugen.special_index)
                ugen_op_name = ugen_op.name
                ugen_name = '{}_{}:{}'.format(
                    ugen_index,
                    ugen_class,
                    ugen_op_name,
                    )
            elif isinstance(ugen, ugentools.UnaryOpUGen):
                ugen_op = synthdeftools.UnaryOperator.from_expr(
                    ugen.special_index)
                ugen_op_name = ugen_op.name
                ugen_name = '{}_{}:{}'.format(
                    ugen_index,
                    ugen_class,
                    ugen_op_name,
                    )
            else:
                ugen_name = '{}_{}'.format(ugen_index, ugen_class)
            return ugen_name
        from supriya.tools import synthdeftools
        from supriya.tools import ugentools
        result = []
        result.append('SynthDef {} {{'.format(self.actual_name))
        for ugen in self._ugens:
            ugen_name = get_ugen_name(ugen)
            for i, input_ in enumerate(ugen.inputs):
                argument_name = None
                if i < len(ugen._ordered_input_names):
                    argument_name = ugen._ordered_input_names[i]
                if isinstance(input_, float):
                    input_index = self._constants.index(input_)
                    input_name = 'const_{}:{}'.format(input_index, input_)
                else:
                    output_index = 0
                    if isinstance(input_, synthdeftools.OutputProxy):
                        output_index = input_.output_index
                        input_ = input_.source
                    input_name = get_ugen_name(input_)
                    input_name += '[{}]'.format(output_index)
                wire = '\t{} -> {}'.format(input_name, ugen_name)
                if argument_name:
                    wire += '[{}:{}]'.format(i, argument_name)
                else:
                    wire += '[{}]'.format(i)
                result.append(wire)
        result.append('}')
        result = '\n'.join(result)
        return result

    ### PRIVATE METHODS ###

    @staticmethod
    def _collect_constants(ugens):
        constants = []
        for ugen in ugens:
            for input_ in ugen._inputs:
                if not isinstance(input_, float):
                    continue
                if input_ not in constants:
                    constants.append(input_)
        return tuple(constants)

    @staticmethod
    def _collect_controls(parameters):
        from supriya.tools import synthdeftools
        from supriya.tools import ugentools
        control_mapping = collections.OrderedDict()
        scalar_parameters = []
        trigger_parameters = []
        audio_parameters = []
        control_parameters = []
        mapping = {
            synthdeftools.ParameterRate.AUDIO: audio_parameters,
            synthdeftools.ParameterRate.CONTROL: control_parameters,
            synthdeftools.ParameterRate.SCALAR: scalar_parameters,
            synthdeftools.ParameterRate.TRIGGER: trigger_parameters,
            }
        for parameter in parameters:
            mapping[parameter.parameter_rate].append(parameter)
        for parameters in mapping.values():
            parameters.sort(key=lambda x: x.name)
        control_ugens = []
        starting_control_index = 0
        if scalar_parameters:
            control = ugentools.Control(
                control_names=scalar_parameters,
                rate=synthdeftools.Rate.SCALAR,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(scalar_parameters)
            for i, parameter in enumerate(scalar_parameters):
                control_mapping[parameter] = control[i]
        if trigger_parameters:
            control = ugentools.TrigControl(
                control_names=trigger_parameters,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(trigger_parameters)
            for i, parameter in enumerate(trigger_parameters):
                control_mapping[parameter] = control[i]
        if audio_parameters:
            control = ugentools.AudioControl(
                control_names=audio_parameters,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(audio_parameters)
            for i, parameter in enumerate(audio_parameters):
                control_mapping[parameter] = control[i]
        if control_parameters:
            control = ugentools.Control(
                control_names=control_parameters,
                rate=synthdeftools.Rate.CONTROL,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(control_parameters)
            for i, parameter in enumerate(control_parameters):
                control_mapping[parameter] = control[i]
        control_ugens = tuple(control_ugens)
        return control_ugens, control_mapping

    def _collect_parameters(self, controls):
        pass

    @staticmethod
    def _extract_parameters(ugens):
        from supriya.tools import synthdeftools
        parameters = set()
        for ugen in ugens:
            if isinstance(ugen, synthdeftools.Parameter):
                parameters.add(ugen)
        ugens = tuple(ugen for ugen in ugens if ugen not in parameters)
        parameters = tuple(sorted(parameters, key=lambda x: x.name))
        return ugens, parameters

    @staticmethod
    def _flatten_ugens(ugens):
        def recurse(ugen):
            flattened_ugens.append(ugen)
            if isinstance(ugen, synthdeftools.Parameter):
                return
            for input_ in ugen.inputs:
                if isinstance(input_, synthdeftools.Parameter):
                    if input_ not in flattened_ugens:
                        flattened_ugens.append(input_)
                elif isinstance(input_, synthdeftools.OutputProxy):
                    if input_.source not in flattened_ugens:
                        flattened_ugens.append(input_.source)
                        recurse(input_.source)
                elif isinstance(input_, synthdeftools.UGen):
                    if input_ not in flattened_ugens:
                        flattened_ugens.append(input_)
                        recurse(input_)
        from supriya.tools import synthdeftools
        flattened_ugens = []
        for ugen in ugens:
            recurse(ugen)
        return flattened_ugens

    @staticmethod
    def _optimize_ugen_graph(ugens):
        return ugens

    @staticmethod
    def _remap_controls(ugens, control_mapping):
        for ugen in ugens:
            inputs = list(ugen.inputs)
            for i, input_ in enumerate(inputs):
                if input_ in control_mapping:
                    output_proxy = control_mapping[input_]
                    inputs[i] = output_proxy
            ugen._inputs = tuple(inputs)

    @staticmethod
    def _sort_ugens_topologically(ugens):
        from supriya.tools import synthdeftools
        ugens = list(ugens)
        available_ugens = []
        sort_bundles = {}
        for ugen in ugens:
            sort_bundles[ugen] = synthdeftools.UGenSortBundle(ugen)
        for ugen in ugens:
            sort_bundle = sort_bundles[ugen]
            sort_bundle._initialize_topological_sort(sort_bundles)
            sort_bundle.descendants[:] = sorted(
                sort_bundles[ugen].descendants,
                key=lambda x: ugens.index(ugen),
                )
        for ugen in reversed(ugens):
            sort_bundles[ugen]._make_available(available_ugens)
        out_stack = []
        while available_ugens:
            available_ugen = available_ugens.pop()
            sort_bundles[available_ugen]._schedule(
                available_ugens, out_stack, sort_bundles)
        return out_stack

    ### PUBLIC METHODS ###

    def allocate(self, server=None):
        from supriya.tools import servertools
        ServerObjectProxy.allocate(self, server=server)
        synthdef_name = self.actual_name
        self.server._synthdefs[synthdef_name] = self
        message = servertools.CommandManager.make_synthdef_receive_message(
            self,
            )
        self.server.send_message(message)

    def compile(self):
        from supriya.tools.synthdeftools import SynthDefCompiler
        synthdefs = [self]
        result = SynthDefCompiler.compile_synthdefs(synthdefs)
        return result

    def free(self):
        from supriya.tools import servertools
        synthdef_name = self.actual_name
        del(self.server._synthdefs[synthdef_name])
        message = servertools.CommandManager.make_synthdef_free_message(
            synthdef=self,
            )
        self.server.send_message(message)
        ServerObjectProxy.free(self)

    ### PUBLIC PROPERTIES ###

    @property
    def actual_name(self):
        return self.name or self.anonymous_name

    @property
    def anonymous_name(self):
        md5 = hashlib.md5()
        md5.update(self._compiled_ugen_graph)
        anonymous_name = md5.hexdigest()
        return anonymous_name

    @property
    def constants(self):
        return self._constants

    @property
    def is_allocated(self):
        if self.server is not None:
            return self in self.server
        return False

    @property
    def name(self):
        return self._name

    @property
    def parameters(self):
        return self._parameters

    @property
    def ugens(self):
        return self._ugens
<commit_msg>Added SynthDef.from_ugens() and .play() convenience methods:<commit_after># -*- encoding: utf-8 -*-
import collections
import copy
import hashlib
from supriya.tools.servertools.ServerObjectProxy import ServerObjectProxy

class SynthDef(ServerObjectProxy):
    r'''A synth definition.

    ::

        >>> from supriya.tools import synthdeftools
        >>> from supriya.tools import ugentools
        >>> builder = synthdeftools.SynthDefBuilder(frequency=440)
        >>> sin_osc = ugentools.SinOsc.ar(frequency=builder['frequency'])
        >>> out = ugentools.Out.ar(bus=0, source=sin_osc)
        >>> builder.add_ugen(out)
        >>> synthdef = builder.build()

    ::

        >>> from supriya import servertools
        >>> server = servertools.Server().boot()

    ::

        >>> synthdef.allocate(server=server)

    ::

        >>> synthdef in server
        True

    ::

        >>> synthdef.free()

    ::

        >>> synthdef in server
        False

    ::

        >>> server.quit()
        <Server: offline>

    '''

    ### CLASS VARIABLES ###

    __slots__ = (
        '_compiled_ugen_graph',
        '_constants',
        '_name',
        '_parameters',
        '_ugens',
        )

    ### INITIALIZER ###

    def __init__(
        self,
        ugens,
        name=None,
        ):
        from supriya.tools import synthdeftools
        ServerObjectProxy.__init__(self)
        compiler = synthdeftools.SynthDefCompiler
        self._name = name
        ugens = copy.deepcopy(ugens)
        ugens = self._flatten_ugens(ugens)
        ugens = self._optimize_ugen_graph(ugens)
        ugens, parameters = self._extract_parameters(ugens)
        control_ugens, control_mapping = self._collect_controls(
            parameters,
            )
        self._parameters = tuple(control_mapping.keys())
        self._remap_controls(ugens, control_mapping)
        ugens = control_ugens + ugens
        ugens = self._sort_ugens_topologically(ugens)
        self._ugens = tuple(ugens)
        self._constants = self._collect_constants(self._ugens)
        self._compiled_ugen_graph = compiler.compile_ugen_graph(self)

    ### SPECIAL METHODS ###

    def __eq__(self, expr):
        if type(expr) != type(self):
            return False
        if expr.name != self.name:
            return False
        if expr._compiled_ugen_graph != self._compiled_ugen_graph:
            return False
        return True

    def __hash__(self):
        hash_values = (
            type(self),
            self._name,
            self._compiled_ugen_graph,
            )
        return hash(hash_values)

    def __str__(self):
        r'''Gets string representation of synth definition.

        ::

            >>> from supriya.tools import synthdeftools
            >>> from supriya.tools import ugentools

        ::

            >>> builder = synthdeftools.SynthDefBuilder()
            >>> sin_one = ugentools.SinOsc.ar()
            >>> sin_two = ugentools.SinOsc.ar(frequency=443)
            >>> sum = sin_one + sin_two
            >>> out = ugentools.Out.ar(bus=0, source=sum)
            >>> builder.add_ugen(out)
            >>> synthdef = builder.build(name='test')

        ::

            >>> print(synthdef)
            SynthDef test {
                const_0:440.0 -> 0_SinOsc[0:frequency]
                const_1:0.0 -> 0_SinOsc[1:phase]
                const_2:443.0 -> 1_SinOsc[0:frequency]
                const_1:0.0 -> 1_SinOsc[1:phase]
                0_SinOsc[0] -> 2_BinaryOpUGen:ADDITION[0:left]
                1_SinOsc[0] -> 2_BinaryOpUGen:ADDITION[1:right]
                const_1:0.0 -> 3_Out[0]
                2_BinaryOpUGen:ADDITION[0] -> 3_Out[1]
            }

        Returns string.
        '''
        def get_ugen_name(ugen):
            ugen_index = self._ugens.index(ugen)
            ugen_class = type(ugen).__name__
            if isinstance(ugen, ugentools.BinaryOpUGen):
                ugen_op = synthdeftools.BinaryOperator.from_expr(
                    ugen.special_index)
                ugen_op_name = ugen_op.name
                ugen_name = '{}_{}:{}'.format(
                    ugen_index,
                    ugen_class,
                    ugen_op_name,
                    )
            elif isinstance(ugen, ugentools.UnaryOpUGen):
                ugen_op = synthdeftools.UnaryOperator.from_expr(
                    ugen.special_index)
                ugen_op_name = ugen_op.name
                ugen_name = '{}_{}:{}'.format(
                    ugen_index,
                    ugen_class,
                    ugen_op_name,
                    )
            else:
                ugen_name = '{}_{}'.format(ugen_index, ugen_class)
            return ugen_name
        from supriya.tools import synthdeftools
        from supriya.tools import ugentools
        result = []
        result.append('SynthDef {} {{'.format(self.actual_name))
        for ugen in self._ugens:
            ugen_name = get_ugen_name(ugen)
            for i, input_ in enumerate(ugen.inputs):
                argument_name = None
                if i < len(ugen._ordered_input_names):
                    argument_name = ugen._ordered_input_names[i]
                if isinstance(input_, float):
                    input_index = self._constants.index(input_)
                    input_name = 'const_{}:{}'.format(input_index, input_)
                else:
                    output_index = 0
                    if isinstance(input_, synthdeftools.OutputProxy):
                        output_index = input_.output_index
                        input_ = input_.source
                    input_name = get_ugen_name(input_)
                    input_name += '[{}]'.format(output_index)
                wire = '\t{} -> {}'.format(input_name, ugen_name)
                if argument_name:
                    wire += '[{}:{}]'.format(i, argument_name)
                else:
                    wire += '[{}]'.format(i)
                result.append(wire)
        result.append('}')
        result = '\n'.join(result)
        return result

    ### PRIVATE METHODS ###

    @staticmethod
    def _collect_constants(ugens):
        constants = []
        for ugen in ugens:
            for input_ in ugen._inputs:
                if not isinstance(input_, float):
                    continue
                if input_ not in constants:
                    constants.append(input_)
        return tuple(constants)

    @staticmethod
    def _collect_controls(parameters):
        from supriya.tools import synthdeftools
        from supriya.tools import ugentools
        control_mapping = collections.OrderedDict()
        scalar_parameters = []
        trigger_parameters = []
        audio_parameters = []
        control_parameters = []
        mapping = {
            synthdeftools.ParameterRate.AUDIO: audio_parameters,
            synthdeftools.ParameterRate.CONTROL: control_parameters,
            synthdeftools.ParameterRate.SCALAR: scalar_parameters,
            synthdeftools.ParameterRate.TRIGGER: trigger_parameters,
            }
        for parameter in parameters:
            mapping[parameter.parameter_rate].append(parameter)
        for parameters in mapping.values():
            parameters.sort(key=lambda x: x.name)
        control_ugens = []
        starting_control_index = 0
        if scalar_parameters:
            control = ugentools.Control(
                control_names=scalar_parameters,
                rate=synthdeftools.Rate.SCALAR,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(scalar_parameters)
            for i, parameter in enumerate(scalar_parameters):
                control_mapping[parameter] = control[i]
        if trigger_parameters:
            control = ugentools.TrigControl(
                control_names=trigger_parameters,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(trigger_parameters)
            for i, parameter in enumerate(trigger_parameters):
                control_mapping[parameter] = control[i]
        if audio_parameters:
            control = ugentools.AudioControl(
                control_names=audio_parameters,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(audio_parameters)
            for i, parameter in enumerate(audio_parameters):
                control_mapping[parameter] = control[i]
        if control_parameters:
            control = ugentools.Control(
                control_names=control_parameters,
                rate=synthdeftools.Rate.CONTROL,
                starting_control_index=starting_control_index,
                )
            control_ugens.append(control)
            starting_control_index += len(control_parameters)
            for i, parameter in enumerate(control_parameters):
                control_mapping[parameter] = control[i]
        control_ugens = tuple(control_ugens)
        return control_ugens, control_mapping

    def _collect_parameters(self, controls):
        pass

    @staticmethod
    def _extract_parameters(ugens):
        from supriya.tools import synthdeftools
        parameters = set()
        for ugen in ugens:
            if isinstance(ugen, synthdeftools.Parameter):
                parameters.add(ugen)
        ugens = tuple(ugen for ugen in ugens if ugen not in parameters)
        parameters = tuple(sorted(parameters, key=lambda x: x.name))
        return ugens, parameters

    @staticmethod
    def _flatten_ugens(ugens):
        def recurse(ugen):
            flattened_ugens.append(ugen)
            if isinstance(ugen, synthdeftools.Parameter):
                return
            for input_ in ugen.inputs:
                if isinstance(input_, synthdeftools.Parameter):
                    if input_ not in flattened_ugens:
                        flattened_ugens.append(input_)
                elif isinstance(input_, synthdeftools.OutputProxy):
                    if input_.source not in flattened_ugens:
                        flattened_ugens.append(input_.source)
                        recurse(input_.source)
                elif isinstance(input_, synthdeftools.UGen):
                    if input_ not in flattened_ugens:
                        flattened_ugens.append(input_)
                        recurse(input_)
        from supriya.tools import synthdeftools
        flattened_ugens = []
        for ugen in ugens:
            recurse(ugen)
        return flattened_ugens

    @staticmethod
    def _optimize_ugen_graph(ugens):
        return ugens

    @staticmethod
    def _remap_controls(ugens, control_mapping):
        for ugen in ugens:
            inputs = list(ugen.inputs)
            for i, input_ in enumerate(inputs):
                if input_ in control_mapping:
                    output_proxy = control_mapping[input_]
                    inputs[i] = output_proxy
            ugen._inputs = tuple(inputs)

    @staticmethod
    def _sort_ugens_topologically(ugens):
        from supriya.tools import synthdeftools
        ugens = list(ugens)
        available_ugens = []
        sort_bundles = {}
        for ugen in ugens:
            sort_bundles[ugen] = synthdeftools.UGenSortBundle(ugen)
        for ugen in ugens:
            sort_bundle = sort_bundles[ugen]
            sort_bundle._initialize_topological_sort(sort_bundles)
            sort_bundle.descendants[:] = sorted(
                sort_bundles[ugen].descendants,
                key=lambda x: ugens.index(ugen),
                )
        for ugen in reversed(ugens):
            sort_bundles[ugen]._make_available(available_ugens)
        out_stack = []
        while available_ugens:
            available_ugen = available_ugens.pop()
            sort_bundles[available_ugen]._schedule(
                available_ugens, out_stack, sort_bundles)
        return out_stack

    ### PUBLIC METHODS ###

    def allocate(self, server=None):
        from supriya.tools import servertools
        ServerObjectProxy.allocate(self, server=server)
        synthdef_name = self.actual_name
        self.server._synthdefs[synthdef_name] = self
        message = servertools.CommandManager.make_synthdef_receive_message(
            self,
            )
        self.server.send_message(message)

    def compile(self):
        from supriya.tools.synthdeftools import SynthDefCompiler
        synthdefs = [self]
        result = SynthDefCompiler.compile_synthdefs(synthdefs)
        return result

    def free(self):
        from supriya.tools import servertools
        synthdef_name = self.actual_name
        del(self.server._synthdefs[synthdef_name])
        message = servertools.CommandManager.make_synthdef_free_message(
            synthdef=self,
            )
        self.server.send_message(message)
        ServerObjectProxy.free(self)

    @staticmethod
    def from_ugens(ugens):
        r'''Makes a synthdef from `ugens`.

        ::

            >>> from supriya.tools import synthdeftools
            >>> from supriya.tools import ugentools
            >>> ugen = ugentools.Out.ar(source=ugentools.SinOsc.ar() * 0.5)
            >>> synthdef = synthdeftools.SynthDef.from_ugens(ugen)

        ::

            >>> print(synthdef)
            SynthDef e0041b72f3e50206267d5f426f7c592f {
                const_0:440.0 -> 0_SinOsc[0:frequency]
                const_1:0.0 -> 0_SinOsc[1:phase]
                0_SinOsc[0] -> 1_BinaryOpUGen:MULTIPLICATION[0:left]
                const_2:0.5 -> 1_BinaryOpUGen:MULTIPLICATION[1:right]
                const_1:0.0 -> 2_Out[0]
                1_BinaryOpUGen:MULTIPLICATION[0] -> 2_Out[1]
            }

        Returns synthdef.
        '''
        from supriya.tools import synthdeftools
        builder = synthdeftools.SynthDefBuilder()
        if isinstance(ugens, collections.Sequence):
            for ugen in ugens:
                builder.add_ugen(ugen)
        elif isinstance(ugens, synthdeftools.UGenMethodMixin):
            builder.add_ugen(ugens)
        synthdef = builder.build()
        return synthdef

    def play(self, server=None, **kwargs):
        r'''Plays the synthdef on the server.

        ::

            >>> from supriya.tools import servertools
            >>> from supriya.tools import synthdeftools
            >>> from supriya.tools import ugentools

        ::

            >>> server = servertools.Server()
            >>> server.boot()
            <Server: udp://127.0.0.1:57751, 8i8o>

        ::

            >>> ugen = ugentools.Out.ar(source=ugentools.SinOsc.ar() * 0.0)
            >>> synthdef = synthdeftools.SynthDef.from_ugens(ugen)
            >>> synth = synthdef.play()

        ::

            >>> server.quit()
            <Server: offline>

        '''
        from supriya.tools import servertools
        if not self.is_allocated:
            self.allocate(server=server)
            self.server.sync()
        synth = servertools.Synth(self, **kwargs)
        synth.allocate(target_node=self.server)
        self.server.sync()
        return synth

    ### PUBLIC PROPERTIES ###

    @property
    def actual_name(self):
        return self.name or self.anonymous_name

    @property
    def anonymous_name(self):
        md5 = hashlib.md5()
        md5.update(self._compiled_ugen_graph)
        anonymous_name = md5.hexdigest()
        return anonymous_name

    @property
    def constants(self):
        return self._constants

    @property
    def is_allocated(self):
        if self.server is not None:
            return self in self.server
        return False

    @property
    def name(self):
        return self._name

    @property
    def parameters(self):
        return self._parameters

    @property
    def ugens(self):
        return self._ugens
<|endoftext|>


<commit_before>import pycqed.analysis_v2.base_analysis as ba
import pycqed.measurement.waveform_control_CC.waveform as wf
import pycqed.analysis.fitting_models as fit_mods
import numpy as np
from numpy.fft import fft, ifft, fftfreq

class DistortionFineAnalysis(ba.BaseDataAnalysis):
    '''
    Analysis for the enhanced cryogenic oscilloscpe.
    '''
    def __init__(self, t_start: str=None, t_stop: str=None,
                 data_file_path: str=None,
                 options_dict: dict=None, extract_only: bool=False,
                 do_fitting: bool=True, auto=True):
        if options_dict is None:
            options_dict = dict()
        options_dict['scan_label'] = 'distortion_scope_fine'
        options_dict['tight_fig'] = False
        options_dict['apply_default_fig_settings'] = False

        super().__init__(t_start=t_start, t_stop=t_stop,
                         data_file_path=data_file_path,
                         options_dict=options_dict,
                         extract_only=extract_only, do_fitting=do_fitting)

        # self.single_timestamp = True
        self.params_dict = {
            # Data
            'xlabel': 'sweep_name',
            'xunit': 'sweep_unit',
            'measurementstring': 'measurementstring',
            'sweep_points': 'sweep_points',
            'value_names': 'value_names',
            'value_units': 'value_units',
            'measured_values': 'measured_values',

            # Pulse parameters
            'gauss_sigma': 'QWG_flux_lutman_QR.S_gauss_sigma',
            'gauss_amp': 'QWG_flux_lutman_QR.S_amp',
            'wave_dict_unit': 'QWG_flux_lutman_QR.wave_dict_unit',
            'sampling_rate': 'QWG_flux_lutman_QR.sampling_rate',
            'QWG_amp': 'QWG.ch3_amp',
            'distortions_amp': 'QWG_flux_lutman_QR.F_amp',

            # Qubit parameters
            'f_max': 'QR.f_max',
            'E_c': 'QR.E_c',
            'V_offset': 'QR.V_offset',
            'V_per_phi0': 'QR.V_per_phi0',
            'asymmetry': 'QR.asymmetry',
        }
        self.numeric_params = [
            'gauss_sigma',
            'gauss_amp',
            'sampling_rate',
            'QWG_amp',
            'distortions_amp',
            'f_max',
            'E_c',
            'V_offset',
            'V_per_phi0',
            'asymmetry',
        ]

        if auto:
            self.run_analysis()

    def process_data(self):
        # Convert lists of duplicated params to floats
        for k in self.numeric_params:
            self.data_dict[k] = self.data_dict[k][1]

        # Background
        cos = (self.data_dict['measured_values'][0][0, ::2] - 0.5) * 2
        sin = (self.data_dict['measured_values'][0][0, 1::2] - 0.5) * 2
        bg_phase = np.unwrap(np.arctan2(sin, cos))
        self.data_dict['background phase'] = np.mean(bg_phase)
        n = fft(bg_phase)
        self.data_dict['noise power spectrum'] = np.conj(n) * n

        # Data
        self.data_dict['cos'] = (
            self.data_dict['measured_values'][1][0, ::2] - 0.5) * 2
        self.data_dict['sin'] = (
            self.data_dict['measured_values'][1][0, 1::2] - 0.5) * 2

        self.data_dict['raw phase'] = np.arctan2(self.data_dict['sin'],
                                                 self.data_dict['cos'])
        self.data_dict['phase'] = (np.unwrap(self.data_dict['raw phase']) -
                                   self.data_dict['background phase'])

        # After unwrapping we can convert phases to degrees
        self.data_dict['raw phase'] = np.rad2deg(self.data_dict['raw phase'])
        self.data_dict['phase'] = np.rad2deg(self.data_dict['phase'])
        self.data_dict['background phase'] = np.rad2deg(
            self.data_dict['background phase'])

        self.data_dict['tau'] = self.data_dict['sweep_points'][1][::2]

        # Calculate integral of the scoping pulse
        data_dt = (self.data_dict['tau'][1] -
                   self.data_dict['tau'][0])
        gauss_pulse = wf.gauss_pulse(
            self.data_dict['gauss_amp'],
            self.data_dict['gauss_sigma'],
            nr_sigma=4,
            sampling_rate=data_dt,
            axis='x', phase=0, motzoi=0, delay=0)[0]  # Returns I and Q

        if self.data_dict['wave_dict_unit'] == 'frac':
            # Convert to volts
            gauss_pulse *= self.data_dict['QWG_amp'] / 2

        # Add zeros to front and back of data because the signal does not
        # start at tau = 0.
        # t0: time from start to center of scoping pulse
        t0 = self.data_dict['gauss_sigma'] * 2
        t0_samples = int(np.round(t0 / data_dt))

        y = np.concatenate((
            np.zeros(t0_samples),
            self.data_dict['phase'],
            np.zeros(t0_samples)))
        N = np.concatenate((
            np.zeros(t0_samples),
            self.data_dict['noise power spectrum'],
            np.zeros(t0_samples)))
        self.y = y

        # Calculate convolution kernel
        # Kernel needs to be sampled at the same rate as data.

        h = 360 * fit_mods.Qubit_dac_sensitivity(
            dac_voltage=gauss_pulse,
            f_max=self.data_dict['f_max'],
            E_c=self.data_dict['E_c'],
            dac_sweet_spot=self.data_dict['V_offset'],
            V_per_phi0=self.data_dict['V_per_phi0'],
            asymmetry=self.data_dict['asymmetry'])
        self.h = h

        # Note that we assume now that the magnitude of the phase error is
        # smaller than 180 degrees, i.e. we assume the phase has not wrapped
        # due to distortions.
        h = np.hstack((h, np.zeros(len(y) - len(h))))
        # Abs of fft, because we want the spectrum of a gauss window centered
        # around zero. Note: f real, symmetric => fft(f) real, symmetric.
        freqs = fftfreq(len(h), d=data_dt)
        H = np.abs(fft(h))
        x = np.real(ifft(fft(y) / H)) * (freqs[1] - freqs[0])
        # Wiener deconvolution
        # x = np.real(ifft(fft(y) * np.conj(H) /
        #                  (np.conj(H) * H + N / (np.conj(y) * y))))

        self.data_dict['tau'] += t0
        self.data_dict['dV'] = x[t0_samples:-t0_samples]

    def prepare_plots(self):
        self.plot_dicts['raw data'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' raw data',
            'xvals': self.data_dict['tau'],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': [self.data_dict['cos'], self.data_dict['sin']],
            'multiple': True,
            'setlabel': ['cos', 'sin'],
            'do_legend': True,
        }

        self.plot_dicts['raw_phase'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' raw phase',
            'xvals': self.data_dict['tau'],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': self.data_dict['raw phase'],
            'ylabel': 'phase',
            'yunit': 'deg',
        }

        self.plot_dicts['phase'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' phase',
            'xvals': self.data_dict['tau'],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': self.data_dict['phase'],
            'ylabel': 'phase',
            'yunit': 'deg',
        }

        self.plot_dicts['dV'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' $\\delta V$',
            'xvals': self.data_dict['tau'][:len(self.data_dict['dV'])],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': self.data_dict['dV'],
            'ylabel': '$\\delta V$',
            'yunit': 'V',
        }

    def run_fitting(self):
        pass
<commit_msg>Unhardcoded some things in cryo-scope analysis<commit_after>import pycqed.analysis_v2.base_analysis as ba
import pycqed.measurement.waveform_control_CC.waveform as wf
import pycqed.analysis.fitting_models as fit_mods
import numpy as np
from numpy.fft import fft, ifft, fftfreq

class DistortionFineAnalysis(ba.BaseDataAnalysis):
    '''
    Analysis for the enhanced cryogenic oscilloscpe.
    '''
    def __init__(self, t_start: str=None, t_stop: str=None,
                 data_file_path: str=None,
                 options_dict: dict=None, extract_only: bool=False,
                 do_fitting: bool=True, auto=True):
        if options_dict is None:
            options_dict = dict()
        options_dict['scan_label'] = 'distortion_scope_fine'
        options_dict['tight_fig'] = False
        options_dict['apply_default_fig_settings'] = False

        super().__init__(t_start=t_start, t_stop=t_stop,
                         data_file_path=data_file_path,
                         options_dict=options_dict,
                         extract_only=extract_only, do_fitting=do_fitting)

        # Extract metadata to know where to extract the parameters from
        self.params_dict = {
            'qubit_name':
                'Experimental Data.Experimental Metadata.qubit_name',
            'QWG_channel':
                'Experimental Data.Experimental Metadata.QWG_channel',
            'flux_LutMan':
                'Experimental Data.Experimental Metadata.flux_LutMan',
        }
        self.numeric_params = []
        self.extract_data()

        # Now actually extract the parameters and the data
        self.params_dict = {
            # Data
            'xlabel': 'sweep_name',
            'xunit': 'sweep_unit',
            'measurementstring': 'measurementstring',
            'sweep_points': 'sweep_points',
            'value_names': 'value_names',
            'value_units': 'value_units',
            'measured_values': 'measured_values',

            # Pulse parameters
            'gauss_sigma':
                '{}.S_gauss_sigma'.format(self.data_dict['flux_LutMan'][0]),
            'gauss_amp': '{}.S_amp'.format(self.data_dict['flux_LutMan'][0]),
            'wave_dict_unit':
                '{}.wave_dict_unit'.format(self.data_dict['flux_LutMan'][0]),
            'sampling_rate':
                '{}.sampling_rate'.format(self.data_dict['flux_LutMan'][0]),
            'QWG_amp':
                'QWG.ch{}_amp'.format(self.data_dict['QWG_channel'][0]),
            'distortions_amp':
                '{}.F_amp'.format(self.data_dict['flux_LutMan'][0]),

            # Qubit parameters
            'f_max': '{}.f_max'.format(self.data_dict['qubit_name'][0]),
            'E_c': '{}.E_c'.format(self.data_dict['qubit_name'][0]),
            'V_offset': '{}.V_offset'.format(self.data_dict['qubit_name'][0]),
            'V_per_phi0':
                '{}.V_per_phi0'.format(self.data_dict['qubit_name'][0]),
            'asymmetry':
                '{}.asymmetry'.format(self.data_dict['qubit_name'][0]),
        }
        self.numeric_params = [
            'gauss_sigma',
            'gauss_amp',
            'sampling_rate',
            'QWG_amp',
            'distortions_amp',
            'f_max',
            'E_c',
            'V_offset',
            'V_per_phi0',
            'asymmetry',
        ]

        if auto:
            self.run_analysis()

    def process_data(self):
        # Convert lists of duplicated params to floats
        for k in self.numeric_params:
            self.data_dict[k] = self.data_dict[k][1]

        # Background
        cos = (self.data_dict['measured_values'][0][0, ::2] - 0.5) * 2
        sin = (self.data_dict['measured_values'][0][0, 1::2] - 0.5) * 2
        bg_phase = np.unwrap(np.arctan2(sin, cos))
        self.data_dict['background phase'] = np.mean(bg_phase)
        n = fft(bg_phase)
        self.data_dict['noise power spectrum'] = np.conj(n) * n

        # Data
        self.data_dict['cos'] = (
            self.data_dict['measured_values'][1][0, ::2] - 0.5) * 2
        self.data_dict['sin'] = (
            self.data_dict['measured_values'][1][0, 1::2] - 0.5) * 2

        self.data_dict['raw phase'] = np.arctan2(self.data_dict['sin'],
                                                 self.data_dict['cos'])
        # Unwrap phase back-to-front, because it makes more sense to assume
        # it's closer to zero at later times.
        self.data_dict['phase'] = (
            np.unwrap(self.data_dict['raw phase'][::-1])[::-1] -
            self.data_dict['background phase'])

        # After unwrapping we can convert phases to degrees
        self.data_dict['raw phase'] = np.rad2deg(self.data_dict['raw phase'])
        self.data_dict['phase'] = np.rad2deg(self.data_dict['phase'])
        self.data_dict['background phase'] = np.rad2deg(
            self.data_dict['background phase'])

        self.data_dict['tau'] = self.data_dict['sweep_points'][1][::2]

        # Calculate integral of the scoping pulse
        data_dt = (self.data_dict['tau'][1] -
                   self.data_dict['tau'][0])
        gauss_pulse = wf.gauss_pulse(
            self.data_dict['gauss_amp'],
            self.data_dict['gauss_sigma'],
            nr_sigma=4,
            sampling_rate=data_dt,
            axis='x', phase=0, motzoi=0, delay=0)[0]  # Returns I and Q

        if self.data_dict['wave_dict_unit'] == 'frac':
            # Convert to volts
            gauss_pulse *= self.data_dict['QWG_amp'] / 2

        # Add zeros to front and back of data because the signal does not
        # start at tau = 0.
        # t0: time from start to center of scoping pulse
        t0 = self.data_dict['gauss_sigma'] * 2
        t0_samples = int(np.round(t0 / data_dt))

        y = np.concatenate((
            np.zeros(t0_samples),
            self.data_dict['phase'],
            np.zeros(t0_samples)))
        N = np.concatenate((
            np.zeros(t0_samples),
            self.data_dict['noise power spectrum'],
            np.zeros(t0_samples)))
        self.y = y

        # Calculate convolution kernel
        # Kernel needs to be sampled at the same rate as data.

        h = 360 * fit_mods.Qubit_dac_sensitivity(
            dac_voltage=gauss_pulse,
            f_max=self.data_dict['f_max'],
            E_c=self.data_dict['E_c'],
            dac_sweet_spot=self.data_dict['V_offset'],
            V_per_phi0=self.data_dict['V_per_phi0'],
            asymmetry=self.data_dict['asymmetry'])
        self.h = h

        # Note that we assume now that the magnitude of the phase error is
        # smaller than 180 degrees, i.e. we assume the phase has not wrapped
        # due to distortions.
        h = np.hstack((h, np.zeros(len(y) - len(h))))
        # Abs of fft, because we want the spectrum of a gauss window centered
        # around zero. Note: f real, symmetric => fft(f) real, symmetric.
        freqs = fftfreq(len(h), d=data_dt)
        H = np.abs(fft(h))
        x = np.real(ifft(fft(y) / H)) * (freqs[1] - freqs[0])
        # Wiener deconvolution
        # x = np.real(ifft(fft(y) * np.conj(H) /
        #                  (np.conj(H) * H + N / (np.conj(y) * y))))

        self.data_dict['tau'] += t0
        self.data_dict['dV'] = x[t0_samples:-t0_samples]

    def prepare_plots(self):
        self.plot_dicts['raw data'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' raw data',
            'xvals': self.data_dict['tau'],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': [self.data_dict['cos'], self.data_dict['sin']],
            'multiple': True,
            'setlabel': ['cos', 'sin'],
            'do_legend': True,
        }

        self.plot_dicts['raw_phase'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' raw phase',
            'xvals': self.data_dict['tau'],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': self.data_dict['raw phase'],
            'ylabel': 'phase',
            'yunit': 'deg',
        }

        self.plot_dicts['phase'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' phase',
            'xvals': self.data_dict['tau'],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': self.data_dict['phase'],
            'ylabel': 'phase',
            'yunit': 'deg',
        }

        self.plot_dicts['dV'] = {
            'plotfn': self.plot_line,
            'title': self.timestamps[0] + ' $\\delta V$',
            'xvals': self.data_dict['tau'][:len(self.data_dict['dV'])],
            'xlabel': self.data_dict['xlabel'][1],
            'xunit': self.data_dict['xunit'][1][0],
            'yvals': self.data_dict['dV'],
            'ylabel': '$\\delta V$',
            'yunit': 'V',
        }

    def run_fitting(self):
        pass
<|endoftext|>


<commit_before>import chainer
from chainer import cuda
import chainer.functions as F
import copy

from chainerrl import agent

def _F_clip(x, x_min, x_max):
    return F.minimum(F.maximum(x, x_min), x_max)

class PPO(agent.AttributeSavingMixin, agent.Agent):
    """Proximal Policy Optimization

    See https://blog.openai.com/openai-baselines-ppo/

    Args:
        model (A3CModel): Model to train.
            state s  |->  (pi(s, _), v(s))
        optimizer (chainer.Optimizer): optimizer used to train the model
        gpu (int): GPU device id if not None nor negative
        gamma (float): Discount factor [0, 1]
        lambd (float): Lambda-return factor [0, 1]
        value_func_coeff (float): Weight coefficient for loss of
            value function (0, inf)
        entropy_coeff (float): Weight coefficient for entropoy bonus [0, inf)
        update_interval (int): Model update interval in step
        minibatch_size (int): Minibatch size
        epochs (int): Training epochs in an update
        clip_eps (float): Epsilon for pessimistic clipping of likelihood ratio
            to update policy
        clip_eps_vf (float): Epsilon for pessimistic clipping of value
            to update value function
    """

    saved_attributes = ['model', 'optimizer']

    def __init__(self, model, optimizer,
                 gpu=None,
                 gamma=0.99,
                 lambd=0.95,
                 value_func_coeff=1.0,
                 entropy_coeff=0.01,
                 update_interval=2048,
                 minibatch_size=64,
                 epochs=10,
                 clip_eps=0.2,
                 clip_eps_vf=0.2,
                 ):
        self.model = model

        if gpu is not None and gpu >= 0:
            cuda.get_device(gpu).use()
            self.model.to_gpu(device=gpu)

        self.optimizer = optimizer
        self.gamma = gamma
        self.lambd = lambd
        self.value_func_coeff = value_func_coeff
        self.entropy_coeff = entropy_coeff
        self.update_interval = update_interval
        self.minibatch_size = minibatch_size
        self.epochs = epochs
        self.clip_eps = clip_eps
        self.clip_eps_vf = clip_eps_vf

        self.xp = self.model.xp
        self.target_model = None
        self.last_state = None

        self.memory = []
        self.last_episode = []

    def _act(self, state, train):
        state = state.astype(self.xp.float32)
        with chainer.using_config('train', train):
            b_state = F.expand_dims(state, axis=0)
            action_distrib, v = self.model(b_state)
            action = action_distrib.sample()
            return action[0].data, v[0].data

    def _train(self):
        if len(self.memory) + len(self.last_episode) >= self.update_interval:
            self.flush_last_episode()
            self.update()
            self.memory = []

    def flush_last_episode(self):
        if self.last_episode:
            self.compute_v_teacher()
            self.memory.extend(self.last_episode)
            self.last_episode = []

    def compute_v_teacher(self):
        adv = 0.0
        for transition in reversed(self.last_episode):
            td_err = (
                transition['reward']
                + (self.gamma * transition['nonterminal']
                   * transition['next_v_pred'])
                - transition['v_pred']
                )
            adv = td_err + self.lambd * adv
            transition['adv'] = adv
            transition['v_teacher'] = adv + transition['v_pred']  # ????

    def lossfun(self, prob_ratio, advs, vs_pred, vs_pred_old, vs_teacher, ent):
        prob_ratio = F.expand_dims(prob_ratio, axis=-1)
        loss_policy = - F.mean(F.minimum(
            prob_ratio * advs,
            F.clip(prob_ratio, 1-self.clip_eps, 1+self.clip_eps) * advs))

        if self.clip_eps_vf is None:
            loss_value_func = F.mean_squared_error(vs_pred, vs_teacher)
        else:
            loss_value_func = F.mean(F.maximum(
                F.square(vs_pred - vs_teacher),
                F.square(_F_clip(vs_pred,
                                 vs_pred_old - self.clip_eps_vf,
                                 vs_pred_old + self.clip_eps_vf)
                         - vs_teacher)
                ))

        loss_entropy = F.mean(ent)

        # print((loss_policy, loss_value_func, loss_entropy))
        return (
            loss_policy
            + self.value_func_coeff * loss_value_func
            + self.entropy_coeff * loss_entropy
            )

    def update(self):
        xp = self.xp
        target_model = copy.deepcopy(self.model)
        dataset_iter = chainer.iterators.SerialIterator(
            self.memory, self.minibatch_size)

        dataset_iter.reset()
        while dataset_iter.epoch < self.epochs:
            batch = dataset_iter.__next__()
            states = xp.array([b['state'] for b in batch], dtype=xp.float32)
            actions = xp.array([b['action'] for b in batch], dtype=xp.int32)
            vs_pred_old = xp.array(
                [b['v_pred'] for b in batch], dtype=xp.float32)

            distribs, vs_pred = self.model(states)
            # print(distribs)
            log_probs = distribs.log_prob(actions)
            target_distribs, _ = target_model(states)
            target_log_probs = target_distribs.log_prob(actions)
            prob_ratio = F.exp(log_probs - target_log_probs)

            advs = xp.array([b['adv'] for b in batch], dtype=xp.float32)
            vs_teacher = xp.array(
                [b['v_teacher'] for b in batch], dtype=xp.float32)
            ent = distribs.entropy

            self.optimizer.update(
                self.lossfun,
                prob_ratio, advs, vs_pred, vs_pred_old, vs_teacher, ent)

    def act_and_train(self, state, reward, done=False):
        action, v = self._act(state, train=True)

        if self.last_state is not None:
            self.last_episode.append({
                'state': self.last_state,
                'action': self.last_action,
                'reward': reward,
                'v_pred': self.last_v,
                'next_state': state,
                'next_v_pred': v,
                'nonterminal': 0.0 if done else 1.0})
        self.last_state = state
        self.last_action = action
        self.last_v = v

        self._train()
        return action

    def act(self, state):
        action, _ = self._act(state, train=False)
        return action

    def stop_episode_and_train(self, state, reward, done=False):
        _, v = self._act(state, train=True)

        if self.last_state is not None:
            self.last_episode.append({
                'state': self.last_state,
                'action': self.last_action,
                'reward': reward,
                'v_pred': self.last_v,
                'next_state': state,
                'next_v_pred': v,
                'nonterminal': 0.0 if done else 1.0})
        self.last_state = None
        del self.last_action
        del self.last_v

        self.flush_last_episode()
        self.stop_episode()

    def stop_episode(self):
        pass

    def get_statistics(self):
        return []
<commit_msg>debug gpu codes<commit_after>import chainer
from chainer import cuda
import chainer.functions as F
import copy

from chainerrl import agent

def _F_clip(x, x_min, x_max):
    return F.minimum(F.maximum(x, x_min), x_max)

class PPO(agent.AttributeSavingMixin, agent.Agent):
    """Proximal Policy Optimization

    See https://blog.openai.com/openai-baselines-ppo/

    Args:
        model (A3CModel): Model to train.
            state s  |->  (pi(s, _), v(s))
        optimizer (chainer.Optimizer): optimizer used to train the model
        gpu (int): GPU device id if not None nor negative
        gamma (float): Discount factor [0, 1]
        lambd (float): Lambda-return factor [0, 1]
        value_func_coeff (float): Weight coefficient for loss of
            value function (0, inf)
        entropy_coeff (float): Weight coefficient for entropoy bonus [0, inf)
        update_interval (int): Model update interval in step
        minibatch_size (int): Minibatch size
        epochs (int): Training epochs in an update
        clip_eps (float): Epsilon for pessimistic clipping of likelihood ratio
            to update policy
        clip_eps_vf (float): Epsilon for pessimistic clipping of value
            to update value function
    """

    saved_attributes = ['model', 'optimizer']

    def __init__(self, model, optimizer,
                 gpu=None,
                 gamma=0.99,
                 lambd=0.95,
                 value_func_coeff=1.0,
                 entropy_coeff=0.01,
                 update_interval=2048,
                 minibatch_size=64,
                 epochs=10,
                 clip_eps=0.2,
                 clip_eps_vf=0.2,
                 ):
        self.model = model

        if gpu is not None and gpu >= 0:
            cuda.get_device_from_id(gpu).use()
            self.model.to_gpu(device=gpu)

        self.optimizer = optimizer
        self.gamma = gamma
        self.lambd = lambd
        self.value_func_coeff = value_func_coeff
        self.entropy_coeff = entropy_coeff
        self.update_interval = update_interval
        self.minibatch_size = minibatch_size
        self.epochs = epochs
        self.clip_eps = clip_eps
        self.clip_eps_vf = clip_eps_vf

        self.xp = self.model.xp
        self.target_model = None
        self.last_state = None

        self.memory = []
        self.last_episode = []

    def _act(self, state, train):
        xp = self.xp
        state = xp.asarray(state, dtype=xp.float32)
        with chainer.using_config('train', train):
            b_state = F.expand_dims(state, axis=0)
            action_distrib, v = self.model(b_state)
            action = action_distrib.sample()
            return action[0].data, v[0].data

    def _train(self):
        if len(self.memory) + len(self.last_episode) >= self.update_interval:
            self.flush_last_episode()
            self.update()
            self.memory = []

    def flush_last_episode(self):
        if self.last_episode:
            self.compute_v_teacher()
            self.memory.extend(self.last_episode)
            self.last_episode = []

    def compute_v_teacher(self):
        adv = 0.0
        for transition in reversed(self.last_episode):
            td_err = (
                transition['reward']
                + (self.gamma * transition['nonterminal']
                   * transition['next_v_pred'])
                - transition['v_pred']
                )
            adv = td_err + self.lambd * adv
            transition['adv'] = adv
            transition['v_teacher'] = adv + transition['v_pred']  # ????

    def lossfun(self, prob_ratio, advs, vs_pred, vs_pred_old, vs_teacher, ent):
        prob_ratio = F.expand_dims(prob_ratio, axis=-1)
        loss_policy = - F.mean(F.minimum(
            prob_ratio * advs,
            F.clip(prob_ratio, 1-self.clip_eps, 1+self.clip_eps) * advs))

        if self.clip_eps_vf is None:
            loss_value_func = F.mean_squared_error(vs_pred, vs_teacher)
        else:
            loss_value_func = F.mean(F.maximum(
                F.square(vs_pred - vs_teacher),
                F.square(_F_clip(vs_pred,
                                 vs_pred_old - self.clip_eps_vf,
                                 vs_pred_old + self.clip_eps_vf)
                         - vs_teacher)
                ))

        loss_entropy = F.mean(ent)

        # print((loss_policy, loss_value_func, loss_entropy))
        return (
            loss_policy
            + self.value_func_coeff * loss_value_func
            + self.entropy_coeff * loss_entropy
            )

    def update(self):
        xp = self.xp
        target_model = copy.deepcopy(self.model)
        dataset_iter = chainer.iterators.SerialIterator(
            self.memory, self.minibatch_size)

        dataset_iter.reset()
        while dataset_iter.epoch < self.epochs:
            batch = dataset_iter.__next__()
            states = xp.array([b['state'] for b in batch], dtype=xp.float32)
            actions = xp.array([b['action'] for b in batch], dtype=xp.int32)
            vs_pred_old = xp.array(
                [b['v_pred'] for b in batch], dtype=xp.float32)

            distribs, vs_pred = self.model(states)
            # print(distribs)
            log_probs = distribs.log_prob(actions)
            target_distribs, _ = target_model(states)
            target_log_probs = target_distribs.log_prob(actions)
            prob_ratio = F.exp(log_probs - target_log_probs)

            advs = xp.array([b['adv'] for b in batch], dtype=xp.float32)
            vs_teacher = xp.array(
                [b['v_teacher'] for b in batch], dtype=xp.float32)
            ent = distribs.entropy

            self.optimizer.update(
                self.lossfun,
                prob_ratio, advs, vs_pred, vs_pred_old, vs_teacher, ent)

    def act_and_train(self, state, reward, done=False):
        action, v = self._act(state, train=True)

        if self.last_state is not None:
            self.last_episode.append({
                'state': self.last_state,
                'action': self.last_action,
                'reward': reward,
                'v_pred': self.last_v,
                'next_state': state,
                'next_v_pred': v,
                'nonterminal': 0.0 if done else 1.0})
        self.last_state = state
        self.last_action = action
        self.last_v = v

        self._train()
        return cuda.to_cpu(action)

    def act(self, state):
        action, _ = self._act(state, train=False)
        return cuda.to_cpu(action)

    def stop_episode_and_train(self, state, reward, done=False):
        _, v = self._act(state, train=True)

        if self.last_state is not None:
            self.last_episode.append({
                'state': self.last_state,
                'action': self.last_action,
                'reward': reward,
                'v_pred': self.last_v,
                'next_state': state,
                'next_v_pred': v,
                'nonterminal': 0.0 if done else 1.0})
        self.last_state = None
        del self.last_action
        del self.last_v

        self.flush_last_episode()
        self.stop_episode()

    def stop_episode(self):
        pass

    def get_statistics(self):
        return []
<|endoftext|>


<commit_before>#!/usr/bin/env python
print "Hello  Word"
<commit_msg>Adding Something else<commit_after>#!/usr/bin/env python
print "Hello  Word"
print " Something else"
<|endoftext|>


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

"""
Core calculator functionality for computing stochastic event sets and ground
motion fields using the 'event-based' method.

Stochastic events sets (which can be thought of as collections of ruptures) are
computed iven a set of seismic sources and investigation time span (in years).

For more information on computing stochastic event sets, see
:mod:`openquake.hazardlib.calc.stochastic`.

One can optionally compute a ground motion field (GMF) given a rupture, a site
collection (which is a collection of geographical points with associated soil
parameters), and a ground shaking intensity model (GSIM).

For more information on computing ground motion fields, see
:mod:`openquake.hazardlib.calc.gmf`.
"""

import random
import collections

import numpy.random

from django.db import transaction
from openquake.hazardlib.calc import filters
from openquake.hazardlib.calc import gmf
from openquake.hazardlib.imt import from_string

from openquake.engine import writer
from openquake.engine.calculators.hazard import general
from openquake.engine.calculators.hazard.classical import (
    post_processing as cls_post_proc)
from openquake.engine.calculators.hazard.event_based import post_processing
from openquake.engine.db import models
from openquake.engine.utils import tasks
from openquake.engine.performance import EnginePerformanceMonitor, LightMonitor

#: Always 1 for the computation of ground motion fields in the event-based
#: hazard calculator.
DEFAULT_GMF_REALIZATIONS = 1

# NB: beware of large caches
inserter = writer.CacheInserter(models.GmfData, 1000)

@tasks.oqtask
def compute_ses_and_gmfs(job_id, src_seeds, gsims_by_rlz, task_no):
    """
    Celery task for the stochastic event set calculator.

    Samples logic trees and calls the stochastic event set calculator.

    Once stochastic event sets are calculated, results will be saved to the
    database. See :class:`openquake.engine.db.models.SESCollection`.

    Optionally (specified in the job configuration using the
    `ground_motion_fields` parameter), GMFs can be computed from each rupture
    in each stochastic event set. GMFs are also saved to the database.

    :param int job_id:
        ID of the currently running job.
    :param src_seeds:
        List of pairs (source, seed)
    :params gsims_by_rlz:
        dictionary of GSIM
    :param task_no:
        an ordinal so that GMV can be collected in a reproducible order
    """
    rlz_ids = [r.id for r in gsims_by_rlz]
    ses_coll = models.SESCollection.objects.get(lt_realization_ids=rlz_ids)

    rnd = random.Random()
    all_ses = models.SES.objects.filter(ses_collection=ses_coll)
    ruptures = []
    ses_ruptures = []
    rupture_seeds = []

    mon1 = LightMonitor('generating ruptures', job_id, compute_ses_and_gmfs)
    mon2 = LightMonitor('sampling ruptures', job_id, compute_ses_and_gmfs)

    # Compute and save stochastic event sets
    for src, seed in src_seeds:
        rnd.seed(seed)
        with mon1:
            rupts = list(src.iter_ruptures())

        for ses in all_ses:
            numpy.random.seed(rnd.randint(0, models.MAX_SINT_32))
            for i, r in enumerate(rupts):
                with mon2:
                    for j in xrange(r.sample_number_of_occurrences()):
                        rup = models.SESRupture(
                            ses=ses,
                            rupture=r,
                            tag='smlt=%02d|ses=%04d|src=%s|i=%04d-%02d' % (
                                ses_coll.ordinal,
                                ses.ordinal, src.source_id, i, j),
                            hypocenter=r.hypocenter.wkt2d,
                            magnitude=r.mag,
                        )
                        rup_seed = rnd.randint(0, models.MAX_SINT_32)
                        ruptures.append(r)
                        ses_ruptures.append(rup)
                        rupture_seeds.append(rup_seed)

    mon1.flush()
    mon2.flush()
    if not ruptures:
        return

    with EnginePerformanceMonitor('saving ses', job_id, compute_ses_and_gmfs):
        rupture_ids = _save_ses_ruptures(ses_ruptures)

    hc = models.HazardCalculation.objects.get(oqjob=job_id)
    if not hc.ground_motion_fields:
        return

    # compute GMFs
    imts = map(from_string, hc.intensity_measure_types)
    params = dict(
        correl_model=general.get_correl_model(hc),
        truncation_level=hc.truncation_level,
        maximum_distance=hc.maximum_distance)

    for rlz, gsims in gsims_by_rlz.items():
        with EnginePerformanceMonitor(
                'computing gmfs', job_id, compute_ses_and_gmfs):
            gmvs_per_site, ruptures_per_site = _compute_gmf(
                params, imts, gsims, hc.site_collection,
                zip(ruptures, rupture_ids, rupture_seeds))

        with EnginePerformanceMonitor(
                'saving gmfs', job_id, compute_ses_and_gmfs):
            gmf_coll = models.Gmf.objects.get(lt_realization=rlz)
            _save_gmfs(gmf_coll, gmvs_per_site, ruptures_per_site,
                       hc.site_collection, task_no)

def _save_ses_ruptures(ruptures):
    """
    Helper function for saving stochastic event set ruptures to the database.
    :param ruptures:
        A list of :class:`openquake.engine.db.models.SESRupture` instances.
    """
    # TODO: Possible future optimization:
    # Refactor this to do bulk insertion of ruptures
    rupture_ids = []
    with transaction.commit_on_success(using='job_init'):
        for r in ruptures:
            r.save()
            rupture_ids.append(r.id)
    return rupture_ids

# NB: I tried to return a single dictionary {site_id: [(gmv, rupt_id),...]}
# but it takes a lot more memory (MS)
def _compute_gmf(params, imts, gsims, site_coll, rupture_id_seed_triples):
    """
    Compute a ground motion field value for each rupture, for all the
    points affected by that rupture, for the given IMT. Returns a
    dictionary with the nonzero contributions to each site id, and a dictionary
    with the ids of the contributing ruptures for each site id.
    assert len(ruptures) == len(rupture_seeds)

    :param params:
        a dictionary containing the keys
        correl_model, truncation_level, maximum_distance
    :param imts:
        a list of hazardlib intensity measure types
    :param gsims:
        a dictionary {tectonic region type -> GSIM instance}
    :param site_coll:
        a SiteCollection instance
    :param rupture_id_seed_triple:
        a list of triples with types
        (:class:`openquake.hazardlib.source.rupture.Rupture`, int, int)
    """
    gmvs_per_site = collections.defaultdict(list)
    ruptures_per_site = collections.defaultdict(list)

    # Compute and save ground motion fields
    for i, (rupture, rup_id, rup_seed) in enumerate(rupture_id_seed_triples):
        gmf_calc_kwargs = {
            'rupture': rupture,
            'sites': site_coll,
            'imts': imts,
            'gsim': gsims[rupture.tectonic_region_type],
            'truncation_level': params['truncation_level'],
            'realizations': DEFAULT_GMF_REALIZATIONS,
            'correlation_model': params['correl_model'],
            'rupture_site_filter': filters.rupture_site_distance_filter(
                params['maximum_distance']),
        }
        numpy.random.seed(rup_seed)
        gmf_dict = gmf.ground_motion_fields(**gmf_calc_kwargs)
        for imt, gmf_1_realiz in gmf_dict.iteritems():
            # since DEFAULT_GMF_REALIZATIONS is 1, gmf_1_realiz is a matrix
            # with n_sites rows and 1 column
            for site, gmv in zip(site_coll, gmf_1_realiz):
                gmv = float(gmv)  # convert a 1x1 matrix into a float
                if gmv:  # nonzero contribution to site
                    gmvs_per_site[imt, site.id].append(gmv)
                    ruptures_per_site[imt, site.id].append(rup_id)
    return gmvs_per_site, ruptures_per_site

@transaction.commit_on_success(using='job_init')
def _save_gmfs(gmf, gmvs_per_site, ruptures_per_site, sites, task_no):
    """
    Helper method to save computed GMF data to the database.

    :param gmf:
        The Gmf instance where to save
    :param gmf_per_site:
        The GMFs per rupture
    :param rupture_per_site:
        The associated rupture ids
    :param sites:
        An :class:`openquake.hazardlib.site.SiteCollection` object,
        representing the sites of interest for a calculation.
    :param task_no:
        The ordinal of the task which generated the current GMFs to save
    """
    for imt, site_id in gmvs_per_site:
        imt_name, sa_period, sa_damping = imt
        inserter.add(models.GmfData(
            gmf=gmf,
            task_no=task_no,
            imt=imt_name,
            sa_period=sa_period,
            sa_damping=sa_damping,
            site_id=site_id,
            gmvs=gmvs_per_site[imt, site_id],
            rupture_ids=ruptures_per_site[imt, site_id]))
    inserter.flush()

class EventBasedHazardCalculator(general.BaseHazardCalculator):
    """
    Probabilistic Event-Based hazard calculator. Computes stochastic event sets
    and (optionally) ground motion fields.
    """
    core_calc_task = compute_ses_and_gmfs

    def task_arg_gen(self, _block_size=None):
        """
        Loop through realizations and sources to generate a sequence of
        task arg tuples. Each tuple of args applies to a single task.
        Yielded results are tuples of the form job_id, sources, ses, seeds
        (seeds will be used to seed numpy for temporal occurence sampling).
        """
        hc = self.hc
        rnd = random.Random()
        rnd.seed(hc.random_seed)
        task_no = 0
        for job_id, block, gsims_by_rlz in super(
                EventBasedHazardCalculator, self).task_arg_gen():
            ss = [(src, rnd.randint(0, models.MAX_SINT_32))
                  for src in block]  # source, seed pairs
            yield job_id, ss, gsims_by_rlz, task_no
            task_no += 1

        # now the source_blocks_per_ltpath dictionary can be cleared
        self.source_blocks_per_ltpath.clear()

    def initialize_ses_db_records(self, ordinal, rlzs):
        """
        Create :class:`~openquake.engine.db.models.Output`,
        :class:`~openquake.engine.db.models.SESCollection` and
        :class:`~openquake.engine.db.models.SES` "container" records for
        a single realization.

        Stochastic event set ruptures computed for this realization will be
        associated to these containers.

        NOTE: Many tasks can contribute ruptures to the same SES.
        """
        rlz_ids = [r.id for r in rlzs]

        output = models.Output.objects.create(
            oq_job=self.job,
            display_name='SES Collection smlt-%d-rlz-%s' % (
                ordinal, ','.join(map(str, rlz_ids))),
            output_type='ses')

        ses_coll = models.SESCollection.objects.create(
            output=output, lt_realization_ids=rlz_ids, ordinal=ordinal)

        for rlz in rlzs:
            if self.job.hazard_calculation.ground_motion_fields:
                output = models.Output.objects.create(
                    oq_job=self.job,
                    display_name='GMF rlz-%s' % rlz.id,
                    output_type='gmf')
                models.Gmf.objects.create(output=output, lt_realization=rlz)

        all_ses = []
        for i in xrange(1, self.hc.ses_per_logic_tree_path + 1):
            all_ses.append(
                models.SES.objects.create(
                    ses_collection=ses_coll,
                    investigation_time=self.hc.investigation_time,
                    ordinal=i))
        return all_ses

    def pre_execute(self):
        """
        Do pre-execution work. At the moment, this work entails:
        parsing and initializing sources, parsing and initializing the
        site model (if there is one), parsing vulnerability and
        exposure files, and generating logic tree realizations. (The
        latter piece basically defines the work to be done in the
        `execute` phase.)
        """
        super(EventBasedHazardCalculator, self).pre_execute()
        for i, rlzs in enumerate(self.rlzs_per_ltpath.itervalues()):
            self.initialize_ses_db_records(i, rlzs)

    def post_process(self):
        """
        If requested, perform additional processing of GMFs to produce hazard
        curves.
        """
        if self.hc.hazard_curves_from_gmfs:
            with EnginePerformanceMonitor('generating hazard curves',
                                          self.job.id):
                self.parallelize(
                    post_processing.gmf_to_hazard_curve_task,
                    post_processing.gmf_to_hazard_curve_arg_gen(self.job),
                    self.log_percent)

            # If `mean_hazard_curves` is True and/or `quantile_hazard_curves`
            # has some value (not an empty list), do this additional
            # post-processing.
            if self.hc.mean_hazard_curves or self.hc.quantile_hazard_curves:
                with EnginePerformanceMonitor(
                        'generating mean/quantile curves', self.job.id):
                    self.do_aggregate_post_proc()

            if self.hc.hazard_maps:
                with EnginePerformanceMonitor(
                        'generating hazard maps', self.job.id):
                    self.parallelize(
                        cls_post_proc.hazard_curves_to_hazard_map_task,
                        cls_post_proc.hazard_curves_to_hazard_map_task_arg_gen(
                            self.job),
                        self.log_percent)
<commit_msg>Added source and rupture filtering<commit_after># Copyright (c) 2010-2013, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core calculator functionality for computing stochastic event sets and ground
motion fields using the 'event-based' method.

Stochastic events sets (which can be thought of as collections of ruptures) are
computed iven a set of seismic sources and investigation time span (in years).

For more information on computing stochastic event sets, see
:mod:`openquake.hazardlib.calc.stochastic`.

One can optionally compute a ground motion field (GMF) given a rupture, a site
collection (which is a collection of geographical points with associated soil
parameters), and a ground shaking intensity model (GSIM).

For more information on computing ground motion fields, see
:mod:`openquake.hazardlib.calc.gmf`.
"""

import random
import collections

import numpy.random

from django.db import transaction
from openquake.hazardlib.calc import filters
from openquake.hazardlib.calc import gmf
from openquake.hazardlib.imt import from_string

from openquake.engine import writer
from openquake.engine.calculators.hazard import general
from openquake.engine.calculators.hazard.classical import (
    post_processing as cls_post_proc)
from openquake.engine.calculators.hazard.event_based import post_processing
from openquake.engine.db import models
from openquake.engine.utils import tasks
from openquake.engine.performance import EnginePerformanceMonitor, LightMonitor

#: Always 1 for the computation of ground motion fields in the event-based
#: hazard calculator.
DEFAULT_GMF_REALIZATIONS = 1

# NB: beware of large caches
inserter = writer.CacheInserter(models.GmfData, 1000)

@tasks.oqtask
def compute_ses_and_gmfs(job_id, src_seeds, gsims_by_rlz, task_no):
    """
    Celery task for the stochastic event set calculator.

    Samples logic trees and calls the stochastic event set calculator.

    Once stochastic event sets are calculated, results will be saved to the
    database. See :class:`openquake.engine.db.models.SESCollection`.

    Optionally (specified in the job configuration using the
    `ground_motion_fields` parameter), GMFs can be computed from each rupture
    in each stochastic event set. GMFs are also saved to the database.

    :param int job_id:
        ID of the currently running job.
    :param src_seeds:
        List of pairs (source, seed)
    :params gsims_by_rlz:
        dictionary of GSIM
    :param task_no:
        an ordinal so that GMV can be collected in a reproducible order
    """
    rlz_ids = [r.id for r in gsims_by_rlz]
    ses_coll = models.SESCollection.objects.get(lt_realization_ids=rlz_ids)

    hc = models.HazardCalculation.objects.get(oqjob=job_id)
    all_ses = models.SES.objects.filter(ses_collection=ses_coll)
    ruptures = []
    ses_ruptures = []
    rupture_seeds = []

    mon1 = LightMonitor('filtering sites', job_id, compute_ses_and_gmfs)
    mon2 = LightMonitor('generating ruptures', job_id, compute_ses_and_gmfs)
    mon3 = LightMonitor('sampling ruptures', job_id, compute_ses_and_gmfs)

    # Compute and save stochastic event sets
    rnd = random.Random()
    for src, seed in src_seeds:
        rnd.seed(seed)

        with mon1:
            s_sites = src.filter_sites_by_distance_to_source(
                hc.maximum_distance, hc.site_collection
            ) if hc.maximum_distance else hc.site_collection
            if s_sites is None:
                continue

        with mon2:
            rupts = []
            for r in src.iter_ruptures():
                r_sites = r.source_typology.\
                    filter_sites_by_distance_to_rupture(
                        r, hc.maximum_distance, s_sites
                    ) if hc.maximum_distance else s_sites
                if r_sites is not None:
                    rupts.append(r)

        for ses in all_ses:
            numpy.random.seed(rnd.randint(0, models.MAX_SINT_32))
            for i, r in enumerate(rupts):
                with mon3:
                    for j in xrange(r.sample_number_of_occurrences()):
                        rup = models.SESRupture(
                            ses=ses,
                            rupture=r,
                            tag='smlt=%02d|ses=%04d|src=%s|i=%04d-%02d' % (
                                ses_coll.ordinal,
                                ses.ordinal, src.source_id, i, j),
                            hypocenter=r.hypocenter.wkt2d,
                            magnitude=r.mag,
                        )
                        rup_seed = rnd.randint(0, models.MAX_SINT_32)
                        ruptures.append(r)
                        ses_ruptures.append(rup)
                        rupture_seeds.append(rup_seed)

    mon1.flush()
    mon2.flush()
    mon3.flush()
    if not ruptures:
        return

    with EnginePerformanceMonitor('saving ses', job_id, compute_ses_and_gmfs):
        rupture_ids = _save_ses_ruptures(ses_ruptures)

    if not hc.ground_motion_fields:
        return

    # compute GMFs
    imts = map(from_string, hc.intensity_measure_types)
    params = dict(
        correl_model=general.get_correl_model(hc),
        truncation_level=hc.truncation_level,
        maximum_distance=hc.maximum_distance)

    for rlz, gsims in gsims_by_rlz.items():
        with EnginePerformanceMonitor(
                'computing gmfs', job_id, compute_ses_and_gmfs):
            gmvs_per_site, ruptures_per_site = _compute_gmf(
                params, imts, gsims, hc.site_collection,
                zip(ruptures, rupture_ids, rupture_seeds))

        with EnginePerformanceMonitor(
                'saving gmfs', job_id, compute_ses_and_gmfs):
            gmf_coll = models.Gmf.objects.get(lt_realization=rlz)
            _save_gmfs(gmf_coll, gmvs_per_site, ruptures_per_site,
                       hc.site_collection, task_no)

def _save_ses_ruptures(ruptures):
    """
    Helper function for saving stochastic event set ruptures to the database.
    :param ruptures:
        A list of :class:`openquake.engine.db.models.SESRupture` instances.
    """
    # TODO: Possible future optimization:
    # Refactor this to do bulk insertion of ruptures
    rupture_ids = []
    with transaction.commit_on_success(using='job_init'):
        for r in ruptures:
            r.save()
            rupture_ids.append(r.id)
    return rupture_ids

# NB: I tried to return a single dictionary {site_id: [(gmv, rupt_id),...]}
# but it takes a lot more memory (MS)
def _compute_gmf(params, imts, gsims, site_coll, rupture_id_seed_triples):
    """
    Compute a ground motion field value for each rupture, for all the
    points affected by that rupture, for the given IMT. Returns a
    dictionary with the nonzero contributions to each site id, and a dictionary
    with the ids of the contributing ruptures for each site id.
    assert len(ruptures) == len(rupture_seeds)

    :param params:
        a dictionary containing the keys
        correl_model, truncation_level, maximum_distance
    :param imts:
        a list of hazardlib intensity measure types
    :param gsims:
        a dictionary {tectonic region type -> GSIM instance}
    :param site_coll:
        a SiteCollection instance
    :param rupture_id_seed_triple:
        a list of triples with types
        (:class:`openquake.hazardlib.source.rupture.Rupture`, int, int)
    """
    gmvs_per_site = collections.defaultdict(list)
    ruptures_per_site = collections.defaultdict(list)

    # Compute and save ground motion fields
    for i, (rupture, rup_id, rup_seed) in enumerate(rupture_id_seed_triples):
        gmf_calc_kwargs = {
            'rupture': rupture,
            'sites': site_coll,
            'imts': imts,
            'gsim': gsims[rupture.tectonic_region_type],
            'truncation_level': params['truncation_level'],
            'realizations': DEFAULT_GMF_REALIZATIONS,
            'correlation_model': params['correl_model'],
        }
        numpy.random.seed(rup_seed)
        gmf_dict = gmf.ground_motion_fields(**gmf_calc_kwargs)
        for imt, gmf_1_realiz in gmf_dict.iteritems():
            # since DEFAULT_GMF_REALIZATIONS is 1, gmf_1_realiz is a matrix
            # with n_sites rows and 1 column
            for site, gmv in zip(site_coll, gmf_1_realiz):
                gmv = float(gmv)  # convert a 1x1 matrix into a float
                if gmv:  # nonzero contribution to site
                    gmvs_per_site[imt, site.id].append(gmv)
                    ruptures_per_site[imt, site.id].append(rup_id)
    return gmvs_per_site, ruptures_per_site

@transaction.commit_on_success(using='job_init')
def _save_gmfs(gmf, gmvs_per_site, ruptures_per_site, sites, task_no):
    """
    Helper method to save computed GMF data to the database.

    :param gmf:
        The Gmf instance where to save
    :param gmf_per_site:
        The GMFs per rupture
    :param rupture_per_site:
        The associated rupture ids
    :param sites:
        An :class:`openquake.hazardlib.site.SiteCollection` object,
        representing the sites of interest for a calculation.
    :param task_no:
        The ordinal of the task which generated the current GMFs to save
    """
    for imt, site_id in gmvs_per_site:
        imt_name, sa_period, sa_damping = imt
        inserter.add(models.GmfData(
            gmf=gmf,
            task_no=task_no,
            imt=imt_name,
            sa_period=sa_period,
            sa_damping=sa_damping,
            site_id=site_id,
            gmvs=gmvs_per_site[imt, site_id],
            rupture_ids=ruptures_per_site[imt, site_id]))
    inserter.flush()

class EventBasedHazardCalculator(general.BaseHazardCalculator):
    """
    Probabilistic Event-Based hazard calculator. Computes stochastic event sets
    and (optionally) ground motion fields.
    """
    core_calc_task = compute_ses_and_gmfs

    def task_arg_gen(self, _block_size=None):
        """
        Loop through realizations and sources to generate a sequence of
        task arg tuples. Each tuple of args applies to a single task.
        Yielded results are tuples of the form job_id, sources, ses, seeds
        (seeds will be used to seed numpy for temporal occurence sampling).
        """
        hc = self.hc
        rnd = random.Random()
        rnd.seed(hc.random_seed)
        task_no = 0
        for job_id, block, gsims_by_rlz in super(
                EventBasedHazardCalculator, self).task_arg_gen():
            ss = [(src, rnd.randint(0, models.MAX_SINT_32))
                  for src in block]  # source, seed pairs
            yield job_id, ss, gsims_by_rlz, task_no
            task_no += 1

        # now the source_blocks_per_ltpath dictionary can be cleared
        self.source_blocks_per_ltpath.clear()

    def initialize_ses_db_records(self, ordinal, rlzs):
        """
        Create :class:`~openquake.engine.db.models.Output`,
        :class:`~openquake.engine.db.models.SESCollection` and
        :class:`~openquake.engine.db.models.SES` "container" records for
        a single realization.

        Stochastic event set ruptures computed for this realization will be
        associated to these containers.

        NOTE: Many tasks can contribute ruptures to the same SES.
        """
        rlz_ids = [r.id for r in rlzs]

        output = models.Output.objects.create(
            oq_job=self.job,
            display_name='SES Collection smlt-%d-rlz-%s' % (
                ordinal, ','.join(map(str, rlz_ids))),
            output_type='ses')

        ses_coll = models.SESCollection.objects.create(
            output=output, lt_realization_ids=rlz_ids, ordinal=ordinal)

        for rlz in rlzs:
            if self.job.hazard_calculation.ground_motion_fields:
                output = models.Output.objects.create(
                    oq_job=self.job,
                    display_name='GMF rlz-%s' % rlz.id,
                    output_type='gmf')
                models.Gmf.objects.create(output=output, lt_realization=rlz)

        all_ses = []
        for i in xrange(1, self.hc.ses_per_logic_tree_path + 1):
            all_ses.append(
                models.SES.objects.create(
                    ses_collection=ses_coll,
                    investigation_time=self.hc.investigation_time,
                    ordinal=i))
        return all_ses

    def pre_execute(self):
        """
        Do pre-execution work. At the moment, this work entails:
        parsing and initializing sources, parsing and initializing the
        site model (if there is one), parsing vulnerability and
        exposure files, and generating logic tree realizations. (The
        latter piece basically defines the work to be done in the
        `execute` phase.)
        """
        super(EventBasedHazardCalculator, self).pre_execute()
        for i, rlzs in enumerate(self.rlzs_per_ltpath.itervalues()):
            self.initialize_ses_db_records(i, rlzs)

    def post_process(self):
        """
        If requested, perform additional processing of GMFs to produce hazard
        curves.
        """
        if self.hc.hazard_curves_from_gmfs:
            with EnginePerformanceMonitor('generating hazard curves',
                                          self.job.id):
                self.parallelize(
                    post_processing.gmf_to_hazard_curve_task,
                    post_processing.gmf_to_hazard_curve_arg_gen(self.job),
                    self.log_percent)

            # If `mean_hazard_curves` is True and/or `quantile_hazard_curves`
            # has some value (not an empty list), do this additional
            # post-processing.
            if self.hc.mean_hazard_curves or self.hc.quantile_hazard_curves:
                with EnginePerformanceMonitor(
                        'generating mean/quantile curves', self.job.id):
                    self.do_aggregate_post_proc()

            if self.hc.hazard_maps:
                with EnginePerformanceMonitor(
                        'generating hazard maps', self.job.id):
                    self.parallelize(
                        cls_post_proc.hazard_curves_to_hazard_map_task,
                        cls_post_proc.hazard_curves_to_hazard_map_task_arg_gen(
                            self.job),
                        self.log_percent)
<|endoftext|>