''' Module to return an axis that has energy-time diagram for CENA.
'''
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.dates import DateFormatter, AutoDateLocator
from matplotlib.colors import LogNorm
import logging
__logger = logging.getLogger('axis_timerange')
__logger.setLevel(logging.DEBUG)
[docs]def axis_time_spectra(xaxis, yaxis, data, rect=None, vmin=None, vmax=None, vscale='linear', voffset=0, colorbar=False, kwds_axes=None, **kwds):
''' Draw time-something image to the current figure.
:param xaxis: Array-like for the time. Time should be matplotlib time.
``datetime.datetime`` should be processed prior by ``date2num()`` function.
:param yaxis: Array-like for y-axis.
:param data: The data to be plotted. len(xaxis) x len(yaxis).
Masked array is also supported.
:param rect: A rectangle to be given to axes. [xmin, ymin, width, height].
If you want to make several axes in one figure, this values look to be used
for identification. See *pylab.axes* for details.
:type data: ``numpy.array`` or ``numpy.ma.masked_array``
:keyword kwds_axes: A dictionaly of the keyword to be given to ``plt.axes()``.
For example, kwds_axes = {"axisbg":"k"} will change the background color.
:keyword dateformat: A format for the x_axis.
:keyword vscale: A scale of the colorbar. 'linear' or 'log'.
:returns: The axis produced.
'''
# fig = kwds.get('figure', plt.gcf())
dateformat = kwds.get('dateformat', "%FT%T")
if vmin is None:
vmin = data.min()
if vmax is None:
vmax = data.max()
if vscale == 'log':
norm = LogNorm(vmin=vmin, vmax=vmax)
data = data + voffset
else:
norm = None
if kwds_axes is None:
kwds_axes={}
X, Y = np.meshgrid(xaxis, yaxis)
__logger.debug('X.shape = %s' % str(X.shape))
__logger.debug('Y.shape = %s' % str(Y.shape))
__logger.debug('V.shape = %s' % str(data.shape))
if rect is not None:
ax = plt.axes(rect, **kwds_axes)
else:
ax = plt.axes(**kwds_axes)
pclr = ax.pcolor(X, Y, data.T, norm=norm, vmax=vmax, vmin=vmin)
formatter = DateFormatter(dateformat)
locator = AutoDateLocator()
ax.xaxis.set_major_formatter(formatter)
ax.xaxis.set_major_locator(locator)
if colorbar:
plt.colorbar(pclr)
return ax
[docs]def insert_datagap(xlist, datalist, gap=1):
''' Insert a datagap into the datalist.
:param xlist: List of ``float``. It should be monotonically increasing.
:param datalist: List of the data. Should be ``ndarray`` or ``masked_array``.
datalist[:, ...] should be the same size as ``datetime_list``.
Use transpose() to make the target axis to the first.
:returns: Gap enbedded data as tuple of array and masked_array.
:rtype: ``tuple of (array, masked_array)``
>>> from datetime import datetime
>>> from datetime import timedelta
xlist has 11 elements. 5 is missing.
>>> x = [0., 1, 2, 3, 4, 6, 7, 8, 9, 10, 11]
Data is defined here. I assume 11 x 3 x 8 x 7 array (4-D).
Note the time list has length of 11.
>>> dat = np.arange(11*3*8*7).reshape(11, 3, 8, 7)
>>> dat.shape
(11, 3, 8, 7)
Now, the data gap will be inserted
Before::
Index: 0 1 2 3 4 5 6 7 8 9 10
X: 0 1 2 3 4 6 7 8 9 10 11
Data: D0 D1 D2 D3 D4 D5 D6 D7 D8 D9 D10
After::
Index: 0 1 2 3 4 5 6 7 8 9 10 11
X: 0 1 2 3 4 5 6 7 8 9 10 11
Data: D0 D1 D2 D3 D4 M D5 D6 D7 D8 D9 D10
>>> x2, dat2 = insert_datagap(x, dat)
>>> print(len(x2))
12
>>> print(dat2.shape)
(12, 3, 8, 7)
>>> print(x2[5])
5
>>> print(dat2.mask[0, 2, 1, 4]) # mask[0, :, :, :] is False
False
>>> print(dat2.mask[5, 1, 7, 2]) # mask[5, :, :, :] is True
True
>>> print(dat2[3, 2, 5, 1] == dat[3, 2, 5, 1])
True
>>> print(dat2[9, 1, 3, 6] == dat[8, 1, 3, 6])
True
'''
if len(xlist) != len(datalist[:, ...]):
raise ValueError("List size mismatch. xlist=%d/datalist.shape[0]=%d"
% (len(xlist), len(datalist[:, 0, 0, 0])))
newlist = [xlist[0]]
newdatalist = [datalist[0]]
nan_datalist = np.ones_like(datalist[0]) * np.nan
for idx in range(1, len(xlist)):
if xlist[idx] > newlist[-1] + gap:
newlist.append(newlist[-1] + gap)
newdatalist.append(nan_datalist)
newlist.append(xlist[idx])
newdatalist.append(datalist[idx])
newdatalist = np.ma.masked_array(newdatalist)
newdatalist = np.ma.masked_invalid(newdatalist)
return newlist, newdatalist