deepdenoiser/data_reader.py

import numpy as np
import pandas as pd
import scipy.signal
import tensorflow as tf

pd.options.mode.chained_assignment = None
import logging
import os
import threading

import obspy
from scipy.interpolate import interp1d

tf.compat.v1.disable_eager_execution()
# from tensorflow.python.ops.linalg_ops import norm
# from tensorflow.python.util import nest


class Config:
    seed = 100
    n_class = 2
    fs = 100
    dt = 1.0 / fs
    freq_range = [0, fs / 2]
    time_range = [0, 30]
    nperseg = 30
    nfft = 60
    plot = False
    nt = 3000
    X_shape = [31, 201, 2]
    Y_shape = [31, 201, n_class]
    signal_shape = [31, 201]
    noise_shape = signal_shape
    use_seed = False
    queue_size = 10
    noise_mean = 2
    noise_std = 1
    # noise_low = 1
    # noise_high = 5
    use_buffer = True
    snr_threshold = 10


# %%
# def normalize(data, window=3000):
#     """
#     data: nsta, chn, nt
#     """
#     shift = window//2
#     nt = len(data)

#     ## std in slide windows
#     data_pad = np.pad(data, ((window//2, window//2)), mode="reflect")
#     t = np.arange(0, nt, shift, dtype="int")
#     # print(f"nt = {nt}, nt+window//2 = {nt+window//2}")
#     std = np.zeros(len(t))
#     mean = np.zeros(len(t))
#     for i in range(len(std)):
#         std[i] = np.std(data_pad[i*shift:i*shift+window])
#         mean[i] = np.mean(data_pad[i*shift:i*shift+window])

#     t = np.append(t, nt)
#     std = np.append(std, [np.std(data_pad[-window:])])
#     mean = np.append(mean, [np.mean(data_pad[-window:])])

#     # print(t)
#     ## normalize data with interplated std
#     t_interp = np.arange(nt, dtype="int")
#     std_interp = interp1d(t, std, kind="slinear")(t_interp)
#     mean_interp = interp1d(t, mean, kind="slinear")(t_interp)
#     data = (data - mean_interp)/(std_interp)
#     return data, std_interp

# %%
def normalize(data, window=200):
    """
    data: nsta, chn, nt
    """
    shift = window // 2
    nt = data.shape[1]

    ## std in slide windows
    data_pad = np.pad(data, ((0, 0), (window // 2, window // 2), (0, 0)), mode="reflect")
    t = np.arange(0, nt, shift, dtype="int")
    # print(f"nt = {nt}, nt+window//2 = {nt+window//2}")
    std = np.zeros(len(t))
    mean = np.zeros(len(t))
    for i in range(len(std)):
        std[i] = np.std(data_pad[:, i * shift : i * shift + window, :])
        mean[i] = np.mean(data_pad[:, i * shift : i * shift + window, :])

    t = np.append(t, nt)
    std = np.append(std, [np.std(data_pad[:, -window:, :])])
    mean = np.append(mean, [np.mean(data_pad[:, -window:, :])])
    # print(t)
    ## normalize data with interplated std
    t_interp = np.arange(nt, dtype="int")
    std_interp = interp1d(t, std, kind="slinear")(t_interp)
    std_interp[std_interp == 0] = 1.0
    mean_interp = interp1d(t, mean, kind="slinear")(t_interp)
    data = (data - mean_interp[np.newaxis, :, np.newaxis]) / std_interp[np.newaxis, :, np.newaxis]
    return data, std_interp


def normalize_batch(data, window=200):
    """
    data: nbn, nf, nt, 2
    """
    assert len(data.shape) == 4
    shift = window // 2
    nbt, nf, nt, nimg = data.shape

    ## std in slide windows
    data_pad = np.pad(data, ((0, 0), (0, 0), (window // 2, window // 2), (0, 0)), mode="reflect")
    t = np.arange(0, nt + shift - 1, shift, dtype="int") # 201 => 0, 100, 200
    std = np.zeros([nbt, len(t)])
    mean = np.zeros([nbt, len(t)])
    for i in range(std.shape[1]):
        std[:, i] = np.std(data_pad[:, :, i * shift : i * shift + window, :], axis=(1, 2, 3))
        mean[:, i] = np.mean(data_pad[:, :, i * shift : i * shift + window, :], axis=(1, 2, 3))
    
    std[:, -1], mean[:, -1] = std[:, -2], mean[:, -2]
    std[:, 0], mean[:, 0] = std[:, 1], mean[:, 1]

    ## normalize data with interplated std
    t_interp = np.arange(nt, dtype="int")
    std_interp = interp1d(t, std, kind="slinear")(t_interp)  ##nbt, nt
    std_interp[std_interp == 0] = 1.0
    mean_interp = interp1d(t, mean, kind="slinear")(t_interp)

    data = (data - mean_interp[:, np.newaxis, :, np.newaxis]) / std_interp[:, np.newaxis, :, np.newaxis]

    if len(t) > 3:  ##need to address this normalization issue in training
        data /= 2.0

    return data


# %%
def py_func_decorator(output_types=None, output_shapes=None, name=None):
    def decorator(func):
        def call(*args, **kwargs):
            nonlocal output_shapes
            # flat_output_types = nest.flatten(output_types)
            flat_output_types = tf.nest.flatten(output_types)
            # flat_values = tf.py_func(
            flat_values = tf.numpy_function(func, inp=args, Tout=flat_output_types, name=name)
            if output_shapes is not None:
                for v, s in zip(flat_values, output_shapes):
                    v.set_shape(s)
            # return nest.pack_sequence_as(output_types, flat_values)
            return tf.nest.pack_sequence_as(output_types, flat_values)

        return call

    return decorator


def dataset_map(iterator, output_types, output_shapes=None, num_parallel_calls=None, name=None):
    dataset = tf.data.Dataset.range(len(iterator))

    @py_func_decorator(output_types, output_shapes, name=name)
    def index_to_entry(idx):
        return iterator[idx]

    return dataset.map(index_to_entry, num_parallel_calls=num_parallel_calls)


class DataReader(object):
    def __init__(
        self,
        signal_dir=None,
        signal_list=None,
        noise_dir=None,
        noise_list=None,
        queue_size=None,
        coord=None,
        config=Config(),
    ):

        self.config = config

        signal_list = pd.read_csv(signal_list, header=0)
        noise_list = pd.read_csv(noise_list, header=0)

        self.signal = signal_list
        self.noise = noise_list
        self.n_signal = len(self.signal)

        self.signal_dir = signal_dir
        self.noise_dir = noise_dir

        self.X_shape = config.X_shape
        self.Y_shape = config.Y_shape
        self.n_class = config.n_class

        self.coord = coord
        self.threads = []
        self.queue_size = queue_size

        self.add_queue()
        self.buffer_signal = {}
        self.buffer_noise = {}
        self.buffer_channels_signal = {}
        self.buffer_channels_noise = {}

    def add_queue(self):
        with tf.device('/cpu:0'):
            self.sample_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None)
            self.target_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None)
            self.queue = tf.queue.PaddingFIFOQueue(
                self.queue_size, ['float32', 'float32'], shapes=[self.config.X_shape, self.config.Y_shape]
            )
            self.enqueue = self.queue.enqueue([self.sample_placeholder, self.target_placeholder])
        return 0

    def dequeue(self, num_elements):
        output = self.queue.dequeue_many(num_elements)
        return output

    def get_snr(self, data, itp, dit=300):
        tmp_std = np.std(data[itp - dit : itp])
        if tmp_std > 0:
            return np.std(data[itp : itp + dit]) / tmp_std
        else:
            return 0

    def add_event(self, sample, channels, j):
        while np.random.uniform(0, 1) < 0.2:
            shift = None
            if channels not in self.buffer_channels_signal:
                self.buffer_channels_signal[channels] = self.signal[self.signal['channels'] == channels]
            fname = os.path.join(self.signal_dir, self.buffer_channels_signal[channels].sample(n=1).iloc[0]['fname'])
            try:
                if fname not in self.buffer_signal:
                    meta = np.load(fname)
                    data_FT = []
                    snr = []
                    for i in range(3):
                        tmp_data = meta['data'][:, i]
                        tmp_itp = meta['itp']
                        snr.append(self.get_snr(tmp_data, tmp_itp))
                        tmp_data -= np.mean(tmp_data)
                        f, t, tmp_FT = scipy.signal.stft(
                            tmp_data,
                            fs=self.config.fs,
                            nperseg=self.config.nperseg,
                            nfft=self.config.nfft,
                            boundary='zeros',
                        )
                        data_FT.append(tmp_FT)
                    data_FT = np.stack(data_FT, axis=-1)
                    self.buffer_signal[fname] = {
                        'data_FT': data_FT,
                        'itp': tmp_itp,
                        'channels': meta['channels'],
                        'snr': snr,
                    }
                meta_signal = self.buffer_signal[fname]
            except:
                logging.error("Failed reading signal: {}".format(fname))
                continue
            if meta_signal['snr'][j] > self.config.snr_threshold:
                tmp_signal = np.zeros([self.X_shape[0], self.X_shape[1]], dtype=np.complex_)
                shift = np.random.randint(-self.X_shape[1], 1, None, 'int')
                tmp_signal[:, -shift:] = meta_signal['data_FT'][:, self.X_shape[1] : 2 * self.X_shape[1] + shift, j]
                if np.isinf(tmp_signal).any() or np.isnan(tmp_signal).any() or (not np.any(tmp_signal)):
                    continue
                tmp_signal = tmp_signal / np.std(tmp_signal)
                sample += tmp_signal / np.random.uniform(1, 5)
        return sample

    def thread_main(self, sess, n_threads=1, start=0):
        stop = False
        while not stop:
            index = list(range(start, self.n_signal, n_threads))
            np.random.shuffle(index)
            for i in index:
                fname_signal = os.path.join(self.signal_dir, self.signal.iloc[i]['fname'])
                try:
                    if fname_signal not in self.buffer_signal:
                        meta = np.load(fname_signal)
                        data_FT = []
                        snr = []
                        for j in range(3):
                            tmp_data = meta['data'][..., j]
                            tmp_itp = meta['itp']
                            snr.append(self.get_snr(tmp_data, tmp_itp))
                            tmp_data -= np.mean(tmp_data)
                            f, t, tmp_FT = scipy.signal.stft(
                                tmp_data,
                                fs=self.config.fs,
                                nperseg=self.config.nperseg,
                                nfft=self.config.nfft,
                                boundary='zeros',
                            )
                            data_FT.append(tmp_FT)
                        data_FT = np.stack(data_FT, axis=-1)
                        self.buffer_signal[fname_signal] = {
                            'data_FT': data_FT,
                            'itp': tmp_itp,
                            'channels': meta['channels'],
                            'snr': snr,
                        }
                    meta_signal = self.buffer_signal[fname_signal]
                except:
                    logging.error("Failed reading signal: {}".format(fname_signal))
                    continue
                channels = meta_signal['channels'].tolist()
                start_tp = meta_signal['itp'].tolist()

                if channels not in self.buffer_channels_noise:
                    self.buffer_channels_noise[channels] = self.noise[self.noise['channels'] == channels]
                fname_noise = os.path.join(
                    self.noise_dir, self.buffer_channels_noise[channels].sample(n=1).iloc[0]['fname']
                )
                try:
                    if fname_noise not in self.buffer_noise:
                        meta = np.load(fname_noise)
                        data_FT = []
                        for i in range(3):
                            tmp_data = meta['data'][: self.config.nt, i]
                            tmp_data -= np.mean(tmp_data)
                            f, t, tmp_FT = scipy.signal.stft(
                                tmp_data,
                                fs=self.config.fs,
                                nperseg=self.config.nperseg,
                                nfft=self.config.nfft,
                                boundary='zeros',
                            )
                            data_FT.append(tmp_FT)
                        data_FT = np.stack(data_FT, axis=-1)
                        self.buffer_noise[fname_noise] = {'data_FT': data_FT, 'channels': meta['channels']}
                    meta_noise = self.buffer_noise[fname_noise]
                except:
                    logging.error("Failed reading noise: {}".format(fname_noise))
                    continue

                if self.coord.should_stop():
                    stop = True
                    break

                j = np.random.choice([0, 1, 2])
                if meta_signal['snr'][j] <= self.config.snr_threshold:
                    continue

                tmp_noise = meta_noise['data_FT'][..., j]
                if np.isinf(tmp_noise).any() or np.isnan(tmp_noise).any() or (not np.any(tmp_noise)):
                    continue
                tmp_noise = tmp_noise / np.std(tmp_noise)

                tmp_signal = np.zeros([self.X_shape[0], self.X_shape[1]], dtype=np.complex_)
                if np.random.random() < 0.9:
                    shift = np.random.randint(-self.X_shape[1], 1, None, 'int')
                    tmp_signal[:, -shift:] = meta_signal['data_FT'][:, self.X_shape[1] : 2 * self.X_shape[1] + shift, j]
                    if np.isinf(tmp_signal).any() or np.isnan(tmp_signal).any() or (not np.any(tmp_signal)):
                        continue
                    tmp_signal = tmp_signal / np.std(tmp_signal)
                    tmp_signal = self.add_event(tmp_signal, channels, j)

                    if np.random.random() < 0.2:
                        tmp_signal = np.fliplr(tmp_signal)

                ratio = 0
                while ratio <= 0:
                    ratio = self.config.noise_mean + np.random.randn() * self.config.noise_std
                # ratio = np.random.uniform(self.config.noise_low, self.config.noise_high)
                tmp_noisy_signal = tmp_signal + ratio * tmp_noise
                noisy_signal = np.stack([tmp_noisy_signal.real, tmp_noisy_signal.imag], axis=-1)
                if np.isnan(noisy_signal).any() or np.isinf(noisy_signal).any():
                    continue
                noisy_signal = noisy_signal / np.std(noisy_signal)
                tmp_mask = np.abs(tmp_signal) / (np.abs(tmp_signal) + np.abs(ratio * tmp_noise) + 1e-4)
                tmp_mask[tmp_mask >= 1] = 1
                tmp_mask[tmp_mask <= 0] = 0
                mask = np.zeros([tmp_mask.shape[0], tmp_mask.shape[1], self.n_class])
                mask[:, :, 0] = tmp_mask
                mask[:, :, 1] = 1 - tmp_mask
                sess.run(self.enqueue, feed_dict={self.sample_placeholder: noisy_signal, self.target_placeholder: mask})

    def start_threads(self, sess, n_threads=8):
        for i in range(n_threads):
            thread = threading.Thread(target=self.thread_main, args=(sess, n_threads, i))
            thread.daemon = True
            thread.start()
            self.threads.append(thread)
        return self.threads


class DataReader_test(DataReader):
    def __init__(
        self,
        signal_dir=None,
        signal_list=None,
        noise_dir=None,
        noise_list=None,
        queue_size=None,
        coord=None,
        config=Config(),
    ):
        self.config = config

        signal_list = pd.read_csv(signal_list, header=0)
        noise_list = pd.read_csv(noise_list, header=0)
        self.signal = signal_list
        self.noise = noise_list
        self.n_signal = len(self.signal)

        self.signal_dir = signal_dir
        self.noise_dir = noise_dir

        self.X_shape = config.X_shape
        self.Y_shape = config.Y_shape
        self.n_class = config.n_class

        self.coord = coord
        self.threads = []
        self.queue_size = queue_size

        self.add_queue()
        self.buffer_signal = {}
        self.buffer_noise = {}
        self.buffer_channels_signal = {}
        self.buffer_channels_noise = {}

    def add_queue(self):
        self.sample_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None)
        self.target_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None)
        self.ratio_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None)
        self.signal_placeholder = tf.compat.v1.placeholder(dtype=tf.complex64, shape=None)
        self.noise_placeholder = tf.compat.v1.placeholder(dtype=tf.complex64, shape=None)
        self.fname_placeholder = tf.compat.v1.placeholder(dtype=tf.string, shape=None)
        self.queue = tf.queue.PaddingFIFOQueue(
            self.queue_size,
            ['float32', 'float32', 'float32', 'complex64', 'complex64', 'string'],
            shapes=[
                self.config.X_shape,
                self.config.Y_shape,
                [],
                self.config.signal_shape,
                self.config.noise_shape,
                [],
            ],
        )
        self.enqueue = self.queue.enqueue(
            [
                self.sample_placeholder,
                self.target_placeholder,
                self.ratio_placeholder,
                self.signal_placeholder,
                self.noise_placeholder,
                self.fname_placeholder,
            ]
        )
        return 0

    def dequeue(self, num_elements):
        output = self.queue.dequeue_up_to(num_elements)
        return output

    def thread_main(self, sess, n_threads=1, start=0):
        index = list(range(start, self.n_signal, n_threads))
        for i in index:
            np.random.seed(i)

            fname = self.signal.iloc[i]['fname']
            fname_signal = os.path.join(self.signal_dir, fname)
            meta = np.load(fname_signal)
            data_FT = []
            snr = []
            for j in range(3):
                tmp_data = meta['data'][..., j]
                tmp_itp = meta['itp']
                snr.append(self.get_snr(tmp_data, tmp_itp))
                tmp_data -= np.mean(tmp_data)
                f, t, tmp_FT = scipy.signal.stft(
                    tmp_data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros'
                )
                data_FT.append(tmp_FT)
            data_FT = np.stack(data_FT, axis=-1)
            meta_signal = {'data_FT': data_FT, 'itp': tmp_itp, 'channels': meta['channels'], 'snr': snr}
            channels = meta['channels'].tolist()
            start_tp = meta['itp'].tolist()

            if channels not in self.buffer_channels_noise:
                self.buffer_channels_noise[channels] = self.noise[self.noise['channels'] == channels]
            fname_noise = os.path.join(
                self.noise_dir, self.buffer_channels_noise[channels].sample(n=1, random_state=i).iloc[0]['fname']
            )
            meta = np.load(fname_noise)
            data_FT = []
            for i in range(3):
                tmp_data = meta['data'][: self.config.nt, i]
                tmp_data -= np.mean(tmp_data)
                f, t, tmp_FT = scipy.signal.stft(
                    tmp_data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros'
                )
                data_FT.append(tmp_FT)
            data_FT = np.stack(data_FT, axis=-1)
            meta_noise = {'data_FT': data_FT, 'channels': meta['channels']}

            if self.coord.should_stop():
                stop = True
                break

            j = np.random.choice([0, 1, 2])
            tmp_noise = meta_noise['data_FT'][..., j]
            if np.isinf(tmp_noise).any() or np.isnan(tmp_noise).any() or (not np.any(tmp_noise)):
                continue
            tmp_noise = tmp_noise / np.std(tmp_noise)

            tmp_signal = np.zeros([self.X_shape[0], self.X_shape[1]], dtype=np.complex_)
            if np.random.random() < 0.9:
                shift = np.random.randint(-self.X_shape[1], 1, None, 'int')
                tmp_signal[:, -shift:] = meta_signal['data_FT'][:, self.X_shape[1] : 2 * self.X_shape[1] + shift, j]
                if np.isinf(tmp_signal).any() or np.isnan(tmp_signal).any() or (not np.any(tmp_signal)):
                    continue
                tmp_signal = tmp_signal / np.std(tmp_signal)
                # tmp_signal = self.add_event(tmp_signal, channels, j)
                # if np.random.random() < 0.2:
                #   tmp_signal = np.fliplr(tmp_signal)

            ratio = 0
            while ratio <= 0:
                ratio = self.config.noise_mean + np.random.randn() * self.config.noise_std
            tmp_noisy_signal = tmp_signal + ratio * tmp_noise
            noisy_signal = np.stack([tmp_noisy_signal.real, tmp_noisy_signal.imag], axis=-1)
            if np.isnan(noisy_signal).any() or np.isinf(noisy_signal).any():
                continue
            std_noisy_signal = np.std(noisy_signal)
            noisy_signal = noisy_signal / std_noisy_signal
            tmp_mask = np.abs(tmp_signal) / (np.abs(tmp_signal) + np.abs(ratio * tmp_noise) + 1e-4)
            tmp_mask[tmp_mask >= 1] = 1
            tmp_mask[tmp_mask <= 0] = 0
            mask = np.zeros([tmp_mask.shape[0], tmp_mask.shape[1], self.n_class])
            mask[:, :, 0] = tmp_mask
            mask[:, :, 1] = 1 - tmp_mask

            sess.run(
                self.enqueue,
                feed_dict={
                    self.sample_placeholder: noisy_signal,
                    self.target_placeholder: mask,
                    self.ratio_placeholder: std_noisy_signal,
                    self.signal_placeholder: tmp_signal,
                    self.noise_placeholder: ratio * tmp_noise,
                    self.fname_placeholder: fname,
                },
            )


class DataReader_pred_queue(DataReader):
    def __init__(self, signal_dir, signal_list, queue_size, coord, config=Config()):
        self.config = config
        signal_list = pd.read_csv(signal_list)
        self.signal = signal_list
        self.n_signal = len(self.signal)
        self.n_class = config.n_class
        self.X_shape = config.X_shape
        self.Y_shape = config.Y_shape
        self.signal_dir = signal_dir

        self.coord = coord
        self.threads = []
        self.queue_size = queue_size
        self.add_placeholder()

    def add_placeholder(self):
        self.sample_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None)
        self.ratio_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None)
        self.fname_placeholder = tf.compat.v1.placeholder(dtype=tf.string, shape=None)
        self.queue = tf.queue.PaddingFIFOQueue(
            self.queue_size, ['float32', 'float32', 'string'], shapes=[self.config.X_shape, [], []]
        )
        self.enqueue = self.queue.enqueue([self.sample_placeholder, self.ratio_placeholder, self.fname_placeholder])

    def dequeue(self, num_elements):
        output = self.queue.dequeue_up_to(num_elements)
        return output

    def thread_main(self, sess, n_threads=1, start=0):
        index = list(range(start, self.n_signal, n_threads))
        shift = 0
        for i in index:
            fname = self.signal.iloc[i]['fname']
            data_signal = np.load(os.path.join(self.signal_dir, fname))
            f, t, tmp_signal = scipy.signal.stft(
                scipy.signal.detrend(np.squeeze(data_signal['data'][shift : self.config.nt + shift])),
                fs=self.config.fs,
                nperseg=self.config.nperseg,
                nfft=self.config.nfft,
                boundary='zeros',
            )
            noisy_signal = np.stack([tmp_signal.real, tmp_signal.imag], axis=-1)
            if np.isnan(noisy_signal).any() or np.isinf(noisy_signal).any() or (not np.any(noisy_signal)):
                continue
            std_noisy_signal = np.std(noisy_signal)
            if std_noisy_signal == 0:
                continue
            noisy_signal = noisy_signal / std_noisy_signal
            sess.run(
                self.enqueue,
                feed_dict={
                    self.sample_placeholder: noisy_signal,
                    self.ratio_placeholder: std_noisy_signal,
                    self.fname_placeholder: fname,
                },
            )


class DataReader_pred:
    def __init__(self, signal_dir, signal_list, format="numpy", sampling_rate=100, config=Config()):
        self.buffer = {}
        self.config = config
        self.format = format
        self.dtype = "float32"
        try:
            signal_list = pd.read_csv(signal_list, sep="\t")["fname"]
        except:
            signal_list = pd.read_csv(signal_list)["fname"]
        self.signal_list = signal_list
        self.n_signal = len(self.signal_list)
        self.signal_dir = signal_dir
        self.sampling_rate = sampling_rate
        self.n_class = config.n_class
        FT_shape = self.get_data_shape()
        self.X_shape = [*FT_shape, 2]
        
    def get_data_shape(self):
        # fname = self.signal_list.iloc[0]['fname']
        # data = np.load(os.path.join(self.signal_dir, fname), allow_pickle=True)["data"]
        # data = np.squeeze(data)
        base_name = self.signal_list[0]
        if self.format == "numpy":
            meta = self.read_numpy(os.path.join(self.signal_dir, base_name))
        elif self.format == "mseed":
            meta = self.read_mseed(os.path.join(self.signal_dir, base_name))
        elif self.format == "hdf5":
            meta = self.read_hdf5(base_name)

        data = meta["data"]
        data = np.transpose(data, [2, 1, 0])

        if self.sampling_rate != 100:
            t = np.linspace(0, 1, data.shape[-1])
            t_interp = np.linspace(0, 1, np.int(np.around(data.shape[-1] * 100.0 / self.sampling_rate)))
            data = interp1d(t, data, kind="slinear")(t_interp)
        f, t, tmp_signal = scipy.signal.stft(
            data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros'
        )
        logging.info(f"Input data shape: {tmp_signal.shape} measured on file {base_name}")
        return tmp_signal.shape

    def __len__(self):
        return self.n_signal

    def read_numpy(self, fname):
        # try:
        if fname not in self.buffer:
            npz = np.load(fname)
            meta = {}
            if len(npz['data'].shape) == 1:
                meta["data"] = npz['data'][:, np.newaxis, np.newaxis]
            elif len(npz['data'].shape) == 2:
                meta["data"] = npz['data'][:, np.newaxis, :]
            else:
                meta["data"] = npz['data']
            if "p_idx" in npz.files:
                if len(npz["p_idx"].shape) == 0:
                    meta["itp"] = [[npz["p_idx"]]]
                else:
                    meta["itp"] = npz["p_idx"]
            if "s_idx" in npz.files:
                if len(npz["s_idx"].shape) == 0:
                    meta["its"] = [[npz["s_idx"]]]
                else:
                    meta["its"] = npz["s_idx"]
            if "t0" in npz.files:
                meta["t0"] = npz["t0"]
            self.buffer[fname] = meta
        else:
            meta = self.buffer[fname]
        return meta
        # except:
        #     logging.error("Failed reading {}".format(fname))
        #     return None

    def read_hdf5(self, fname):
        data = self.h5_data[fname][()]
        attrs = self.h5_data[fname].attrs
        meta = {}
        if len(data.shape) == 2:
            meta["data"] = data[:, np.newaxis, :]
        else:
            meta["data"] = data
        if "p_idx" in attrs:
            if len(attrs["p_idx"].shape) == 0:
                meta["itp"] = [[attrs["p_idx"]]]
            else:
                meta["itp"] = attrs["p_idx"]
        if "s_idx" in attrs:
            if len(attrs["s_idx"].shape) == 0:
                meta["its"] = [[attrs["s_idx"]]]
            else:
                meta["its"] = attrs["s_idx"]
        if "t0" in attrs:
            meta["t0"] = attrs["t0"]
        return meta

    def read_s3(self, format, fname, bucket, key, secret, s3_url, use_ssl):
        with self.s3fs.open(bucket + "/" + fname, 'rb') as fp:
            if format == "numpy":
                meta = self.read_numpy(fp)
            elif format == "mseed":
                meta = self.read_mseed(fp)
            else:
                raise (f"Format {format} not supported")
        return meta

    def read_mseed(self, fname):

        mseed = obspy.read(fname)
        mseed = mseed.detrend("spline", order=2, dspline=5 * mseed[0].stats.sampling_rate)
        mseed = mseed.merge(fill_value=0)
        starttime = min([st.stats.starttime for st in mseed])
        endtime = max([st.stats.endtime for st in mseed])
        mseed = mseed.trim(starttime, endtime, pad=True, fill_value=0)
        if mseed[0].stats.sampling_rate != self.sampling_rate:
            logging.warning(f"Sampling rate {mseed[0].stats.sampling_rate} != {self.sampling_rate} Hz")

        order = ['3', '2', '1', 'E', 'N', 'Z']
        order = {key: i for i, key in enumerate(order)}
        comp2idx = {"3": 0, "2": 1, "1": 2, "E": 0, "N": 1, "Z": 2}

        t0 = starttime.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
        nt = len(mseed[0].data)
        data = np.zeros([nt, 3], dtype=self.dtype)
        ids = [x.get_id() for x in mseed]
        if len(ids) == 3:
            for j, id in enumerate(sorted(ids, key=lambda x: order[x[-1]])):
                data[:, j] = mseed.select(id=id)[0].data.astype(self.dtype)
        else:
            if len(ids) > 3:
                logging.warning(f"More than 3 channels {ids}!")
            for jj, id in enumerate(ids):
                j = comp2idx[id[-1]]
                data[:, j] = mseed.select(id=id)[0].data.astype(self.dtype)

        data = data[:, np.newaxis, :]
        meta = {"data": data, "t0": t0}
        return meta

    def __getitem__(self, i):
        # fname = self.signal.iloc[i]['fname']
        # data = np.load(os.path.join(self.signal_dir, fname), allow_pickle=True)["data"]
        # data = np.squeeze(data)
        base_name = self.signal_list[i]

        if self.format == "numpy":
            meta = self.read_numpy(os.path.join(self.signal_dir, base_name))
        elif self.format == "mseed":
            meta = self.read_mseed(os.path.join(self.signal_dir, base_name))
        elif self.format == "hdf5":
            meta = self.read_hdf5(base_name)

        data = meta["data"]  # nt, 1, nch
        data = np.transpose(data, [2, 1, 0])  # nch, 1, nt
        if np.mod(data.shape[-1], 3000) == 1:  # 3001=>3000
            data = data[..., :-1]
        if "t0" in meta:
            t0 = meta["t0"]
        else:
            t0 = "1970-01-01T00:00:00.000"

        if self.sampling_rate != 100:
            logging.warning(f"Resample from {self.sampling_rate} to 100!")
            t = np.linspace(0, 1, data.shape[-1])
            t_interp = np.linspace(0, 1, np.int(np.around(data.shape[-1] * 100.0 / self.sampling_rate)))
            data = interp1d(t, data, kind="slinear")(t_interp)
        # sos = scipy.signal.butter(4, 0.1, 'high', fs=100, output='sos')  ## for stability of long sequence
        # data = scipy.signal.sosfilt(sos, data)
        f, t, tmp_signal = scipy.signal.stft(
            data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros'
        )  # nch, 1, nf, nt
        noisy_signal = np.stack([tmp_signal.real, tmp_signal.imag], axis=-1)  # nch, 1, nf, nt, 2
        noisy_signal[np.isnan(noisy_signal)] = 0
        noisy_signal[np.isinf(noisy_signal)] = 0
        # noisy_signal, std_noisy_signal = normalize(noisy_signal)
        # return noisy_signal.astype(self.dtype), std_noisy_signal.astype(self.dtype), fname

        return noisy_signal.astype(self.dtype), base_name, t0

    def dataset(self, batch_size, num_parallel_calls=4):
        dataset = dataset_map(
            self,
            output_types=(self.dtype, "string", "string"),
            output_shapes=(self.X_shape, None, None),
            num_parallel_calls=num_parallel_calls,
        )
        dataset = tf.compat.v1.data.make_one_shot_iterator(
            dataset.batch(batch_size).prefetch(batch_size * 3)
        ).get_next()
        return dataset


if __name__ == "__main__":

    # %%
    data_reader = DataReader_pred(signal_dir="./Dataset/yixiao/", signal_list="./Dataset/yixiao.csv")
    noisy_signal, std_noisy_signal, fname = data_reader[0]
    print(noisy_signal.shape, std_noisy_signal.shape, fname)
    batch = data_reader.dataset(10)
    init = tf.compat.v1.initialize_all_variables()
    sess = tf.compat.v1.Session()
    sess.run(init)
    print(sess.run(batch))