text:
"""
 Copyright (C) 2018-2020 Intel Corporation

 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.
"""

import unittest

import numpy as np

from mo.middle.passes.fusing.decomposition import convert_scale_shift_to_mul_add, convert_batch_norm
from mo.utils.unittest.graph import build_graph
from mo.utils.ir_engine.compare_graphs import compare_graphs

nodes_attributes = {
    'placeholder_1': {'shape': None, 'type': 'Parameter', 'kind': 'op', 'op': 'Parameter'},
    'placeholder_1_data': {'value': None, 'shape': None, 'kind': 'data', 'data_type': None},
    'placeholder_2': {'shape': None, 'type': 'Parameter', 'kind': 'op', 'op': 'Parameter'},
    'placeholder_2_data': {'value': None, 'shape': None, 'kind': 'data', 'data_type': None},
    # ScaleShift layer
    'scaleshift_1': {'type': 'ScaleShift', 'kind': 'op', 'op': 'ScaleShift', 'axis': 0},
    'const_scaleshift_1_w': {'value': None, 'shape': None, 'kind': 'op'},
    'scaleshift_1_w': {'value': None, 'shape': None, 'kind': 'data'},
    'const_scaleshift_1_b': {'value': None, 'shape': None, 'kind': 'op'},
    'scaleshift_1_b': {'value': None, 'shape': None, 'kind': 'data'},
    'scaleshift_1_data': {'value': None, 'shape': None, 'kind': 'data'},
    # Mul and Add operations
    'mul_1': {'type': None, 'value': None, 'kind': 'op', 'op': 'Mul'},
    'const_mul_1_w': {'value': None, 'shape': None, 'kind': 'op'},
    'mul_1_w': {'value': None, 'shape': None, 'kind': 'data'},
    'mul_1_data': {'value': None, 'shape': None, 'kind': 'data'},
    'add_1': {'type': None, 'kind': 'op', 'op': 'Add'},
    'const_add_1_w': {'value': None, 'shape': None, 'kind': 'op'},
    'add_1_w': {'value': None, 'shape': None, 'kind': 'data'},
    'add_1_data': {'value': None, 'shape': None, 'kind': 'data'},
    # Mul and Add operations
    'mul_2': {'type': None, 'kind': 'op', 'op': 'Mul'},
    'const_mul_2_w': {'value': None, 'shape': None, 'kind': 'op'},
    'mul_2_w': {'value': None, 'shape': None, 'kind': 'data'},
    'mul_2_data': {'value': None, 'shape': None, 'kind': 'data'},
    'add_2': {'type': None, 'kind': 'op', 'op': 'Add'},
    'const_add_2_w': {'value': None, 'shape': None, 'kind': 'op'},
    'add_2_w': {'value': None, 'shape': None, 'kind': 'data'},
    'add_2_data': {'value': None, 'shape': None, 'kind': 'data'},
    # Reshape
    'placeholder_2/Reshape_': {'type': 'Reshape', 'kind': 'op', 'op': 'Reshape'},
    'placeholder_2/Reshape_data': {'value': None, 'shape': None, 'kind': 'data'},
    'placeholder_2/Reshape_const': {'type': 'Const', 'kind': 'op', 'op': 'Const', 'value': None},
    'placeholder_2/Reshape_const_data': {'kind': 'data', 'value': None, 'shape': None},
    # BatchNorm operation
    'bn_op': {'type': None, 'kind': 'op', 'op': 'BatchNorm', 'can_be_fused': True},
    'const_bn_const': {'value': None, 'shape': None, 'kind': 'op'},
    'bn_const': {'value': None, 'shape': None, 'kind': 'data'},
    'const_bn_beta': {'value': None, 'shape': None, 'kind': 'op'},
    'bn_beta': {'value': None, 'shape': None, 'kind': 'data'},
    'const_bn_mean': {'value': None, 'shape': None, 'kind': 'op'},
    'bn_mean': {'value': None, 'shape': None, 'kind': 'data'},
    'const_bn_var': {'value': None, 'shape': None, 'kind': 'op'},
    'bn_var': {'value': None, 'shape': None, 'kind': 'data'},
    'bn_data': {'value': None, 'shape': None, 'kind': 'data'},
    # Concat1 operation
    'concat': {'type': 'Concat', 'kind': 'op', 'op': 'Concat'},
    'concat_data': {'value': None, 'shape': None, 'kind': 'data'},
    'op_output': {'kind': 'op', 'op': 'Result'}
}

class ScaleShiftToMulAdd(unittest.TestCase):
    # ScaleShift -> Mul
    def test_scaleshift_to_mul_1(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_1_data', 'scaleshift_1'),
                             ('const_scaleshift_1_w', 'scaleshift_1_w'),
                             ('scaleshift_1_w', 'scaleshift_1'),
                             ('scaleshift_1', 'scaleshift_1_data'),
                             ('scaleshift_1_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'scaleshift_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'scaleshift_1_data': {}
                             })

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_1_data', 'mul_1'),
                                 ('const_mul_1_w', 'mul_1_w'),
                                 ('mul_1_w', 'mul_1'),
                                 ('mul_1', 'scaleshift_1_data'),
                                 ('scaleshift_1_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'const_mul_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'mul_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'mul_1': {'can_be_fused': True},
                                 'scaleshift_1_data': {}
                                 })

        graph.graph['layout'] = 'NHWC'
        convert_scale_shift_to_mul_add(graph)
        graph.clean_up()
        (flag, resp) = compare_graphs(graph, graph_ref, 'placeholder_1')
        self.assertTrue(flag, resp)

    # ScaleShift  2 inputs-> Mul
    def test_scaleshift2_to_mul(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_2', 'placeholder_2_data'),
                             ('placeholder_1_data', 'scaleshift_1'),
                             ('placeholder_2_data', 'scaleshift_1'),
                             ('scaleshift_1', 'scaleshift_1_data'),
                             ('scaleshift_1_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'placeholder_2_data': {'shape': np.array([1, 227])},
                             'scaleshift_1_data': {}
                             })

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_2', 'placeholder_2_data'),
                                 ('placeholder_2_data', 'placeholder_2/Reshape_'),
                                 ('placeholder_2/Reshape_const', 'placeholder_2/Reshape_const_data'),
                                 ('placeholder_2/Reshape_const_data', 'placeholder_2/Reshape_'),
                                 ('placeholder_2/Reshape_', 'placeholder_2/Reshape_data'),
                                 ('placeholder_1_data', 'mul_1'),
                                 ('placeholder_2/Reshape_data', 'mul_1'),
                                 ('mul_1', 'scaleshift_1_data'),
                                 ('scaleshift_1_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'placeholder_2_data': {'shape': np.array([1, 227])},
                                 'placeholder_2/Reshape_const': {'value': np.array([1, 227, 1, 1]), 'shape': [4]},
                                 'placeholder_2/Reshape_const_data': {'value': np.array([1, 227, 1, 1]), 'shape': [4]},
                                 'placeholder_2/Reshape_data': {'shape': np.array([1, 227, 1, 1])},
                                 'mul_1': {'can_be_fused': True},
                                 'scaleshift_1_data': {}
                                 })

        graph.graph['layout'] = 'NHWC'
        convert_scale_shift_to_mul_add(graph)
        graph.clean_up()
        (flag, resp) = compare_graphs(graph, graph_ref, 'placeholder_1')
        self.assertTrue(flag, resp)

    # ScaleShift  2 inputs-> Mul (axis = 1)
    def test_scaleshift2_axis1_to_mul(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_2', 'placeholder_2_data'),
                             ('placeholder_1_data', 'scaleshift_1'),
                             ('placeholder_2_data', 'scaleshift_1'),
                             ('scaleshift_1', 'scaleshift_1_data'),
                             ('scaleshift_1_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'placeholder_2_data': {'shape': np.array([227])},
                             'scaleshift_1': {'axis': 1},
                             'scaleshift_1_data': {}
                             })

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_2', 'placeholder_2_data'),
                                 ('placeholder_2_data', 'placeholder_2/Reshape_'),
                                 ('placeholder_2/Reshape_const', 'placeholder_2/Reshape_const_data'),
                                 ('placeholder_2/Reshape_const_data', 'placeholder_2/Reshape_'),
                                 ('placeholder_2/Reshape_', 'placeholder_2/Reshape_data'),
                                 ('placeholder_1_data', 'mul_1'),
                                 ('placeholder_2/Reshape_data', 'mul_1'),
                                 ('mul_1', 'scaleshift_1_data'),
                                 ('scaleshift_1_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'placeholder_2_data': {'shape': np.array([227])},
                                 'placeholder_2/Reshape_const': {'value': np.array([1, 227, 1, 1]), 'shape': [4]},
                                 'placeholder_2/Reshape_const_data': {'value': np.array([1, 227, 1, 1]), 'shape': [4]},
                                 'placeholder_2/Reshape_data': {'shape': np.array([1, 227, 1, 1])},
                                 'mul_1': {'can_be_fused': True},
                                 'scaleshift_1_data': {}
                                 })

        graph.graph['layout'] = 'NHWC'
        convert_scale_shift_to_mul_add(graph)
        graph.clean_up()
        (flag, resp) = compare_graphs(graph, graph_ref, 'placeholder_1')
        self.assertTrue(flag, resp)

    # ScaleShift -> Mul (Zero biases)
    def test_scaleshift_to_mul_2(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_1_data', 'scaleshift_1'),
                             ('const_scaleshift_1_w', 'scaleshift_1_w'),
                             ('const_scaleshift_1_b', 'scaleshift_1_b'),
                             ('scaleshift_1_w', 'scaleshift_1'),
                             ('scaleshift_1_b', 'scaleshift_1'),
                             ('scaleshift_1', 'scaleshift_1_data'),
                             ('scaleshift_1_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'scaleshift_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'scaleshift_1_b': {'shape': np.array([3]), 'value': np.array([0, 0, 0])},
                             'scaleshift_1_data': {}
                             })

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_1_data', 'mul_1'),
                                 ('const_mul_1_w', 'mul_1_w'),
                                 ('mul_1_w', 'mul_1'),
                                 ('mul_1', 'scaleshift_1_data'),
                                 ('scaleshift_1_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'const_mul_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'mul_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'mul_1': {'can_be_fused': True},
                                 'scaleshift_1_data': {}
                                 })

        graph.graph['layout'] = 'NHWC'
        convert_scale_shift_to_mul_add(graph)
        graph.clean_up()
        (flag, resp) = compare_graphs(graph, graph_ref, 'placeholder_1')
        self.assertTrue(flag, resp)

    # ScaleShift -> Mul->Add
    def test_scaleshift_to_mul_add(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_1_data', 'scaleshift_1'),
                             ('const_scaleshift_1_w', 'scaleshift_1_w'),
                             ('const_scaleshift_1_b', 'scaleshift_1_b'),
                             ('scaleshift_1_w', 'scaleshift_1'),
                             ('scaleshift_1_b', 'scaleshift_1'),
                             ('scaleshift_1', 'scaleshift_1_data'),
                             ('scaleshift_1_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'scaleshift_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'scaleshift_1_b': {'shape': np.array([3]), 'value': np.array([3, 2, 1])},
                             'scaleshift_1_data': {}
                             })

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_1_data', 'mul_1'),
                                 ('const_mul_1_w', 'mul_1_w'),
                                 ('mul_1_w', 'mul_1'),
                                 ('mul_1', 'mul_1_data'),
                                 ('mul_1_data', 'add_1'),
                                 ('const_add_1_w', 'add_1_w'),
                                 ('add_1_w', 'add_1'),
                                 ('add_1', 'scaleshift_1_data'),
                                 ('scaleshift_1_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'const_mul_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'mul_1_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'const_add_1_w': {'shape': np.array([3]), 'value': np.array([3, 2, 1])},
                                 'add_1_w': {'shape': np.array([3]), 'value': np.array([3, 2, 1])},
                                 'mul_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'add_1': {'can_be_fused': True},
                                 'mul_1': {'can_be_fused': True},
                                 'scaleshift_1_data': {}
                                 })

        graph.graph['layout'] = 'NHWC'
        convert_scale_shift_to_mul_add(graph)
        graph.clean_up()
        (flag, resp) = compare_graphs(graph, graph_ref, 'placeholder_1')
        self.assertTrue(flag, resp)

    # ScaleShift -> None (Zero weights and biases)
    def test_scaleshift_to_nothing(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_1_data', 'scaleshift_1'),
                             ('const_scaleshift_1_w', 'scaleshift_1_w'),
                             ('const_scaleshift_1_b', 'scaleshift_1_b'),
                             ('scaleshift_1_w', 'scaleshift_1'),
                             ('scaleshift_1_b', 'scaleshift_1'),
                             ('scaleshift_1', 'scaleshift_1_data'),
                             ('scaleshift_1_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'scaleshift_1_w': {'shape': np.array([3]), 'value': np.array([1, 1, 1])},
                             'scaleshift_1_b': {'shape': np.array([3]), 'value': np.array([0, 0, 0])},
                             'scaleshift_1_data': {'shape': np.array([1, 227, 227, 3])}
                             }, nodes_with_edges_only=True)

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_1_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])}}
                                ,nodes_with_edges_only=True)

        graph.graph['layout'] = 'NHWC'
        convert_scale_shift_to_mul_add(graph)
        graph.clean_up()
        (flag, resp) = compare_graphs(graph, graph_ref, 'placeholder_1')
        self.assertTrue(flag, resp)

    # ScaleShift -> ScaleShift (can_be_fused=False)
    def test_scaleshift_can_be_fused(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_1_data', 'scaleshift_1'),
                             ('const_scaleshift_1_w', 'scaleshift_1_w'),
                             ('const_scaleshift_1_b', 'scaleshift_1_b'),
                             ('scaleshift_1_w', 'scaleshift_1'),
                             ('scaleshift_1_b', 'scaleshift_1'),
                             ('scaleshift_1', 'scaleshift_1_data'),
                             ('scaleshift_1_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'scaleshift_1_w': {'shape': np.array([3]), 'value': np.array([1, 1, 1])},
                             'scaleshift_1_b': {'shape': np.array([3]), 'value': np.array([0, 0, 0])},
                             'scaleshift_1': {'can_be_fused': False},
                             'scaleshift_1_data': {'shape': np.array([1, 227, 227, 3])}
                             })

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_1_data', 'scaleshift_1'),
                                 ('const_scaleshift_1_w', 'scaleshift_1_w'),
                                 ('const_scaleshift_1_b', 'scaleshift_1_b'),
                                 ('scaleshift_1_w', 'scaleshift_1'),
                                 ('scaleshift_1_b', 'scaleshift_1'),
                                 ('scaleshift_1', 'scaleshift_1_data'),
                                 ('scaleshift_1_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'const_scaleshift_1_w': {'shape': np.array([3]), 'value': np.array([1, 1, 1])},
                                 'scaleshift_1_w': {'shape': np.array([3]), 'value': np.array([1, 1, 1])},
                                 'const_scaleshift_1_b': {'shape': np.array([3]), 'value': np.array([0, 0, 0])},
                                 'scaleshift_1_b': {'shape': np.array([3]), 'value': np.array([0, 0, 0])},
                                 'scaleshift_1': {'can_be_fused': False},
                                 'scaleshift_1_data': {'shape': np.array([1, 227, 227, 3])}
                                 })

        convert_scale_shift_to_mul_add(graph)
        graph.clean_up()

        (flag, resp) = compare_graphs(graph, graph_ref, 'scaleshift_1_data')
        self.assertTrue(flag, resp)

class BatchNormDecomposition(unittest.TestCase):
    def test_bn_decomposition_1(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_1_data', 'bn_op'),
                             ('const_bn_const', 'bn_const'),
                             ('const_bn_beta', 'bn_beta'),
                             ('const_bn_mean', 'bn_mean'),
                             ('const_bn_var', 'bn_var'),
                             ('bn_const', 'bn_op'),
                             ('bn_beta', 'bn_op'),
                             ('bn_mean', 'bn_op'),
                             ('bn_var', 'bn_op'),
                             ('bn_op', 'bn_data'),
                             ('concat', 'concat_data'),
                             ('bn_data', 'concat'),
                             ('concat_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'bn_op': {'eps': 1.2},
                             'bn_const': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_beta': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_mean': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_var': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_data': {'shape': np.array([1, 227, 227, 3])},
                             'concat_data': {}
                             }, nodes_with_edges_only=True)

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_1_data', 'mul_1'),
                                 ('const_mul_1_w', 'mul_1_w'),
                                 ('mul_1_w', 'mul_1'),
                                 ('mul_1', 'mul_1_data'),
                                 ('mul_1_data', 'add_1'),
                                 ('const_add_1_w', 'add_1_w'),
                                 ('add_1_w', 'add_1'),
                                 ('add_1', 'add_1_data'),
                                 ('add_1_data', 'mul_2'),
                                 ('const_mul_2_w', 'mul_2_w'),
                                 ('mul_2_w', 'mul_2'),
                                 ('mul_2', 'mul_2_data'),
                                 ('mul_2_data', 'add_2'),
                                 ('const_add_2_w', 'add_2_w'),
                                 ('add_2_w', 'add_2'),
                                 ('add_2', 'add_2_data'),
                                 ('concat', 'concat_data'),
                                 ('add_2_data', 'concat'),
                                 ('concat_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'const_mul_1_w': {'shape': np.array([3]),
                                             'value': np.array([0.67419986, 0.55901699, 0.48795004])},
                                 'mul_1_w': {'shape': np.array([3]),
                                             'value': np.array([0.67419986, 0.55901699, 0.48795004])},
                                 'const_mul_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'mul_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'const_add_1_w': {'shape': np.array([3]),
                                             'value': np.array([-0.67419986, -1.11803399, -1.46385011])},
                                 'add_1_w': {'shape': np.array([3]),
                                             'value': np.array([-0.67419986, -1.11803399, -1.46385011])},
                                 'const_add_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'add_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'add_2_data': {'shape': np.array([1, 227, 227, 3])},
                                 'mul_1': {'can_be_fused': True},
                                 'mul_2': {'can_be_fused': True},
                                 'add_1': {'can_be_fused': True},
                                 'add_2': {'can_be_fused': True},
                                 'concat_data': {}
                                 }, nodes_with_edges_only=True)

        graph.graph['layout'] = 'NHWC'
        convert_batch_norm(graph)
        graph.clean_up()

        (flag, resp) = compare_graphs(graph, graph_ref, 'concat_data')
        self.assertTrue(flag, resp)

    # 'can_be_fused': False for BatchNorm
    def test_bn_decomposition_2(self):
        graph = build_graph(nodes_attributes,
                            [('placeholder_1', 'placeholder_1_data'),
                             ('placeholder_1_data', 'bn_op'),
                             ('const_bn_const', 'bn_const'),
                             ('const_bn_beta', 'bn_beta'),
                             ('const_bn_mean', 'bn_mean'),
                             ('const_bn_var', 'bn_var'),
                             ('bn_const', 'bn_op'),
                             ('bn_beta', 'bn_op'),
                             ('bn_mean', 'bn_op'),
                             ('bn_var', 'bn_op'),
                             ('bn_op', 'bn_data'),
                             ('concat', 'concat_data'),
                             ('bn_data', 'concat'),
                             ('concat_data', 'op_output')
                             ],
                            {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                             'bn_op': {'eps': 1.2, 'can_be_fused': False},
                             'bn_const': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_beta': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_mean': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_var': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                             'bn_data': {'shape': np.array([1, 227, 227, 3])},
                             'concat_data': {}
                             })

        graph_ref = build_graph(nodes_attributes,
                                [('placeholder_1', 'placeholder_1_data'),
                                 ('placeholder_1_data', 'mul_1'),
                                 ('const_mul_1_w', 'mul_1_w'),
                                 ('mul_1_w', 'mul_1'),
                                 ('mul_1', 'mul_1_data'),
                                 ('mul_1_data', 'add_1'),
                                 ('const_add_1_w', 'add_1_w'),
                                 ('add_1_w', 'add_1'),
                                 ('add_1', 'add_1_data'),
                                 ('add_1_data', 'mul_2'),
                                 ('const_mul_2_w', 'mul_2_w'),
                                 ('mul_2_w', 'mul_2'),
                                 ('mul_2', 'mul_2_data'),
                                 ('mul_2_data', 'add_2'),
                                 ('const_add_2_w', 'add_2_w'),
                                 ('add_2_w', 'add_2'),
                                 ('add_2', 'add_2_data'),
                                 ('concat', 'concat_data'),
                                 ('add_2_data', 'concat'),
                                 ('concat_data', 'op_output')
                                 ],
                                {'placeholder_1_data': {'shape': np.array([1, 227, 227, 3])},
                                 'const_mul_1_w': {'shape': np.array([3]),
                                             'value': np.array([0.67419986, 0.55901699, 0.48795004])},
                                 'mul_1_w': {'shape': np.array([3]),
                                             'value': np.array([0.67419986, 0.55901699, 0.48795004])},
                                 'const_mul_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'mul_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'const_add_1_w': {'shape': np.array([3]),
                                             'value': np.array([-0.67419986, -1.11803399, -1.46385011])},
                                 'add_1_w': {'shape': np.array([3]),
                                             'value': np.array([-0.67419986, -1.11803399, -1.46385011])},
                                 'const_add_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'add_2_w': {'shape': np.array([3]), 'value': np.array([1, 2, 3])},
                                 'add_2_data': {'shape': np.array([1, 227, 227, 3])},
                                 'mul_1': {'can_be_fused': False},
                                 'mul_2': {'can_be_fused': False},
                                 'add_1': {'can_be_fused': False},
                                 'add_2': {'can_be_fused': False},
                                 'concat_data': {}
                                 })

        graph.graph['layout'] = 'NHWC'
        convert_batch_norm(graph)
        graph.clean_up()

        (flag, resp) = compare_graphs(graph, graph_ref, 'concat_data')
        self.assertTrue(flag, resp)

max_stars_repo_name: fujunwei/dldt

max_issues_repo_name: fujunwei/dldt

max_forks_repo_name: fujunwei/dldt


text:
import time
import pandas as pd
import numpy as np

from sklearn.pipeline import Pipeline
from sklearn.metrics import pairwise
from sklearn.decomposition import TruncatedSVD
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.preprocessing import Normalizer, normalize

from modules.utils.similarity_measures import SimilarityMeasure
from modules.utils.tokenizers import WordNetBased_LemmaTokenizer

from modules.models.generic_model import GenericModel
from modules.models.model_hyperps import LSI_Model_Hyperp

"""
params_dict = {
    'lsi__similarity_measure' : SimilarityMeasure.COSINE,
    'lsi__name' : 'LSI',
    'lsi__vectorizer' : TfidfVectorizer(),
    'lsi__vectorizer__stop_words' : 'english',
    'lsi__vectorizer__tokenizer' : Tokenizer(),
    'lsi__vectorizer__use_idf' : True,          # optional if type(Vectorizer) == TfidfVectorizer
    'lsi__vectorizer__smooth_idf' : True,       # optional if type(Vectorizer) == TfidfVectorizer
    'lsi__vectorizer__ngram_range' : (1,2),
    'lsi__svd_model' : TruncatedSVD(),
    'lsi__svd_model__n_components' : 5
}
"""
class LSI(GenericModel):
    def __init__(self, **kwargs):
        self._svd_matrix = None
        self._query_vector = None

        self.vectorizer = None
        self.svd_model = None

        super().__init__()

        self.similarity_measure = None

        self.set_basic_params(**kwargs)
        self.set_vectorizer(**kwargs)
        self.set_svd_model(**kwargs)

    def set_name(self, name):
        super().set_name(name)

    def set_model_gen_name(self, gen_name):
        super().set_model_gen_name(gen_name)

    def set_basic_params(self, **kwargs):
        self.set_name('LSI' if LSI_Model_Hyperp.NAME.value not in kwargs.keys() else kwargs[LSI_Model_Hyperp.NAME.value])
        self.set_similarity_measure(SimilarityMeasure.COSINE)
        self.set_model_gen_name('lsi')

    def set_similarity_measure(self, sim_measure):
        self.similarity_measure = sim_measure

    def set_vectorizer(self, **kwargs):
        self.vectorizer = TfidfVectorizer(stop_words='english',
                                             use_idf=True,
                                             smooth_idf=True) if LSI_Model_Hyperp.VECTORIZER.value not in kwargs.keys() else kwargs[LSI_Model_Hyperp.VECTORIZER.value]

        vec_params = {key.split('__')[2]:kwargs[key] for key,val in kwargs.items() if '__vectorizer__' in key}
        self.vectorizer.set_params(**vec_params)

    def set_svd_model(self, **kwargs):
        self.svd_model = TruncatedSVD(n_components = 100,
                                         algorithm = 'randomized',
                                         n_iter = 10,
                                         random_state = 42) if LSI_Model_Hyperp.SVD_MODEL.value not in kwargs.keys() else kwargs[LSI_Model_Hyperp.SVD_MODEL.value]

        svd_model_params = {key.split('__')[2]:kwargs[key] for key,val in kwargs.items() if '__svd_model__' in key}
        self.svd_model.set_params(**svd_model_params)

    def recover_links(self, corpus, query, test_cases_names, bug_reports_names):
        starttime = time.time()

        if self.similarity_measure == SimilarityMeasure.COSINE:
            self._recover_links_cosine(corpus, query, test_cases_names, bug_reports_names)

        elif self.similarity_measure == SimilarityMeasure.JACCARD_INDEX:
            self._recover_links_jaccard(corpus, query, test_cases_names, bug_reports_names)

        elif self.similarity_measure == SimilarityMeasure.EDIT_DISTANCE:
            self._recover_links_edit(corpus, query, test_cases_names, bug_reports_names)

        self._record_docs_feats(corpus, query, test_cases_names, bug_reports_names)

        endtime = time.time()

        print(f' ..Total processing time: {round(endtime-starttime, 2)} seconds', )

    def _record_docs_feats(self, corpus, query, test_cases_names, bug_reports_names):
        self.mrw_tcs = self._recover_mrw_list(test_cases_names, corpus)
        self.mrw_brs = self._recover_mrw_list(bug_reports_names, query)

        self.dl_tcs = self._recover_dl_list(test_cases_names, corpus)
        self.dl_brs = self._recover_dl_list(bug_reports_names, query)

        index = list(test_cases_names) + list(bug_reports_names)
        self.docs_feats_df = pd.DataFrame(index=index,
                                         columns=['mrw','dl'])

        for tc_name, mrw in self.mrw_tcs:
            self.docs_feats_df.at[tc_name, 'mrw'] = mrw

        for tc_name, dl in self.dl_tcs:
            self.docs_feats_df.at[tc_name, 'dl'] = dl

        for br_name, mrw in  self.mrw_brs:
            self.docs_feats_df.at[br_name, 'mrw'] = mrw

        for br_name, dl in self.dl_brs:
            self.docs_feats_df.at[br_name, 'dl'] = dl

    def _recover_dl_list(self, artf_names, artf_descs):
        tokenizer = WordNetBased_LemmaTokenizer()
        dl_list = []
        for artf_name, artf_desc in zip(artf_names, artf_descs):
            dl_list.append((artf_name, len(tokenizer.__call__(artf_desc))))
        return dl_list

    def _recover_mrw_list(self, artf_names, artf_descs):
        N_REL_WORDS = 6
        mrw_list = [] # list of tuples (artf_name, mrw_list={})

        for artf_name, artf_desc in zip(artf_names, artf_descs):
            X = self.vectorizer.transform([artf_desc])
            df1 = pd.DataFrame(X.T.toarray())
            df1['token'] = self.vectorizer.get_feature_names()
            df1.sort_values(by=0, ascending=False, inplace=True)
            mrw = list(df1.iloc[0:N_REL_WORDS,1].values)
            mrw_list.append((artf_name, mrw))

        return mrw_list

    def _recover_links_cosine(self, corpus, query, test_cases_names, bug_reports_names):
        svd_transformer = Pipeline([('vec', self.vectorizer),
                            ('svd', self.svd_model)])

        self._svd_matrix = svd_transformer.fit_transform(corpus)
        self._query_vector = svd_transformer.transform(query)
        self._sim_matrix = pairwise.cosine_similarity(X=self._svd_matrix, Y=self._query_vector)

        #self._sim_matrix =  super().normalize_sim_matrix(self._sim_matrix)
        self._sim_matrix = pd.DataFrame(data=self._sim_matrix, index=test_cases_names, columns=bug_reports_names)

    def _recover_links_jaccard(self, corpus, query, test_cases_names, bug_reports_names):
        tokenizer = self.vectorizer.tokenizer

        corpus_tokens = [tokenizer.__call__(doc) for doc in corpus]
        query_tokens = [tokenizer.__call__(doc) for doc in query]

        self._sim_matrix = pd.DataFrame(index = test_cases_names,
                                       columns = bug_reports_names,
                                       data = np.zeros(shape=(len(test_cases_names), len(bug_reports_names)), dtype='int8'))

        for br_id, doc_query_tset in zip(bug_reports_names, query_tokens):
            for tc_id, doc_corpus_tset in zip(test_cases_names, corpus_tokens):
                self._sim_matrix.at[tc_id, br_id] = nltk.jaccard_distance(set(doc_corpus_tset), set(doc_query_tset))

    def _recover_links_edit(self, corpus, query, test_cases_names, bug_reports_names):
        self._sim_matrix = pd.DataFrame(index = test_cases_names,
                                       columns = bug_reports_names,
                                       data = np.zeros(shape=(len(test_cases_names), len(bug_reports_names)), dtype='int8'))

        for br_id, doc_query in zip(bug_reports_names, query):
            for tc_id, doc_corpus in zip(test_cases_names, corpus):
                self._sim_matrix.at[tc_id, br_id] = nltk.edit_distance(doc_corpus, doc_query)

        normalizer = Normalizer(copy=False).fit(self._sim_matrix.values)
        self._sim_matrix = pd.DataFrame(data=normalizer.transform(self._sim_matrix.values), index=test_cases_names, columns=bug_reports_names)

    def model_setup(self):
        return {"Setup" :
                  [
                      {"Name" : self.get_name()},
                      {"Similarity Measure" : self.get_similarity_measure()},
                      {"SVD Model" : self.svd_model.get_params()},
                      {"Vectorizer" : self.vectorizer.get_params()},
                      {"Vectorizer Type" : type(self.vectorizer)}
                  ]
               }

    def get_query_vector(self):
        return self._query_vector

    def get_svd_matrix(self):
        return self._svd_matrix

    def get_vectorizer_type(self):
        return type(self.vectorizer)

    def get_tokenizer_type(self):
        return type(self.vectorizer.tokenizer)

    def get_name(self):
        return super().get_name()

    def get_model_gen_name(self):
        return super().get_model_gen_name()

    def get_similarity_measure(self):
        return self.similarity_measure

    def get_sim_matrix(self):
        return super().get_sim_matrix()

    def save_sim_matrix(self):
        super().save_sim_matrix()

ext: py

max_stars_repo_name: guilhermemg/trace-links-tc-br

max_issues_repo_name: guilhermemg/trace-links-tc-br

max_forks_repo_name: guilhermemg/trace-links-tc-br

reason: import numpy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Last update, 16/01/2018
Continuous Replicator Dynamic in Asexual model
check the stability
@author: Shota Shibasaki
"""

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

a1=1
e1=0.0001
a2=0.5
e2=0.0001
h=0.5
e0=0.000001

def Matrix(a1,b1,a2,b2,theta):
    """
   Defining patoff matrixes
   A1: payoff matrix for fruiting body formation
   A2: payoff matrix for macrocyst formation
   A: patoff matrix generated by the mixture of A1 and A2.
      see Eq [2] in the main text for more detail.
   Parameters
   a2: benefit of cmutual ooperation in macrocyst formation
   b2: benefit of exploitation for defectors in macrocyst formation
   theta: prob. of macrocyst fromation

   """
    A1=np.array([[a1,h*a1,0],\
                 [e0,e0,e0],\
                 [b1,h*e1,e1]])

    A2=np.array([[a2,0,a2],\
                 [b2,e2,b2],\
                 [a2,0,a2]])
    A=(1-theta)*A1+(theta)*A2
    return A

def func(x, t, a1,b1,a2,b2,theta):
    """
    Continuous replicator dynamics
    x: array of the frequency of the three strategy
    t: time
    a2, b2, theta: defined in function Matrix
    """
    #x is array and A is matrix
    A=Matrix(a1,b1,a2,b2,theta)
    phi=np.dot(x,A)
    phi=np.dot(phi,x.T)#mean fitness in the population
    #calculate replicator dynamics
    f0=np.dot(A[0,:],x.T)-phi
    f1=np.dot(A[1,:],x.T)-phi
    f2=np.dot(A[2,:],x.T)-phi
    f=np.array([x[0]*f0,x[1]*f1,x[2]*f2])
    return(f)

def main():
    # condition of the linear stability in the discrete replicator
    THETA=np.array([])#theta
    DEL1=np.array([])#b1-a1
    DEL2=np.array([])#b2-a2
    step=26#number of points in each coordinate axe
    the=np.linspace(0,1,step)
    del1=np.linspace(0,1,step)
    del2=np.linspace(0,1,step)
    for t in range(np.size(the)):
        theta=the[t]
        #print(theta)
        for d1 in range(np.size(del1)):
            b1=a1+del1[d1]
            #print(b1)
            for d2 in range(np.size(del2)):
                b2=a2+del2[d2]
                #check there is a global attractor in replicator dynamics
                #print(b2)
                #original payoff matrix CC,DC,CD
                A=Matrix(a1,b1,a2,b2,theta)
                #arranged payoff matrix for cntinuous replicator equation
                Ac=np.array([[A[0,0]-A[0,0],A[0,1]-A[1,1],A[0,2]-A[2,2]],\
                           [A[1,0]-A[0,0],A[1,1]-A[1,1],A[1,2]-A[2,2]],\
                           [A[2,0]-A[0,0],A[2,1]-A[1,1],A[2,2]-A[2,2]]])
                odds=(1-theta)/theta
                if (b2-a2)/(a1-e1)<odds and odds<e2/(h*e1-e0):
                 #condition for RSP game
                    if np.linalg.det(Ac)>0:
                    #unique global attractor in continuous replicator dynamics
                        #analytical solution
                        p1 = A[2,1]*A[0,2]+A[1,2]*A[0,1]+A[2,1]*A[1,2]
                        p2 = A[1,0]*A[0,2]+A[1,2]*A[2,0]+A[0,2]*A[2,0]
                        p3 = A[2,1]*A[1,0]+A[0,1]*A[2,0]+A[0,1]*A[1,0]
                        pa=np.array([p1,p2,p3])
                        p=pa/sum(pa)
                        #calculate the linearized discrete replicator
                        Ave=(p.dot(A)).dot(p.T)
                        ev1=np.array([1,0,0]).T
                        ev2=np.array([0,1,0]).T
                        ev3=np.array([0,0,1]).T
                        L=np.zeros([3,3])
                        L[0,0]=1+p[0]*(A[0,0]-(p.dot(A)).dot(ev1))/Ave
                        L[0,1]=0+p[0]*(A[0,1]-(p.dot(A)).dot(ev2))/Ave
                        L[0,2]=0+p[0]*(A[0,2]-(p.dot(A)).dot(ev3))/Ave
                        L[1,0]=0+p[1]*(A[1,0]-(p.dot(A)).dot(ev1))/Ave
                        L[1,1]=1+p[1]*(A[1,1]-(p.dot(A)).dot(ev2))/Ave
                        L[1,2]=0+p[1]*(A[1,2]-(p.dot(A)).dot(ev3))/Ave
                        L[2,0]=0+p[2]*(A[2,0]-(p.dot(A)).dot(ev1))/Ave
                        L[2,1]=0+p[2]*(A[2,1]-(p.dot(A)).dot(ev2))/Ave
                        L[2,2]=1+p[2]*(A[2,2]-(p.dot(A)).dot(ev3))/Ave
                        #conditions where |eigenvalues of L| < 1
                        (eig,v)=np.linalg.eig(L)
                        meig=abs(eig)
                        if max(meig)<1:
                            THETA=np.append(THETA,theta)
                            DEL1=np.append(DEL1,del1[d1])
                            DEL2=np.append(DEL2,del2[d2])
    #3D scatter plot
    fig=plt.figure(figsize=(6,6))
    ax  = Axes3D(fig)
    ax.plot(DEL1, DEL2, THETA, "o", color="b", ms= 8, mew=0.5)
    ax.set_xlim(0,1.0)
    ax.set_ylim(0,1.0)
    ax.set_zlim(0,1.0)
    ax.invert_xaxis()
    ax.invert_yaxis()
    #ax.grid(False)
    ax.patch.set_facecolor('white')
    ax.set_xlabel('$b_1-a_1$',fontsize=12)
    ax.set_ylabel('$b_2-a_2$',fontsize=12)
    ax.set_zlabel(r'$\theta$',fontsize=12)
    #plt.savefig("CondDisRep_asexual.pdf")
    #add the boundary of conditions for continuous replicator dyanmics
    z1=a1/(a1+del2)
    z2=(h*e1-e0)/(h*e1-e0+e2+del2*0)
    ax.plot(del2,z1,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z1,zs=1,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z2,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z2,zs=1,zdir='x',color="k",linestyle="--")
    ax.plot(del2,z1,zs=0,zdir='x',color="k",linestyle="--")
    z3=a1/(a1+1+0*del1)
    z4=a1/(a1+0+0*del1)
    z5=(h*e1-e0)/(h*e1-e0+e2+1*+del1*0)
    z6=(h*e1-e0)/(h*e1-e0+e2+0+del1*0)
    ax.plot(del1,z3,zs=1,zdir='y',color="k",linestyle="--")
    ax.plot(del1,z4,zs=0,zdir='y',color="k",linestyle="--")
    ax.plot(del1,z5,zs=0,zdir='y',color="k",linestyle="--")
    ax.plot(del1,z6,zs=1,zdir='y',color="k",linestyle="--")
    zz=np.linspace((h*e1-e0)/(h*e1-e0+e2),1,step)
    y1=zz*0+0
    ax.plot(y1,zz,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(y1,zz,zs=1,zdir='x',color="k",linestyle="--")
    zz=np.linspace((h*e1-e0)/(h*e1-e0+e2),a1/(a1+1),step)
    y1=zz*0+1
    ax.plot(y1,zz,zs=0,zdir='x',color="k",linestyle="--")
    ax.plot(y1,zz,zs=1,zdir='x',color="k",linestyle="--")
    plt.savefig("CondRSPandDisc_Merge_asexual.pdf")
    #plt.show()

if __name__ == '__main__':
    main()

max_stars_repo_name: ShotaSHIBASAKI/Cyclic-dominance-emerges

max_issues_repo_name: ShotaSHIBASAKI/Cyclic-dominance-emerges

max_forks_repo_name: ShotaSHIBASAKI/Cyclic-dominance-emerges

reason: import numpy,from scipy


text:
from __future__ import absolute_import, division, print_function

import os.path
import torch
import numpy as np
import skimage.io as io
import png
import cv2
import random
import math
from . import exp_aug as flow_transforms
from PIL import Image
import torch.nn.functional as tf

width_to_date = dict()
width_to_date[1242] = '2011_09_26'
width_to_date[1224] = '2011_09_28'
width_to_date[1238] = '2011_09_29'
width_to_date[1226] = '2011_09_30'
width_to_date[1241] = '2011_10_03'

def get_date_from_width(width):
	return width_to_date[width]

def list_flatten(input_list):
	return [img for sub_list in input_list for img in sub_list]

def intrinsic_scale(mat, sy, sx):
	out = mat.clone()
	out[0, 0] *= sx
	out[0, 2] *= sx
	out[1, 1] *= sy
	out[1, 2] *= sy
	return out

def kitti_adjust_intrinsic(k_l1, k_r1, crop_info):
	str_x = crop_info[0]
	str_y = crop_info[1]
	k_l1[0, 2] -= str_x
	k_l1[1, 2] -= str_y
	k_r1[0, 2] -= str_x
	k_r1[1, 2] -= str_y
	return k_l1, k_r1

def kitti_crop_image_list(img_list, crop_info):
	str_x = crop_info[0]
	str_y = crop_info[1]
	end_x = crop_info[2]
	end_y = crop_info[3]

	transformed = [img[str_y:end_y, str_x:end_x, :] for img in img_list]

	return transformed

def numpy2torch(array):
	assert(isinstance(array, np.ndarray))
	if array.ndim == 3:
		array = np.transpose(array, (2, 0, 1))
	else:
		array = np.expand_dims(array, axis=0)
	return torch.from_numpy(array.copy()).float()

def read_image_as_byte(filename):
	return io.imread(filename)

def read_png_flow(flow_file):
	flow_object = png.Reader(filename=flow_file)
	flow_direct = flow_object.asDirect()
	flow_data = list(flow_direct[2])
	(w, h) = flow_direct[3]['size']
	flow = np.zeros((h, w, 3), dtype=np.float64)
	for i in range(len(flow_data)):
		flow[i, :, 0] = flow_data[i][0::3]
		flow[i, :, 1] = flow_data[i][1::3]
		flow[i, :, 2] = flow_data[i][2::3]

	invalid_idx = (flow[:, :, 2] == 0)
	flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0
	flow[invalid_idx, 0] = 0
	flow[invalid_idx, 1] = 0
	return flow[:, :, 0:2], (1 - invalid_idx * 1)[:, :, None]

def read_png_file(flow_file):
	"""
	Read from KITTI .png file
	:param flow_file: name of the flow file
	:return: optical flow data in matrix
	"""
	flow = cv2.imread(flow_file,-1)[:,:,::-1].astype(np.float64)
 #   flow_object = png.Reader(filename=flow_file)
 #   flow_direct = flow_object.asDirect()
 #   flow_data = list(flow_direct[2])
 #   (w, h) = flow_direct[3]['size']
 #   #print("Reading %d x %d flow file in .png format" % (h, w))
 #   flow = np.zeros((h, w, 3), dtype=np.float64)
 #   for i in range(len(flow_data)):
 #       flow[i, :, 0] = flow_data[i][0::3]
 #       flow[i, :, 1] = flow_data[i][1::3]
 #       flow[i, :, 2] = flow_data[i][2::3]

	invalid_idx = (flow[:, :, 2] == 0)
	flow[:, :, 0:2] = (flow[:, :, 0:2] - 2 ** 15) / 64.0
	flow[invalid_idx, 0] = 0
	flow[invalid_idx, 1] = 0
	return flow

def read_png_disp(disp_file):
	disp_np = io.imread(disp_file).astype(np.uint16) / 256.0
	disp_np = np.expand_dims(disp_np, axis=2)
	mask_disp = (disp_np > 0).astype(np.float64)
	return disp_np, mask_disp

def read_full_disp(path):
	data = Image.open(path)
	data = np.ascontiguousarray(data,dtype=np.float32)/256
	return data

def read_raw_calib_file(filepath):
	# From https://github.com/utiasSTARS/pykitti/blob/master/pykitti/utils.py
	"""Read in a calibration file and parse into a dictionary."""
	data = {}
	#print(filepath)
	with open(filepath, 'r') as f:
		for line in f.readlines():
			key, value = line.split(':', 1)
			# The only non-float values in these files are dates, which
			# we don't care about anyway
			try:
				data[key] = np.array([float(x) for x in value.split()])
			except ValueError:
				pass
	return data

def get_r_matrix(rx=0,ry=0,rz=0):
	theta_x = 2*math.pi * rx / 360
	theta_y = 2*math.pi * ry / 360
	theta_z = 2*math.pi * rz / 360

	Ry = np.array([[1,0,0],
				[0,math.cos(theta_x),-math.sin(theta_x)],
				[0,math.sin(theta_x),math.cos(theta_x)]])
	Rx = np.array([[math.cos(theta_y),0,math.sin(theta_y)],
				[0,1,0],
				[-math.sin(theta_y),0,math.cos(theta_y)]]
				)
	Rz = np.array([[math.cos(theta_z),-math.sin(theta_z),0],
				[math.sin(theta_z),math.cos(theta_z),0],
				[0,0,1]])

	total_R = np.matmul(np.matmul(Rz,Ry),Rx)

	return total_R

def read_calib_into_dict(path_dir):

	calibration_file_list = ['2011_09_26', '2011_09_28', '2011_09_29', '2011_09_30', '2011_10_03']
	intrinsic_dict_l = {}
	intrinsic_dict_r = {}

	for ii, date in enumerate(calibration_file_list):
		file_name = "cam_intrinsics/calib_cam_to_cam_" + date + '.txt'
		file_name_full = os.path.join(path_dir, file_name)
		file_data = read_raw_calib_file(file_name_full)
		P_rect_02 = np.reshape(file_data['P_rect_02'], (3, 4))
		P_rect_03 = np.reshape(file_data['P_rect_03'], (3, 4))
		intrinsic_dict_l[date] = P_rect_02[:3, :3]
		intrinsic_dict_r[date] = P_rect_03[:3, :3]

	return intrinsic_dict_l, intrinsic_dict_r

def read_full_calib_into_dict(path_dir):

	calibration_file_list = ['2011_09_26', '2011_09_28', '2011_09_29', '2011_09_30', '2011_10_03']
	intrinsic_dict_l = {}
	intrinsic_dict_r = {}

	for ii, date in enumerate(calibration_file_list):
		file_name = "cam_intrinsics/calib_cam_to_cam_" + date + '.txt'
		file_name_full = os.path.join(path_dir, file_name)
		file_data = read_raw_calib_file(file_name_full)
		P_rect_02 = np.reshape(file_data['P_rect_02'], (3, 4))
		P_rect_03 = np.reshape(file_data['P_rect_03'], (3, 4))
		intrinsic_dict_l[date] = P_rect_02
		intrinsic_dict_r[date] = P_rect_03

	return intrinsic_dict_l, intrinsic_dict_r

def process_calib_into_dict(P_rect_20, P_rect_30):
	data = {}
	data['K_cam2'] = P_rect_20[0:3, 0:3]
	data['K_cam3'] = P_rect_30[0:3, 0:3]
	data['b20'] = P_rect_20[0, 3] / P_rect_20[0, 0]
	data['b30'] = P_rect_30[0, 3] / P_rect_30[0, 0]
	return data

def pixel_coord_np(width, height):
	x = np.linspace(0, width - 1, width).astype(np.int)
	y = np.linspace(0, height - 1, height).astype(np.int)
	[x, y] = np.meshgrid(x, y)
	return np.vstack((x.flatten(), y.flatten(), np.ones_like(x.flatten())))

def get_shear_x(sy=0,sz=0):
	Shx = np.array([[1,0,0,0],
				[sy,1,0,0],
				[sz,0,1,0],
				[0,0,0,1]])

	return Shx

def get_shear_y(sx=0,sz=0):
	Shy = np.array([[1,sx,0,0],
				[0,1,0,0],
				[0,sz,1,0],
				[0,0,0,1]])
	return Shy

def get_shear_z(sx=0,sy=0):
	Shz = np.array([[1,0,sx,0],
				[0,1,sy,0],
				[0,0,1,0],
				[0,0,0,1]])
	return Shz

def get_valid_map(new_coords,h,w):
	return (new_coords[:,0] < w) * (new_coords[:,0] >= 0) * (new_coords[:,1] < h) * (new_coords[:,1] >= 0)

def perform_shearing(cam_coords_l,cam_coords_r):
	sh = get_shear_y(random.uniform(-0.05,0.05),random.uniform(-0.05,0.05))
	#expand input dim
	ones_l = np.ones((cam_coords_l.shape[0],1))
	cam_coords_l = np.concatenate((cam_coords_l,ones_l),axis=1)
	ones_r = np.ones((cam_coords_r.shape[0],1))
	cam_coords_r = np.concatenate((cam_coords_r,ones_r),axis=1)

	shed_cam_coords_l = np.matmul(sh,cam_coords_l.T).T
	shed_cam_coords_r = np.matmul(sh,cam_coords_r.T).T

	return shed_cam_coords_l[:,:3],shed_cam_coords_r[:,:3]

def perform_rotate(cam_coords_l,cam_coords_r):
	R = get_r_matrix(0,random.uniform(-8,8),0)
	return np.matmul(cam_coords_l,R),np.matmul(cam_coords_r,R)

def perform_trans(cam_coords_l,cam_coords_r,depth_min):
	tz = random.uniform(-depth_min,0)
	T = [0,0,tz]
	valid_area_l = cam_coords_l[:,2] >= -T[-1]
	cam_coords_l[valid_area_l,-1] = cam_coords_l[valid_area_l,-1] + T[-1]
	valid_area_r = cam_coords_r[:,2] >= -T[-1]
	cam_coords_r[valid_area_r,-1] = cam_coords_r[valid_area_r,-1] + T[-1]
	return cam_coords_l,cam_coords_r

def perform_scaling(cam_coords_l,cam_coords_r):
	scaling_x = random.uniform(0,0.1)
	scaling_y = random.uniform(0,0.1)
	scaling_z = random.uniform(0,0.1)
	cam_coords_l[:,0] = cam_coords_l[:,0] * scaling_x
	cam_coords_l[:,1] = cam_coords_l[:,1] * scaling_y
	cam_coords_l[:,2] = cam_coords_l[:,2] * scaling_z
	cam_coords_r[:,0] = cam_coords_r[:,0] * scaling_x
	cam_coords_r[:,1] = cam_coords_r[:,1] * scaling_y
	cam_coords_r[:,2] = cam_coords_r[:,2] * scaling_z
	return cam_coords_l,cam_coords_r

def threed_warp(image_l,depth_l,image_r,depth_r,intrinsic_l,intrinsic_r):
	height_l = depth_l.shape[0]
	width_l = depth_l.shape[1]
	height_r = depth_r.shape[0]
	width_r = depth_r.shape[1]
	depth_min = np.min(np.vstack((depth_l,depth_r)))

	inv_intrinsic_l = np.linalg.inv(intrinsic_l)
	inv_intrinsic_r = np.linalg.inv(intrinsic_r)

	image_coord_l = pixel_coord_np(width_l,height_l)
	cam_coords_l = (inv_intrinsic_l[:3, :3] @ image_coord_l * depth_l.flatten()).T
	image_coord_r = pixel_coord_np(width_r,height_r)
	cam_coords_r = (inv_intrinsic_l[:3, :3] @ image_coord_r * depth_r.flatten()).T

	# rotation
	if torch.rand(1) > 0.1:
		aug_cam_coords_l,aug_cam_coords_r = perform_rotate(cam_coords_l,cam_coords_r)
	else:
		aug_cam_coords_l,aug_cam_coords_r = perform_scaling(cam_coords_l,cam_coords_r)
	# shearing
	# aug_cam_coords_l,aug_cam_coords_r = perform_shearing(aug_cam_coords_l,aug_cam_coords_r)
	#translation
	aug_cam_coords_l,aug_cam_coords_r = perform_trans(aug_cam_coords_l,aug_cam_coords_r,depth_min)
	# generate sceneflow
	sf_l = (aug_cam_coords_l - cam_coords_l).reshape(height_l,width_l,3).astype(np.float32)
	sf_r = (aug_cam_coords_r - cam_coords_r).reshape(height_r,width_r,3).astype(np.float32)
	#get warpped image
	new_image_coords_l = (intrinsic_l @ aug_cam_coords_l.T / (aug_cam_coords_l[:,2]+ 1e-6)).T
	new_image_coords_r = (intrinsic_r @ aug_cam_coords_r.T / (aug_cam_coords_r[:,2]+ 1e-6)).T
	new_image_coords_tensor_l = torch.from_numpy(new_image_coords_l[:,:2]).view(height_l,width_l,2)
	new_image_coords_l_x = new_image_coords_tensor_l[:,:,0:1] / (width_l - 1) * 2 - 1
	new_image_coords_l_y = new_image_coords_tensor_l[:,:,1:2] / (height_l - 1) * 2 - 1
	normed_image_coords_l = torch.cat((new_image_coords_l_x, new_image_coords_l_y),dim=2)
	new_image_coords_tensor_r = torch.from_numpy(new_image_coords_r[:,:2]).view(height_r,width_r,2)
	new_image_coords_r_x = new_image_coords_tensor_r[:,:,0:1] / (width_r - 1) * 2 - 1
	new_image_coords_r_y = new_image_coords_tensor_r[:,:,1:2] / (height_r - 1) * 2 - 1
	normed_image_coords_r = torch.cat((new_image_coords_r_x, new_image_coords_r_y),dim=2)
	# transfer to tensor:
	image_l = torch.from_numpy(image_l.astype(float).transpose(2,0,1))
	image_r = torch.from_numpy(image_r.astype(float).transpose(2,0,1))
	image_l_mask = torch.ones_like(image_l)
	image_r_mask = torch.ones_like(image_r)
	# warpping
	warpped_im_l = tf.grid_sample(image_l.unsqueeze(0),normed_image_coords_l.unsqueeze(0))
	warpped_im_r = tf.grid_sample(image_r.unsqueeze(0),normed_image_coords_r.unsqueeze(0))
	warpped_mask_l = (tf.grid_sample(image_l_mask.unsqueeze(0), normed_image_coords_l.unsqueeze(0)) >= 1)
	warpped_mask_r = (tf.grid_sample(image_r_mask.unsqueeze(0), normed_image_coords_r.unsqueeze(0)) >= 1)

	warpped_im_l = (warpped_im_l.float() * warpped_mask_l.float()).squeeze(0).numpy().transpose(1,2,0)
	warpped_im_r = (warpped_im_r.float() * warpped_mask_r.float()).squeeze(0).numpy().transpose(1,2,0)
	warpped_mask_l = warpped_mask_l.squeeze(0).numpy().transpose(1,2,0).all(-1)[:,:,np.newaxis]
	warpped_mask_r = warpped_mask_r.squeeze(0).numpy().transpose(1,2,0).all(-1)[:,:,np.newaxis]

	# get new image coordinate

	#sf_bl = -cv2.remap(sf_l,new_img_xl,new_img_yl,cv2.INTER_LINEAR)
	#sf_br = -cv2.remap(sf_r,new_img_xr,new_img_yr,cv2.INTER_LINEAR)
	#cv2.imwrite("l.png", warpped_im_l)
	#cv2.imwrite("r.png", warpped_im_r)
	#cv2.imwrite("l_og.png", image_l)
	#cv2.imwrite("r_og.png", image_r)

	out_dict = {
		"sf_l":sf_l,
		"sf_r":sf_r,
		"warpped_im_l":warpped_im_l,
		"warpped_im_r":warpped_im_r,
		"valid_l": warpped_mask_l,
		"valid_r": warpped_mask_r
	}
	return out_dict

def readPFM(file):
	import re
	file = open(file, 'rb')

	color = None
	width = None
	height = None
	scale = None
	endian = None

	header = file.readline().rstrip()
	if header == b'PF':
		color = True
	elif header == b'Pf':
		color = False
	else:
		raise Exception('Not a PFM file.')

	dim_match = re.match(b'^(\d+)\s(\d+)\s$', file.readline())
	if dim_match:
		width, height = map(int, dim_match.groups())
	else:
		raise Exception('Malformed PFM header.')

	scale = float(file.readline().rstrip())
	if scale < 0: # little-endian
		endian = '<'
		scale = -scale
	else:
		endian = '>' # big-endian

	data = np.fromfile(file, endian + 'f')
	shape = (height, width, 3) if color else (height, width)

	data = np.reshape(data, shape)
	data = np.flipud(data)
	return data, scale

def generate_gt_expansion(iml0,iml1,shape,flowl0,d1, d2, bl, fl, cx, cy, count_path=None, order=1, prob=1):
	np.ascontiguousarray(flowl0,dtype=np.float32)
	flowl0[np.isnan(flowl0)] = 1e6
	#print(flowl0.shape)
	th,tw = shape

	flowl0[:,:,2] = np.logical_and(np.logical_and(flowl0[:,:,2]==1, d1!=0), d2!=0).astype(float)
	shape = d1.shape
	mesh = np.meshgrid(range(shape[1]),range(shape[0]))
	xcoord = mesh[0].astype(float)
	ycoord = mesh[1].astype(float)

	P0 = triangulation(d1, xcoord, ycoord, bl=1, fl = fl, cx = cx, cy = cy)
	P1 = triangulation(d2, xcoord + flowl0[:,:,0], ycoord + flowl0[:,:,1], bl=1, fl = fl, cx = cx, cy = cy)
	dis0 = P0[2]
	dis1 = P1[2]

	change_size =  dis0.reshape(shape).astype(np.float32)
	flow3d = (P1-P0)[:3].reshape((3,)+shape).transpose((1,2,0))

	gt_normal = np.concatenate((d1[:,:,np.newaxis],d2[:,:,np.newaxis],d2[:,:,np.newaxis]),-1)
	change_size = np.concatenate((change_size[:,:,np.newaxis],gt_normal,flow3d),2)

	iml1 = (iml1)/255.
	iml0 = (iml0)/255.
	iml0 = iml0[:,:,::-1].copy()
	iml1 = iml1[:,:,::-1].copy()

	try:
		with open(count_path, 'r') as f:
			iter_counts = int(f.readline())
	except:
		iter_counts = 0

	schedule = [0.5, 1., 50000.]  # initial coeff, final_coeff, half life
	schedule_coeff = schedule[0] + (schedule[1] - schedule[0]) * (2/(1+np.exp(-1.0986*iter_counts/schedule[2])) - 1)

	if np.random.binomial(1,prob):
		co_transform1 = flow_transforms.Compose([
						flow_transforms.SpatialAug([th,tw],
										scale=[0.2,0.,0.1],
										rot=[0.4,0.],
										trans=[0.4,0.],
										squeeze=[0.3,0.], schedule_coeff=schedule_coeff, order=order),
		])
	else:
		co_transform1 = flow_transforms.Compose([
		flow_transforms.RandomCrop([th,tw]),
		])

	co_transform2 = flow_transforms.Compose([
		flow_transforms.pseudoPCAAug( schedule_coeff=schedule_coeff),
		#flow_transforms.PCAAug(schedule_coeff=schedule_coeff),
		flow_transforms.ChromaticAug( schedule_coeff=schedule_coeff, noise=0.06),
		])

	flowl0 = np.concatenate([flowl0,change_size],-1)
	augmented,flowl0,intr = co_transform1([iml0, iml1], flowl0, [fl,cx,cy,bl])
	imol0 = augmented[0]
	imol1 = augmented[1]
	augmented,flowl0,intr = co_transform2(augmented, flowl0, intr)

	iml0 = augmented[0]
	iml1 = augmented[1]
	flowl0 = flowl0.astype(np.float32)
	change_size = flowl0[:,:,3:]
	flowl0 = flowl0[:,:,:3]

	# randomly cover a region
	sx=0;sy=0;cx=0;cy=0
	if np.random.binomial(1,0.5):
		sx = int(np.random.uniform(25,100))
		sy = int(np.random.uniform(25,100))
		#sx = int(np.random.uniform(50,150))
		#sy = int(np.random.uniform(50,150))
		cx = int(np.random.uniform(sx,iml1.shape[0]-sx))
		cy = int(np.random.uniform(sy,iml1.shape[1]-sy))
		iml1[cx-sx:cx+sx,cy-sy:cy+sy] = np.mean(np.mean(iml1,0),0)[np.newaxis,np.newaxis]

	#iml0_bgr = iml0[...,::-1].copy()
	#iml1_bgr = iml1[...,::-1].copy()

	return iml0, iml1, flowl0, change_size, np.asarray([cx-sx,cx+sx,cy-sy,cy+sy])

def triangulation(disp, xcoord, ycoord, bl=1, fl = 450, cx = 479.5, cy = 269.5):
	#mask = disp != 0
	depth = bl*fl / (disp) # 450px->15mm focal length
	#depth = depth * mask
	#print(depth.shape, type(depth))
	X = (xcoord - cx) * depth / fl
	Y = (ycoord - cy) * depth / fl
	Z = depth
	P = np.concatenate((X[np.newaxis],Y[np.newaxis],Z[np.newaxis]),0).reshape(3,-1)
	P = np.concatenate((P,np.ones((1,P.shape[-1]))),0)
	return P

def threed_warp_flow(image_l,depth_l,image_r,depth_r,intrinsic_l,intrinsic_r):
	height_l = depth_l.shape[0]
	width_l = depth_l.shape[1]
	height_r = depth_r.shape[0]
	width_r = depth_r.shape[1]
	depth_min = np.min(np.vstack((depth_l,depth_r)))

	inv_intrinsic_l = np.linalg.inv(intrinsic_l)
	inv_intrinsic_r = np.linalg.inv(intrinsic_r)

	image_coord_l = pixel_coord_np(width_l,height_l)
	cam_coords_l = (inv_intrinsic_l[:3, :3] @ image_coord_l * depth_l.flatten()).T
	image_coord_r = pixel_coord_np(width_r,height_r)
	cam_coords_r = (inv_intrinsic_l[:3, :3] @ image_coord_r * depth_r.flatten()).T

	# rotation
	if torch.rand(1) > 0.5:
		aug_cam_coords_l,aug_cam_coords_r = perform_rotate(cam_coords_l,cam_coords_r)
	else:
		aug_cam_coords_l,aug_cam_coords_r = perform_scaling(cam_coords_l,cam_coords_r)
	# shearing
	# aug_cam_coords_l,aug_cam_coords_r = perform_shearing(aug_cam_coords_l,aug_cam_coords_r)
	#translation
	aug_cam_coords_l,aug_cam_coords_r = perform_trans(aug_cam_coords_l,aug_cam_coords_r,depth_min)
	#generate flow
	new_image_coords_l = (intrinsic_l @ aug_cam_coords_l.T / (aug_cam_coords_l[:,2]+ 1e-6)).T
	new_image_coords_r = (intrinsic_r @ aug_cam_coords_r.T / (aug_cam_coords_r[:,2]+ 1e-6)).T
	#print(image_coord_l.shape)
	flow_l = (new_image_coords_l.T - image_coord_l).T[:,:2].reshape(height_l,width_l,2)
	flow_r = (new_image_coords_r.T - image_coord_r).T[:,:2].reshape(height_r,width_r,2)

	#get warpped image
	new_image_coords_tensor_l = torch.from_numpy(new_image_coords_l[:,:2]).view(height_l,width_l,2)
	new_image_coords_l_x = new_image_coords_tensor_l[:,:,0:1] / (width_l - 1) * 2 - 1
	new_image_coords_l_y = new_image_coords_tensor_l[:,:,1:2] / (height_l - 1) * 2 - 1
	normed_image_coords_l = torch.cat((new_image_coords_l_x, new_image_coords_l_y),dim=2)
	new_image_coords_tensor_r = torch.from_numpy(new_image_coords_r[:,:2]).view(height_r,width_r,2)
	new_image_coords_r_x = new_image_coords_tensor_r[:,:,0:1] / (width_r - 1) * 2 - 1
	new_image_coords_r_y = new_image_coords_tensor_r[:,:,1:2] / (height_r - 1) * 2 - 1
	normed_image_coords_r = torch.cat((new_image_coords_r_x, new_image_coords_r_y),dim=2)
	# transfer to tensor:
	image_l = torch.from_numpy(image_l.astype(float).transpose(2,0,1))
	image_r = torch.from_numpy(image_r.astype(float).transpose(2,0,1))
	image_l_mask = torch.ones_like(image_l)
	image_r_mask = torch.ones_like(image_r)
	# warpping
	warpped_im_l = tf.grid_sample(image_l.unsqueeze(0),normed_image_coords_l.unsqueeze(0))
	warpped_im_r = tf.grid_sample(image_r.unsqueeze(0),normed_image_coords_r.unsqueeze(0))
	warpped_mask_l = (tf.grid_sample(image_l_mask.unsqueeze(0), normed_image_coords_l.unsqueeze(0)) >= 1)
	warpped_mask_r = (tf.grid_sample(image_r_mask.unsqueeze(0), normed_image_coords_r.unsqueeze(0)) >= 1)

	warpped_im_l = (warpped_im_l.float() * warpped_mask_l.float()).squeeze(0).numpy().transpose(1,2,0)
	warpped_im_r = (warpped_im_r.float() * warpped_mask_r.float()).squeeze(0).numpy().transpose(1,2,0)
	warpped_mask_l = warpped_mask_l.squeeze(0).numpy().transpose(1,2,0).all(-1)[:,:,np.newaxis]
	warpped_mask_r = warpped_mask_r.squeeze(0).numpy().transpose(1,2,0).all(-1)[:,:,np.newaxis]

	# get new image coordinate

	#sf_bl = -cv2.remap(sf_l,new_img_xl,new_img_yl,cv2.INTER_LINEAR)
	#sf_br = -cv2.remap(sf_r,new_img_xr,new_img_yr,cv2.INTER_LINEAR)
	#cv2.imwrite("l.png", warpped_im_l)
	#cv2.imwrite("r.png", warpped_im_r)
	#cv2.imwrite("l_og.png", image_l)
	#cv2.imwrite("r_og.png", image_r)

	out_dict = {
		"sf_l":flow_l,
		"sf_r":flow_r,
		"warpped_im_l":warpped_im_l,
		"warpped_im_r":warpped_im_r,
		"valid_l": warpped_mask_l,
		"valid_r": warpped_mask_r
	}
	return out_dict


text:
import glob
import os
import logging.config
import numpy as np
from PIL import Image as Pil_Image
from matplotlib import pyplot as plt
from albumentations import *
from util import logger_init

log = logging.getLogger(__name__)

def save_aug_image(image_name, img_path, pig_img_aug1, prefix):
    """
    Saves the augmented image

    """
    log.info('Saving augmented image...')
    img_aug_opencv = np.array(pig_img_aug1)
    pil_img = Pil_Image.fromarray(img_aug_opencv)
    aug_img_name = prefix + image_name
    pil_img.save(os.path.join(img_path, aug_img_name))

def plot_image(pig_img_aug, label, rows):
    """Plots the augmented image"""
    log.info('Plotting augmented image...')
    global num
    num = num + 1
    plt.subplot(rows, 5, num)
    plt.title(label)
    plt.imshow(pig_img_aug)

def image_resize(image, width=None, height=None, inter=cv2.INTER_AREA):
    """
    Initialize the dimensions of the image to be resized and grab the
    image size.
    """
    log.info('Resizing image...')
    dim = None
    (h, w) = image.shape[:2]
    # if both the width and height are None, then return the
    # original image
    if width is None and height is None:
        return image
    # check to see if the width is None
    if width is None:
        # calculate the ratio of the height and construct the
        # dimensions
        r = height / float(h)
        dim = (int(w * r), height)
    # otherwise, the height is None
    else:
        # calculate the ratio of the width and construct the dimensions
        r = width / float(w)
        dim = (width, int(h * r))
    # resize the image
    resized = cv2.resize(image, dim, interpolation=inter)
    return resized

def clean_augmented_images(path, pattern):
    """
    Deletes all files in sub folders depending on the pattern.
    @param:
        - path : str Path of the to delete file in sub folders
        - pattern: str Pattern of the file beginning
    """
    log.info('Start deleting process of images with the pattern ' + pattern + '...')
    all_files = [y
                 for x in os.walk(path)
                 for y in glob.glob(os.path.join(x[0], pattern))
                 ]
    if len(all_files) == 0:
        log.warning('No files for this pattern ' + pattern + ' in the folders')
    for file in all_files:
        os.remove(file)
    log.info('Removed ' + str(len(all_files)) + ' Elements')
    log.info('Finished deleting process')

def counts_file_in_sub_folder_with_specific_pattern(path, pattern):
    """
       Counts all files in sub folders depending on the pattern.
       @param:
           - path : str Path of the to delete file in sub folders
           - pattern: str Pattern of the file beginning
       """
    log.info('Start counting process of images with the pattern ' + pattern + '...')
    all_files = [y
                 for x in os.walk(path)
                 for y in glob.glob(os.path.join(x[0], pattern))
                 ]
    log.info('Finished counting process - Total elements: ' + str(len(all_files)))
    return len(all_files)

max_stars_repo_name: hurschler/pig-face-recognition

max_issues_repo_name: hurschler/pig-face-recognition

max_forks_repo_name: hurschler/pig-face-recognition


text:
import sys
import os
import glob
import argparse
import numpy as np

sys.path.insert(0, '/gpfs/mira-home/ekovacs/.local/lib/python2.7/site-packages')
sys.path.insert(0, '/gpfs/mira-home/ekovacs/cosmology/cosmodc2')
sys.path.insert(0, '/gpfs/mira-home/ekovacs/cosmology/galsampler/build/lib.linux-x86_64-2.7')
sys.path.insert(0, '/gpfs/mira-home/ekovacs/cosmology/halotools/build/lib.linux-x86_64-2.7')

from cosmodc2.write_umachine_color_mocks_to_disk import write_snapshot_mocks_to_disk
from cosmodc2.get_fof_info import get_fof_info

parser = argparse.ArgumentParser()

parser.add_argument("commit_hash",
    help="Commit hash to save in output files")
parser.add_argument("-nsnap",
    help="Number of snapshots to loop over. Default is 29.",
        default=29, type=int)
parser.add_argument("-verbose",
    help="Turn on extra printing",
        action='store_true', default=False)

args = parser.parse_args()

pkldirname = "/home/ekovacs/cosmology/cosmodc2/cosmodc2"

#get list of simulation halo files
redshift_strings, snapshots, alphaQ_halos_fname_list = get_fof_info(pkldirname, nsnapshot=args.nsnap)
expansion_factors = [1./(1+float(z)) for z in redshift_strings]
if(args.verbose):
    print("target z's and a's:", redshift_strings, expansion_factors)

umachine_mstar_ssfr_mock_dirname = (
    '/projects/DarkUniverse_esp/kovacs/AlphaQ/protoDC2_v4_um_sfr_catalogs_and_halos')
sfr_files = sorted([os.path.basename(f) for f in glob.glob(umachine_mstar_ssfr_mock_dirname+'/sfr*')])
um_expansion_factors = np.asarray([float(f.split('sfr_catalog_')[-1].split('_value_added.hdf5')[0]) for f in sfr_files])
closest_snapshots = [np.abs(um_expansion_factors - a).argmin() for a in expansion_factors]
if(args.verbose):
    print('index of closest snapshots:',closest_snapshots)

umachine_mstar_ssfr_mock_basename_list = [sfr_files[n] for n in closest_snapshots]
umachine_mstar_ssfr_mock_fname_list = list(
    (os.path.join(umachine_mstar_ssfr_mock_dirname, basename)
    for basename in umachine_mstar_ssfr_mock_basename_list))
if(args.verbose):
    print('umachine_mstar_ssfr_mock_basename_list:',umachine_mstar_ssfr_mock_basename_list)

umachine_host_halo_dirname = (
    '/projects/DarkUniverse_esp/kovacs/AlphaQ/protoDC2_v4_um_sfr_catalogs_and_halos')
umachine_host_halo_basename_list = [sfr_files[n].replace('sfr', 'halo') for n in closest_snapshots]
umachine_host_halo_fname_list = list(
    (os.path.join(umachine_host_halo_dirname, basename)
    for basename in umachine_host_halo_basename_list))
if(args.verbose):
    print('umachine_host_halo_fname_list:',umachine_host_halo_basename_list)

target_halo_dirname = "/projects/DarkUniverse_esp/heitmann/OuterRim/M000/L360/HACC001/analysis/Halos/b0168/fofp_new"
target_halo_basename_list = [os.path.basename(q) for q in alphaQ_halos_fname_list]
target_halo_fname_list = list(
    (os.path.join(target_halo_dirname, basename)
    for basename in target_halo_basename_list))
if(args.verbose):
    print('target_halo_basename_list:',target_halo_basename_list)

output_mock_dirname = (
        "/projects/DarkUniverse_esp/kovacs/AlphaQ/galsampler_alphaq_outputs_v4_6")

output_mock_basename_list = [''.join(["umachine_color_mock_v4_", t.replace('fofproperties', ''), 'hdf5']) for t in target_halo_basename_list]
output_color_mock_fname_list = list(
    (os.path.join(output_mock_dirname, basename)
    for basename in output_mock_basename_list))
if(args.verbose):
    print('output_mock_basename_list:', output_mock_basename_list)

redshift_list = [float(z) for z in redshift_strings]
commit_hash = args.commit_hash
target_halo_loader = "gio"
Lbox_target_halos = 256.

write_snapshot_mocks_to_disk(
            umachine_mstar_ssfr_mock_fname_list, umachine_host_halo_fname_list,
            target_halo_fname_list, output_color_mock_fname_list,
            redshift_list, commit_hash, target_halo_loader, Lbox_target_halos)

max_stars_repo_name: ArgonneCPAC/skysim

max_issues_repo_name: ArgonneCPAC/skysim


text:
import numpy as np
from numpy.testing import assert_allclose
from astroML.filters import savitzky_golay, wiener_filter

def test_savitzky_golay():
    y = np.zeros(100)
    y[::2] = 1
    f = savitzky_golay(y, window_size=3, order=1)
    assert_allclose(f, (2 - y) / 3.)

def test_savitzky_golay_fft():
    y = np.random.normal(size=100)

    for width in [3, 5]:
        for order in range(width - 1):
            f1 = savitzky_golay(y, width, order, use_fft=False)
            f2 = savitzky_golay(y, width, order, use_fft=True)
            assert_allclose(f1, f2)

def test_wiener_filter_simple():
    t = np.linspace(0, 1, 256)
    h = np.zeros_like(t)
    h[::2] = 1000
    s = wiener_filter(t, h)
    assert_allclose(s, np.mean(h))

def test_wienter_filter_spike():
    np.random.seed(0)
    N = 2048
    dt = 0.05

    t = dt * np.arange(N)
    h = np.exp(-0.5 * ((t - 20.) / 1.0) ** 2) + 10
    hN = h + np.random.normal(0, 0.05, size=h.shape)
    h_smooth = wiener_filter(t, hN)

    assert_allclose(h, h_smooth, atol=0.03)

max_stars_repo_name: aragilar/astroML

max_issues_repo_name: aragilar/astroML

max_forks_repo_name: aragilar/astroML

reason: import numpy,from numpy


text:
from typing import Sequence, Tuple

import numpy as np
from exrec.core.interface import Logger

def als(
    logger: Logger,
    users: Sequence[int],
    items: Sequence[int],
    ratings: Sequence[float],
    U: np.ndarray,
    V: np.ndarray,
    n_samples: int,
    dim: int,
    reg: float,
    lr: float,
    epochs: int,
    seed: int,
) -> Tuple[np.ndarray, np.ndarray]:

    np.random.seed(seed=seed)
    index = np.arange(n_samples)
    np.random.shuffle(index)

    for e in range(epochs):
        logger.info(f"Epoch: {e}/{epochs}")

        for idx in index:
            i = users[idx]
            j = items[idx]
            diff = ratings[idx] - np.dot(U[i, :], V[j, :])
            U[i, :] = U[i, :] + lr * (2 * diff * V[j, :] - reg * U[i, :])
            V[j, :] = V[j, :] + lr * (2 * diff * U[i, :] - reg * V[j, :])

    return (U, V)

max_stars_repo_name: InkToYou/Ex-Rec

max_issues_repo_name: InkToYou/Ex-Rec

max_forks_repo_name: InkToYou/Ex-Rec


text:
import numpy as np
#import gym
import gym.envs
import torch
from HalfCheetah.PyBullUtils import ZFilter, get_action
import pybullet as p

seed=123
env = gym.make("HalfCheetahPyBulletEnv-v0")
env.seed(seed)
torch.manual_seed(seed)
env.render()
env.reset()

for i in range(p.getNumBodies()):
    print(p.getBodyInfo(i))
    if p.getBodyInfo(i)[1].decode() == "cheetah":
        torsoId = i

num_inputs = env.observation_space.shape[0]
num_actions = env.action_space.shape[0]

running_state = ZFilter((num_inputs,), clip=5)
running_state.load('HC')

actor = torch.load('D:/Models/HCActor1.pt')
critic = torch.load('D:/Models/HCCritic1.pt')

def test(iterations = 10):
    actor.eval(), critic.eval()
    scores = []
    distance = 5
    yaw = 0
    for iter in range(iterations):
        steps = 0
        rewards = []
        state = env.reset()
        env.render()
        state = running_state(state)
        score = 0
        done = False
        while not done:
            mu, std, _ = actor(torch.Tensor(state).unsqueeze(0))
            action = get_action(mu, std)[0]
            next_state, reward, done, _ = env.step(action)
            next_state = running_state(next_state)
            rewards.append(reward)
            score += reward
            steps +=1
            state = next_state
            humanPos = p.getLinkState(torsoId,4)[0]
            p.resetDebugVisualizerCamera(distance,yaw,-20,humanPos)
            env.render()
        print('Episode {} Finished After {} Steps With Score {}'.format(iter, steps, score))
        scores.append(score)
    print('Testing Completed With Average Score {}'.format(np.mean(scores)))

test()


text:
import numpy as np
import pandas as pd
import matplotlib as plt
from matplotlib import pyplot as pllt
import math
from scipy.special import gamma

def main():
    data = pd.read_csv('wyniki.csv')
    data = np.array(data)
    single_time = np.array([data[i][3]/data[i][4] for i in range(len(data))])
    data = np.column_stack((data, single_time))
    n = np.array([data[i][1].count(' ') for i in range(len(data))])
    data = np.column_stack((data, n))
    x = [data[i][6] for i in range(len(data))]
    y = [data[i][5] for i in range(len(data))]
    pllt.plot(x, y, label='TSP brute force',
              linewidth=3.0, marker='.', markersize=15)
    x_e = np.linspace(min([data[i][6] for i in range(len(data))]), max(
        [data[i][6] for i in range(len(data))]), 100)
    param = max(y)/max(gamma(x_e))
    pllt.plot(x_e, gamma(x_e)*param, label='n!', linestyle='--', linewidth=2.0)
    pllt.xlabel('Liczba wierzchołków w grafie')
    pllt.ylabel('Czas wykonywania algorytmu[s]')
    pllt.legend()
    pllt.show()

if __name__ == "__main__":
    main()

max_stars_repo_name: tomekrzymyszkiewicz/TSP-brute-force

max_issues_repo_name: tomekrzymyszkiewicz/TSP-brute-force

max_forks_repo_name: tomekrzymyszkiewicz/TSP-brute-force


text:
#=========================================================
# IMPORT STATEMENTS
#=========================================================

import caesar
from readgadget import *
import sys
import pylab as plt
import os
os.environ["OMP_NUM_THREADS"] = "24"
import numpy as np
import fsps
import h5py
from scipy import interpolate
from astropy import constants as const
from astropy.cosmology import FlatLambdaCDM
from caesar.utils import memlog
from caesar.property_manager import MY_DTYPE
from yt.funcs import mylog

#from pygas import *
#from auxloser import t_elapsed,parse_args,progress_bar,hubble_z

#from scipy.ndimage.filters import gaussian_filter

# start overall timer
CLIGHT_AA = const.c.to('AA/s').value
Solar = {'total':0.0134, 'He':0.2485, 'C':2.38e-3, 'N':0.70e-3, 'O':5.79e-3, 'Ne':1.26e-3, 'Mg':7.14e-4, 'Si':6.71e-4, 'S':3.12e-4, 'Ca':0.65e-4, 'Fe':1.31e-3} # Asplund abundances used in Simba

#=========================================================
# ROUTINES TO COMPUTE SPECTRA AND MAGNITUDES
#=========================================================

# photometry class
class photometry:

    def __init__(self, obj, group_list, ds=None, band_names='v', ssp_table_file='SSP_Chab_EL.hdf5', view_dir='x', use_dust=True, use_cosmic_ext=True, kernel_type='cubic', nproc=-1):

        from caesar.property_manager import ptype_ints
        self.obj = obj  # caesar object
        self.groups = group_list  # list of objects to process

        # optional arguments
        self.band_names = band_names
        if hasattr(self.obj,'_kwargs') and 'fsps_bands' in self.obj._kwargs:
            self.band_names = self.obj._kwargs['fsps_bands']
        self.ssp_table_file = os.path.expanduser('~/caesar/%s'%ssp_table_file)
        if hasattr(self.obj,'_kwargs') and 'ssp_table_file' in self.obj._kwargs:
            self.ssp_table_file = self.obj._kwargs['ssp_table_file']
        self.ext_law = 'mw'
        if hasattr(self.obj,'_kwargs') and 'ext_law' in self.obj._kwargs:
            self.ext_law = self.obj._kwargs['ext_law'].lower()
        if hasattr(self.obj,'_kwargs') and 'view_dir' in self.obj._kwargs:
            view_dir = self.obj._kwargs['view_dir'].lower()
        if view_dir is 'x': self.viewdir = 0
        if view_dir is 'y': self.viewdir = 1
        if view_dir is 'z': self.viewdir = 2
        self.use_dust = use_dust  # if False, will use metals plus an assumed dust-to-metal ratio
        if hasattr(self.obj,'_kwargs') and 'use_dust' in self.obj._kwargs:
            use_dust = self.obj._kwargs['use_dust'].lower()
        self.use_cosmic_ext = use_cosmic_ext
        if hasattr(self.obj,'_kwargs') and 'use_cosmic_ext' in self.obj._kwargs:
            use_cosmic_ext = self.obj._kwargs['use_cosmic_ext'].lower()
        self.kernel_type = kernel_type
        self.nkerntab = 2000
        if nproc == -1:
            try:
                self.nproc = obj.nproc
            except:
                self.nproc = 1
        else:
            self.nproc = nproc

        # useful quantities
        self.boxsize = self.obj.simulation.boxsize
        self.solar_abund = Solar
        self.lumtoflux_abs = const.L_sun.to('erg/s').value/(4* np.pi * 10.**2 * const.pc.to('cm').value**2)
        cosmo = FlatLambdaCDM(H0=100.*self.obj.simulation.hubble_constant, Om0=self.obj.simulation.omega_matter, Tcmb0=2.73)
        lumdist = cosmo.luminosity_distance(self.obj.simulation.redshift).to('pc').value
        self.lumtoflux = const.L_sun.to('erg/s').value/(4* np.pi * lumdist**2 * const.pc.to('cm').value**2)
        self.lumtoflux *= 1.+self.obj.simulation.redshift  # we compute apparent mags by blueshifting the band, which reduces the flux by (1+z); correct for this here

        # if there is no data_manager, assume we're running interactively
        # this means we have to load in the particle info, and set some other info we need
        if not hasattr(self.obj,'data_manager'):
            from caesar.data_manager import DataManager
            from caesar.property_manager import DatasetType
            self.obj.data_manager = DataManager(self.obj)
            self.obj._ds_type = DatasetType(ds)
            self.obj.yt_dataset = ds
            self.obj.units = dict(
                mass='Msun',
                length='kpccm',
                velocity='km/s',
                time='yr',
                temperature='K'
            )
            self.obj.data_manager._photometry_init()

    def run_pyloser(self):

        from caesar.cyloser import compute_AV, compute_mags

        self.init_pyloser()
        #computes AV for all stars in snapshot
        self.obj.AV_star = compute_AV(self)
        #find the AV for stars belonging to the groups that were asked for
        self.Av_per_group()
        spect_dust, spect_nodust = compute_mags(self)

        return spect_dust, spect_nodust

    def init_pyloser(self):
        from caesar.cyloser import init_kerntab
        self.init_ssp_table()
        self.init_extinction()
        self.init_bands()
        init_kerntab(self)
        self.init_stars_to_process()

    #separate AV_all_stars by group
    def Av_per_group(self):
        memlog('Finding LOS A_V values for %d objects'%(len(self.groups)))
        try:
            import tqdm
            for obj_ in tqdm.tqdm(self.groups):
                Av_per_star = self.obj.AV_star[obj_.slist]
                obj_.group_Av = Av_per_star

        except:
            for obj_ in self.groups:
                Av_per_star = self.obj.AV_star[obj_.slist]
                obj_.group_Av = Av_per_star

    # initialize extinction curves.  order: 0=Calzetti, 1=Chevallard, 2=Conroy, 3=Cardelli(MW), 4=SMC, 5=LMC, 6=Mix Calz/MW, 7=Composite Calz/MW/SMC; see atten_laws.py for details (these return optical depths)
    def init_extinction(self):
        from caesar.pyloser.atten_laws import calzetti,chevallard,conroy,cardelli,smc,lmc
        wave = self.ssp_wavelengths.astype(np.float64)
        self.ext_curves = []
        self.ext_curves.append(calzetti(wave))
        self.ext_curves.append(chevallard(wave))
        self.ext_curves.append(conroy(wave))
        self.ext_curves.append(cardelli(wave))
        self.ext_curves.append(smc(wave))
        self.ext_curves.append(lmc(wave))
        self.ext_curves = np.asarray(self.ext_curves)

        memlog('Starting photometry using %s extinction law'%self.ext_law)
        if 'calzetti' in self.ext_law: self.ext_law = 0
        elif 'chevallard' in self.ext_law: self.ext_law = 1
        elif 'conroy' in self.ext_law: self.ext_law = 2
        elif self.ext_law == 'mw' or self.ext_law == 'cardelli' or 'CCM' in self.ext_law: self.ext_law = 3
        elif 'smc' in self.ext_law: self.ext_law = 4
        elif 'lmc' in self.ext_law: self.ext_law = 5
        elif self.ext_law == 'mix_calz_MW': self.ext_law = 6
        elif self.ext_law == 'composite': self.ext_law = 7
        else:
            mylog.warning('Extinction law %s not recognized, assuming composite'%self.ext_law)
            self.ext_law = 7

    # set up star and gas lists in each object
    def init_stars_to_process(self):
        from caesar.group import Group, collate_group_ids
        from caesar.property_manager import ptype_ints
        from caesar.cyloser import smass_at_formation

        #if isinstance(self.groups[0],Group):
        self.ngroup, self.gasids, self.gid_bins = collate_group_ids(self.groups,'gas',self.obj.simulation.ngas)
        self.ngroup, self.starids, self.sid_bins = collate_group_ids(self.groups,'star',self.obj.simulation.nstar)
        #else:
        #    sys.exit('Must provide a list of Caesar groups.')
        self.scount = sum([len(i.slist) for i in self.groups])
        self.gcount = sum([len(i.glist) for i in self.groups])
        self.Nobjs = len(self.groups)

        # get original stellar mass at time of formation
        self.obj.smass_orig = smass_at_formation(self.obj,self.groups,self.ssp_mass,self.ssp_ages,nproc=self.nproc)

        memlog('Loaded %d stars and %d gas in %d objects to process'%(self.scount,self.gcount,self.Nobjs))
        return

    # initialize band transmission data interpolated to FSPS wavelengths
    def init_bands(self):
        import fsps
        if isinstance(self.band_names,str):
            self.band_names = [self.band_names]
        if self.band_names[0] == 'all':
            self.band_names = fsps.list_filters()
        elif self.band_names[0] == 'uvoir':
            self.band_names = []
            for ib,b in enumerate(fsps.list_filters()):
                band = fsps.filters.get_filter(b)  # look up characteristics of desired band
                band_wave = band.transmission[0]   # filter wavelengths
                band_trans = band.transmission[1]  # filter response function
                meanwave = np.sum(band.transmission[0]*band.transmission[1])/np.sum(band.transmission[1])
                if meanwave < 50000: self.band_names.append(b)
        else:
            # collect all filters containing the input string(s)
            allfilters = fsps.list_filters()
            mybands = []
            for b in self.band_names:  # check that requested bands are actually available
                for b_all in allfilters:
                    if b in b_all:
                        if b == b_all: mybands.append(b_all)  # if exact match, add
                        elif len(b)>3: mybands.append(b_all)  # avoid adding matching short band names (e.g. 'u')
            if len(mybands) == 0:
                assert b in allfilters, 'Band %s not found among available FSPS filters! Call fsps.list_filters() to list filters.'%self.band_names
            self.band_names = mybands
        # V band is always computed, so that one has A_V (= V_dust - V_nodust)
        if 'v' not in self.band_names:
            self.band_names.append('v')

        # Madau IGM attenuation is applied directly to rest-frame bandpasses only when computing apparent magnitudes; compute this curve here for specific redshift
        redshift = self.obj.simulation.redshift
        if self.use_cosmic_ext:
            from synphot import etau_madau  # see synphot.readthedocs.io/en/latest/synphot/tutorials.html
            extcurve = etau_madau(self.ssp_wavelengths*(1.+redshift), redshift)
            cosmic_ext = extcurve(self.ssp_wavelengths)
        else: cosmic_ext = np.ones(len(self.ssp_wavelengths))

        # set up band information
        nbands = len(self.band_names)
        self.band_meanwave = np.zeros(nbands,dtype=MY_DTYPE)
        self.band_indexes = np.zeros(nbands+1,dtype=np.int32)
        self.band_ftrans = np.empty(0,dtype=MY_DTYPE)
        self.band_iwave0 = np.zeros(nbands,dtype=np.int32)
        self.band_iwave1 = np.zeros(nbands,dtype=np.int32)
        self.band_indz = np.zeros(nbands+1,dtype=np.int32)
        self.band_ztrans = np.empty(0,dtype=MY_DTYPE)
        self.band_iwz0 = np.zeros(nbands,dtype=np.int32)
        self.band_iwz1 = np.zeros(nbands,dtype=np.int32)
        for ib,b in enumerate(self.band_names):
            band = fsps.filters.get_filter(b)  # look up characteristics of desired band
            band_wave = band.transmission[0]   # filter wavelengths
            band_trans = band.transmission[1]  # filter response function
            self.band_meanwave[ib] = np.sum(band.transmission[0]*band.transmission[1])/np.sum(band.transmission[1])
            # Set up transmission curve in region probed by rest-frame band
            ind = np.where((self.ssp_wavelengths > band_wave[0]) & (self.ssp_wavelengths < band_wave[-1]))[0] # indices of wavelengths in the band
            self.band_iwave0[ib] = ind[0]
            self.band_iwave1[ib] = ind[-1]+1
            ftrans = np.interp(self.ssp_wavelengths[ind],band_wave,band_trans)  # transmission at those wavelengths
            dnu = CLIGHT_AA/self.ssp_wavelengths[ind[0]:ind[-1]+1] - CLIGHT_AA/self.ssp_wavelengths[ind[0]+1:ind[-1]+2]  # convert to delta-nu
            self.band_ftrans = np.append(self.band_ftrans, ftrans*dnu)
            self.band_indexes[ib+1] = len(self.band_ftrans)
            # Now set up band for apparent mag computation
            # We will blueshift the band, corresponding to redshifting the intrinsic spectrum
            ind = np.where((self.ssp_wavelengths > band_wave[0]*self.obj.simulation.scale_factor) & (self.ssp_wavelengths < band_wave[-1]*self.obj.simulation.scale_factor))[0] # indices of wavelengths for redshifted rest-frame spectrum (i.e. blueshifted band)
            self.band_iwz0[ib] = ind[0]
            self.band_iwz1[ib] = ind[-1]+1
            ftrans = np.interp(self.ssp_wavelengths[ind],band_wave*self.obj.simulation.scale_factor,band_trans)  # transmission at those wavelengths
            dnu = CLIGHT_AA/self.ssp_wavelengths[ind[0]:ind[-1]+1] - CLIGHT_AA/self.ssp_wavelengths[ind[0]+1:ind[-1]+2]  # convert to delta-nu
            self.band_ztrans = np.append(self.band_ztrans, np.array(ftrans*dnu*cosmic_ext[ind]))
            self.band_indz[ib+1] = len(self.band_ztrans)

        memlog('Computing %d bands: %s'%(len(self.band_names),self.band_names))

    # initialize SSP table, by either generating it if it doesn't exist or reading it in
    def init_ssp_table(self):
        import os
        read_flag = False
        if os.path.exists(self.ssp_table_file):
            try:
                self.read_ssp_table(self.ssp_table_file)
                memlog('Read SSP table %s'%self.ssp_table_file)
                read_flag = True
            except:
                memlog('Error reading SSP table %s, will generate...'%self.ssp_table_file)
        if not read_flag:  # generate table with Caesar default options
            ssp_ages, ssp_logZ, mass_remaining, wavelengths, ssp_spectra = generate_ssp_table(self.ssp_table_file, return_table=True, imf_type=1,add_neb_emission=True,sfh=0,zcontinuous=1)  # note Caesar default FSPS options; run generate_ssp_table() separately to set desired FSPS options
            self.ssp_ages = np.array(ssp_ages,dtype=MY_DTYPE)
            self.ssp_logZ = np.array(ssp_logZ,dtype=MY_DTYPE)
            self.ssp_mass = np.array(mass_remaining,dtype=MY_DTYPE)
            self.ssp_wavelengths = np.array(wavelengths,dtype=MY_DTYPE)
            self.ssp_spectra = np.array(ssp_spectra,dtype=MY_DTYPE)

    def read_ssp_table(self,ssp_lookup_file):
        hf = h5py.File(ssp_lookup_file,'r')
        for i in hf.keys():
            if i=='wavelengths': wavelengths = list(hf[i])
            if i=='mass_remaining': mass_remaining = list(hf[i])
            if i=='ages': ssp_ages = list(hf[i])
            if i=='logZ': ssp_logZ = list(hf[i])
            if i=='spectra': ssp_spectra = list(hf[i])
        self.ssp_ages = np.array(ssp_ages,dtype=MY_DTYPE)
        self.ssp_logZ = np.array(ssp_logZ,dtype=MY_DTYPE)
        self.ssp_mass = np.array(mass_remaining,dtype=MY_DTYPE)
        self.ssp_wavelengths = np.array(wavelengths,dtype=MY_DTYPE)
        self.ssp_spectra = np.array(ssp_spectra,dtype=MY_DTYPE)

def generate_ssp_table(ssp_lookup_file,Zsol=Solar['total'],oversample=[2,2],return_table=False,**fsps_options):
        '''
        Generates an SPS lookup table, oversampling in [age,metallicity] by oversample
        '''
        import fsps
        mylog.info('Generating SSP lookup table %s'%(ssp_lookup_file))
        mylog.info('with FSPS options: %s'%(fsps_options))
        fsps_opts = ''
        for key, value in fsps_options.items():
            fsps_opts = fsps_opts + ("{0} = {1}, ".format(key, value))
        fsps_opts = np.string_(fsps_opts)
        fsps_ssp = fsps.StellarPopulation(**fsps_options)
        wavelengths = fsps_ssp.wavelengths
        ssp_ages = []
        mass_remaining = []
        ssp_ages.append(fsps_ssp.ssp_ages[0])
        mass_remaining.append(fsps_ssp.stellar_mass[0])
        for i in range(len(fsps_ssp.ssp_ages)-1):
            for j in range(i+1,i+oversample[0]):
                ssp_ages.append((fsps_ssp.ssp_ages[j]-fsps_ssp.ssp_ages[j-1])*(j-i)/oversample[0]+fsps_ssp.ssp_ages[j-1])
                mass_remaining.append((fsps_ssp.stellar_mass[j]-fsps_ssp.stellar_mass[j-1])*(j-i)/oversample[0]+fsps_ssp.stellar_mass[j-1])
            ssp_ages.append(fsps_ssp.ssp_ages[j])
            mass_remaining.append(fsps_ssp.stellar_mass[j])
        ssp_logZ = []
        ssp_logZ.append(fsps_ssp.zlegend[0])
        for i in range(len(fsps_ssp.zlegend)-1):
            for j in range(i+1,i+oversample[1]):
                ssp_logZ.append((fsps_ssp.zlegend[j]-fsps_ssp.zlegend[j-1])*(j-i)/oversample[1]+fsps_ssp.zlegend[j-1])
            ssp_logZ.append(fsps_ssp.zlegend[j])
        ssp_logZ = np.log10(ssp_logZ)
        ssp_spectra = []
        for age in ssp_ages:
            for Zmet in ssp_logZ:
                fsps_ssp.params["logzsol"] = Zmet-np.log10(Zsol)
                spectrum = fsps_ssp.get_spectrum(tage=10**(age-9))[1]
                ssp_spectra.append(spectrum)
        with h5py.File(ssp_lookup_file, 'w') as hf:
            hf.create_dataset('fsps_options',data=fsps_opts)
            hf.create_dataset('ages',data=ssp_ages)
            hf.create_dataset('logZ',data=ssp_logZ)
            hf.create_dataset('mass_remaining',data=mass_remaining)
            hf.create_dataset('wavelengths',data=wavelengths)
            hf.create_dataset('spectra',data=ssp_spectra)
        memlog('Generated lookup table with %d ages and %d metallicities'%(len(ssp_ages),len(ssp_logZ)))

        if return_table:
            return ssp_ages, ssp_logZ, mass_remaining, wavelengths, ssp_spectra

reason: import numpy,from scipy,from astropy