<gh_stars>100-1000
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from . import _unittest as unittest
import cmath
import decimal
import math
import re

try:
    import numpy
except ImportError:
    numpy = False

from datatest._vendor.predicate import (
    _check_type,
    _check_callable,
    _check_wildcard,
    _check_truthy,
    _check_falsy,
    _check_nan,
    _check_regex,
    _check_set,
    _get_matcher_parts,
    get_matcher,
    MatcherBase,
    MatcherObject,
    MatcherTuple,
    Predicate,
    PredicateIntersectionType,
    PredicateUnionType,
)

class TestCheckType(unittest.TestCase):
    def test_isinstance(self):
        function = lambda x: _check_type(int, x)
        self.assertTrue(function(0))
        self.assertTrue(function(1))
        self.assertFalse(function(0.0))
        self.assertFalse(function(1.0))

    def test_is_type(self):
        self.assertTrue(_check_type(int, int))

class TestCheckCallable(unittest.TestCase):
    def test_function(self):
        def divisible3or5(x):  # <- Helper function.
            return (x % 3 == 0) or (x % 5 == 0)

        function = lambda x: _check_callable(divisible3or5, x)
        self.assertFalse(function(1))
        self.assertFalse(function(2))
        self.assertTrue(function(3))
        self.assertFalse(function(4))
        self.assertTrue(function(5))
        self.assertTrue(function(6))

    def test_error(self):
        def fails_internally(x):  # <- Helper function.
            raise TypeError('raising an error')

        function = lambda x: _check_callable(fails_internally, x)
        with self.assertRaises(TypeError):
            self.assertFalse(function('abc'))

    def test_identity(self):
        def always_false(x):  # <- Helper function.
            return False

        function = lambda x: _check_callable(always_false, x)
        self.assertTrue(function(always_false))

    def test_identity_with_error(self):
        def fails_internally(x):  # <- Helper function.
            raise TypeError('raising an error')

        function = lambda x: _check_callable(fails_internally, x)
        self.assertTrue(function(fails_internally))

class TestCheckWildcard(unittest.TestCase):
    def test_always_true(self):
        self.assertTrue(_check_wildcard(1))
        self.assertTrue(_check_wildcard(object()))
        self.assertTrue(_check_wildcard(None))

class TestCheckTruthy(unittest.TestCase):
    def test_matches(self):
        self.assertTrue(_check_truthy('x'))
        self.assertTrue(_check_truthy(1.0))
        self.assertTrue(_check_truthy([1]))
        self.assertTrue(_check_truthy(range(1)))

    def test_nonmatches(self):
        self.assertFalse(_check_truthy(''))
        self.assertFalse(_check_truthy(0.0))
        self.assertFalse(_check_truthy([]))
        self.assertFalse(_check_truthy(range(0)))

class TestCheckFalsy(unittest.TestCase):
    def test_matches(self):
        self.assertTrue(_check_falsy(''))
        self.assertTrue(_check_falsy(0.0))
        self.assertTrue(_check_falsy([]))
        self.assertTrue(_check_falsy(range(0)))

    def test_nonmatches(self):
        self.assertFalse(_check_falsy('x'))
        self.assertFalse(_check_falsy(1.0))
        self.assertFalse(_check_falsy([1]))
        self.assertFalse(_check_falsy(range(1)))

class TestCheckNaN(unittest.TestCase):
    def test_matches(self):
        self.assertTrue(_check_nan(float('NaN')))
        self.assertTrue(_check_nan(complex(float('NaN'))))
        self.assertTrue(_check_nan(decimal.Decimal('NaN')))
        if hasattr(math, 'nan'):  # New in version 3.5
            self.assertTrue(_check_nan(math.nan))
        if hasattr(cmath, 'nan'):  # New in version 3.6
            self.assertTrue(_check_nan(cmath.nan))

    def test_nonmatches(self):
        self.assertFalse(_check_nan('x'))
        self.assertFalse(_check_nan(1))
        self.assertFalse(_check_nan(1.0))
        self.assertFalse(_check_nan(complex(1)))
        self.assertFalse(_check_nan(decimal.Decimal('1.123')))

    @unittest.skipUnless(numpy, 'requires numpy')
    def test_numpy_cases(self):
        self.assertTrue(_check_nan(numpy.nan))
        self.assertFalse(_check_nan(numpy.int64(123)))

class TestCheckRegex(unittest.TestCase):
    def test_function(self):
        regex = re.compile('(Ch|H)ann?ukk?ah?')
        function = lambda x: _check_regex(regex, x)

        self.assertTrue(function('Happy Hanukkah'))
        self.assertTrue(function('Happy Chanukah'))
        self.assertFalse(function('Merry Christmas'))

    def test_incompatible_types(self):
        regex = re.compile('abc')
        self.assertFalse(_check_regex(regex, 123))
        self.assertFalse(_check_regex(regex, ('a', 'b')))

    def test_identity(self):
        regex = re.compile('abc')
        self.assertTrue(_check_regex(regex, regex))

class TestCheckSet(unittest.TestCase):
    def test_function(self):
        function = lambda x: _check_set(set(['abc', 'def']), x)
        self.assertTrue(function('abc'))
        self.assertFalse(function('xyz'))

    def test_whole_set_equality(self):
        function = lambda x: _check_set(set(['abc', 'def']), x)
        self.assertTrue(function(set(['abc', 'def'])))

    def test_unhashable_check(self):
        function = lambda x: _check_set(set(['abc', 'def']), x)
        self.assertFalse(function(['abc']))
        self.assertFalse(function((1, ['xyz'])))

class TestGetMatcherParts(unittest.TestCase):
    def test_type(self):
        pred_handler, repr_string = _get_matcher_parts(int)
        self.assertTrue(pred_handler(1))
        self.assertFalse(pred_handler(1.0))
        self.assertEqual(repr_string, 'int')

    def test_callable(self):
        def userfunc(x):
            return x == 1
        pred_handler, repr_string = _get_matcher_parts(userfunc)
        self.assertTrue(pred_handler(1))
        self.assertFalse(pred_handler(2))
        self.assertEqual(repr_string, 'userfunc')

        userlambda = lambda x: x == 1
        pred_handler, repr_string = _get_matcher_parts(userlambda)
        self.assertTrue(pred_handler(1))
        self.assertFalse(pred_handler(2))
        self.assertEqual(repr_string, '<lambda>')

    def test_ellipsis_wildcard(self):
        pred_handler, repr_string = _get_matcher_parts(Ellipsis)
        self.assertIs(pred_handler, _check_wildcard)
        self.assertEqual(repr_string, '...')

    def test_truthy(self):
        pred_handler, repr_string = _get_matcher_parts(True)
        self.assertIs(pred_handler, _check_truthy)
        self.assertEqual(repr_string, 'True')

    def test_falsy(self):
        pred_handler, repr_string = _get_matcher_parts(False)
        self.assertIs(pred_handler, _check_falsy)
        self.assertEqual(repr_string, 'False')

    def test_nan(self):
        pred_handler, repr_string = _get_matcher_parts(float('nan'))
        self.assertIs(pred_handler, _check_nan)
        self.assertEqual(repr_string, 'NaN')

    def test_regex(self):
        regex = re.compile('ab[cd]')

        pred_handler, repr_string = _get_matcher_parts(regex)
        self.assertTrue(pred_handler('abc'))
        self.assertFalse(pred_handler('abe'))
        self.assertEqual(repr_string, "re.compile('ab[cd]')")

    def test_set(self):
        myset = set(['a'])
        pred_handler, repr_string = _get_matcher_parts(myset)
        self.assertTrue(pred_handler('a'))
        self.assertFalse(pred_handler('b'))
        self.assertEqual(repr_string, repr(myset))

    def test_no_special_handling(self):
        self.assertIsNone(_get_matcher_parts(1))
        self.assertIsNone(_get_matcher_parts(0))

class TestMatcherInheritance(unittest.TestCase):
    def test_inheritance(self):
        self.assertTrue(issubclass(MatcherTuple, MatcherBase))
        self.assertTrue(issubclass(MatcherObject, MatcherBase))

class TestGetMatcher(unittest.TestCase):
    def assertIsInstance(self, obj, cls, msg=None):  # New in Python 3.2.
        if not isinstance(obj, cls):
            standardMsg = '%s is not an instance of %r' % (safe_repr(obj), cls)
            self.fail(self._formatMessage(msg, standardMsg))

    def test_single_value(self):
        # Check for MatcherObject wrapping.
        def isodd(x):  # <- Helper function.
            return x % 2 == 1
        matcher = get_matcher(isodd)
        self.assertIsInstance(matcher, MatcherObject)

        # When original is adequate, it should be returned unchanged.
        original = object()
        matcher = get_matcher(original)
        self.assertIs(matcher, original)

    def test_tuple_of_values(self):
        # Check for MatcherTuple wrapping.
        def isodd(x):  # <- Helper function.
            return x % 2 == 1
        matcher = get_matcher((1, isodd))
        self.assertIsInstance(matcher, MatcherTuple)

        # When tuple contains no MatcherObject objects,
        # the original should be returned unchanged.
        original = ('abc', 123)
        matcher = get_matcher(original)
        self.assertIs(matcher, original)

    def test_get_matcher_from_matcher(self):
        original = get_matcher((1, 'abc'))
        matcher = get_matcher(original)
        self.assertIs(matcher, original)

    def test_get_matcher_from_predicate(self):
        predicate = Predicate('abc')
        matcher = get_matcher(predicate)
        self.assertIs(matcher, predicate.matcher)

    def test_integration(self):
        """A small integration test that checks a tuple containing all
        of the different special handling cases.
        """
        def mycallable(x):  # <- Helper function.
            return x == '_'

        myregex = re.compile('_')

        myset = set(['_'])

        matcher = get_matcher(
            (mycallable,  myregex, myset, '_', Ellipsis)
        )

        self.assertTrue(matcher == ('_', '_', '_', '_', '_'))   # <- Passes all conditions.
        self.assertFalse(matcher == ('X', '_', '_', '_', '_'))  # <- Callable returns False.
        self.assertFalse(matcher == ('_', 'X', '_', '_', '_'))  # <- Regex has no match.
        self.assertFalse(matcher == ('_', '_', 'X', '_', '_'))  # <- Not in set.
        self.assertFalse(matcher == ('_', '_', '_', 'X', '_'))  # <- Does not equal string.
        self.assertTrue(matcher == ('_', '_', '_', '_', 'X'))   # <- Passes all conditions (wildcard).

        expected = "(mycallable, re.compile('_'), {0!r}, '_', ...)".format(myset)
        self.assertEqual(repr(matcher), expected)

class TestPredicate(unittest.TestCase):
    def test_predicate_function(self):
        pred = Predicate('abc')
        self.assertTrue(pred('abc'))
        self.assertFalse(pred('def'))
        self.assertFalse(pred(123))

        pred = Predicate(re.compile('^abc$'))
        self.assertTrue(pred('abc'))
        self.assertFalse(pred('def'))
        self.assertFalse(pred(123))

        pred = Predicate(1)
        self.assertTrue(pred(1))
        self.assertFalse(pred(2))
        self.assertFalse(pred('abc'))

        pred = Predicate(('abc', int))
        self.assertTrue(pred(('abc', 1)))
        self.assertFalse(pred(('abc', 1.0)))

        pred = Predicate((str, float('nan')))
        self.assertTrue(pred(('abc', float('nan'))))
        self.assertFalse(pred(('abc', 1.0)))
        self.assertFalse(pred(('abc', 'xyz')))

    @unittest.skipUnless(numpy, 'requires numpy')
    def test_numpy_types(self):
        """Check that built-in types can match numpy types."""
        # Match numpy.character sub-types.
        pred = Predicate(str)
        self.assertTrue(pred(numpy.str_('abc')))
        self.assertTrue(pred(numpy.unicode_('def')))

        # Match numpy.integer sub-types.
        pred = Predicate(int)
        self.assertTrue(pred(numpy.int8(123)))
        self.assertTrue(pred(numpy.uint64(456)))

        # Match numpy.floating sub-types.
        pred = Predicate(float)
        self.assertTrue(pred(numpy.float32(1.0)))
        self.assertTrue(pred(numpy.float64(2.0)))

        # Match numpy.complexfloating sub-types.
        pred = Predicate(complex)
        self.assertTrue(pred(numpy.complex64(1.0)))
        self.assertTrue(pred(numpy.complex128(2.0)))

    @unittest.skipUnless(numpy, 'requires numpy')
    def test_numpy_equality_error(self):
        """Doing `numpy.dtype(float) == 1` raises a TypeError.
        Comparisons that error-out should be caught and considered
        non-matches/False.
        """
        pred = Predicate(numpy.dtype(float))
        self.assertFalse(pred(1))

        pred = Predicate(1)
        self.assertFalse(pred(numpy.dtype(float)))

    def test_inverted_logic(self):
        pred = ~Predicate('abc')
        self.assertFalse(pred('abc'))
        self.assertTrue(pred('def'))

    def test_repr(self):
        pred = Predicate('abc')
        self.assertEqual(repr(pred), "Predicate('abc')")

        pred = ~Predicate('abc')
        self.assertEqual(repr(pred), "~Predicate('abc')")

        pred = Predicate('abc', name='custom_name')
        self.assertEqual(repr(pred), "Predicate('abc', name='custom_name')")

    def test_optional_name(self):
        pred1 = Predicate('abc')  # <- No name arg provided.
        self.assertFalse(hasattr(pred1, '__name__'))

        pred2 = Predicate(pred1)  # <- No name arg inherited from pred1.
        self.assertFalse(hasattr(pred1, '__name__'))

        pred3 = Predicate('abc', name='pred3_name')  # <- Provides name.
        self.assertEqual(pred3.__name__, 'pred3_name')

        pred4 = Predicate(pred3)  # <- Inherits name from pred3.
        self.assertEqual(pred4.__name__, 'pred3_name')

        pred5 = Predicate(pred3, name='pred5_name')  # <- Overrides pred3 name.
        self.assertEqual(pred5.__name__, 'pred5_name')

        # Test bad name values.
        with self.assertRaises(ValueError):
            Predicate('abc', name='1foo')

        with self.assertRaises(ValueError):
            Predicate('abc', name='foo!')

        with self.assertRaises(ValueError):
            Predicate('abc', name='foo bar')

        with self.assertRaises(ValueError):
            Predicate('abc', name='foo()')

        with self.assertRaises(ValueError):
            Predicate('abc', name='foo ')

        with self.assertRaises(ValueError):
            Predicate('abc', name='')

    def test_str(self):
        pred = Predicate('abc')
        self.assertEqual(str(pred), "'abc'")

        inverted = ~Predicate('abc')
        self.assertEqual(str(inverted), "not 'abc'")

    def test_predicate_from_predicate(self):
        pred1 = Predicate('abc')
        pred2 = Predicate(pred1)
        self.assertIsNot(pred1, pred2, msg='should be different object')
        self.assertIs(pred1.obj, pred2.obj, msg='should keep original reference')
        self.assertEqual(pred1.matcher, pred2.matcher)
        self.assertEqual(pred1._inverted, pred2._inverted)

        from_inverted = Predicate(~Predicate('abc'))
        self.assertTrue(from_inverted._inverted)

    @unittest.skipUnless(numpy, 'requires numpy')
    def test_predicate_from_predicate_numpy(self):
        pred1 = Predicate(numpy.int64)
        pred2 = Predicate(pred1)
        self.assertIsNot(pred1, pred2, msg='should be different object')
        self.assertIs(pred1.obj, pred2.obj, msg='should keep original reference')
        self.assertIs(pred1.matcher, pred2.matcher)
        self.assertEqual(pred1._inverted, pred2._inverted)

    def test_passthrough(self):
        """Callable predicates should return the values provided by
        the given function as-is--the values should not be converted
        to True or False.
        """
        TOKEN = object()

        def divisible_or_token(x):  # <- Helper function.
            if x % 3 == 0:
                return True
            if x % 5 == 0:
                return TOKEN
            return False

        predicate = Predicate(divisible_or_token)
        self.assertEqual(predicate(1), False)
        self.assertEqual(predicate(3), True)
        self.assertIs(predicate(5), TOKEN, msg='TOKEN should be returned, not True.')

    def test_equality_inconsistency(self):
        """Badly behaved objects could have inconsistent EQ and NE
        behavior. To make sure that Predicate objects are consistent
        they should only use`==` internally, not `!=`.
        """
        class Inconsistent(object):
            def __init__(self_, value):
                self_.value = value

            def __eq__(self_, other):
                return self_.value == other

            def __ne__(self_, other):
                """Badly behaved not-equals method."""
                return self_.__eq__(other)  # <- DECEPTIVE RESULT!

        # Test `Inconsistent` class, itself.
        obj = Inconsistent(1)
        self.assertTrue(obj == 1, msg='expected behavior')
        self.assertTrue(obj != 1, msg='badly behaved comparison result')
        self.assertFalse(obj == 2, msg='expected behavior')
        self.assertFalse(obj != 2, msg='badly behaved comparison result')

        # Test predicate matching.
        pred = Predicate(Inconsistent(1))
        self.assertTrue(pred(1))
        self.assertFalse(pred(2))

        # Test inverted predicate matching.
        pred = ~Predicate(Inconsistent(1))
        self.assertFalse(pred(1))
        self.assertTrue(pred(2))

    def test_equality_failure(self):
        class BadObj(object):
            def __eq__(self, other):
                if isinstance(other, BadObj):
                    return True
                raise TypeError('Sudden but inevitable betrayal!')

        pred = Predicate(BadObj())
        self.assertFalse(pred(1))

        pred = Predicate(1)
        self.assertFalse(pred(BadObj()))

        pred = ~Predicate(BadObj())  # Check inverted case.
        self.assertTrue(pred(1))

        pred = ~Predicate(1)  # Check inverted case.
        self.assertTrue(pred(BadObj()))

class TestPredicateIntersectionType(unittest.TestCase):
    def setUp(self):
        """Define simple predicates to use for testing."""
        self.pred_gt3 = Predicate(lambda x: x > 3)  # Greater-than three.
        self.pred_even = Predicate(lambda x: x % 2 == 0)  # Is even.

    def test_basics(self):
        pred = PredicateIntersectionType(self.pred_gt3, self.pred_even)
        self.assertFalse(pred(1))
        self.assertFalse(pred(2))
        self.assertFalse(pred(3))
        self.assertTrue(pred(4))
        self.assertFalse(pred(5))
        self.assertTrue(pred(6))
        self.assertFalse(pred(7))

    def test_inverted(self):
        # Using the inversion operator (~).
        inv_pred = ~PredicateIntersectionType(self.pred_gt3, self.pred_even)
        self.assertTrue(inv_pred(1))
        self.assertTrue(inv_pred(2))
        self.assertTrue(inv_pred(3))
        self.assertFalse(inv_pred(4))
        self.assertTrue(inv_pred(5))
        self.assertFalse(inv_pred(6))
        self.assertTrue(inv_pred(7))

    def test_repr(self):
        pred = PredicateIntersectionType(self.pred_gt3, self.pred_even)
        self.assertEqual(repr(pred), 'Predicate(<lambda>) & Predicate(<lambda>)')

        inv_pred = ~pred
        self.assertEqual(repr(inv_pred), '~(Predicate(<lambda>) & Predicate(<lambda>))')

    def test_bad_type(self):
        with self.assertRaises(TypeError):
            PredicateIntersectionType(self.pred_gt3, 'foobarbaz')

        with self.assertRaises(TypeError):
            PredicateIntersectionType('foobarbaz', self.pred_gt3)

    def test_bitwise_operator(self):
        pred = self.pred_gt3 & self.pred_even  # <- Bitwise operator.
        self.assertIsInstance(pred, PredicateIntersectionType)

    def test_bitwise_operator_bad_type(self):
        with self.assertRaises(TypeError):
            self.pred_gt3 & 'foobarbaz'

        with self.assertRaises(TypeError):
            'foobarbaz' & self.pred_gt3

    def test_intersection_method(self):
        pred = self.pred_gt3.intersection(self.pred_even)  # <- Intersection method.
        self.assertIsInstance(pred, PredicateIntersectionType)

        with self.assertRaises(TypeError):
            self.pred_gt3.intersection('foobarbaz')

class TestPredicateUnionType(unittest.TestCase):
    def setUp(self):
        """Define simple predicates to use for testing."""
        self.pred_foo = Predicate('foo')  # Is "foo".
        self.pred_bar = Predicate('bar')  # Is "bar".

    def test_basics(self):
        pred = PredicateUnionType(self.pred_foo, self.pred_bar)
        self.assertTrue(pred('foo'))
        self.assertTrue(pred('bar'))
        self.assertFalse(pred('baz'))

    def test_inverted(self):
        # Using the inversion operator (~).
        inv_pred = ~PredicateUnionType(self.pred_foo, self.pred_bar)
        self.assertFalse(inv_pred('foo'))
        self.assertFalse(inv_pred('bar'))
        self.assertTrue(inv_pred('baz'))

    def test_repr(self):
        pred = PredicateUnionType(self.pred_foo, self.pred_bar)
        self.assertEqual(repr(pred), "Predicate('foo') | Predicate('bar')")

        inv_pred = ~pred
        self.assertEqual(repr(inv_pred), "~(Predicate('foo') | Predicate('bar'))")

    def test_bad_type(self):
        with self.assertRaises(TypeError):
            PredicateIntersectionType(self.pred_foo, 'foobarbaz')

        with self.assertRaises(TypeError):
            PredicateIntersectionType('foobarbaz', self.pred_foo)

    def test_bitwise_operator(self):
        pred = self.pred_foo | self.pred_bar
        self.assertIsInstance(pred, PredicateUnionType)

    def test_bitwise_operator_bad_type(self):
        with self.assertRaises(TypeError):
            self.pred_foo | 'foobarbaz'

        with self.assertRaises(TypeError):
            'foobarbaz' | self.pred_foo

    def test_union_method(self):
        pred = self.pred_foo.union(self.pred_bar)  # <- Union method.
        self.assertIsInstance(pred, PredicateUnionType)

        with self.assertRaises(TypeError):
            self.pred_foo.union('foobarbaz')

class TesCombinedTypeNesting(unittest.TestCase):
    def setUp(self):
        """Define simple predicates to use for testing."""
        self.pred_foo = Predicate('foo')  # Is "foo".
        self.pred_bar = Predicate('bar')  # Is "bar".
        self.pred_baz = Predicate('baz')  # Is "baz".
        self.pred_gt3 = Predicate(lambda x: x > 3)  # Greater-than three.
        self.pred_even = Predicate(lambda x: x % 2 == 0)  # Is even.

    def test_behavior(self):
        pred = self.pred_foo | self.pred_bar | self.pred_baz | (self.pred_gt3 & self.pred_even)
        self.assertFalse(pred(1))
        self.assertFalse(pred(2))
        self.assertFalse(pred(3))
        self.assertTrue(pred(4))
        self.assertFalse(pred(5))
        self.assertTrue(pred(6))
        self.assertFalse(pred(7))
        self.assertTrue(pred('foo'))
        self.assertTrue(pred('bar'))
        self.assertTrue(pred('baz'))
        self.assertFalse(pred('qux'))

    def test_repr(self):
        """Several of the following predicates are logically nonsensical
        but since we're only testing the repr behavior, this is OK.
        """
        pred = self.pred_foo | self.pred_bar | self.pred_baz
        expected = "Predicate('foo') | Predicate('bar') | Predicate('baz')"
        self.assertEqual(repr(pred), expected)

        pred = self.pred_foo & self.pred_bar & self.pred_baz
        expected = "Predicate('foo') & Predicate('bar') & Predicate('baz')"
        self.assertEqual(repr(pred), expected)

        pred = self.pred_foo | self.pred_bar & self.pred_baz  # "&" takes precedence
        expected = "Predicate('foo') | (Predicate('bar') & Predicate('baz'))"
        self.assertEqual(repr(pred), expected)

        pred = (self.pred_foo | self.pred_bar) & self.pred_baz  # Change order
        expected = "(Predicate('foo') | Predicate('bar')) & Predicate('baz')"
        self.assertEqual(repr(pred), expected)

        pred = self.pred_foo & self.pred_bar | self.pred_baz  # "&" takes precedence
        expected = "(Predicate('foo') & Predicate('bar')) | Predicate('baz')"
        self.assertEqual(repr(pred), expected)

        pred = self.pred_foo & (self.pred_bar | self.pred_baz)  # Change order
        expected = "Predicate('foo') & (Predicate('bar') | Predicate('baz'))"
        self.assertEqual(repr(pred), expected)

        pred = self.pred_foo | self.pred_bar | self.pred_gt3 & self.pred_even | self.pred_baz
        expected = "Predicate('foo') | Predicate('bar') | (Predicate(<lambda>) & Predicate(<lambda>)) | Predicate('baz')"
        self.assertEqual(repr(pred), expected)

        pred = self.pred_foo | (self.pred_bar | self.pred_gt3) & self.pred_even | self.pred_baz
        expected = "Predicate('foo') | ((Predicate('bar') | Predicate(<lambda>)) & Predicate(<lambda>)) | Predicate('baz')"
        self.assertEqual(repr(pred), expected)

if __name__ == '__main__':
    unittest.main()


"""
performance test
"""

import numpy as np
import matplotlib.pyplot as plt
import miepy
from tqdm import tqdm
from functools import partial
from scipy.sparse.linalg import bicg, bicgstab
from miepy.interactions import solve_linear_system
from topics.photonic_clusters.create_lattice import hexagonal_lattice_particles
from timer import time_function

nm = 1e-9

Ag = miepy.materials.Ag()
radius = 75*nm
source = miepy.sources.plane_wave.from_string(polarization='rhc')
source = miepy.sources.gaussian_beam(2500*nm, [1, 1j], power=.002)

x = np.linspace(-4000*nm, 4000*nm, 100)
y = np.linspace(-4000*nm, 4000*nm, 100)
source = miepy.sources.grid_interpolate_source(source, [x,y,0])
separation = 600*nm

x = np.arange(-1500*nm, 1500*nm, 50*nm)
y = np.arange(-1500*nm, 1500*nm, 50*nm)
# source = miepy.sources.grid_interp_source(source, grid=(x,y,0))

def tests(Nmax, step=1):
    Nparticles = np.arange(1, Nmax+1, step)
    t_force, t_flux, t_build, t_solve, t_source = [np.zeros_like(Nparticles, dtype=float) for i in range(5)]
    for i,N in enumerate(Nparticles):
        print(N, Nmax)
        # positions = [[n*separation, 0, 0] for n in range(N)]
        positions = hexagonal_lattice_particles(N)*separation
        mie = miepy.sphere_cluster(position=positions,
                                   radius=radius,
                                   material=Ag,
                                   source=source,
                                   wavelength=600*nm,
                                   lmax=2)

        t_force[i] = time_function(mie.force)
        t_flux[i] = time_function(mie.cross_sections)
        t_build[i] = time_function(partial(miepy.interactions.sphere_aggregate_tmatrix,
                              mie.position, mie.mie_scat, mie.material_data.k_b))

        A = miepy.interactions.sphere_aggregate_tmatrix(mie.position, mie.mie_scat, k=mie.material_data.k_b)
        t_solve[i] = time_function(partial(solve_linear_system, A, mie.p_src, method=miepy.solver.bicgstab))

        x = np.linspace(0, N*separation, 1)
        y = 2*radius*np.ones_like(x)
        z = np.zeros_like(x)

        t_source[i] = time_function(mie._solve_source_decomposition)

    fig, ax = plt.subplots()

    ax.plot(Nparticles, t_force*1e3, '-o', label='force')
    ax.plot(Nparticles, t_flux*1e3,'-o',  label='flux')
    ax.plot(Nparticles, t_build*1e3, '-o', label='build')
    ax.plot(Nparticles, t_solve*1e3, '-o', label='solve')
    ax.plot(Nparticles, t_source*1e3, '-o', label='source')

    ax.legend()
    ax.set(xlabel='number of particles', ylabel='runtime (ms)')

    plt.show()

tests(10, step=1)


from lightweaver.fal import Falc82
from lightweaver.rh_atoms import H_6_atom, H_6_CRD_atom, H_3_atom, C_atom, O_atom, OI_ord_atom, Si_atom, Al_atom, CaII_atom, Fe_atom, FeI_atom, He_9_atom, He_atom, He_large_atom, MgII_atom, N_atom, Na_atom, S_atom
import lightweaver as lw
import numpy as np
import scipy.interpolate as interp
from astropy.convolution import Box1DKernel
from astropy.convolution import convolve
from astropy.io import fits
from enum import IntEnum
from mpi4py import MPI
from tqdm import tqdm
import pickle
import argparse

class tags(IntEnum):
    READY = 0
    DONE = 1
    EXIT = 2
    START = 3

def smooth(sig, width):
    return convolve(sig, Box1DKernel(width))

def iterate_ctx_crd(ctx, Nscatter=10, NmaxIter=500):
    for i in range(NmaxIter):
        dJ = ctx.formal_sol_gamma_matrices(verbose=False)
        if i < Nscatter:
            continue
        delta = ctx.stat_equil(printUpdate=False)

        if dJ < 3e-3 and delta < 1e-3:
            # print(i, flush=True)
            # print('----------')
            return

def synth_spectrum(atmos, depthData=False, Nthreads=1, conserveCharge=False):
    atmos.quadrature(5)
    aSet = lw.RadiativeSet([H_6_atom(),
    C_atom(),
     OI_ord_atom(), Si_atom(), Al_atom(),
    CaII_atom(),
    Fe_atom(),
    He_9_atom(),
    MgII_atom(), N_atom(), Na_atom(), S_atom()
    ])
    # aSet.set_active('H', 'Ca')
    aSet.set_active('Ca')
    spect = aSet.compute_wavelength_grid()

    eqPops = aSet.compute_eq_pops(atmos)

    ctx = lw.Context(atmos, spect, eqPops, Nthreads=Nthreads, conserveCharge=conserveCharge)
    if depthData:
        ctx.depthData.fill = True

    iterate_ctx_crd(ctx)
    eqPops.update_lte_atoms_Hmin_pops(atmos, quiet=True)
    ctx.formal_sol_gamma_matrices(verbose=False)
    return ctx

def master_work(filename, write_frequency=1):
    task_index = 0
    num_workers = size - 1
    closed_workers = 0

    fmodel = fits.open('/net/drogon/scratch1/aasensio/3dcubes/Enhanced_network_385_tau_from_RH_01_tau8.fits')
    bifrost = fmodel[0].data[:].astype('<f8').reshape((11, 96, -1))

    n = bifrost.shape[-1]

    log_departure_list = [None] * n
    T_list = [None] * n
    tau_list = [None] * n
    vturb_list = [None] * n
    cmass = [None]
    success = True

    tasks = [0] * n
    pointer = 0

    print("*** Master starting with {0} workers".format(num_workers))
    with tqdm(initial=pointer, total=n, ncols=140) as pbar:
        while closed_workers < num_workers:
            dataReceived = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            source = status.Get_source()
            tag = status.Get_tag()

            if tag == tags.READY:
                # Worker is ready, so send it a task
                try:
                    task_index = tasks.index(0)

                    tau500 = np.ascontiguousarray(10.0**bifrost[0, ::-1, task_index])
                    T = np.ascontiguousarray(bifrost[1, ::-1, task_index])
                    vlos = np.ascontiguousarray(bifrost[5, ::-1, task_index]) / 100.0  # m/s
                    vturb = np.ascontiguousarray(bifrost[3, ::-1, task_index]) / 100.0 # m/s

                    dataToSend = {'index': task_index, 'tau500': tau500, 'T': T, 'vlos': vlos, 'vturb': vturb}
                    comm.send(dataToSend, dest=source, tag=tags.START)

                    tasks[task_index] = 1

                    pbar.set_postfix(sent=f'{task_index}->{source}')

                except:
                    comm.send(None, dest=source, tag=tags.EXIT)

            elif tag == tags.DONE:
                index = dataReceived['index']
                success = dataReceived['success']

                if (not success):
                    tasks[index] = -1
                else:
                    log_departure_list[index] = dataReceived['log_departure']
                    T_list[index] = dataReceived['T']
                    tau_list[index] = dataReceived['tau']
                    vturb_list[index] = dataReceived['vturb']
                    cmass = dataReceived['cmass']

                pbar.update(1)

            elif tag == tags.EXIT:
                print(" * MASTER : worker {0} exited.".format(source))
                closed_workers += 1

            if (pbar.n / write_frequency == pbar.n // write_frequency):

                with open(f'{filename}_logdeparture.pk', 'wb') as filehandle:
                    pickle.dump(log_departure_list[0:task_index], filehandle)

                with open(f'{filename}_T.pk', 'wb') as filehandle:
                    pickle.dump(T_list[0:task_index], filehandle)

                with open(f'{filename}_vturb.pk', 'wb') as filehandle:
                    pickle.dump(vturb_list[0:task_index], filehandle)

                with open(f'{filename}_tau.pk', 'wb') as filehandle:
                    pickle.dump(tau_list[0:task_index], filehandle)

                with open(f'{filename}_cmass.pk', 'wb') as filehandle:
                    pickle.dump(cmass, filehandle)

    print("Master finishing")

    with open(f'{filename}_cmass.pk', 'wb') as filehandle:
        pickle.dump(cmass, filehandle)

    with open(f'{filename}_logdeparture.pk', 'wb') as filehandle:
        pickle.dump(log_departure_list, filehandle)

    with open(f'{filename}_T.pk', 'wb') as filehandle:
        pickle.dump(T_list, filehandle)

    with open(f'{filename}_vturb.pk', 'wb') as filehandle:
        pickle.dump(vturb_list, filehandle)

    with open(f'{filename}_tau.pk', 'wb') as filehandle:
        pickle.dump(tau_list, filehandle)

def slave_work(rank):

    while True:
        comm.send(None, dest=0, tag=tags.READY)
        dataReceived = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)

        tag = status.Get_tag()

        if tag == tags.START:
            # Do the work here
            task_index = dataReceived['index']
            tau500 = dataReceived['tau500']
            T = dataReceived['T']
            vlos = dataReceived['vlos']
            vturb = dataReceived['vturb']

            success = 1

            try:
                atmos = lw.Atmosphere.make_1d(scale=lw.ScaleType.Tau500, depthScale=tau500, temperature=T, vlos=vlos, vturb=vturb, verbose=False)
                ctx = synth_spectrum(atmos, depthData=True, conserveCharge=False)
                tau = atmos.tauRef
                cmass = atmos.cmass
                temperature = atmos.temperature
                vturb = atmos.vturb
                log_departure = np.log10(ctx.activeAtoms[0].n / ctx.activeAtoms[0].nStar)
            except:
                success = 0

            dataToSend = {'index': task_index, 'T': temperature, 'log_departure': log_departure, 'tau': tau, 'cmass': cmass, 'vturb': vturb, 'success': success}

            comm.send(dataToSend, dest=0, tag=tags.DONE)

        elif tag == tags.EXIT:
            break

    comm.send(None, dest=0, tag=tags.EXIT)

if (__name__ == '__main__'):

    # Initializations and preliminaries
    comm = MPI.COMM_WORLD   # get MPI communicator object
    size = comm.size        # total number of processes
    rank = comm.rank        # rank of this process
    status = MPI.Status()   # get MPI status object

    print(f"Node {rank}/{size} active", flush=True)

    if rank == 0:
        parser = argparse.ArgumentParser(description='Generate synthetic models and solve NLTE problem')
        parser.add_argument('--f', '--freq', default=1, type=int, metavar='FREQ', help='Frequency of model write')

        parsed = vars(parser.parse_args())

        master_work('bifrost', write_frequency=parsed['f'])
    else:
        slave_work(rank)


<reponame>jeti182/tvatoj-power<gh_stars>1-10
# MIT License

# Copyright (c) 2020 <NAME>

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from numpy import random, exp, array, zeros, log, clip
from theano.tensor import mean, cast
from theano.tensor import exp as ttexp
from math import floor
from numpy.random import random as runif
from tqdm import tqdm
import pymc3
import pandas as pd
from scipy.optimize import fmin
from scipy.stats import *
from tqdm import tqdm
from scipy.stats import beta
import sys, logging
import warnings

logger = logging.getLogger(__name__)
try:
    import coloredlogs
    coloredlogs.install(level='DEBUG')
except ImportError:
    logging.info('If you like the terminal output colored.' +
                 'Install colored coloredlogs (e.g., pip install coloredlogs)')

# The TVAOJ psychometric function, see Tünnermann, Petersen, & Scharlau (2015):
def tvatoj_psychometric_function(SOA, C, wp, vp=None, vr=None):
    """ Takes SOAs in ms and either C in 1/ms and w or vp and vr in 1/ms """
    if vp is None or vr is None:
        vp = C * wp
        vr = C * (1 - wp)
    SOA = array(SOA)
    left = (1-exp(-vp*abs(SOA))) + exp(-vp*abs(SOA)) * vp/(vp+vr)
    right = exp(-vr*abs(SOA))*vp/(vp+vr)
    return ((SOA<=0)*left  + (SOA>0)*right)

# A generative simulation of the process
def simulate_subject_toj(SOA, reps, C, wp):
    v1 = C * wp                             # attentional weights and overall rate C ...
    v2 = C * (1 -wp)                        # ... determine the individual rates
    probe_first_count = 0                   # Our counter each SOA starts with zero
    for i in range(0, reps):                # For every repetition
        tS = -log(1 - runif(1)) / v2        # let stimulus 2 race and record its VSTM arrival
        tC =  SOA - log(1 - runif(1)) / v1  # sane for stimulus 1, offset by the SOA
        if tC < tS:                         # Did 1 arrive before 2?
            probe_first_count += 1          # Count as a "probe first judment"
    return probe_first_count                # Return the result across all SOAs

# Simulate TOJs for a group of participants, by drawing
# their individual parameters from distributions
def simulate_tojs(simulation_setup):

    s = simulation_setup # For convenient access ...
    single_wp=False

    # Get the paras per individual
    if 'C_sd_within' in s: # within subject design
        logging.info('[SIM] Simulating two different (but correlated) C parameters.')
        C_sub_mu = clip(random.normal(s['C_mu'], s['C_sd_between'], size=s['num_participants']), 0, None)
        C_a_sub =  clip(random.normal(C_sub_mu, s['C_sd_within'], size=s['num_participants']), 0, None)
        C_n_sub =  clip(random.normal(C_sub_mu, s['C_sd_within'], size=s['num_participants']), 0,None)
        wp_a_sub =  clip(random.normal(s['wp_a_mu'], s['wp_a_sd_between'], size=s['num_participants']), 0, None)
        wp_n_sub =  clip(random.normal(s['wp_n_mu'], s['wp_n_sd_between'], size=s['num_participants']), 0, None)
    elif 'C_a_mu' in s: # between design
        logging.info('[SIM] Simulating two independent C parameters.')
        C_a_sub =  clip(random.normal(s['C_a_mu'], s['C_a_sd_between'], size=s['num_participants']), 0, None)
        C_n_sub =  clip(random.normal(s['C_n_mu'], s['C_n_sd_between'], size=s['num_participants']), 0,None)
        wp_a_sub =  clip(random.normal(s['wp_a_mu'], s['wp_a_sd_between'], size=s['num_participants']), 0, None)
        wp_n_sub =  clip(random.normal(s['wp_n_mu'], s['wp_n_sd_between'], size=s['num_participants']), 0, None)
    elif 'C_single_mu' in s:
        logging.info('[SIM] Simulating a single C parameter for both conditions.')
        C_a_sub =  clip(random.normal(s['C_single_mu'], s['C_single_sd_between'], size=s['num_participants']), 0, None)
        C_n_sub = C_a_sub
        wp_a_sub =  clip(random.normal(s['wp_a_mu'], s['wp_a_sd_between'], size=s['num_participants']), 0, None)
        wp_n_sub =  clip(random.normal(s['wp_n_mu'], s['wp_n_sd_between'], size=s['num_participants']), 0, None)
    elif 'wp_mu' in s: # A single wp ==> Single condition experient
        logging.info('[SIM] Simulating a single condition.')
        C_sub =  clip(random.normal(s['C_mu'], s['C_sd_between'], size=s['num_participants']), 0, None)
        wp_sub =  clip(random.normal(s['wp_mu'], s['wp_sd_between'], size=s['num_participants']), 0, None)
        single_wp=True
    else:
        logger.error('Could not infer the design from the simulation parameters provided. Please refer to the exmaples')
        sys.exit('Aborting')

    # Get the TOJs
    participant_id = []
    condition_id = []
    probe_first_count = []
    repetitions = []
    SOA = []

    if single_wp:
        condition_nums = [0]
    else:
        condition_nums = [0, 1]

    for p in range(0, s['num_participants']):
        for i,soa in enumerate(s['SOAs']):
            for c in condition_nums:
                participant_id.append(p)
                condition_id.append(c)
                SOA.append(soa)
                repetitions.append(s['repetitions'][i])
                if c == 0 and not single_wp: # simulate a neutral condition TOJ
                    probe_first_count.append(simulate_subject_toj(soa, s['repetitions'][i], C_n_sub[p], wp_n_sub[p]))
                if c == 1 and not single_wp: # simulate an attention condition TOJ
                    probe_first_count.append(simulate_subject_toj(soa, s['repetitions'][i], C_a_sub[p], wp_a_sub[p]))
                if c == 0 and single_wp:
                    probe_first_count.append(simulate_subject_toj(soa, s['repetitions'][i], C_sub[p], wp_sub[p]))

    df = pd.DataFrame()
    df['participant_id'] = participant_id
    df['condition_id'] = condition_id
    df['SOA'] = SOA
    df['repetitions'] = repetitions
    df['probe_first_count'] = probe_first_count

    return df

# Using the non-centered reparamtrization to reduce divergenses
# See here for the rationale: https://twiecki.io/blog/2017/02/08/bayesian-hierchical-non-centered/
def hierarchical_model_noncentered(data, single_C=False, single_wp=False):
    '''Sets up a pymc3 model based on TVATOJ.

    :param data: A TOJ dataframe as return by the simulations
    :param single_C: Whether to use single C (for both conditions)
    :param single_wp: Whether to use a single wp (implies single C and produces a model for a single condition only)

    :returns: Model
    :rtype: pymc3.Model
    '''

    model = pymc3.Model()
    with model:

        p_id = data['participant_id']
        c_id = data['condition_id']

        if single_wp:
            wp_c_id = len(data['condition_id']) * [0]
            single_C = True
        else:
            wp_c_id = c_id

        if single_C:
            C_c_id = len(data['condition_id']) * [0]
        else:
            C_c_id = c_id

        pfc =  pymc3.Data('probe_first_count',data['probe_first_count'])

        C_mu = pymc3.Normal('C_mu', 0.080, 0.050, shape=len(set(C_c_id)))
        C_sd = pymc3.HalfCauchy('C_sd', 0.1, shape=len(set(C_c_id)))

        wp_mu = pymc3.Normal('wp_mu', 0.5,0.2, shape=len(set(wp_c_id)))
        wp_sd = pymc3.HalfCauchy('wp_sd', 0.2, shape=len(set(wp_c_id)))

        wp_e = pymc3.Normal('wp_e', 0,1, shape=(len(set(p_id)), len(set(wp_c_id))))
        C_e = pymc3.Normal('C_e', 0,1, shape=(len(set(p_id)), len(set(C_c_id))))

        C = pymc3.Deterministic('C', (C_mu + C_e * C_sd).clip(0.0001, 0.9999))
        wp = pymc3.Deterministic('wp',  (wp_mu + wp_e * wp_sd).clip(0.0001, 0.9999))

        theta = pymc3.Deterministic('theta', tvatoj_psychometric_function(
            data['SOA'], C[(p_id, C_c_id)], wp[(p_id, wp_c_id)]))

        y = pymc3.Binomial('y', n=cast(data['repetitions'], 'int64'),
                                p=theta, observed=pfc,
                                dtype='int64')

        # The deterministic transformation could probably be externalized
        # However, here the calculation is most safe to produce prober within subject estimates
        vp = pymc3.Deterministic('vp', wp * C)
        vr = pymc3.Deterministic('vr', (1 - wp) * C)

        vp_mean = pymc3.Deterministic('vp_mean', mean(vp, axis=0))
        vr_mean = pymc3.Deterministic('vr_mean', mean(vr, axis=0))
        if not single_wp:
            va_diff_mean = pymc3.Deterministic('va_diff_mean', mean(vp[:,1] - vr[:,1])) # Diff of probe and ref rate in the attention cond
            vp_diff_mean = pymc3.Deterministic('vp_diff_mean', mean(vp[:,1] - vp[:,0])) # Diff of attention and neutral condition probe rates
            vr_diff_mean = pymc3.Deterministic('vr_diff_mean', mean(vr[:,1] - vr[:,0])) # Diff of attention and neutral condition probe rates
            wpa_mean = pymc3.Deterministic('wpa_mean', mean(wp[:,1]))
            wp_diff_mean = pymc3.Deterministic('wp_diff_mean', mean(wp[:,1] - wp[:,0]))
        else:
            wp_vs_point5_mean = pymc3.Deterministic('wp_mean', mean(wp))
        return(model)

# This function is borrowed from @aloctavodia, who ported it from <NAME>'s scripts
# https://github.com/aloctavodia/Doing_bayesian_data_analysis/blob/master/HDIofICDF.py
def HDIofICDF(dist_name, credMass=0.95, **args):
    # freeze distribution with given arguments
    distri = dist_name(**args)
    # initial guess for HDIlowTailPr
    incredMass =  1.0 - credMass

    def intervalWidth(lowTailPr):
        return distri.ppf(credMass + lowTailPr) - distri.ppf(lowTailPr)

    # find lowTailPr that minimizes intervalWidth
    HDIlowTailPr = fmin(intervalWidth, incredMass, ftol=1e-8, disp=False)[0]
    # return interval as array([low, high])
    return distri.ppf([HDIlowTailPr, credMass + HDIlowTailPr])

def sim_and_fit(setup, model_func, iterations, condition_func,
                goal_var_names=None, log_var_names=['C_mu', 'wp_mu'],
                single_C=False, single_wp=False, outfile='out.csv',
                turn_off_warnings=True,
                tune=1000,
                target_accept=0.85,
                init='adapt_diag'):

    if (turn_off_warnings):
        warnings.filterwarnings("ignore")
        logging.warning('Attention: Warnings turned off. ') # There is so much from pymc3 and theano ..

    if log_var_names==None or len(log_var_names) < 1:
        sys.exit('log_var_names should not be empty or None! Log at least one variable!')
    num_success=0
    model = None
    for i in tqdm(range(iterations), desc='Overall progress'):
        data = simulate_tojs(setup)
        if model is None:
            model = model_func(data, single_C=single_C, single_wp=single_wp)
        with model:
            pymc3.set_data({'probe_first_count': data['probe_first_count']})
            trace = pymc3.sample(2000, tune=tune, cores=4, init=init, target_accept=target_accept)
            summary_stats = pymc3.summary(trace, var_names=goal_var_names, hdi_prob=0.95)
            print(summary_stats)
        success = condition_func(summary_stats) * 1 # Either 0 or 1, depending on reaching our goals.
        num_success += success
        attempts = (i+1)
        success_rate = num_success / attempts
        hdi = HDIofICDF(beta,a=1+num_success, b=1+(attempts-num_success))
        logging.info(('[ESTIMATE] Success rate: %.2f' % success_rate +
                     ' [95 %% HDI: %.2f to %.2f]' % (hdi[0],hdi[1]) +
                     '\n' + '-'* 20))

        out_df = pymc3.summary(trace, var_names=log_var_names, hdi_prob=0.95)
        out_df.insert(0, 'iteration', attempts)
        out_df.insert(1, 'success', success)
        out_df.insert(2, 'power_est', success_rate)
        out_df.insert(3, 'power_hdi_2.5%', hdi[0])
        out_df.insert(4, 'power_hdi_97.5%', hdi[1])
        if attempts == 1:
            out_df.to_csv(outfile)
        else:
            out_df.to_csv(outfile, mode='a', header=False)

'''
Convenience function to fit with logging.
'''
def fit(model, outfile='fit.csv'):
    with model:
        trace = pymc3.sample(2000, tune=1000, cores=4, init='adapt_diag') #, target_accept=.85)
        summary_stats = pymc3.summary(trace, hdi_prob=0.95)
        summary_stats.to_csv(outfile)
    logger.info('The model was fitted and a summary was written to: ' + outfile)
    logger.info('You can analyze the returned trace with help of the Arviz library (https://arviz-devs.github.io/arviz/)')
    logger.info('For instance, plot parameter posteriors with arviz.plot_posterior(trace, var_names=["C", "wp"])')
    return trace


# Hysteresis model
# http://eprints.lancs.ac.uk/1375/1/MFI_10c.pdf
# Identification of Hysteresis Functions Using a Multiple Model Approach
# Mihaylova, Lampaert et al

import numpy as npy
from scipy.optimize import root
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import copy

#%%

plt.close('all')

class HysteresisModel:
    """
    Hysteresis model comprising a number of elementary Maxwell-slip models
    refer http://eprints.lancs.ac.uk/1375/1/MFI_10c.pdf
    """

    def __init__(self,N,K,W=None,delta=None):

        self.N = N
        """
        Integer number of elementary models
        """

        self.K = npy.ravel(npy.abs(K))
        """
        Array of stiffness values for each elementary model
        """

        if delta is None :
            # K and W specified

            if W is None:
                raise ValueError("Error: either delta or W arguments "+
                                 "must be provided!")
            else:
                W = npy.ravel(npy.abs(W))          # limiting friction values

        else:
            # K and delta specified
            # W to be inferred, given this input
            delta = npy.abs(npy.ravel(delta))
            W = self.K * delta

        self.W = W
        """
        Array of limiting friction values for each elementary model
        """

        # Initialise matrices F and C, which do not vary with input
        self.F = npy.asmatrix(npy.identity(self.N))
        self.C = npy.asmatrix(npy.diag(-self.K))

        # Initialise matrices G and D, as empty
        self.G = npy.asmatrix(npy.empty((self.N,1)))
        self.D = npy.asmatrix(npy.empty((self.N,1)))

        # Initialise array to contain case indexs
        self.case = npy.zeros((self.N,),dtype=int)

    @property
    def x0(self):
        return self._x0

    @x0.setter
    def x0(self,x0):
        """
        Set initial states
        """
        self._x0 = npy.asmatrix(npy.ravel(x0)).T
        self.x =self.x0

        if self.x.shape[0] != self.N:
            raise ValueError("Error: x0 wrong shape!")

    @property
    def x(self):
        return self._x

    @x.setter
    def x(self,val):
        #print("States updated")
        self._x = npy.asmatrix(val)

    def update(self,u,save_states=True):
        """
        Function to advance state-space description of model dynamics
        by a single time step, returning next state and output vectors
        """

        x = copy.deepcopy(self.x)

        # Define G and D matrix entries
        for i in range(self.N): # loop over all elementary models

            Wi = self.W[i]
            Ki = self.K[i]

            # Evaluate switching parameter
            fi = Ki * (u - x[i])

            if fi > Wi:
                # Case 2
                self.case[i] = 2
                self.G[i] = 1
                self.D[i] = 0
                x[i] = -Wi/Ki

            elif fi < -Wi:
                # Case 3
                self.case[i] = 3
                self.G[i] = 1
                self.D[i] = 0
                x[i] = +Wi/Ki

            else:
                # Case 1
                self.case[i] = 1
                self.G[i] = 0
                self.D[i] = Ki

        # Compute next states and output
        # using eqns (10) and (11) in Mihaylova's paper
        x_next = self.F * x + self.G * u
        y_k = self.C * x + self.D * u

        Fh_k = y_k.sum() # total hysteresis force

        # Update states
        if save_states:
            self.x = x_next

        return x_next, y_k, Fh_k

    def run(self,x0,uVals):
        """
        Run simulation from initial conditions, given inputs u
            x0    : column vector [Nx1]
            u     : list or vector of length (nSteps,1)
        """

        # Convert and check shape of u
        uVals = npy.ravel(uVals)
        nSteps = uVals.shape[0]

        # Initialise state space eqns
        self.x0 = x0

        # Step through state space eqns
        xVals = npy.zeros((nSteps,self.N))
        yVals = npy.zeros((nSteps,self.N))
        Fh_vals = npy.zeros((nSteps,))

        for k, u_k in enumerate(uVals):

            # Get next states and output
            x_k, y_k, Fh_k = self.update(u_k)

            # Store
            xVals[k,:] = npy.ravel(x_k)
            yVals[k,:] = y_k.T
            Fh_vals[k] = Fh_k

        # Store results
        self.uVals = uVals
        self.xVals = xVals
        self.yVals = yVals
        self.FhVals = Fh_vals

        # Return states and output for each step
        return xVals, yVals, Fh_vals

    def write_results(self,
                      fname='results.csv',
                      delimiter=','):

        arr = npy.asmatrix(self.uVals).T
        titles = ["u"]
        N = self.N

        arr = npy.hstack((arr,self.xVals))
        titles += ["x%d" % (i+1) for i in range(N)]

        arr = npy.hstack((arr,self.yVals))
        titles += ["y%d" % (i+1) for i in range(N)]

        arr = npy.hstack((arr,npy.asmatrix(self.FhVals).T))
        titles += ["Fh"]

        npy.savetxt(fname=fname,
                    X=arr,
                    delimiter=delimiter,
                    header=delimiter.join(str(x) for x in titles))

    def PlotResults_timeSeries(self,tVals):
        """
        Plot results as time series
        [t,u], [t,x], [t,y], [t,Fh]
        """

        fig, axarr = plt.subplots(4,sharex=True)
        fig.set_size_inches(16,9,forward=True)

        ax1 = axarr[0]
        ax1.plot(tVals,self.uVals)
        ax1.xaxis.set_visible(False)
        ax1.set_ylabel("u")
        ax1.set_xlabel("Input displacement, u(t)")

        ax2 = axarr[1]
        ax2.plot(tVals,self.xVals)
        ax2.xaxis.set_visible(False)
        ax2.set_ylabel("x")
        ax2.set_title("States of\nelementary models, x(t)")

        ax3 = axarr[2]
        ax3.plot(tVals,self.yVals)
        ax3.xaxis.set_visible(False)
        ax3.set_ylabel("y")
        ax3.set_title("Outputs from\nelementary models, y(t)")

        ax4 = axarr[3]
        ax4.plot(tVals,self.FhVals)
        ax4.set_xlabel("Time (seconds)")
        ax4.set_ylabel("F$_h$")
        ax4.set_title("Net output F$_h$")

    def PlotResults(self):
        """
        Plot results as [u,x], [u,y], [u,Fh] plots
        """

        fig, axarr = plt.subplots(1,3,sharex=True)
        fig.set_size_inches(16,9,forward=True)

        ax1 = axarr[0]
        ax1.plot(self.uVals,self.xVals)
        ax1.set_xlabel("Input u")
        ax1.set_title("States of\nelementary models, x")

        ax2 = axarr[1]

        ax2.plot(self.uVals,self.yVals)
        ax2.set_xlabel("Slip (u-x)")
        ax2.set_title("Outputs from\nelementary models, y")

        ax3 = axarr[2]
        ax3.plot(self.uVals,self.FhVals)
        ax3.set_xlabel("Input u")
        ax3.set_title("Net output F$_h$")

class static_response():
    """
    Class used to compute response to forcing input
    """

    def __init__(self,hys_obj,K1, K2):
        self.hys_obj = hys_obj
        self.K1 = K1
        self.K2 = K2

    def net_force(self,d,F_ext,verbose=False):
        """
        Function which defines net force
        given position 'u' and external force 'F_ext'
        """

        u = d[0] - d[1]  # relative displacement at friction interface
        F_hys = self.hys_obj.update(u=u,save_states=False)[2]

        F_net_1 = self.K1 * d[0] + F_hys - F_ext
        F_net_2 = self.K2 * d[1] - F_hys
        F_net = npy.array([F_net_1,F_net_2])

        if verbose:
            print("u = %.3e" % u)
            print("x = {0}".format(self.hys_obj.x))
            print("F_hys = {0}".format(F_hys))
            print("F_net = {0}".format(F_net))

        return F_net

    def run(self,F_vals,x0=None,d0=None):

        # Define function to solve for next u
        def solve(d_last,F_k,hys_obj):

            # Determine next u to satify equilibrium - i.e. zero net force
            sol = root(fun=self.net_force,x0=d_last,args=(F_k,))
            d_k = sol.x
            u_k = d_k[0]-d_k[1]

            F_net = self.net_force(d_k,F_k)

            if not sol.success:
                pass#print(sol.message)

            x_k, y_k, F_hys_k = hys_obj.update(u=u_k,save_states=True)

            return F_hys_k, d_k, u_k, x_k, y_k, F_net

        # Set initial conditions
        if x0 is None:
            x0 = npy.zeros((self.hys_obj.N,))
        self.hys_obj.x0 = x0

        if d0 is None:
            d0 = npy.array([0.0,0.0])
        d_j = d0 # initial guess

        # Run step by step
        F_hys_vals = []
        x_vals = []
        u_vals = []
        y_vals = []
        F_net_vals = []

        for j, F_j in enumerate(F_vals):

            #print("--- Step #%d ---" % j)
            F_hys_j, d_j, u_j, x_j, y_j, F_net = solve(d_j,F_j,self.hys_obj)

            F_hys_vals.append(F_hys_j)
            x_vals.append(npy.ravel(x_j))
            y_vals.append(npy.ravel(y_j))
            u_vals.append(u_j)
            F_net_vals.append(F_net)

        self.x_vals = x_vals
        self.y_vals = y_vals
        self.u_vals = u_vals
        self.F_hys_vals = F_hys_vals
        self.F_vals = F_vals
        self.F_net_vals = F_net_vals

    def plot(self):

        fig, axarr = plt.subplots(3,2,sharex='col')
        fig.set_size_inches(14,8)

        ax = axarr[0,0]
        ax.plot(self.F_vals,label='$F_{external}$')
        ax.plot(self.F_hys_vals,label='$F_{hysteresis}$')
        ax.legend()
        ax.set_ylabel("Forces")

        ax = axarr[1,0]
        ax.plot(self.u_vals)
        ax.set_ylabel("Displacement, u")

        ax = axarr[2,0]
        ax.plot(self.x_vals)
        ax.set_xlabel("Step index")
        ax.set_ylabel("States, x")

        ax = axarr[0,1]
        ax.plot(self.u_vals,self.y_vals)
        ax.set_ylabel("Outputs, y")

        ax = axarr[1,1]
        ax.plot(self.u_vals,self.F_hys_vals)
        ax.set_ylabel("$F_{hysteresis}$")

        ax = axarr[2,1]
        ax.plot(self.u_vals,self.F_vals)
        ax.set_xlabel("Displacement, u")
        ax.set_ylabel("$F_{external}$")

        return fig

# -------- TEST ROUTINE ----------

if __name__ == "__main__":

    test_routine = 1

    if test_routine == 0:

        # Define hysteresis model
        K = [1000,2000,3000]
        delta = [1,2,3]
        Ne = len(K)
        hys = HysteresisModel(Ne,K,delta=delta)

        # Define displacement inputs
        dt = 0.02
        tmax = 10
        u0 = 10

        import random

        def randomWalk(N,normalise=True):

            x= [0]

            for j in range(N-1):
                step_x = random.randint(0,1)
                if step_x == 1:
                    x.append(x[j] + 1 + 0.05*npy.random.normal())
                else:
                    x.append(x[j] - 1 + 0.05*npy.random.normal())

            x = npy.asarray(x)

            if normalise:
                absmaxVal = npy.max([npy.max(x),-npy.min(x)])
                x = x / absmaxVal

            return x

        tVals = npy.arange(0,tmax,dt)
        uVals = u0*randomWalk(tVals.shape[0])

        #uVals = 4.5*npy.sin(2*npy.pi*0.5*tVals)

        # Obtain states and outputs by state space stepping
        hys.run(npy.zeros((Ne,)),uVals)

        # Plot results
        hys.PlotResults()
        hys.PlotResults_timeSeries(tVals)

        #hys.write_results()

    elif test_routine==1:

        # Define hysteresis model
        K = [1000,2000,3000]
        W = [1000,1000,1000]
        Ne = len(K)
        hys = HysteresisModel(Ne,K,W=W)

        # Define force function
        # Define displacement inputs
        dt = 0.02
        tmax = 10
        u0 = 10
        F0 = 3000

        t_vals = npy.arange(0,tmax,dt)
        F_vals = F0 * (npy.sin(2*npy.pi*t_vals) + npy.sin(2*npy.pi*3.2*t_vals))

        # Define spring
        K_spring = 1500

        # Define and run analysis
        analysis = static_response(hys_obj=hys,K_spring=K_spring)
        analysis.run(F_vals=F_vals)
        analysis.plot()

    else:
        raise ValueError("No test selected!")

#%%

#


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue May 11 09:37:22 2021

@author: <NAME>
"""
'''
Complexity = Hard
'''
'''
Given two sorted arrays nums1 and nums2 of size m and n respectively, return the median of the two sorted arrays.

The overall run time complexity should be O(log (m+n)).

Example 1:

Input: nums1 = [1,3], nums2 = [2]
Output: 2.00000
Explanation: merged array = [1,2,3] and median is 2.
Example 2:

Input: nums1 = [1,2], nums2 = [3,4]
Output: 2.50000
Explanation: merged array = [1,2,3,4] and median is (2 + 3) / 2 = 2.5.
Example 3:

Input: nums1 = [0,0], nums2 = [0,0]
Output: 0.00000
Example 4:

Input: nums1 = [], nums2 = [1]
Output: 1.00000
Example 5:

Input: nums1 = [2], nums2 = []
Output: 2.00000

Constraints:

nums1.length == m
nums2.length == n
0 <= m <= 1000
0 <= n <= 1000
1 <= m + n <= 2000
-106 <= nums1[i], nums2[i] <= 106
Accepted
937,027
Submissions
2,961,335
'''

'''
Results of below solution:

Runtime: 92 ms, faster than 64.31% of Python3 online submissions for Median of Two Sorted Arrays.
Memory Usage: 14.4 MB, less than 78.54% of Python3 online submissions for Median of Two Sorted Arrays.
'''
class Solution:
    def findMedianSortedArrays(self, nums1, nums2) -> float:
        array=sorted(nums1+nums2)
        import statistics
        return statistics.median(array)

y=Solution()
nums1 = [1,3]
nums2 = [2]
print(y.findMedianSortedArrays(nums1, nums2))
nums1 = [1,2]
nums2 = [3,4]
print(y.findMedianSortedArrays(nums1, nums2))
nums1 = [0,0]
nums2 = [0,0]
print(y.findMedianSortedArrays(nums1, nums2))
nums1 = []
nums2 = [1]
print(y.findMedianSortedArrays(nums1, nums2))
nums1 = [2]
nums2 = []
print(y.findMedianSortedArrays(nums1, nums2))


# -*- coding:utf-8 -*-
__author__ = 'yangjian'
"""

"""

from collections import defaultdict

import numpy as np
import pandas as pd
from scipy.cluster import hierarchy
from scipy.stats import spearmanr
from sklearn.impute import SimpleImputer

from hypernets.core import randint
from hypernets.tabular import sklearn_ex as skex, dask_ex as dex, column_selector as cs
from hypernets.tabular.cfg import TabularCfg as cfg
from hypernets.utils import logging

logger = logging.get_logger(__name__)

def select_by_multicollinearity(X, method=None):
    """
    Adapted from https://scikit-learn.org/stable/auto_examples/inspection/plot_permutation_importance_multicollinear.html
    handling multicollinearity is by performing hierarchical clustering on the features’ Spearman
    rank-order correlations, picking a threshold, and keeping a single feature from each cluster.
    """
    X_shape = X.shape
    if dex.is_dask_dataframe(X):
        X_shape = dex.compute(X_shape)[0]
    sample_limit = cfg.multi_collinearity_sample_limit
    if X_shape[0] > sample_limit:
        logger.info(f'{X_shape[0]} rows data found, sample to {sample_limit}')
        frac = sample_limit / X_shape[0]
        X, _, = dex.train_test_split(X, train_size=frac, random_state=randint())

    n_values = [X[c].value_counts() for c in X.columns]
    if dex.is_dask_dataframe(X):
        n_values = dex.compute(*n_values)
    one_values = [n.name for n in n_values if len(n) <= 1]
    if len(one_values) > 0:
        X = X[[c for c in X.columns if c not in one_values]]

    logger.info('computing correlation')
    if (method is None or method == 'spearman') and isinstance(X, pd.DataFrame):
        Xt = SimpleImputer(missing_values=np.nan, strategy='most_frequent').fit_transform(X)
        corr = spearmanr(Xt).correlation
    elif isinstance(X, pd.DataFrame):
        Xt = X.copy()
        cols = cs.column_number_exclude_timedelta(X)
        if cols:
            Xt[cols] = SimpleImputer(missing_values=np.nan, strategy='most_frequent').fit_transform(Xt[cols])
        Xt = skex.SafeOrdinalEncoder().fit_transform(Xt)
        corr = Xt.corr(method=method).values
    else:  # dask
        Xt = dex.SafeOrdinalEncoder().fit_transform(X)
        corr = Xt.corr(method='pearson' if method is None else method).compute().values

    logger.info('computing cluster')
    corr_linkage = hierarchy.ward(corr)
    cluster_ids = hierarchy.fcluster(corr_linkage, 1, criterion='distance')
    cluster_id_to_feature_ids = defaultdict(list)
    for idx, cluster_id in enumerate(cluster_ids):
        cluster_id_to_feature_ids[cluster_id].append(idx)
    selected = [X.columns[v[0]] for v in cluster_id_to_feature_ids.values()]
    unselected = list(set(X.columns.to_list()) - set(selected)) + one_values
    feature_clusters = [[X.columns[i] for i in v] for v in cluster_id_to_feature_ids.values()]
    return feature_clusters, selected, unselected


<reponame>DeepK/hoDMD-experiments
from scipy.fftpack import dct
from traintestutil import train_test
import numpy
from functools import partial
import argparse

def dct_keep(to_keep, vectors):
  return numpy.hstack(dct(vectors)[:to_keep, :])

if __name__ == "__main__":

  parser = argparse.ArgumentParser(description='Evaluate DCT based sentence embedding')
  parser.add_argument("pickled_training_data_path", help="pickled train path")
  parser.add_argument("pickled_test_data_path", help="pickled test path")
  parser.add_argument("DCT_components_to_keep", help="DCT components to keep", type = int)

  args = parser.parse_args()

  pickled_training_data_path = args.pickled_training_data_path
  pickled_test_data_path = args.pickled_test_data_path
  DCT_components_to_keep = args.DCT_components_to_keep

  print ("DCT components to keep -> %s"%DCT_components_to_keep)

  func_to_apply = partial(dct_keep, DCT_components_to_keep)

  filter_by_len = DCT_components_to_keep
  results = train_test(pickled_training_data_path, pickled_test_data_path, func_to_apply, filter_by_len)
  results = results.split("\n")[-2]
  print (results)


'''
Copyright(C) 2016 Engineering Department, University of Cambridge, UK.

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

Author
    <NAME> <<EMAIL>>
'''

import os
import subprocess
import numpy as np
import scipy.interpolate
from scipy import signal as sig

from lib import sigproc as sp
from . import fileio

def resample(wav, fs, trgfs, method=2, deterministic=True):
    '''
        deterministic [True] : Try to make it deterministic.
                               (e.g. sox (mehtod=2) is not deterministic by default)
                    ATTENTION This option has been tested only for method==2
    '''
    if method==1:
        # sndfile-resample (libresample)
        # 'c' argument
        #0 : Best Sinc Interpolator
        #1 : Medium Sinc Interpolator (default)
        #2 : Fastest Sinc Interpolator TO AVOID
        #3 : ZOH Interpolator TO AVOID
        #4 : Linear Interpolator TO AVOID
        # sndfile-resample _seems_ to be always deterministic

        tmpinfname = sp.gentmpfile('sndfile-resample-in.wav')
        tmpoutfname = sp.gentmpfile('sndfile-resample-out.wav')

        try:
            wavwrite(tmpinfname, wav, fs)

            cmd = 'sndfile-resample -c 0 -to '+str(trgfs)+' '+tmpinfname+' '+tmpoutfname
            out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
            #print(out)

            syn, synfs, synenc = wavread(tmpoutfname)
        except:
            if os.path.exists(tmpinfname): os.remove(tmpinfname)
            if os.path.exists(tmpoutfname): os.remove(tmpoutfname)
            raise

        if os.path.exists(tmpinfname): os.remove(tmpinfname)
        if os.path.exists(tmpoutfname): os.remove(tmpoutfname)

    elif method==2:
        # SOX
        # VHQ: -v -s: The fastest with the results among the bests
        # ATTENTION:If deterministic=False, sox is NOT deterministic!
        #           I.e. it does NOT produce the same samples for each run!

        tmpinfname = sp.gentmpfile('sox-resample-in.wav')
        tmpoutfname = sp.gentmpfile('sox-resample-out.wav')

        try:
            fileio.wavwrite(tmpinfname, wav, fs)

            cmd = 'sox '
            if deterministic: cmd += '--no-dither '
            cmd += tmpinfname+' '+tmpoutfname+' rate -v -s '+str(trgfs)
            out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
            #print(out)

            syn, synfs, synenc = fileio.wavread(tmpoutfname)
        except:
            if os.path.exists(tmpinfname): os.remove(tmpinfname)
            if os.path.exists(tmpoutfname): os.remove(tmpoutfname)
            raise

        if os.path.exists(tmpinfname): os.remove(tmpinfname)
        if os.path.exists(tmpoutfname): os.remove(tmpoutfname)

    elif method==3:
        '''
        Resample using FFT and power of 2
        Create sometimes a significant peak at Nyquist
        '''
        syn = wav.copy()
        wavlen = syn.shape[0]
        wavlenpow2  = int(np.power(2, np.floor(np.log2(wavlen))+1))
        syn = np.pad(syn, (0, wavlenpow2-wavlen), constant_values=(0,0), mode='constant')
        syn = scipy.signal.resample(syn, np.round(len(syn)*float(trgfs)/fs))
        syn = syn[:np.round(wavlen*float(trgfs)/fs)]

    if 0:
        import matplotlib.pyplot as plt
        plt.ion()
        plt.plot(np.arange(len(wav))/float(fs), wav, 'k')
        plt.plot(np.arange(len(syn))/float(trgfs), syn, 'b')
        from IPython.core.debugger import  Pdb; Pdb().set_trace()

    return syn

# Resample feature using the nearest
def featureresample(ts, X, nts):
    if len(X.shape)>1:
        Y = np.zeros((len(nts), X.shape[1]))
    else:
        Y = np.zeros(len(nts))
    for n, t in enumerate(nts):
        idx = np.where(ts>=t)[0]
        if len(idx)==0:
            idx = X.shape[0]-1
        else:
            idx = np.min(idx) # Nearest
        idx = np.clip(idx, 0, X.shape[0]-1)
        if len(X.shape)>1:
            Y[n,:] = X[idx,:]
        else:
            Y[n] = X[idx]
    return Y

def f0s_resample_pitchsync(f0s, nbperperiod, f0min=20.0, f0max=5000.0):
    f0s = f0s.copy()

    # Interpolate where there is zero values
    f0s[:,1] = np.interp(f0s[:,0], f0s[f0s[:,1]>0,0], f0s[f0s[:,1]>0,1])

    f0s[:,1] = np.clip(f0s[:,1], f0min, f0max)

    ts = [0.0]
    while ts[-1]<f0s[-1,0]:
        cf0 = np.interp(ts[-1], f0s[:,0], f0s[:,1])
        ts.append(ts[-1]+(1.0/nbperperiod)/cf0)
    f0s = np.vstack((ts, np.interp(ts, f0s[:,0], f0s[:,1]))).T

    return f0s

def f0s_resample_cst(f0s, timeshift):
    f0s = f0s.copy()

    vcs = f0s.copy()
    vcs[vcs[:,1]>0,1] = 1

    nts = np.arange(f0s[0,0], f0s[-1,0], timeshift)

    # The voicing resampling has to be done using nearest ...
    vcsfn = scipy.interpolate.interp1d(vcs[:,0], vcs[:,1], kind='nearest', bounds_error=False, fill_value=0)

    # ... whereas the frequency resampling need linear interpolation, while ignoring the voicing
    f0s = np.interp(nts, f0s[f0s[:,1]>0,0], f0s[f0s[:,1]>0,1])

    # Put back the voicing
    f0s[vcsfn(nts)==0] = 0.0

    f0s = np.vstack((nts, f0s)).T

    if 0:
        plt.plot(f0s[:,0], f0s[:,1])

    return f0s

def f0s_rmsteps(f0s):
    '''
    Description
        Removes steps in the F0 curve.

        Steps can come from some F0 estimator (e.g. those based on GCI
        detection are likely to exhibits these).

        For pulse synthesis, it avoids some glitches around the main lobes

        It might be bad for creaky voice (oversmoothing the f0 curve),
        though F0 estimate in creaky voice is quite likely to be wrong anyway.
    '''
    f0sori = f0s.copy()
    f0s = f0s.copy()
    voicedi = np.where(f0s[:,1]>0)[0]
    shift = np.mean(np.diff(f0s[:,0]))
    fc = (1.0/shift)/4.0  # The cut-off frequency
    hshift = (1.0/fc)/8.0 # The high sampling rate for resampling the original curve
    data = np.interp(np.arange(0.0, f0s[-1,0], hshift), f0s[voicedi,0], f0s[voicedi,1])
    b, a = sig.butter(8, fc/(0.5/hshift), btype='low')
    f0ss = sig.filtfilt(b, a, data)
    f0s[voicedi,1] = np.interp(f0s[voicedi,0], np.arange(0.0, f0s[-1,0], hshift), f0ss)

    if 0:
        plt.plot(f0sori[:,0], f0sori[:,1], 'k')
        plt.plot(f0s[:,0], f0s[:,1], 'b')
        from IPython.core.debugger import  Pdb; Pdb().set_trace()

    return f0s


from pathlib import Path
import os
import matplotlib as mpl
if os.environ.get('DISPLAY') is None:  # NOQA
    mpl.use('Agg')  # NOQA
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib
import itertools
import numpy as np

# from pylab import rcParams
# rcParams['figure.figsize'] = (12, 12)
fontsize = 12
params = {
    'figure.figsize': (12, 12),
    'axes.labelsize': fontsize,
    'axes.titlesize': fontsize,
    'legend.fontsize': fontsize,
    'xtick.labelsize': fontsize - 1,
    'ytick.labelsize': fontsize - 1,
}
mpl.rcParams.update(params)

def histogram(
        values,
        *,
        name,
        filetype="png",
        plt_kwargs=None,
        hist_kwargs=None):
    n, bins, patches = mpl_plot(
        plt_fn=lambda: plt.hist(values, linewidth=10, **(hist_kwargs or {})),
        name=name,
        filetype=filetype,
        **(plt_kwargs or {}))
    return n, bins

def scatter(
        x,
        y,
        *,
        name,
        filetype="png",
        plt_kwargs=None,
        scatter_kwargs=None,
        fit_y=None,
        fit_plot_kwargs=None,
        legend=False):
    def plot_fit():
        if fit_y is not None:
            plt.plot(x, fit_y, **(fit_plot_kwargs or {}))

    def legend_fn():
        if legend:
            if isinstance(legend, dict):
                plt.legend(**legend)
            else:
                plt.legend()

    if isinstance(y, dict):
        def plt_fn():
            for label, _y in y.items():
                plt.scatter(x, _y, label=label, **(scatter_kwargs or {}))
            plot_fit()
            legend_fn()
    else:
        def plt_fn():
            plt.scatter(x, y, **(scatter_kwargs or {}))
            plot_fit()
            legend_fn()
    mpl_plot(
        plt_fn=plt_fn,
        name=name,
        filetype=filetype,
        **plt_kwargs)

def mpl_plot(
        plt_fn,
        *,
        name,
        filetype="png",
        **plt_kwargs):
    Figure.file_types = [filetype]
    Figure.set_defaults(**plt_kwargs)
    with Figure(name):
        return plt_fn()

# http://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html#the-seq2seq-model  # NOQA
def plot_attention(
        input_labels, output_labels, attentions,
        out_colors=None, filepath=None):
    # Set up figure with colorbar
    fig = plt.figure()
    ax = fig.add_subplot(111)
    cax = ax.matshow(attentions, cmap='bone')
    fig.colorbar(cax)

    # Set up axes
    ax.tick_params(axis='both', which='major', labelsize=10)
    ax.tick_params(axis='both', which='minor', labelsize=8)
    ax.set_xticklabels([''] + output_labels, rotation=90)
    ax.set_yticklabels([''] + input_labels)

    # Show label at every tick
    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))

    if out_colors:
        out_colors = ["k"] + out_colors
        assert len(out_colors) == 1 + len(output_labels), \
            f"got {len(out_colors)} colors for {len(output_labels)} labels"
        for xtick, color in zip(ax.get_xticklabels(), out_colors):
            xtick.set_color(color)

    if filepath:
        plt.savefig(filepath)
    else:
        plt.show()
    plt.close()

# http://scikit-learn.org/stable/auto_examples/model_selection/plot_confusion_matrix.html  # NOQA
def plot_confusion_matrix(
        cm, classes,
        normalize=False, title='Confusion matrix', cmap=plt.cm.Blues,
        filepath=None):
    """This function prints and plots the confusion matrix.
    Normalization can be applied by setting `normalize=True`.
    """
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        plt.imshow(
            cm, interpolation='nearest', cmap=cmap, vmin=0, vmax=1)
    else:
        plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=90)
    plt.yticks(tick_marks, classes)

    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(
                j, i, f"{cm[i, j]:.2f}",
                horizontalalignment="center",
                color="white" if cm[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    if filepath:
        plt.savefig(filepath)
    else:
        plt.show()
    plt.close()

# https://stackoverflow.com/questions/18195758/set-matplotlib-colorbar-size-to-match-graph  # NOQA
def add_colorbar(im, aspect=20, pad_fraction=0.5, **kwargs):
    """Add a vertical color bar to an image plot."""
    from mpl_toolkits import axes_grid1
    divider = axes_grid1.make_axes_locatable(im.axes)
    width = axes_grid1.axes_size.AxesY(im.axes, aspect=1./aspect)
    pad = axes_grid1.axes_size.Fraction(pad_fraction, width)
    current_ax = plt.gca()
    cax = divider.append_axes("right", size=width, pad=pad)
    plt.sca(current_ax)
    return im.axes.figure.colorbar(im, cax=cax, **kwargs)

def simple_imshow(
        matrix,
        cmap="viridis", figsize=(10, 10), aspect_equal=True, outfile=None,
        title=None, xlabel=None, ylabel=None,
        xticks=True,
        yticks=True,
        xtick_labels=None,
        ytick_labels=None,
        xtick_locs_labels=None,
        ytick_locs_labels=None,
        xtick_label_rotation='vertical',
        xgrid=None,
        ygrid=None,
        colorbar=True, scale="lin", cbar_title=None,
        bad_color=None,
        origin='upper'):
    if aspect_equal and figsize[1] is None:
        matrix_aspect = matrix.shape[0] / matrix.shape[1]
        width = figsize[0]
        height = max(3, width * matrix_aspect)
        figsize = (width, height)
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(1, 1, 1)
    if aspect_equal:
        ax.set_aspect('equal')
    if title:
        plt.title(title)
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)
    norm = matplotlib.colors.SymLogNorm(1) if scale == "log" else None
    cmap = mpl.cm.get_cmap(cmap)
    if bad_color is not None:
        cmap.set_bad(bad_color)
    im = plt.imshow(
        matrix, interpolation='nearest', cmap=cmap, norm=norm, origin=origin)
    if xtick_labels is not None:
        assert xtick_locs_labels is None
        locs = np.arange(0, len(xtick_labels))
        xtick_locs_labels = locs, xtick_labels
    if ytick_labels is not None:
        assert ytick_locs_labels is None
        locs = np.arange(0, len(ytick_labels))
        ytick_locs_labels = locs, ytick_labels
    if xtick_locs_labels is not None:
        plt.xticks(*xtick_locs_labels, rotation=xtick_label_rotation)
    if ytick_locs_labels is not None:
        plt.yticks(*ytick_locs_labels)
    if xgrid is not None or ygrid is not None:
        if xgrid is not None:
            ax.set_xticks(xgrid, minor=True)
        if ygrid is not None:
            ax.set_yticks(ygrid, minor=True)
        ax.grid(which="minor")
    if xticks is not True:
        plt.tick_params(
            axis='x',          # changes apply to the x-axis
            which='both',      # both major and minor ticks are affected
            bottom=False,      # ticks along the bottom edge are off
            top=False,)         # ticks along the top edge are off
    if yticks is not True:
        plt.tick_params(
            axis='y',          # changes apply to the y-axis
            which='both',      # both major and minor ticks are affected
            left=False,      # ticks along the bottom edge are off
            right=False,)         # ticks along the top edge are off
    if colorbar:
        cbar = add_colorbar(im)
        if cbar_title:
            cbar.ax.set_ylabel(cbar_title, rotation=270)
    plt.tight_layout()
    if outfile:
        plt.savefig(outfile)
    else:
        plt.show()
    plt.clf()

def embed_2d(
        emb, emb_method="UMAP", umap_n_neighbors=15, umap_min_dist=0.1,
        return_proj=False):
    if hasattr(emb_method, 'fit_transform'):
        proj = emb_method
    elif emb_method.lower() == "umap":
        try:
            from umap import UMAP
        except ImportError:
            print("Please install umap to use emb_method='UMAP'")
            print("pip install umap-learn (NOT pip install umap)")
            print("https://github.com/lmcinnes/umap")
            raise
        proj = UMAP(
            init="random",
            n_neighbors=umap_n_neighbors,
            min_dist=umap_min_dist)
    else:
        import sklearn.manifold
        proj = getattr(sklearn.manifold, emb_method)()
    emb_2d = proj.fit_transform(emb)
    if return_proj:
        return emb_2d, proj
    return emb_2d

def plot_embeddings(
        emb, emb_method=None,
        labels=None, color=None, classes=None, class2color=None, title=None,
        outfile=None, cmap="viridis", max_labels=100,
        colorbar_ticks=None, reverse_colorbar=False, colorbar_label=None,
        label_fontpath=None,
        **scatter_kwargs):
    """
    Plot a scatterplot of the embeddings contained in emb.

    emb: an array with dim (n_embeddings x 2) or (n_embeddings x emb_dim).
    In the latter case an embedding method emb_method should be supplied
    to project from emb_dim to dim=2.

    emb_method: "UMAP", "TSNE", or any other algorithm in sklearn.manifold
    labels: Optional text labels for each embedding
    color: Optional color for each embedding, according to which it will be
    colored in the plot.
    classes:  Optional class for each embedding, according to which it will
    be colored in the plot.
    class2color: A map which determines the color assigned to each class
    outfile: If provided, save plot to this file instead of showing it
    cmap: colormap
    max_labels: maximum number of labels to be displayed
    """
    from matplotlib.ticker import NullFormatter
    if emb_method:
        x, y = embed_2d(emb, emb_method).T
    else:
        x, y = emb.T
    figsize = (14, 12) if color is not None else (12, 12)
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111)
    ax.set_aspect('equal')
    ax.xaxis.set_major_formatter(NullFormatter())
    ax.yaxis.set_major_formatter(NullFormatter())
    if not scatter_kwargs:
        scatter_kwargs = dict(marker="o", s=1, alpha=1)
    if classes is not None:
        for cls in set(classes):
            i = (classes == cls).nonzero()
            ax.scatter(x[i], y[i], label=cls, **scatter_kwargs)
    elif color is not None:
        sc = ax.scatter(x, y, c=color, cmap=cmap, **scatter_kwargs)
        cb = fig.colorbar(sc, ticks=colorbar_ticks)
        if reverse_colorbar:
            cb.ax.invert_yaxis()
        if colorbar_label:
            cb.set_label(colorbar_label)
    else:
        ax.scatter(x, y, **scatter_kwargs)

    if labels is not None:
        if label_fontpath:
            import matplotlib.font_manager as fm
            fontproperties = fm.FontProperties(fname=label_fontpath)
        else:
            fontproperties = None
        n_labels = len(labels)
        for i in range(len(emb)):
            if (
                    max_labels < 0 or
                    n_labels <= max_labels or
                    not i % (n_labels // max_labels)):
                ax.annotate(
                    labels[i], (x[i], y[i]), alpha=0.76, size=10,
                    fontproperties=fontproperties)
    if title:
        plt.title(title)
    plt.axis('tight')
    if classes is not None:
        plt.legend(loc='best', scatterpoints=1, markerscale=5, fontsize=10)
    plt.tight_layout()
    if outfile:
        plt.savefig(str(outfile))
    else:
        plt.show()

def plot_dendrogram(
        dist,
        labels,
        outfile=None,
        method="centroid",
        figsize=(50, 45),
        font_size=10,
        cmap='magma_r',
        ):
    from scipy.cluster import hierarchy
    fig = plt.figure(figsize=figsize)
    # dendrogram
    axdendro = fig.add_axes([0.09, 0.1, 0.2, 0.8])
    axdendro.set_xticks([])
    axdendro.set_yticks([])
    Y = hierarchy.linkage(dist, method=method)
    Z = hierarchy.dendrogram(
        Y, orientation='right', labels=labels, leaf_font_size=font_size)
    # distance matrix
    index = Z['leaves']
    D = dist[index, :]
    D = D[:, index]
    axmatrix = fig.add_axes([0.3, 0.1, 0.6, 0.8])
    im = axmatrix.matshow(D, aspect='auto', origin='lower', cmap=cmap)
    axmatrix.set_xticks([])
    axmatrix.set_yticks([])
    # colorbar
    axcolor = fig.add_axes([0.91, 0.1, 0.02, 0.8])
    plt.colorbar(im, cax=axcolor)

    if outfile:
        fig.savefig(str(outfile))
    else:
        fig.show()
    plt.close(fig)

def get_palette(categories, cmap=None):
    from bokeh.palettes import (
        Category20,
        Category20b,
        Category20c,
        viridis,
        )
    n_cat = len(set(categories))
    if cmap is not None:
        return cmap[n_cat]
    if n_cat <= 20:
        if n_cat <= 2:
            palette = Category20[3]
            return [palette[0], palette[-1]]
        else:
            return Category20[n_cat]
    if n_cat <= 40:
        return Category20[20] + Category20b[20]
    if n_cat <= 60:
        return Category20[20] + Category20b[20] + Category20c[20]
    return viridis(n_cat)

def plot_embeddings_bokeh(
        emb,
        emb_method='UMAP',
        classes=None,
        class_category=None,
        labels=None,
        color=None,
        raw_colors=None,
        color_category=None,
        color_categorical=False,
        cmap=None,
        cmap_reverse=False,
        colorbar=False,
        colorbar_ticks=None,
        outfile=None,
        title=None,
        scatter_labels=False,
        tooltip_fields=None,
        figure_kwargs=None,
        write_png=False,
        return_plot=False,
        plot_width=None,
        plot_height=None,
        reuse_figure=None,
        **circle_kwargs,
        ):
    """
    Creates an interactive scatterplot of the embeddings contained in emb,
    using the bokeh library.

    emb: an array with dim (n_embeddings x 2) or (n_embeddings x emb_dim).
    In the latter case an embedding method emb_method should be supplied
    to project from emb_dim to dim=2.

    emb_method: "UMAP", "TSNE", or any other algorithm in sklearn.manifold
    labels: Optional text labels for each embedding
    color: Optional color for each embedding, according to which it will be
    colored in the plot.
    classes:  Optional class for each embedding, according to which it will
    be colored in the plot.
    outfile: If provided, save plot to this file instead of showing it
    cmap: colormap
    title: optional title of the plot
    """
    from bokeh.plotting import figure, output_file, show, save
    from bokeh.models import (
        ColumnDataSource, CategoricalColorMapper, LinearColorMapper,
        ColorBar, FixedTicker, Text)
    from bokeh.palettes import Viridis256

    if emb_method:
        emb = embed_2d(emb, emb_method)
    else:
        assert emb.shape[1] == 2

    if outfile:
        output_file(outfile)

    if cmap is not None:
        if isinstance(cmap, str):
            import bokeh.palettes
            # matplotib suffix for reverse color maps
            if cmap.endswith("_r"):
                cmap_reverse = True
                cmap = cmap[:-2]
            cmap = getattr(bokeh.palettes, cmap)
        elif isinstance(cmap, dict):
            cmap = cmap[max(cmap.keys())]
        if cmap_reverse:
            if isinstance(cmap, dict):
                new_cmap = {}
                for k, v in cmap.items():
                    v = list(v)
                    v.reverse()
                    new_cmap[k] = v
                cmap = new_cmap
            else:
                cmap = list(cmap)
                cmap.reverse()

    source_dict = dict(x=emb[:, 0], y=emb[:, 1])
    if labels is not None:
        source_dict["label"] = labels

    if raw_colors is not None:
        assert color is None
        if any(isinstance(c, str) for c in raw_colors):
            assert all(isinstance(c, str) for c in raw_colors)
        else:
            assert all(len(c) == 3 for c in raw_colors)
            assert cmap is None
            from bokeh.colors import RGB
            raw_colors = [RGB(*c) for c in raw_colors]
        source_dict["color"] = raw_colors
        color_conf = {"field": "color"}
    elif color is not None:
        if any(isinstance(c, str) for c in color):
            assert all(isinstance(c, str) for c in color)
            palette = get_palette(color, cmap=cmap)
            color_mapper = CategoricalColorMapper(
                    factors=sorted(set(color)),
                    palette=palette)
        else:
            if cmap is None:
                cmap = Viridis256
            elif isinstance(cmap, dict):
                cmap = cmap[max(cmap.keys())]
            color_mapper = LinearColorMapper(cmap)
        color_conf = {
            "field": "color",
            "transform": color_mapper}
        source_dict["color"] = color
    else:
        color_conf = "red"

    if classes is not None:
        source_dict["class"] = classes
    if tooltip_fields:
        for k, v in tooltip_fields.items():
            source_dict[k] = v
    source = ColumnDataSource(source_dict)

    tools = "crosshair,pan,wheel_zoom,box_zoom,reset,hover"
    figure_kwargs = figure_kwargs or {}
    if plot_width is not None:
        figure_kwargs['plot_width'] = plot_width
    if plot_height is not None:
        figure_kwargs['plot_height'] = plot_height
    if reuse_figure is None:
        p = figure(tools=tools, sizing_mode='stretch_both', **figure_kwargs)
    else:
        p = reuse_figure
    if title:
        p.title.text = title

    if colorbar:
        if colorbar_ticks:
            ticker = FixedTicker(ticks=colorbar_ticks)
        else:
            ticker = None
        colorbar = ColorBar(
            color_mapper=color_mapper, ticker=ticker)

    if labels is not None and scatter_labels:
        glyph = Text(
            x="x", y="y", text="label", angle=0.0,
            text_color=color_conf, text_alpha=0.95, text_font_size="8pt",
            **circle_kwargs)
        p.add_glyph(source, glyph)
    else:
        plot_kwargs = dict(
                x='x', y='y',
                source=source,
                color=color_conf,
                **circle_kwargs
                )
        if classes is not None:
            legend_field = 'class'
        elif color is not None and not raw_colors:
            legend_field = 'color'
        else:
            legend_field = None
        if legend_field:
            plot_kwargs['legend_field'] = legend_field
            # sort by color field to order the legend entries nicely
            sorted_source = source.to_df().sort_values(legend_field)
            plot_kwargs['source'] = source.from_df(sorted_source)

        p.circle(**plot_kwargs)

    if labels is not None:
        from bokeh.models import HoverTool
        from collections import OrderedDict
        hover = p.select(dict(type=HoverTool))
        hover_entries = [
            ("label", "@label{safe}"),
            ("(x, y)", "(@x, @y)"),
            ]
        if color is not None and color_category:
            hover_entries.append((color_category, "@color"))
        if classes is not None and class_category:
            hover_entries.append((class_category, "@class"))
        if tooltip_fields:
            for field in tooltip_fields:
                hover_entries.append((field, "@" + field))
        hover.tooltips = OrderedDict(hover_entries)
    if colorbar:
        assert color is not None
        p.add_layout(colorbar, 'right')
    if return_plot:
        return p
    if outfile:
        save(p)
        if write_png:
            from bokeh.io import export_png
            png_file = outfile.with_suffix('.png')
            export_png(p, filename=png_file)
    else:
        show(p)

class Figure():
    """Provides a context manager that automatically saves and closes
    a matplotlib plot.

    >>> with Figure("figure_name"):
    >>>     plt.plot(x, y)
    >>> # saves plot to {Figure.fig_dir}/{figure_name}.{Figure.file_type}

    When creating many figures with the same settings, e.g. plt.xlim(0, 100)
    and plt.ylim(0, 1.0), defaults can be set with:

    >>> Figure.set_defaults(xlim=(0, 100), ylim=(0, 1.0))
    >>> # context manager will call plt.xlim(0, 100) and plt.ylim(0, 1.0)
    """
    fig_dir = Path("out/fig")
    file_types = ["png", "pdf"]
    default_plt_calls = {}
    late_calls = ["xscale", "xlim", "yscale", "ylim"]  # order is important

    def __init__(
            self, name,
            figwidth=6, figheight=None, fontsize=12,
            invert_xaxis=False, invert_yaxis=False,
            **kwargs):
        self.fig = plt.figure()
        self.fig.set_figwidth(figwidth)
        phi = 1.6180
        self.fig.set_figheight(figheight or figwidth / phi)
        # params = {
        #     'figure.figsize': (figwidth, figheight or figwidth / phi),
        #     'axes.labelsize': fontsize,
        #     'axes.titlesize': fontsize,
        #     'legend.fontsize': fontsize,
        #     'xtick.labelsize': fontsize - 1,
        #     'ytick.labelsize': fontsize - 1,
        # }
        # mpl.rcParams.update(params)
        self.name = name
        self.plt_calls = {**kwargs}
        self.invert_xaxis = invert_xaxis
        self.invert_yaxis = invert_yaxis
        for attr, val in self.default_plt_calls.items():
            if attr not in self.plt_calls:
                self.plt_calls[attr] = val

    def __enter__(self):
        for attr, val in self.plt_calls.items():
            # if attr in self.late_calls:
            #     continue
            try:
                getattr(plt, attr)(val)
            except:
                getattr(plt, attr)(*val)

        return self.fig

    def __exit__(self, exc_type, exc_val, exc_tb):
        # for attr in self.late_calls:
        #     if attr in self.plt_calls:
        #         print(attr, self.plt_calls[attr])
        #         getattr(plt, attr)(self.plt_calls[attr])
        if self.invert_xaxis:
            plt.gca().invert_xaxis()
        if self.invert_yaxis:
            plt.gca().invert_yaxis()
        plt.tight_layout()
        for file_type in self.file_types:
            outfile = self.fig_dir / f"{self.name}.{file_type}"
            plt.savefig(outfile)
        plt.clf()

    @classmethod
    def set_defaults(cls, **kwargs):
        cls.default_plt_calls = kwargs
        for attr, val in kwargs.items():
            setattr(cls, attr, val)

    @classmethod
    def reset_defaults(cls):
        cls.default_plt_calls = {}

linestyles = [
    "-", "--", "-.", ":",
    "-", "--", "-.", ":",
    "-", "--", "-.", ":",
    "-", "--", "-.", ":",
    "-", "--", "-.", ":",
    "-", "--", "-.", ":"]

try:
    from bokeh.palettes import Category20
    colors = Category20[20]
except ImportError:
    try:
        import seaborn as sns
        colors = sns.color_palette("muted")
    except ImportError:
        # https://gist.github.com/huyng/816622
        colors = [
            "348ABD", "7A68A6", "A60628",
            "467821", "CF4457", "188487", "E24A33",
            "348ABD", "7A68A6", "A60628",
            "467821", "CF4457", "188487", "E24A33",
            "348ABD", "7A68A6", "A60628",
            "467821", "CF4457", "188487", "E24A33",
            ]

# https://matplotlib.org/api/markers_api.html
markers = [
    ".",   # point
    ",",   # pixel
    "o",   # circle
    "v",   # triangle_down
    "^",   # triangle_up
    "<",   # triangle_left
    ">",   # triangle_right
    "1",   # tri_down
    "2",   # tri_up
    "3",   # tri_left
    "4",   # tri_right
    "8",   # octagon
    "s",   # square
    "p",   # pentagon
    "P",   # plus (filled)
    "*",   # star
    "h",   # hexagon1
    "H",   # hexagon2
    "+",   # plus
    "x",   # x
    "X",   # x (filled)
    "D",   # diamond
    "d",   # thin_diamond
    "|",   # vline
    "_",   # hline
    ]

if __name__ == "__main__":
    plot_attention(
        "1 2 3 4".split(),
        "a b c d".split(),
        np.random.rand(4, 4),
        out_colors="r g b r".split())