# -*-Python-*-
# Created by bgrierson at 30 Aug 2017 16:11
"""
This script shows the difference between the native xarray reindexing and the OMFIT provided reindex_interp function.
"""
x = linspace(0, 2 * np.pi, 31)
t1 = linspace(0, 1, 11)
tt, xx = meshgrid(t1, x)
f1 = np.cos(xx - tt)
t2 = linspace(0.123704, 1.123, 41)
tt, xx = meshgrid(t2, x)
f2 = np.sin(xx - tt)
da1 = DataArray(f1, coords=[('x', x), ('time', t1)], name='some_data')
da2 = DataArray(f2, coords=[('x', x), ('time', t2)], name='other_data')
# we can interpolate each DataArray to a common time base
t = np.unique(np.hstack((t1, t2))) # use a combination of the two time bases
kws = dict(fill_value='extrapolate', bounds_error=False) # enable linear extrapolation beyond the original bounds
da1i = reindex_interp(da1, time=t, interpolate_kws=kws).rename('some_reindex_interp')
da2i = reindex_interp(da2, time=t, interpolate_kws=kws).rename('other_reindex_interp')
# and once they have matched dimensions, they can be gathered in a common "Dataset"
ds = Dataset({'da1': da1i, 'da2': da2i})
# finally, we check what we did visually
fig, ax = subplots(3, 2, sharex=True)
da1.plot(ax=ax[0, 0])
da1.reindex(time=t).rename('some_reindex').plot(ax=ax[1, 0])
da1i.plot(ax=ax[2, 0])
da2.plot(ax=ax[0, 1])
da2.reindex(time=t).rename('other_reindex').plot(ax=ax[1, 1])
da2i.plot(ax=ax[2, 1])
# make it pretty
fig.tight_layout()
