<gh_stars>0
import unittest
from parameterized import parameterized
import numpy as np
import pandas as pd

import neuroseries as nts
import os

class TsTestCase(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_create_ts(self):
        """
        calling convention
        ts = nts.Ts(a)
        ts is an instance of pd.DataFrame
        """
        a = np.random.randint(0, 10000000, 100)
        a.sort()
        ts = nts.Ts(a)
        self.assertIsInstance(ts, pd.Series, msg="ts doesn't return DataFrame")
        self.assertIsInstance(ts.index.values, np.ndarray,
                              msg="ts doesn't return array values")
        self.assertIs(ts.index.dtype, np.dtype(np.int64), msg='index type is not int64')
        np.testing.assert_array_almost_equal_nulp(a, ts.index.values)

    def test_create_ts_from_double(self):
        """
        data get converted to ts and back fine
        """
        a = np.floor(np.random.rand(100)*1000000)
        a.sort()
        ts = nts.Ts(a)
        self.assertIs(ts.index.dtype, np.dtype(np.int64), msg='index type is not int64')
        np.testing.assert_array_almost_equal_nulp(a, ts.index.values)

    def test_create_ts_time_units(self):
        """
        internally data are stored as us
        """
        a = np.random.randint(0, 1000, 100)
        a.sort()
        ts = nts.Ts(a/1000, time_units='ms')
        np.testing.assert_array_almost_equal_nulp(ts.index.values, a)
        ts = nts.Ts(a/1000000, time_units='s')
        # np.testing.assert_array_almost_equal_nulp(ts.index.values, a.astype(np.int64), nulp=100)
        self.assertTrue(np.all(ts.index.values >= a.astype(np.int64)-1))
        self.assertTrue(np.all(ts.index.values <= a.astype(np.int64)+1))
    @staticmethod
    def test_create_ts_time_units_double():
        """
        conversion of time units from floating point type
        """
        a = np.floor(np.random.rand(100) * 1000000)
        a.sort()
        ts = nts.Ts(a.copy(), time_units='ms')
        # noinspection PyTypeChecker
        np.testing.assert_array_almost_equal_nulp(ts.index.values/1000, a)

    @staticmethod
    def test_create_ts_from_non_sorted():
        """
        if ts are not sorted, a warning should be returned and the timestamps sorted for you
        """
        a = np.random.randint(0, 1000, 100)
        # with self.assertWarns(UserWarning):
        #     ts = nts.Ts(a)
        ts = nts.Ts(a)

        np.testing.assert_array_almost_equal_nulp(np.sort(a), ts.index.values)

    def test_create_ts_wrong_units(self):
        """
        if the units are unsupported it should raise ValueError
        """
        a = np.random.randint(0, 10000000, 100)
        a.sort()
        # noinspection PyUnusedLocal
        ts = 1
        with self.assertRaises(ValueError):
            ts = nts.Ts(a, time_units='min')
        self.assertTrue(ts)

    def test_times_data(self):
        """
        tests the times and data properties
        """
        a = np.random.randint(0, 10000000, 100)
        a.sort()
        b = np.random.randn(100)
        t = nts.Tsd(a, b)
        np.testing.assert_array_almost_equal_nulp(b, t.data())
        np.testing.assert_array_almost_equal_nulp(a, t.times())
        np.testing.assert_array_almost_equal_nulp(a/1000., t.times(units='ms'))
        np.testing.assert_array_almost_equal_nulp(a/1.0e6, t.times(units='s'))
        with self.assertRaises(ValueError):
            t.times(units='banana')

class TsRestrictTestCase(unittest.TestCase):
    def setUp(self):
        from scipy.io import loadmat
        self.mat_data1 = \
            loadmat(os.path.join(nts.get_test_data_dir(), 'restrict_ts_data_1.mat'))
        self.mat_data_left = \
            loadmat(os.path.join(nts.get_test_data_dir(), 'restrict_ts_data_2.mat'))
        self.mat_data_right = \
            loadmat(os.path.join(nts.get_test_data_dir(), 'restrict_ts_data_3.mat'))

    def tearDown(self):
        pass

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_realign(self, data_class):
        """
        first simple realign case
        """

        d_a = self.mat_data1['d_a']
        d_a = d_a.reshape((len(d_a),))
        t_a = data_class(self.mat_data1['t_a'].astype(np.int64), d_a)
        t_b = nts.Ts(self.mat_data1['t_b'].astype(np.int64))
        t_closest = t_a.realign(t_b, align='closest')
        dt = self.mat_data1['d_closest'].reshape((len(self.mat_data1['d_closest'],)))
        self.assertTrue((t_closest.values.ravel() != dt).sum() < 10)
        np.testing.assert_array_almost_equal_nulp(t_closest.values.ravel(), dt)

        t_next = t_a.realign(t_b, align='next')
        dt = self.mat_data1['d_next'].reshape((len(self.mat_data1['d_next'],)))
        np.testing.assert_array_almost_equal_nulp(t_next.values.ravel(), dt)

        t_prev = t_a.realign(t_b, align='prev')
        dt = self.mat_data1['d_prev'].reshape((len(self.mat_data1['d_prev'],)))
        np.testing.assert_array_almost_equal_nulp(t_prev.values.ravel(), dt)

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_realign_left(self, data_class):
        d_a = self.mat_data_left['d_a']
        d_a = d_a.reshape((len(d_a),))
        t_a = data_class(self.mat_data_left['t_a'].astype(np.int64), d_a)
        t_b = nts.Ts(self.mat_data_left['t_b'].astype(np.int64))
        t_closest = t_a.realign(t_b)
        dt = self.mat_data_left['d_closest'].ravel()
        self.assertTrue((t_closest.values.ravel() != dt).sum() < 10)
        np.testing.assert_array_almost_equal_nulp(t_closest.values.ravel(), dt)

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_realign_right(self, data_class):
        d_a = self.mat_data_right['d_a']
        d_a = d_a.reshape((len(d_a),))
        t_a = data_class(self.mat_data_right['t_a'].astype(np.int64), d_a)
        t_b = nts.Ts(self.mat_data_right['t_b'].astype(np.int64))
        t_closest = t_a.realign(t_b)
        dt = self.mat_data_right['d_closest'].ravel()
        self.assertTrue((t_closest.values.ravel() != dt).sum() < 10)
        np.testing.assert_array_almost_equal_nulp(t_closest.values.ravel(), dt)

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_realign_wrong_units(self, data_class):
        d_a = self.mat_data1['d_a']
        d_a = d_a.reshape((len(d_a),))
        t_a = data_class(self.mat_data1['t_a'].astype(np.int64), d_a)
        t_b = nts.Ts(self.mat_data1['t_b'].astype(np.int64))
        # noinspection PyUnusedLocal
        t_closest = 1
        with self.assertRaises(ValueError):
            t_closest = t_a.realign(t_b, align='banana')
        self.assertTrue(t_closest)

class IntervalSetOpsTestCase(unittest.TestCase):
    def setUp(self):
        from scipy.io import loadmat
        self.mat_data1 = loadmat(os.path.join(nts.get_test_data_dir(), 'interval_set_data_1.mat'))

        # note that data are n x 1 here, need to be converted to 1-D in constructor
        self.a1 = self.mat_data1['a1']
        self.b1 = self.mat_data1['b1']
        self.int1 = nts.IntervalSet(self.a1, self.b1, expect_fix=True)

        self.a2 = self.mat_data1['a2']
        self.b2 = self.mat_data1['b2']
        self.int2 = nts.IntervalSet(self.a2, self.b2, expect_fix=True)

    def tearDown(self):
        pass

    def test_create_interval_set(self):
        """
        create an interval set from start and end points "to be fixed"
        """
        a_i1 = self.mat_data1['a_i1'].ravel()
        b_i1 = self.mat_data1['b_i1'].ravel()
        np.testing.assert_array_almost_equal_nulp(a_i1, self.int1['start'])
        np.testing.assert_array_almost_equal_nulp(b_i1, self.int1['end'])

        a_i2 = self.mat_data1['a_i2'].ravel()
        b_i2 = self.mat_data1['b_i2'].ravel()
        np.testing.assert_array_almost_equal_nulp(a_i2, self.int2['start'])
        np.testing.assert_array_almost_equal_nulp(b_i2, self.int2['end'])

    def test_create_interval_set_from_scalars(self):
        int1 = nts.IntervalSet(100, 2100)
        self.assertIsInstance(int1, nts.IntervalSet)
        np.testing.assert_array_almost_equal_nulp(np.array((100,)), int1['start'])
        np.testing.assert_array_almost_equal_nulp(np.array((2100,)), int1['end'])

    def test_iterator(self):
        a_i1 = self.mat_data1['a_i1'].ravel().astype(np.int64)
        b_i1 = self.mat_data1['b_i1'].ravel().astype(np.int64)
        for i, int_i in self.int1.iterrows():
            self.assertEqual(a_i1[i], int_i['start'])
            self.assertEqual(b_i1[i], int_i['end'])

    def test_timespan_tot_length(self):
        """
        return the total length and the timespan of the interval set
        """
        a_i1 = self.mat_data1['a_i1'].ravel()
        b_i1 = self.mat_data1['b_i1'].ravel()
        time_span = self.int1.time_span()
        self.assertIsInstance(time_span, nts.IntervalSet)
        self.assertEqual(time_span['start'][0], a_i1[0])
        self.assertEqual(time_span['end'][0], b_i1[-1])
        tot_l = (b_i1 - a_i1).sum()
        self.assertAlmostEqual(self.int1.tot_length(), tot_l)

    def test_intersect(self):
        """
        intersection of the interval sets
        """
        a_intersect = self.mat_data1['a_intersect'].ravel()
        b_intersect = self.mat_data1['b_intersect'].ravel()
        int_intersect = self.int1.intersect(self.int2)
        self.assertIsInstance(int_intersect, nts.IntervalSet)

        np.testing.assert_array_almost_equal_nulp(int_intersect['start'], a_intersect)
        np.testing.assert_array_almost_equal_nulp(int_intersect['end'], b_intersect)

        int_intersect = self.int2.intersect(self.int1)

        np.testing.assert_array_almost_equal_nulp(int_intersect['start'], a_intersect)
        np.testing.assert_array_almost_equal_nulp(int_intersect['end'], b_intersect)

    def test_union(self):
        """
        union of the interval sets
        """
        a_union = self.mat_data1['a_union'].ravel()
        b_union = self.mat_data1['b_union'].ravel()
        int_union = self.int1.union(self.int2)
        self.assertIsInstance(int_union, nts.IntervalSet)

        np.testing.assert_array_almost_equal_nulp(int_union['start'], a_union)
        np.testing.assert_array_almost_equal_nulp(int_union['end'], b_union)

        int_union = self.int2.union(self.int1)

        np.testing.assert_array_almost_equal_nulp(int_union['start'], a_union)
        np.testing.assert_array_almost_equal_nulp(int_union['end'], b_union)

    def test_setdiff(self):
        """
        diffs of the interval sets
        """
        a_diff1 = self.mat_data1['a_diff1'].ravel()
        b_diff1 = self.mat_data1['b_diff1'].ravel()
        a_diff2 = self.mat_data1['a_diff2'].ravel()
        b_diff2 = self.mat_data1['b_diff2'].ravel()

        int_diff1 = self.int1.set_diff(self.int2)
        self.assertIsInstance(int_diff1, nts.IntervalSet)

        np.testing.assert_array_almost_equal_nulp(int_diff1['start'], a_diff1)
        np.testing.assert_array_almost_equal_nulp(int_diff1['end'], b_diff1)

        int_diff2 = self.int2.set_diff(self.int1)

        np.testing.assert_array_almost_equal_nulp(int_diff2['start'], a_diff2)
        np.testing.assert_array_almost_equal_nulp(int_diff2['end'], b_diff2)

class IntervalSetDropMergeTestCase(unittest.TestCase):
    def setUp(self):
        from scipy.io import loadmat
        self.mat_data1 = loadmat(os.path.join(nts.get_test_data_dir(),
            'interval_set_data_dropmerge1.mat'))

        # note that data are n x 1 here, need to be converted to 1-D in constructor
        self.a1 = self.mat_data1['a1'].ravel()
        self.b1 = self.mat_data1['b1'].ravel()
        self.int1 = nts.IntervalSet(self.a1, self.b1, expect_fix=True)

        self.a1_drop = self.mat_data1['a_i1_drop'].ravel()
        self.b1_drop = self.mat_data1['b_i1_drop'].ravel()
        self.int1_drop = nts.IntervalSet(self.a1_drop, self.b1_drop, expect_fix=True)

        self.a1_merge = self.mat_data1['a_i1_merge'].ravel()
        self.b1_merge = self.mat_data1['b_i1_merge'].ravel()
        self.int1_merge = nts.IntervalSet(self.a1_merge, self.b1_merge, expect_fix=True)

    def tearDown(self):
        pass

    def test_drop_short(self):
        i_drop = self.int1.drop_short_intervals(threshold=2.e3)
        self.assertIsInstance(i_drop, nts.IntervalSet)
        np.testing.assert_array_almost_equal_nulp(i_drop['start'], self.a1_drop)
        np.testing.assert_array_almost_equal_nulp(i_drop['end'], self.b1_drop)

    def test_merge_close(self):
        i_merge = self.int1.merge_close_intervals(threshold=3.e3)
        self.assertIsInstance(i_merge, nts.IntervalSet)
        np.testing.assert_array_almost_equal_nulp(i_merge['start'], self.a1_merge)
        np.testing.assert_array_almost_equal_nulp(i_merge['end'], self.b1_merge)

class TsdUnitsTestCase(unittest.TestCase):
    def setUp(self):
        from scipy.io import loadmat
        self.mat_data1 = loadmat(os.path.join(nts.get_test_data_dir(), 'interval_set_data_1.mat'))
        self.tsd_t = self.mat_data1['t'].ravel()
        self.tsd_d = self.mat_data1['d'].ravel()

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_as_dataframe(self, data_class):
        self.tsd = data_class(self.tsd_t, self.tsd_d)
        if data_class == nts.Tsd:
            df = self.tsd.as_series()
        else:
            df = self.tsd.as_dataframe()
        np.testing.assert_array_almost_equal_nulp(df.values, self.tsd.values)
        np.testing.assert_array_almost_equal_nulp(df.index.values, self.tsd.index.values)

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_units_context(self, data_class):
        self.tsd = data_class(self.tsd_t, self.tsd_d)
        with nts.TimeUnits('ms'):
            t = self.tsd.times()
            # noinspection PyTypeChecker
            np.testing.assert_array_almost_equal_nulp(self.tsd_t/1000., t)

    @staticmethod
    def test_times_units_ts():
        """
        tests the units calling of times
        """
        a = np.random.randint(0, 10000000, 100)
        a.sort()
        ts = nts.Ts(a)

        np.testing.assert_array_almost_equal_nulp(a, ts.times('us'))
        np.testing.assert_array_almost_equal_nulp(a/1000., ts.times('ms'))
        np.testing.assert_array_almost_equal_nulp(a/1.e6, ts.times('s'))

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_as_units_ts(self, data_class):
        """
        as_units returns tsd dataframe
        :return:
        """
        self.tsd = data_class(self.tsd_t, self.tsd_d)
        tsd_ms = self.tsd.as_units('ms')
        if data_class == nts.Tsd:
            self.assertIsInstance(tsd_ms, pd.Series)
        else:
            self.assertIsInstance(tsd_ms, pd.DataFrame)
        # noinspection PyTypeChecker
        np.testing.assert_array_almost_equal_nulp(self.tsd_t/1000., tsd_ms.index.values)

        tsd_s = self.tsd.as_units('s')
        # noinspection PyTypeChecker
        np.testing.assert_array_almost_equal_nulp(self.tsd_t/1.e6, tsd_s.index.values)

class TsdSupportTestCase(unittest.TestCase):
    def setUp(self):
        a1 = np.arange(1, 500000, 100)
        a2 = np.arange(800000, 2300000, 100)
        a3 = np.arange(5200000, 8900000, 100)
        self.t = np.hstack((a1, a2, a3))
        self.d = np.random.randn(*self.t.shape)

    def tearDown(self):
        pass

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_gaps(self, data_class):
        self.t1 = data_class(self.t, self.d)
        gaps = self.t1.gaps(500, method='absolute')
        st = gaps['start']
        en = gaps['end']
        np.testing.assert_array_almost_equal_nulp(st, np.array((499902, 2299901)))
        np.testing.assert_array_almost_equal_nulp(en, np.array((799999, 5199999)))

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def test_support(self, data_class):
        self.t1 = data_class(self.t, self.d)
        support = self.t1.support(min_gap=500, method='absolute')
        np.testing.assert_array_almost_equal_nulp(support['start'], np.array((0, 799999, 5199999)))
        np.testing.assert_array_almost_equal_nulp(support['end'], np.array((499902, 2299901, 8899901)))

class TsdIntervalSetRestrictTestCase(unittest.TestCase):
    def setUp(self):
        from scipy.io import loadmat
        self.mat_data1 = loadmat(os.path.join(nts.get_test_data_dir(),
            'interval_set_data_1.mat'))

        self.a1 = self.mat_data1['a1'].ravel()
        self.b1 = self.mat_data1['b1'].ravel()
        self.int1 = nts.IntervalSet(self.a1, self.b1, expect_fix=True)

        self.a2 = self.mat_data1['a2'].ravel()
        self.b2 = self.mat_data1['b2'].ravel()
        self.int2 = nts.IntervalSet(self.a2, self.b2, expect_fix=True)

        self.tsd_t = self.mat_data1['t'].ravel()
        self.tsd_d = self.mat_data1['d'].ravel()
        self.tsd = nts.Tsd(self.tsd_t, self.tsd_d)

    def tearDown(self):
        pass

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def testRestrict(self, data_class):
        """
        IntervalSet restrict of tsd
        """
        self.tsd = data_class(self.tsd_t, self.tsd_d)
        t_r1 = self.mat_data1['t_r1'].ravel()
        d_r1 = self.mat_data1['d_r1'].ravel()
        tsd_r1 = self.tsd.restrict(self.int1)
        self.assertIsInstance(tsd_r1, data_class)
        self.assertEqual(tsd_r1.index.name, "Time (us)")
        np.testing.assert_array_almost_equal_nulp(t_r1, tsd_r1.times())
        np.testing.assert_array_almost_equal_nulp(d_r1, tsd_r1.values.ravel())

        t_r2 = self.mat_data1['t_r2'].ravel()
        d_r2 = self.mat_data1['d_r2'].ravel()
        tsd_r2 = self.tsd.restrict(self.int2)
        np.testing.assert_array_almost_equal_nulp(t_r2, tsd_r2.times())
        np.testing.assert_array_almost_equal_nulp(d_r2, tsd_r2.values.ravel())

    @parameterized.expand([
        (nts.Tsd,),
        (nts.TsdFrame,)
    ])
    def testRange(self, data_class):
        self.tsd = data_class(self.tsd_t, self.tsd_d)
        range_interval = nts.IntervalSet(9.e8, 3.e9)

        int1_r = self.int1.intersect(range_interval)
        tsd_r = self.tsd.restrict(range_interval)
        tsd_r_r1 = self.tsd.restrict(int1_r)

        with nts.Range(range_interval):
            np.testing.assert_array_almost_equal_nulp(self.tsd.r.times(), tsd_r.times())
            np.testing.assert_array_almost_equal_nulp(self.int1.r['start'], int1_r['start'])
            np.testing.assert_array_almost_equal_nulp(self.int1.r['end'], int1_r['end'])

            np.testing.assert_array_almost_equal_nulp(self.tsd.r.restrict(self.int1.r).times(),
                                                      tsd_r_r1.times())

            # testing caching
            self.assertIsNotNone(self.tsd.r_cache)
            self.assertIsNotNone(self.int1.r_cache)
            np.testing.assert_array_almost_equal_nulp(self.tsd.r.times(), tsd_r.times())
            np.testing.assert_array_almost_equal_nulp(self.int1.r['start'], int1_r['start'])
            np.testing.assert_array_almost_equal_nulp(self.int1.r['end'], int1_r['end'])
            np.testing.assert_array_almost_equal_nulp(self.tsd.r.restrict(self.int1.r).times(),
                                                      tsd_r_r1.times())

        self.assertIsNone(self.tsd.r_cache)
        self.assertIsNone(self.int1.r_cache)

        with nts.Range(9.e8, 3.e9):
            np.testing.assert_array_almost_equal_nulp(self.tsd.r.times(), tsd_r.times())

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


<gh_stars>1-10
import pysubs2
import numpy as np
import tensorflow as tf
from scipy.io import wavfile

def read_audio(audio_filepath, target_fs=1000, mono=True):
    """Read an audio file and optionnaly resample it into a mono channel.

    Parameters
    ----------
        audio_filepath : string, required
            filepath to the audio WAV file
        target_fs : float
            target sampling rate (default: 1 kHz)
        mono : bool
            convert the audio into a mono-channel (default: True)

    Returns
    -------
        `np.array` [int16/int32/float32] : audio samples
    """
    fs, audio = wavfile.read(audio_filepath)
    channels = audio.shape[-1] if len(audio.shape) > 1 else 1

    # sampling the signal
    if target_fs !=  fs:
        n_samples = int(audio.shape[0]*(target_fs/fs))
        resample_idx = np.insert((fs/target_fs)*np.ones(n_samples - 1), 0, 0.)
        resample_idx = np.array(np.round(np.cumsum(resample_idx)), dtype=int)
        # getting mono channel
        audio_resampled = audio[resample_idx]
    else:
        audio_resampled = audio
    if mono & (channels > 1):
        audio_resampled = audio_resampled[:, 0]

    return audio_resampled, fs

def save_audio(audio, audio_filepath, fs):
    """Save an audio file.

    Parameters
    ----------
        audio : `np.array` [int16/int32/float32], required
            audio samples
        audio_filepath : string, required
            output filepath to the audio WAV file
        fs : float, required
            sampling rate in Hz
    """
    wavfile.write(audio_filepath, fs, audio)

def read_subs(input_sub_filepath, target_fs=1000, encoding="utf-8"):
    """Read a subtitle file and convert it into a numerical signal.

    Parameters
    ----------
        input_sub_filepath : string, required
            filepath to the input subtitle SRT file
        target_fs : float
            target sampling rate in Hz (default: 1 kHz)
        encoding : string
            encoding to choose (default: utf-8)

    Returns
    -------
        subs_signal : `np.array` [np.int32] of size [last subtitle end event in ms + 1]
            subtitle signal (takes 0. or 1. value)
        subs : `pysubs2.ssafile.SSAFile`
            subtitle object with all the subtitle events
        starts : `np.array` [np.int32]
            subtitles starts events in ms
        ends : `np.array` [np.int32]
            subtitles ends events in ms
    """
    subs = pysubs2.load(input_sub_filepath, encoding=encoding)
    duration_sub_in_ms = subs[-1].end

    # box signal correspond to the subs
    target_ratio = target_fs / 1000.
    subs_signal = np.zeros(( np.int32(duration_sub_in_ms*target_ratio + 1),), dtype=np.float32)
    starts = np.array([], dtype=np.int32)
    ends = np.array([], dtype=np.int32)
    for sub in subs:
        start = np.int32(sub.start)
        start_ix = np.int32(sub.start*target_ratio)
        end = np.int32(sub.end)
        end_ix = np.int32(sub.end*target_ratio)
        starts = np.append(starts, start)
        ends = np.append(ends, end)
        subs_signal[start_ix:end_ix] = 1.

    return subs_signal, subs, starts, ends

def save_subs(subs, output_sub_filepath, encoding="utf-8"):
    """Save a subtitle file.

    Parameters
    ----------
        subs : `pysubs2.ssafile.SSAFile`, required
            subtitle object with all the subtitle events
        output_sub_filepath : string, required
            output filepath to the subtitle SRT file
        encoding : string
            encoding to choose (default: utf-8)
    """
    subs.save(output_sub_filepath, encoding=encoding)

def save_labels(subs, output_label_filepath):
    """Save a subtitle file into an audacity label TXT file.

    Parameters
    ----------
        subs : `pysubs2.ssafile.SSAFile`, required
            subtitle object with all the subtitle events
        output_label_filepath : string, required
            output filepath to the label TXT file
    """
    with open(output_label_filepath, "w") as fst:
        for sub in subs:
            start = sub.start
            end = sub.end
            text = sub.text
            fst.write("{}\t{}\t{}\n".format(start/1000., end/1000., text))
    fst.close()

def gaussian_kernel_1d(filter_length):
    """Get a 1D gaussian kernel.

    Parameters
    ----------
        filter_length : int, required
            width of the 1D gaussian kernel

    Returns
    -------
        `np.array` [float] of size [filter_length] : 1D gaussian kernel
    """
    #99% of the values
    sigma = (filter_length/2)/2.33
    width = int(filter_length/2)

    norm = 1.0 / (np.sqrt(2*np.pi) * sigma)
    kernel = [norm * np.exp((-1)*(x**2)/(2 * sigma**2)) for x in range(-width, width + 1)]

    return np.float32(kernel / np.sum(kernel))

def tf_1d_gaussian_filtering(signal, kernel_size):
    """Tensorflow convolution of the input signal with a gaussian kernel.

    Parameters
    ----------
        signal : `np.array` or `tf.tensor` [float], required
            input signal to be convolved with a gaussian kernel
        kernel_size : int. required
            width of the 1D gaussian kernel

    Returns
    -------
        `np.array` [float] : convolved input signal with a gaussian kernel
    """
    f = tf.reshape(gaussian_kernel_1d(kernel_size), [-1, 1, 1])
    signal_filtered = tf.reshape(tf.constant(signal, dtype=tf.float32), [1, -1, 1])

    return tf.reshape(tf.nn.conv1d(signal_filtered, filters=f, stride=1, padding='SAME'), [-1])

def filter_audio(audio, threshold=None, kernel_size=None):
    """Audio pre-processing for a vocal audio input (absolute, convolution and thresholding).

    Parameters
    ----------
        audio : `np.array` [float], required
            input audio signal
        threshold : float
            input audio will be thresholded by this amount (default: np.quantile(audio, 1/4))
        kernel_size : int
            width of the 1D gaussian kernel for the input audio convolution (default: 100)

    Returns
    -------
        `np.array` [float] : input audio pre-processed
    """
    # get and convolve the audio magnitude spectrum
    audio_mag = np.abs(audio)
    audio_filtered = tf_1d_gaussian_filtering(audio_mag, kernel_size=kernel_size)
    default_threshold = np.quantile(audio_filtered[audio_mag>0], 1/4)
    # manual threshold to filter noise
    if threshold is None:
        threshold = default_threshold
    else:
        threshold = max(threshold, default_threshold)
    audio_preprocessed = np.array(audio_filtered > threshold, dtype=np.float32)

    return audio_preprocessed

def rescale_audio_subs(audio, subs_signal):
    """Resize both the input audio and subtitle signal so they have the same length.

    Parameters
    ----------
        audio : `np.array` [float], required
            input audio signal
        subs_signal : `np.array` [float], required
            input subtitle signal

    Returns
    -------
        audio : `np.array` [float]
            input audio signal resized
        subs_signal : `np.array` [float]
            input subtitle signal resized
    """
    # rescaling sub or audio so they have the same length
    duration_sub = subs_signal.shape[0]
    duration_audio = audio.shape[0]

    if duration_sub < duration_audio:
        duration = duration_audio - duration_sub
        subs_signal = tf.concat([subs_signal, tf.zeros((duration), dtype=tf.float32)], axis=0)
    else:
        duration = duration_sub - duration_audio
        audio = tf.concat([audio, tf.zeros((duration), dtype=tf.float32)], axis=0)

    return audio, subs_signal

def get_subs_mask(subs_signal, starts, ends, max_duration_allowed=10000, fs=1000):
    """Get a mask from a subtitle signal and starts/ends event, where each cluster is identified by its int value.

    Parameters
    ----------
        subs_signal : `np.array` [float], required
            input subtitle signal
        starts : `np.array` [np.int32], required
            subtitles starts events in ms
        ends : `np.array` [np.int32], required
            subtitles ends events in ms
        max_duration_allowed : int
            maximum duration in ms allowed between two subtitle events to form one cluster (default: 10000 ms)
        fs : float
            Sampling rate in Hz (default: 1 kHz)

    Returns
    -------
        `np.array` [float] : mask with cluster id values for each sample in the input subtitle signal
    """
    mask_subs_signal = np.zeros_like(subs_signal)
    num_sub = len(starts)
    middle_subs = np.array([], dtype=np.int32)
    fs_ratio = fs / 1000.
    # clustering the subs into n groups
    for i in range(num_sub - 1):
        dura = starts[i+1]*fs_ratio - ends[i]*fs_ratio
        if(dura > max_duration_allowed*fs_ratio):
            middle_sub = np.int32((ends[i]*fs_ratio + starts[i+1]*fs_ratio)/2)
            middle_subs = np.append(middle_subs, [middle_sub])
    mask_subs_signal[middle_subs] = 1

    return np.cumsum(mask_subs_signal)

def resync_subs(params, subs, mask=None, fs=1000, cut_start=0, cut_end=90060990):
    """Re-synchronize the subtitle object given the transformation parameters.

    Parameters
    ----------
        params : `list` [float], required
            transformation parameters
        subs : `pysubs2.ssafile.SSAFile`, required
            subtitle object with all the subtitle events
        mask : `np.array` [float]
            mask with cluster id values for each sample in the input subtitle signal
        fs : float
            Sampling rate in Hz (default: 1 kHz)
        cut_start : int
            Minimum timestamp value (ms) to include (default: 0)
        cut_end : int
            Maximum timestamp value (ms) to include (default: 90060990)

    Returns
    -------
        `pysubs2.ssafile.SSAFile` : subtitle object with all the re-synchronized subtitle events
    """
    target_ratio = fs / 1000.
    max_duration = int(subs[-1].end + 3600000)
    for sub, i in zip(subs, range(len(subs))):
        input_start = int(sub.start * target_ratio)
        input_end = int(sub.end * target_ratio)
        start =  params[0]*input_start + params[1]
        end = params[0]*input_end + params[1]
        # non-rigid
        if len(params) != 2:
            if mask is None:
                idx_start =  input_start
                idx_end = input_end
            else:
                idx_start = int(mask[input_start])
                idx_end = int(mask[input_end])
            start = start + params[2][idx_start]
            end = end + params[2][idx_end]
        if (cut_start < sub.start) & (sub.end < cut_end):
            # updating subtitle file
            subs[i].start = min(max(0, int(start/target_ratio)), max_duration)
            subs[i].end = min(max(0, int(end/target_ratio)), max_duration)

    return subs

def add_credits(subs):
    """Add credits to the software at the end of the subtitle SRT file.

    Parameters
    ----------
        subs : `pysubs2.ssafile.SSAFile`, required
            subtitle object with all the subtitle events

    Returns
    -------
        `pysubs2.ssafile.SSAFile` : subtitle object with credits appended at the end
    """

    start = int(subs[-1].end) + 2000
    end = start + 2000
    text = "Processed with ReSuber.\nCheck the github page  https://github.com/polak0v/ReSuber  !"
    event = pysubs2.SSAEvent(start=start, end=end, text=text)
    subs += [event]

    return subs


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

import numpy as np
import scipy.linalg

import qutip
from qutip.core import data as _data
import pytest

def f(t, args):
    return t * (1 - 0.5j)

def liouvillian_ref(H, c_ops=()):
    L = -1.0j * (qutip.spre(H) - qutip.spost(H)) if H else 0
    for c in c_ops:
        if c.issuper:
            L += c
        else:
            cdc = c.dag() * c
            L += qutip.sprepost(c, c.dag())
            L -= 0.5 * (qutip.spre(cdc) + qutip.spost(cdc))
    return L

class TestMatVec:
    """
    A test class for the QuTiP function for matrix/vector conversion.
    """
    def testOperatorVector(self):
        """
        Superoperator: Operator - vector - operator conversion.
        """
        N = 3
        rho1 = qutip.rand_dm(N)
        rho2 = qutip.vector_to_operator(qutip.operator_to_vector(rho1))
        np.testing.assert_allclose(rho1.full(), rho2.full(), 1e-8)

    def testsuperrep(self):
        N = 3
        rho1 = qutip.rand_dm(N)
        as_vec = qutip.operator_to_vector(rho1)
        assert as_vec.superrep == 'super'

        with pytest.raises(TypeError) as err:
            bad_vec = as_vec.copy()
            bad_vec.superrep = ""
            qutip.vector_to_operator(bad_vec)
        assert err.value.args[0] == ("only defined for operator-kets "
                                     "in super format")

        with pytest.raises(TypeError) as err:
            qutip.operator_to_vector(as_vec)
        assert err.value.args[0] == ("Cannot convert object already "
                                     "in super representation")

    def testOperatorVectorTensor(self):
        """
        Superoperator: Operator - vector - operator conversion with a tensor product state.
        """
        Na = 3
        Nb = 2
        rhoa = qutip.rand_dm(Na)
        rhob = qutip.rand_dm(Nb)
        rho1 = qutip.tensor(rhoa, rhob)
        rho2 = qutip.vector_to_operator(qutip.operator_to_vector(rho1))
        np.testing.assert_allclose(rho1.full(), rho2.full(), 1e-8)

    def testOperatorVectorNotSquare(self):
        """
        Superoperator: Operator - vector - operator conversion for non-square matrix.
        """
        op1 = qutip.Qobj(np.random.rand(6).reshape((3, 2)))
        op2 = qutip.vector_to_operator(qutip.operator_to_vector(op1))
        np.testing.assert_allclose(op1.full(), op2.full(), 1e-8)

    def testOperatorSpreAppl(self):
        """
        Superoperator: apply operator and superoperator from left (spre)
        """
        N = 3
        rho = qutip.rand_dm(N)
        U = qutip.rand_unitary(N)
        rho1 = U * rho
        rho2_vec = qutip.spre(U) * qutip.operator_to_vector(rho)
        rho2 = qutip.vector_to_operator(rho2_vec)
        np.testing.assert_allclose(rho1.full(), rho2.full(), 1e-8)

    def testOperatorSpostAppl(self):
        """
        Superoperator: apply operator and superoperator from right (spost)
        """
        N = 3
        rho = qutip.rand_dm(N)
        U = qutip.rand_unitary(N)
        rho1 = rho * U
        rho2_vec = qutip.spost(U) * qutip.operator_to_vector(rho)
        rho2 = qutip.vector_to_operator(rho2_vec)
        np.testing.assert_allclose(rho1.full(), rho2.full(), 1e-8)

    def testOperatorUnitaryTransform(self):
        """
        Superoperator: Unitary transformation with operators and superoperators
        """
        N = 3
        rho = qutip.rand_dm(N)
        U = qutip.rand_unitary(N)
        rho1 = U * rho * U.dag()
        rho2_vec = qutip.sprepost(U, U.dag()) * qutip.operator_to_vector(rho)
        rho2 = qutip.vector_to_operator(rho2_vec)
        np.testing.assert_allclose(rho1.full(), rho2.full(), 1e-8)

    def testMatrixVecMat(self):
        """
        Superoperator: Conversion matrix to vector to matrix
        """
        M = _data.create(np.random.rand(10, 10))
        V = qutip.stack_columns(M)
        M2 = qutip.unstack_columns(V)
        np.testing.assert_allclose(M.to_array(), M2.to_array(), 1e-8)

    def testVecMatVec(self):
        """
        Superoperator: Conversion vector to matrix to vector
        """
        V = _data.create(np.random.rand(100, 1))
        M = qutip.unstack_columns(V)
        V2 = qutip.stack_columns(M)
        np.testing.assert_allclose(V.to_array(), V2.to_array(), 1e-8)

    def testVecMatIndexConversion(self):
        """
        Superoperator: Conversion between matrix and vector indices
        """
        N = 10
        for i in range(N * N):
            assert i == qutip.stacked_index(N, *qutip.unstacked_index(N, i))

    def testVecMatIndexCompability(self):
        """
        Superoperator: Compatibility between matrix/vector and
        corresponding index conversions.
        """
        N = 10
        M = _data.create(np.random.rand(N, N))
        V = qutip.stack_columns(M)
        for idx in range(N * N):
            i, j = qutip.unstacked_index(N, idx)
            assert V.to_array()[idx, 0] == M.to_array()[i, j]

    def test_reshuffle(self):
        U1 = qutip.rand_unitary(2)
        U2 = qutip.rand_unitary(3)
        U3 = qutip.rand_unitary(4)
        U = qutip.tensor(U1, U2, U3)
        S = qutip.to_super(U)
        S_col = qutip.reshuffle(S)
        assert S_col.dims[0] == [[2, 2], [3, 3], [4, 4]]
        assert qutip.reshuffle(S_col) == S

    def test_sprepost(self):
        U1 = qutip.rand_unitary(3)
        U2 = qutip.rand_unitary(3)
        S1 = qutip.spre(U1) * qutip.spost(U2)
        S2 = qutip.sprepost(U1, U2)
        assert S1 == S2

    def testLiouvillianImplem(self):
        """
        Superoperator: Randomized comparison of standard and reference
        Liouvillian functions.
        """
        N1, N2, N3 = 3, 4, 5
        a1 = qutip.tensor(qutip.rand_dm(N1, density=0.75),
                          qutip.qeye([N2, N3]))
        a2 = qutip.tensor(qutip.qeye(N1),
                          qutip.rand_dm(N2, density=0.75),
                          qutip.qeye(N3))
        a3 = qutip.tensor(qutip.qeye([N1, N2]),
                          qutip.rand_dm(N3, density=0.75))
        H = a1.dag()*a1 + a2.dag()*a2 + a3.dag()*a3
        c_ops = [np.sqrt(0.01) * a1, np.sqrt(0.025) * a2, np.sqrt(0.05) * a3]
        L1 = qutip.liouvillian(H, c_ops)
        L2 = liouvillian_ref(H, c_ops)
        np.testing.assert_allclose(L1.full(), L2.full(), 1e-8)

class TestSuper_td:
    """
    A test class for the QuTiP superoperator functions.
    """
    N = 3
    t1 = qutip.QobjEvo([qutip.qeye(N)*(1 + 0.1j),
                        [qutip.create(N) * (1 - 0.1j), f]])
    t2 = qutip.QobjEvo([qutip.destroy(N) * (1 - 0.2j)])
    t3 = qutip.QobjEvo([[qutip.num(N) * (1 + 0.2j), f]])
    q1 = qutip.qeye(N) * (1 + 0.3j)
    q2 = qutip.destroy(N) * (1 - 0.3j)
    q3 = qutip.num(N) * (1 + 0.4j)

    def test_spre_td(self):
        "Superoperator: spre, time-dependent"
        assert qutip.spre(self.t1)(0.5) == qutip.spre(self.t1(0.5))

    def test_spost_td(self):
        "Superoperator: spre, time-dependent"
        assert qutip.spost(self.t1)(0.5) == qutip.spost(self.t1(0.5))

    def test_sprepost_td(self):
        "Superoperator: sprepost, time-dependent"
        # left QobjEvo
        assert (qutip.sprepost(self.t1, self.q2)(0.5)
                == qutip.sprepost(self.t1(0.5), self.q2))
        # left QobjEvo
        assert (qutip.sprepost(self.q2, self.t1)(0.5)
                == qutip.sprepost(self.q2, self.t1(0.5)))
        # left 2 QobjEvo, one cte
        assert (qutip.sprepost(self.t1, self.t2)(0.5)
                == qutip.sprepost(self.t1(0.5), self.t2(0.5)))

    def test_operator_vector_td(self):
        "Superoperator: operator_to_vector, time-dependent"
        assert (qutip.operator_to_vector(self.t1)(0.5)
                == qutip.operator_to_vector(self.t1(0.5)))
        vec = qutip.operator_to_vector(self.t1)
        assert (qutip.vector_to_operator(vec)(0.5)
                == qutip.vector_to_operator(vec(0.5)))

    def test_liouvillian_td(self):
        "Superoperator: liouvillian, time-dependent"
        assert (qutip.liouvillian(self.t1)(0.5)
                == qutip.liouvillian(self.t1(0.5)))
        assert (qutip.liouvillian(None, [self.t2])(0.5)
                == qutip.liouvillian(None, [self.t2(0.5)]))
        assert (qutip.liouvillian(self.t1, [self.t2, self.q1, self.t3],
                                  chi=[1, 2, 3])(0.5)
                == qutip.liouvillian(self.t1(0.5),
                                     [self.t2(0.5), self.q1, self.t3(0.5)],
                                     chi=[1, 2, 3]))

    def test_lindblad_dissipator_td(self):
        "Superoperator: lindblad_dissipator, time-dependent"
        assert (qutip.lindblad_dissipator(self.t2)(0.5)
                == qutip.lindblad_dissipator(self.t2(0.5)))
        assert (qutip.lindblad_dissipator(self.t2, self.q1)(0.5)
                == qutip.lindblad_dissipator(self.t2(0.5), self.q1))
        assert (qutip.lindblad_dissipator(self.q1, self.t2)(0.5)
                == qutip.lindblad_dissipator(self.q1, self.t2(0.5)))


<reponame>mostamndi/Merra2BC<filename>merra2_module.py
import config
import re
import os
from netCDF4 import Dataset
from datetime import datetime, timedelta
import numpy as np
from scipy import interpolate

from multiprocessing import Pool
from functools import partial

Ptop_mera=1 #Pa   (=0.01 hPa)
mera_lat=0
mera_lon=0

shifted_lons=False
shift_index=0

merra_files=[]
mera_times={}                           #map between time and index in file
mera_times_files={}                     #map between time and file index
merra_vars=[]
mer_number_of_x_points=0
mer_number_of_y_points=0
mer_number_of_z_points=0

numbers = re.compile(r'(\d+)')
def numericalSort(value):
    parts = numbers.split(value)
    return parts[9]

def get_file_index_by_time(time):
    return mera_times_files.get(time)

def get_index_in_file_by_time(time):
    return mera_times.get(time)

def get_file_name_by_index(index):
    return merra_files[index]

#********************************
#Horizontal interpolation of 3d Merra field on WRF boundary
def hor_interpolate_3dfield_on_wrf_boubdary(FIELD, wrf_length, wrf_lon, wrf_lat):
    FIELD_BND=np.zeros([mer_number_of_z_points, wrf_length])
    for z_level in range(mer_number_of_z_points):
        f = interpolate.RectBivariateSpline(mera_lat, mera_lon, FIELD[z_level,:,:])
        FIELD_BND[z_level,:]=f(wrf_lat,wrf_lon,grid=False)
    return FIELD_BND

#Vertical interpolation of Merra boundary on WRF boundary
def ver_interpolate_3dfield_on_wrf_boubdary(MER_HOR_SPECIE_BND,MER_HOR_PRES_BND,WRF_PRES_BND,wrf_nz, wrf_length):
    WRF_SPECIE_BND = np.zeros([wrf_nz,wrf_length])  # Required SPEC on WRF boundary
    for i in range(0,wrf_length):
        f = interpolate.interp1d(MER_HOR_PRES_BND[:,i], MER_HOR_SPECIE_BND[:,i], kind='linear',bounds_error=False,fill_value=0)
        WRF_SPECIE_BND[:,i]=f(WRF_PRES_BND[:,i])
    return WRF_SPECIE_BND

#Horizontal interpolation of 3d Merra field on WRF horizontal grid
def hor_interpolate_3dfield_on_wrf_grid(FIELD, wrf_ny, wrf_nx, wrf_lon, wrf_lat):
    FIELD_HOR=np.zeros([mer_number_of_z_points, wrf_ny, wrf_nx])

    for z_level in range(mer_number_of_z_points):
        f = interpolate.RectBivariateSpline(mera_lat, mera_lon, FIELD[z_level,:,:])
        FIELD_HOR[z_level,:,:]=f(wrf_lat,wrf_lon,grid=False).reshape(wrf_ny, wrf_nx)

    return FIELD_HOR

#Vertical interpolation on WRF grid
def ver_interpolate_3dfield_on_wrf_grid(MER_HOR_SPECIE, MER_HOR_PRES,WRF_PRES,wrf_nz, wrf_ny, wrf_nx):
    WRF_SPECIE = np.zeros([wrf_nz,wrf_ny,wrf_nx])  # Required SPEC on WRF grid
    for x in range(0,wrf_nx,1):
        for y in range(0,wrf_ny,1):
            f = interpolate.interp1d(MER_HOR_PRES[:,y,x], MER_HOR_SPECIE[:,y,x], kind='linear',bounds_error=False,fill_value=0)
            WRF_SPECIE[:,y,x]=f(WRF_PRES[:,y,x])
    return WRF_SPECIE
#********************************

#extracts 3d field from merra2 file from given time
def get_3dfield_by_time(time,merra_file,field_name):
    mera_time_idx=get_index_in_file_by_time(time)
    field=merra_file.variables[field_name][mera_time_idx,:]

    if shifted_lons:
        field=np.roll(field,shift_index,axis=2)

    return np.flipud(field)

def get_pressure_by_time(time,merra_file):
    global Ptop_mera
    #MER_Pres will be restored on 73 edges
    MER_Pres = np.zeros([mer_number_of_z_points+1,mer_number_of_y_points,mer_number_of_x_points])
    #filling top edge with Ptop_mera
    MER_Pres[0,:,:]=Ptop_mera

    # Extract deltaP from NetCDF file at index defined by time
    mera_time_idx=get_index_in_file_by_time(time)
    DELP = merra_file.variables['DELP'][mera_time_idx,:]  #Pa

    for z_level in range(mer_number_of_z_points):
        MER_Pres[z_level+1]=MER_Pres[z_level]+DELP[z_level]

    #BUT! we need pressure on 72 levels
    #=> averaging pressure values on adjacent edges
    MER_Pres = (MER_Pres[0:mer_number_of_z_points:1][:,:] + MER_Pres[1::1][:,:]) / 2

    if shifted_lons:
        MER_Pres=np.roll(MER_Pres,shift_index,axis=2)

    MER_Pres=np.flipud(MER_Pres)
    return MER_Pres

def initialise():
    global merra_files,mer_number_of_x_points,mer_number_of_y_points,mer_number_of_z_points,mera_lon,mera_lat,merra_vars,shifted_lons,shift_index

    merra_files=sorted([f for f in os.listdir(config.mera_dir) if re.match(config.mera_files, f)], key=numericalSort)
    #print "Open "+config.mera_dir+"/"+merra_files[0]
    merra_f = Dataset(config.mera_dir+"/"+merra_files[0],'r')
    mer_number_of_x_points=merra_f.variables['lon'].size
    mer_number_of_y_points=merra_f.variables['lat'].size
    #not all merra2 files (loading diagnostic) have 'lev' variable
    try:
        mer_number_of_z_points=merra_f.variables['lev'].size
    except Exception:
        pass

    print "MERRA2 dimensions: [bottom_top]="+str(mer_number_of_z_points)+" [south_north]="+str(mer_number_of_y_points)+" [west_east]="+str(mer_number_of_x_points)

    merra_vars = [var for var in merra_f.variables]

    mera_lon  = merra_f.variables['lon'][:]
    mera_lat  = merra_f.variables['lat'][:]

    #if data is given in range of 0_360, then we need to shift lons and data to the -180_180
    if(max(mera_lon)>180):
        print "###########################"
        print "ATTENTION!!!:"
        print "SHIFTING LONGITUDES"
        index=0
        for lon in mera_lon:
            if lon >180:
                mera_lon[index]=mera_lon[index]-360.0
            index=index+1
        shift_index=len(mera_lon)/2
        mera_lon=np.roll(mera_lon,shift_index)
        shifted_lons=True
        print "###########################"

    print "Lower left corner: lat="+str(min(mera_lat))+" long="+str(min(mera_lon))
    print "Upper right corner: lat="+str(max(mera_lat))+" long="+str(max(mera_lon))

    #number of times in  mera file
    times_per_file=merra_f.variables['time'].size
    merra_f.close()

    index=0
    for merra_file in merra_files:
        date=numbers.split(merra_file)[9]
        for i in range(0,times_per_file,1):
            t=datetime.strptime(date, '%Y%m%d')+timedelta(minutes =(i*(24/times_per_file)*60))
            mera_times_files.update({t.strftime("%Y-%m-%d_%H:%M:%S"):index})
            mera_times.update({t.strftime("%Y-%m-%d_%H:%M:%S"):i})
        index=index+1


<gh_stars>10-100
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Oct 12 15:10:49 2019
This is a library of tensor network method for Ising models
@author: Bruce
"""

import numpy as np
import scipy.integrate as integrate
from ncon import ncon
from tntools.initialtensors import get_initial_tensor

def exactSol(J):
    """
    g = exactSol(J).
    ---------------------
    Return: g is defined as ln(Z) / N = -beta * f, where Z is the partition
    function, N is the total number of spin, f the free energy per site.
    """

    g = np.zeros(len(J))
    for i in range(len(J)):
        k = 1 / (np.sinh(2*J[i])**2)
        integrand = lambda theta: np.log((np.cosh(2*J[i]))**2
        + 1 / k * np.sqrt(1 + k**2 - 2*k * np.cos(2*theta)))
        g[i] = np.log(2) / 2 + 1 / (2 * np.pi) * integrate.quad(integrand, 0,np.pi)[0]
    return g

def Ising2dT(beta = 0.4, h = 0, isSym = False):
    """
    T = Ising2dT(J,h).
    -------------------------
    Set up the initial tensor for 2d classical Ising model on a square lattice.
    Argument: J is defined to be beta * J = J / kT, and h is
    defined to be beta*h = h / kT, where J and h are conventional coupling constants.
    Return: a rank 4 tensor T[i,j,k,l]. Each index of the tensor represents
    physical classical spin, and the tensor T represents the Boltzmann weight
    for interaction on one plaquettes.
    """
    pars = {"model":"ising", "dtype":"float64",
            "J":1, "H":h, "beta":beta, "symmetry_tensors":isSym}
    T0 = get_initial_tensor(pars)

    return T0

def Ising2dTwall(J = 0.1):
    """
    Set up the initial tensor for 2d classical Ising model on a square lattice
    according to the method described in Xiang's HOTRG paper and Hauru's Gilt
    paper. The number of initial tensor is the same as the number of initial
    spin variables.

    Parameters
    ----------
    J : double, optional
        Inverse temperature J/T. The default is 0.1.

    Returns
    -------
    T: four leg tensor

    """
    raise Warning("This function is deprecated now!")
    # The sqaure root of weight matrix (See Hauru's Gilt paper for details)
    M = np.array([[np.sqrt(np.cosh(J)), np.sqrt(np.sinh(J))],
         [np.sqrt(np.cosh(J)), - np.sqrt(np.sinh(J))]])
    T = np.einsum("ai,aj,ak,al->ijkl",M,M,M,M)
    return T

def checkSym(T):
    """
    Check whether a tensor T has horizontal and verticle symmetry

    Parameters
    ----------
    T : four-leg tensor
       j|
     i--T--k.
       l|

    Returns
    -------
    None.

    """
    Tarr = T.to_ndarray()
    ishSym = np.allclose(Tarr,Tarr.transpose([2,1,0,3]))
    isvSym = np.allclose(Tarr,Tarr.transpose([0,3,2,1]))
    print("Horizontal symmetry: {}.".format(ishSym))
    print("Vertical symmetry: {}.".format(isvSym))

def calg(A,Anorm,initN = 4, traceord = [1,1,2,2]):
    numlevels = len(A)-1
    FreeEnergy = np.zeros(numlevels)
    Anorm = np.array(Anorm)
    for k in range(1,numlevels+1):
        FreeEnergy[k-1] = (np.sum(4.0 ** (-np.arange(0,k+1)) *
                                  np.log(Anorm[:(k+1)])) +
                           (1/4)**k * np.log(ncon(
                               [A[k]],[traceord]).norm())) /(2*initN)
    return FreeEnergy

if __name__ == "__main__":
    import argparse
    from trg import mainTRG, calg

    def compareg(n_iter = 12,Dcut = 8, isSym = False):
        Tc = 2 / np.log(1 + np.sqrt(2))
        gext = exactSol([1/Tc])[0]
        A, Anorm = mainTRG(1.0, Dcut, n_iter, isSym = isSym)
        gapp = calg(A,Anorm, initN = 1,  traceord = [1,2,1,2])
        err = np.abs(gapp - gext) / gext
        return err, A
    parser = argparse.ArgumentParser(
        "Test TRG implementation")
    parser.add_argument("--chi", dest = "chi", type = int,
                    help = "bound dimension (default: 10)",
                    default = 10)
    parser.add_argument("--iterN", dest = "iterN", type = int,
                    help = "Number of TRG iterations (default: 25)",
                    default = 25)
    parser.add_argument("--isSym", help = "whether to use Z2 symmetry",
                        action = "store_true")
    args = parser.parse_args()
    chi = args.chi
    iterN = args.iterN
    isSym = args.isSym
    err,A = compareg(iterN, chi, isSym)
    print("Error of free energy at Tc is ")
    print(err)


"""Utility functions for ray computation."""
import math
import numpy as np
from scipy.spatial.transform import Rotation as R
import box_utils
import torch

def apply_batched_transformations(inputs, transformations):
    """Batched transformation of inputs.

    Args:
        inputs: List of [R, S, 3]
        transformations: [R, 4, 4]

    Returns:
        transformed_inputs: List of [R, S, 3]
    """
    # if rotation_only:
    #    transformations[:, :3, 3] = torch.zeros((3,), dtype=torch.float)

    transformed_inputs = []
    for x in inputs:
        N_samples = x.size()[1]
        homog_transformations = transformations.unsqueeze(1)  # [R, 1, 4, 4]
        homog_transformations = torch.tile(homog_transformations, (1, N_samples, 1, 1))  # [R, S, 4, 4]
        homog_component = torch.ones_like(x)[..., 0:1]  # [R, S, 1]
        homog_x = torch.cat((x, homog_component), axis=-1)  # [R, S, 4]
        homog_x = homog_x.unsqueeze(2)
        transformed_x = torch.matmul(
            homog_x,
            torch.transpose(homog_transformations, 2, 3))  # [R, S, 1, 4]
        transformed_x = transformed_x[..., 0, :3]  # [R, S, 3]
        transformed_inputs.append(transformed_x)
    return transformed_inputs

def get_transformation_from_params(params):
    translation, rotation = [0, 0, 0], [0, 0, 0]
    if 'translation' in params:
        translation = params['translation']
    if 'rotation' in params:
        rotation = params['rotation']
    translation = torch.tensor(translation, dtype=torch.float)
    rotmat = torch.tensor(R.from_euler('xyz', rotation, degrees=True).as_matrix(), dtype=torch.float)
    return translation, rotmat

def rotate_dirs(dirs, rotmat):
    """
    Args:
        dirs: [R, 3] float tensor.
        rotmat: [3, 3]
    """
    if type(dirs) == np.ndarray:
        dirs = torch.tensor(dirs).float()
    #rotmat = rotmat.unsqueeze(0)
    rotmat = torch.broadcast_to(rotmat, (dirs.shape[0], 3, 3))  # [R, 3, 3]
    dirs_obj = torch.matmul(dirs.unsqueeze(1), torch.transpose(rotmat, 1, 2))  # [R, 1, 3]
    dirs_obj = dirs_obj.squeeze(1)  # [R, 3]
    return dirs_obj

def transform_dirs(dirs, params, inverse=False):
    _, rotmat = get_transformation_from_params(params)  # [3,], [3, 3]
    if inverse:
        rotmat = torch.transpose(rotmat, 0, 1)  # [3, 3]
    dirs_transformed = rotate_dirs(dirs, rotmat)
    return dirs_transformed

def transform_rays(ray_batch, params, use_viewdirs, inverse=False):
    """Transform rays into object coordinate frame given o2w transformation params.

    Note: do not assume viewdirs is always the normalized version of rays_d (e.g., in staticcam case).

    Args:
        ray_batch: [R, M] float tensor. Batch of rays.
        params: Dictionary containing transformation parameters:
            'translation': List of 3 elements. xyz translation.
            'rotation': List of 3 euler angles in xyz.
        use_viewdirs: bool. Whether to we are using viewdirs.
        inverse: bool. Whether to apply inverse of the transformations provided in 'params'.

    Returns:
        ray_batch_obj: [R, M] float tensor. The ray batch, in object coordinate frame.
    """
    rays_o, rays_d = ray_batch[:, 0:3], ray_batch[:, 3:6]
    translation, rotmat = get_transformation_from_params(params)  # [3,], [3, 3]

    if inverse:
        translation = -1 * translation  # [3,]
        rotmat = torch.transpose(rotmat, 1, 0)  # [3, 3]

    translation_inverse = -1 * translation
    rotmat_inverse = torch.transpose(rotmat, 1, 0)

    # Transform the ray origin.
    rays_o_obj, _ = box_utils.ray_to_box_coordinate_frame_pairwise(
        box_center=translation_inverse,
        box_rotation_matrix=rotmat_inverse,
        rays_start_point=rays_o,
        rays_end_point=rays_d)

    # Only apply rotation to rays_d.
    rays_d_obj = rotate_dirs(rays_d, rotmat)

    ray_batch_obj = update_ray_batch_slice(ray_batch, rays_o_obj, 0, 3)
    ray_batch_obj = update_ray_batch_slice(ray_batch_obj, rays_d_obj, 3, 6)
    if use_viewdirs:
        # Grab viewdirs from the ray batch itself. Because it may be different from rays_d
        # (as in the staticcam case).
        viewdirs = ray_batch[:, 8:11]
        viewdirs_obj = rotate_dirs(viewdirs, rotmat)
        ray_batch_obj = update_ray_batch_slice(ray_batch_obj, viewdirs_obj, 8, 11)
    return ray_batch_obj

def transform_points_into_world_coordinate_frame(pts, params, check_numerics=False):
    translation, rotmat = get_transformation_from_params(params)  # [3,], [3, 3]

    # pts_flat = pts.view(-1, 3)  # [RS, 3]
    # num_examples = pts_flat.size()[0]  # RS

    # translation = translation.unsqueeze(0)
    # translation = torch.tile(translation, (num_examples, 1))  # [RS, 3]
    # rotmat = rotmat.unsqueeze(0)
    # rotmat = torch.tile(rotmat, (num_examples, 1, 1))

    # # pts_flat_transformed = torch.matmul(pts_flat[:, None, :], torch.transpose(rotmat, 2, 1))  # [RS, 1, 3]
    # pts_flat_transformed = pts_flat[:, None, :]  # [RS, 1, 3]
    # pts_flat_transformed += translation[:, None, :]  # [RS, 1, 3]
    # pts_transformed = pts_flat_transformed.view(pts.size())  # [R, S, 3]
    chunk = 256
    # Check batch transformations works without rotation.
    if check_numerics:
        transformations = np.eye(4)
        transformations[:3, 3] = translation
        transformations = torch.tensor(transformations, dtype=torch.float)  # [4, 4]
        transformations = torch.tile(transformations[None, ...], (pts.size()[0], 1, 1))  # [R, 4, 4]
        pts_transformed1 = []
        for i in range(0, pts.size()[0], chunk):
            pts_transformed1_chunk = apply_batched_transformations(
                inputs=[pts[i:i+chunk]], transformations=transformations[i:i+chunk])[0]
            pts_transformed1.append(pts_transformed1_chunk)
        pts_transformed1 = torch.cat(pts_transformed1, dim=0)

        pts_transformed2 = pts + translation[None, None, :]

    # Now add rotation
    transformations = np.eye(4)
    transformations = torch.tensor(transformations, dtype=torch.float)
    transformations[:3, :3] = rotmat
    transformations[:3, 3] = translation
    #transformations = torch.tensor(transformations, dtype=torch.float)  # [4, 4]
    transformations = torch.tile(transformations[None, ...], (pts.size()[0], 1, 1))  # [R, 4, 4]
    pts_transformed = []
    for i in range(0, pts.size()[0], chunk):
        pts_transformed_chunk = apply_batched_transformations(
            inputs=[pts[i:i+chunk]], transformations=transformations[i:i+chunk])[0]
        pts_transformed.append(pts_transformed_chunk)
    pts_transformed = torch.cat(pts_transformed, dim=0)
    return pts_transformed

# def transform_rays(ray_batch, translation, use_viewdirs):
#    """Apply transformation to rays.

#    Args:
#        ray_batch: [R, M] float tensor. All information necessary
#            for sampling along a ray, including: ray origin, ray direction, min
#            dist, max dist, and unit-magnitude viewing direction.
#        translation: [3,] float tensor. The (x, y, z) translation to apply.
#        use_viewdirs: Whether to use view directions.

#    Returns:
#        ray_batch: [R, M] float tensor. Transformed ray batch.
#    """
#   assert translation.size()[0] == 3, "translation.size()[0] must be 3..."

#    # Since we are only supporting translation for now, only ray origins need to be
#    # modified. Ray directions do not need to change.
#    rays_o = ray_batch[:, 0:3] + translation
#    rays_remaining = ray_batch[:, 3:]
#    ray_batch = torch.cat((rays_o, rays_remaining), dim=1)
#    return ray_batch

def compute_rays_length(rays_d):
    """Compute ray length.

    Args:
        rays_d: [R, 3] float tensor. Ray directions.

    Returns:
        rays_length: [R, 1] float tensor. Ray lengths.
    """
    rays_length = torch.norm(rays_d, dim=-1, keepdim=True)  # [N_rays, 1]
    return rays_length

def normalize_rays(rays):
    """Normalize ray directions.

    Args:
        rays: [R, 3] float tensor. Ray directions.

    Returns:
        normalized_rays: [R, 3] float tensor. Normalized ray directions.
    """
    normalized_rays = rays / compute_rays_length(rays_d=rays)
    return normalized_rays

def compute_ray_dirs_and_length(rays_o, rays_dst):
    """Compute ray directions.

    Args:
        rays_o: [R, 3] float tensor. Ray origins.
        rays_dst: [R, 3] float tensor. Ray destinations.

    Returns:
        rays_d: [R, 3] float tensor. Normalized ray directions.
    """
    # The ray directions are the difference between the ray destinations and the
    # ray origins.
    rays_d = rays_dst - rays_o  # [R, 3]  # Direction out of light source

    # Compute the length of the rays.
    rays_length = compute_rays_length(rays_d=rays_d)

    # Normalized the ray directions.
    rays_d = rays_d / rays_length  # [R, 3]  # Normalize direction
    return rays_d, rays_length

def update_ray_batch_slice(ray_batch, x, start, end):
    left = ray_batch[:, :start]  # [R, ?]
    right = ray_batch[:, end:]  # [R, ?]
    updated_ray_batch = torch.cat((left, x, right), dim=-1)
    return updated_ray_batch

def update_ray_batch_bounds(ray_batch, bounds):
    updated_ray_batch = update_ray_batch_slice(ray_batch=ray_batch, x=bounds,
                                               start=6, end=8)
    return updated_ray_batch

def create_ray_batch(
    rays_o, rays_dst, rays_i, use_viewdirs, rays_near=None, rays_far=None, epsilon=1e-10):
    # Compute the ray directions.
    rays_d = rays_dst - rays_o  # [R,3]  # Direction out of light source
    rays_length = compute_rays_length(rays_d=rays_d)  # [R, 1]
    rays_d = rays_d / rays_length  # [R, 3]  # Normalize direction
    viewdirs = rays_d  # [R, 3]

    # If bounds are not provided, set the beginning and end of ray as sampling bounds.
    if rays_near is None:
        rays_near = torch.zeros((rays_o.size()[0], 1), dtype=torch.float) + epsilon  # [R, 1]
    if rays_far is None:
        rays_far = rays_length  # [R, 1]

    ray_batch = torch.cat((rays_o, rays_d, rays_near, rays_far), dim=-1)
    if use_viewdirs:
        ray_batch = torch.cat((ray_batch, viewdirs), dim=-1)
    ray_batch = torch.cat((ray_batch, rays_i), dim=-1)
    return ray_batch

def sample_random_lightdirs(num_rays, num_samples, upper_only=False):
    """Randomly sample directions in the unit sphere.

    Args:
        num_rays: int or tensor shape dimension. Number of rays.
        num_samples: int or tensor shape dimension. Number of samples per ray.
        upper_only: bool. Whether to sample only on the upper hemisphere.

    Returns:
        lightdirs: [R, S, 3] float tensor. Random light directions sampled from the unit
            sphere for each sampled point.
    """
    if upper_only:
        min_z = 0
    else:
        min_z = -1

    phi = torch.rand(num_rays, num_samples) * (2 * math.pi)  # [R, S]
    cos_theta = torch.rand(num_rays, num_samples) * (1 - min_z) + min_z  # [R, S]
    theta = torch.acos(cos_theta)  # [R, S]

    x = torch.sin(theta) * torch.cos(phi)
    y = torch.sin(theta) * torch.sin(phi)
    z = torch.cos(theta)

    lightdirs = torch.cat((x[..., None], y[..., None], z[..., None]), dim=-1)  # [R, S, 3]
    return lightdirs

def get_light_positions(rays_i, img_light_pos):
    """Extracts light positions given scene IDs.

    Args:
        rays_i: [R, 1] float tensor. Per-ray image IDs.
        img_light_pos: [N, 3] float tensor. Per-image light positions.

    Returns:
        rays_light_pos: [R, 3] float tensor. Per-ray light positions.
    """
    rays_light_pos = img_light_pos[rays_i.long()].squeeze()  # [R, 3]
    return rays_light_pos

def get_lightdirs(lightdirs_method, num_rays=None, num_samples=None, rays_i=None,
    metadata=None, ray_batch=None, use_viewdirs=False, normalize=False):
    """Compute lightdirs.

    Args:
        lightdirs_method: str. Method to use for computing lightdirs.
        num_rays: int or tensor shape dimension. Number of rays.
        num_samples: int or tensor shape dimension. Number of samples per ray.
        rays_i: [R, 1] float tensor. Ray image IDs.
        metadata: [N, 3] float tensor. Metadata about each image. Currently only light
            position is provided.
        ray_batch: [R, M] float tensor. Ray batch.
        use_viewdirs: bool. Whether to use viewdirs.
        normalize: bool. Whether to normalize lightdirs.

    Returns;
        lightdirs: [R, S, 3] float tensor. Light directions for each sample.
    """
    if lightdirs_method == 'viewdirs':
        raise NotImplementedError
        assert use_viewdirs
        lightdirs = ray_batch[:, 8:11]  # [R, 3]
        lightdirs *= 1.5
        lightdirs = torch.tile(lightdirs[:, None, :], (1, num_samples, 1))
    elif lightdirs_method == 'metadata':
        lightdirs = get_light_positions(rays_i, metadata)  # [R, 3]
        lightdirs = torch.tile(lightdirs[:, None, :], (1, num_samples, 1))  # [R, S, 3]
    elif lightdirs_method == 'random':
        lightdirs = sample_random_lightdirs(num_rays, num_samples)  # [R, S, 3]
    elif lightdirs_method == 'random_upper':
        lightdirs = sample_random_lightdirs(num_rays, num_samples, upper_only=True)  # [R, S, 3]
    else:
        raise ValueError(f'Invalid lightdirs_method: {lightdirs_method}.')
    if normalize:
        lightdirs_flat = lightdirs.view(-1, 3)  # [RS, 3]
        lightdirs_flat = normalize_rays(lightdirs_flat)  # [RS, 3]
        lightdirs = lightdirs_flat.view(lightdirs.size())  # [R, S, 3]
    return lightdirs


<gh_stars>0
import argparse
import cv2
import numpy as np
import os
import pickle
import sys
from gaussian2d import gaussian2d
from hashkey import hashkey
from math import floor, pi
from matplotlib import pyplot as plt
from scipy import interpolate
from skimage import transform
from numba import njit, prange

def bicubic2x(mat):
    mch=False
    if len(mat.shape)==3:
        h,w,ch=mat.shape
        mch=True
    if mch:
        ret = np.zeros((h*2,w*2,ch))
        for i in range(ch):
            ret[:,:,i]=_bicubic2x(mat[:,:,i])
        return ret
    else:
        return _bicubic2x(mat)

#@njit(parallel=True)
def _bicubic2x(mat):
    h,w=mat.shape
    heightgridLR = np.linspace(0,h-1,h)
    widthgridLR = np.linspace(0,w-1,w)
    heightgridHR = np.linspace(0,h-0.5,h*2)
    widthgridHR = np.linspace(0,w-0.5,w*2)
    interp=interpolate.interp2d(widthgridLR, heightgridLR, mat, kind='cubic')
    return interp(widthgridHR,heightgridHR)

def bicubic0_5x(mat):
    mch=False
    if len(mat.shape)==3:
        h,w,ch=mat.shape
        mch=True
    if mch:
        ret = np.zeros((int(h/2),int(w/2),ch))
        for i in range(ch):
            ret[:,:,i]=_bicubic0_5x(mat[:,:,i])
        return ret
    else:
        return _bicubic0_5x(mat)

#@njit(parallel=True)
def _bicubic0_5x(mat):
    h,w=mat.shape
    heightgridHR = np.linspace(0,h-1,h)
    widthgridHR = np.linspace(0,w-1,w)
    heightgridLR = np.linspace(0,h-2,h/2)
    widthgridLR = np.linspace(0,w-2,w/2)
    interp=interpolate.interp2d(widthgridHR, heightgridHR, mat, kind='cubic')
    return interp(widthgridLR,heightgridLR)


from statistics import mean, median
from time import perf_counter

class Timer:
    def __init__(self, func=None):
        self._start = None
        self._finish = None
        self.runs = list()
        self._func = func

    def __enter__(self):
        self._start = perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._finish = perf_counter()
        self.runs.append(self._finish - self._start)

    def __call__(self, *args, **kwargs):
        with self:
            return self._func(*args, **kwargs)

    @property
    def elapsed(self):
        return sum(self.runs)

    @property
    def min(self):
        return min(self.runs)

    @property
    def max(self):
        return max(self.runs)

    @property
    def mean(self):
        return mean(self.runs)

    @property
    def median(self):
        return median(self.runs)