from IPython.core.display import display, HTML
display(HTML("<style>.container { width:80% !important; }</style>")) 
%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.gridspec import GridSpec
import datetime
import scipy.fftpack as fft
import numpy as np
import pandas as pd
from obspy.clients.fdsn import Client
from obspy.signal.spectral_estimation import PPSD
from obspy import UTCDateTime

from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()

import matplotlib as mpl
mpl.rcParams['date.autoformatter.day'] = "%Y-%m-%d"
mpl.rcParams['date.autoformatter.hour'] = "%Y-%m-%d %Hh"
mpl.rcParams['date.autoformatter.minute'] = "%Y-%m-%d %H:%M"

mpl.rcParams['figure.figsize'] = "12,6"
mpl.rcParams['figure.dpi'] = 100

s = UTCDateTime("2012-02-03")
c = Client("IRIS")
st = c.get_waveforms("BE","UCC","","HHZ", s, s+86400, attach_response=True)
inv = c.get_stations(network="BE", station="UCC", level="response")

acceleration = st.copy().remove_response(output="ACC", pre_filt=(0.01,0.02,30.0,35.0))

velocity = st.copy().remove_response(output="VEL", pre_filt=(0.01,0.02,30.0,35.0))

displacement = st.copy().remove_response(output="DISP", pre_filt=(0.01,0.02,30.0,35.0))

st.plot();
acceleration.plot();
velocity.plot();
displacement.plot();

ppsd = PPSD(st[0].stats, metadata=inv,
            ppsd_length=600.0,
            overlap=0.0,
            period_smoothing_width_octaves=0.010,
            period_step_octaves=0.0125,
            period_limits=(0.02, 50),
            db_bins=(-200, 20, 0.5))
# ppsd._nfft *= 32
# ppsd.nfft

ppsd.add(st)

d:\pythonforsource\msnoise_stack\obspy-git\obspy\signal\spectral_estimation.py:1000: RuntimeWarning: Mean of empty slice.
  smoothed_psd.append(specs.mean())
C:\ProgramData\Anaconda3\envs\py37\lib\site-packages\numpy\core\_methods.py:85: RuntimeWarning: invalid value encountered in double_scalars
  ret = ret.dtype.type(ret / rcount)

True

ppsd.plot()
ppsd.plot_spectrogram()

d:\pythonforsource\msnoise_stack\obspy-git\obspy\signal\spectral_estimation.py:2098: MatplotlibDeprecationWarning: 
The set_clim function was deprecated in Matplotlib 3.1 and will be removed in 3.3. Use ScalarMappable.set_clim instead.
  cb.set_clim(*fig.ppsd.color_limits)

10**(1/2.) == np.sqrt(10**1)

True

velocity_spectrum = []
displacement_spectrum = []
acc_spectrum = []


permin = 1/20.
permax = 1/1.

per = ppsd.period_bin_centers
ix = np.where((per>=permin) & (per<=permax))

for psd in ppsd.psd_values:
    # acceleration power spectrum in Hz
    spec = psd[ix][::-1]
    f = 1.0/per[ix][::-1]

    # remove NaNs from the list
    valid = np.where(np.isfinite(spec))[0]
    spec = spec[valid]
    f = f[valid]
    
    # The acceleration amplitude spectrum (dB to Power! = divide by 10 and not 20!)
    amp = 10.0**(spec/10.) 
    acc_spectrum.append(amp)
    
    # velocity spectrum
    vamp = amp / (2.0 * np.pi * f)**2
    velocity_spectrum.append(vamp)
    
    # displacement spectrum
    damp =  vamp / (2.0 * np.pi * f)**2
    displacement_spectrum.append(damp)
#     break
    

index = pd.DatetimeIndex(pd.date_range(st[0].stats.starttime.datetime, periods=len(velocity_spectrum), freq="600S"))    

velocity_spectrum = pd.DataFrame(velocity_spectrum, index=index, columns= f )
displacement_spectrum = pd.DataFrame(displacement_spectrum, index=index, columns= f )
acc_spectrum = pd.DataFrame(acc_spectrum, index=index, columns= f )

velocity_spectrum = velocity_spectrum.resample("600S").mean()
displacement_spectrum = displacement_spectrum.resample("600S").mean()
acc_spectrum = acc_spectrum.resample("600S").mean()

fig = plt.figure()
plt.pcolormesh(acc_spectrum.index, acc_spectrum.columns, acc_spectrum.T)
plt.colorbar().set_label("Spectral Amplitude")
plt.title("Acceleration, max = %.2e m/s²" % acc_spectrum.values.max())
plt.ylabel("Frequency (Hz)")
fig.autofmt_xdate()
plt.show()

fig = plt.figure()
plt.pcolormesh(velocity_spectrum.index, velocity_spectrum.columns, velocity_spectrum.T)
plt.colorbar().set_label("Spectral Amplitude")
plt.title("Velocity, max = %.2e m/s" % velocity_spectrum.values.max())
plt.ylabel("Frequency (Hz)")
fig.autofmt_xdate()
plt.show()

fig = plt.figure()
plt.pcolormesh(displacement_spectrum.index, displacement_spectrum.columns, displacement_spectrum.T)
plt.colorbar().set_label("Spectral Amplitude")
plt.title("Displacement, max = %.2e m" % displacement_spectrum.values.max())
plt.ylabel("Frequency (Hz)")
fig.autofmt_xdate()
plt.show()

def rms(a):
    return np.sqrt(np.trapz(a.values, a.index))

# converting ObsPy Stream to dataframe for easy resampling (could do st.slide() etc... but it's easy like this)
index = pd.DatetimeIndex(pd.date_range(acceleration[0].stats.starttime.datetime, periods=acceleration[0].stats.npts, freq="10ms"))
acc = acceleration.copy().filter("bandpass", freqmin=1./permax, freqmax=1./permin, corners=16)
acc = pd.DataFrame(acc[0].data, index=index).dropna()

# converting ObsPy Stream to dataframe for easy resampling (could do st.slide() etc... but it's easy like this)
index = pd.DatetimeIndex(pd.date_range(velocity[0].stats.starttime.datetime, periods=velocity[0].stats.npts, freq="10ms"))
vel = velocity.copy().filter("bandpass", freqmin=1./permax, freqmax=1./permin, corners=16)
vel = pd.DataFrame(vel[0].data, index=index).dropna()

# converting ObsPy Stream to dataframe for easy resampling (could do st.slide() etc... but it's easy like this)
index = pd.DatetimeIndex(pd.date_range(displacement[0].stats.starttime.datetime, periods=displacement[0].stats.npts, freq="10ms"))
dis = displacement.copy().filter("bandpass", freqmin=1./permax, freqmax=1./permin, corners=16)
dis = pd.DataFrame(dis[0].data, index=index)

scale = 1e6
ticks = ticker.FuncFormatter(lambda x, pos: '{0:g}'.format(x*scale))

for label, item, spectrum in [("Acceleration µm/s²", acc, acc_spectrum), ("Velocity µm/s", vel, velocity_spectrum), ("Displacement µm", dis, displacement_spectrum)]:
    # Getting the RMS from the PSD
    a = spectrum.apply(rms, axis=1)
    
    # Getting the RMS from the time series
    r = item.resample("600S")
    xSTD = r.std()
    
    # Plot Title
    title = "Filtered data [%.2f-%.2f] Hz" % (1./permax, 1./permin)
    
    # Plotting
    plt.figure(figsize=(16,10,))
    gs = GridSpec(2,1, height_ratios=[1,4])

    ax = plt.subplot(gs[0])
    plt.plot(item.index, item, lw=1, c="k", label=title)
    plt.plot(xSTD.index, xSTD, label="RMS", c='orange')
    plt.plot(xSTD.index, xSTD*3, label="3*RMS", c='r')
    plt.gca().yaxis.set_major_formatter(ticks)
    plt.ylabel(label)
    plt.legend(loc=4, ncol=3, fontsize=14)

    
    ax2 = plt.subplot(gs[1], sharex=ax)
    plt.plot(a.index, a, label="RMS (PSD)")
    plt.plot(xSTD.index, xSTD, label="RMS (time domain std)")
    plt.gca().yaxis.set_major_formatter(ticks)
    plt.ylabel(label)

    plt.axvline(datetime.datetime(2012,2,3,14,20), label="Snowing in Brussels", ls="--", c="silver")

    plt.legend(ncol=1, fontsize=14)
    plt.grid()
    fig.autofmt_xdate()
    plt.xlim(xSTD.index[0], xSTD.index[-1])
    plt.tight_layout()
    plt.savefig("%s - %s.png" % (title, label.split(" ")[0]), dpi=300)
    plt.show()
#     break