From 7f4571fe116bc54ec5f714e9bad864d876943531 Mon Sep 17 00:00:00 2001
From: "Michael Hirsch, Ph.D" <scivision@users.noreply.github.com>
Date: Thu, 29 Mar 2018 13:12:23 -0400
Subject: [PATCH] 3 color movie generation plots working

---
 dascutils/plots.py | 81 ++++++++++++++++++++++++++--------------------
 1 file changed, 46 insertions(+), 35 deletions(-)

diff --git a/dascutils/plots.py b/dascutils/plots.py
index 1f83925..6844954 100644
--- a/dascutils/plots.py
+++ b/dascutils/plots.py
@@ -1,12 +1,10 @@
 from pathlib import Path
-from os import devnull
-from numpy import arange
-from pytz import UTC
-from datetime import datetime
+import xarray
+import numpy as np
+from datetime import datetime, timedelta
 from scipy.interpolate import interp1d
 from matplotlib.pyplot import draw,pause,figure
 from matplotlib.colors import LogNorm
-import matplotlib.animation as anim
 #
 try:
     import themisasi.plots as themisplot
@@ -40,44 +38,57 @@ def histogram_dasc(imgs, odir=None):
         fg.savefig(ofn, bbox_inches='tight')
 
 
-def moviedasc(img,odir,cadence,rows=None,cols=None):
+def moviedasc(imgs:xarray.Dataset, odir:Path, cadence:float, rows=None, cols=None):
+
+    if odir:
+        print('writing to',odir)
+        odir = Path(odir).expanduser()
 
     fg = figure(figsize=(15,5))
-    axs = fg.subplots(1,3)
+
+    wavelength = [d for d in imgs.data_vars if not isinstance(d,str)]
+    axs = np.atleast_1d(fg.subplots(1,len(wavelength)))
+
     hi = []; ht=[]
-    for a,w,x,mm,c in zip(axs,wavelength,(0.225,0.5,0.775),
-                     ((350,800),(350,9000),(350,900)),('b','g','r')):
-        #a.axis('off') #this also removes xlabel,ylabel
-        a.set_xticks([]); a.set_yticks([])
-        a.set_xlabel('{} nm'.format(w),color=c)
-        hi.append(a.imshow(img[0][0],vmin=mm[0],vmax=mm[1],
+    time = imgs.time.values
+# %% setup figures
+    for ax,w,mm,c in zip(axs,
+                       wavelength,
+                       ((350,800),(350,9000),(350,900)),
+                       ('b','g','r')):
+        #ax.axis('off') #this also removes xlabel,ylabel
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.set_xlabel(f'{w} nm', color=c)
+
+        hi.append(ax.imshow(imgs[w][0],
+                           vmin=mm[0],vmax=mm[1],
                            origin='lower',
                            norm=LogNorm(),cmap='gray'))
-        ht.append(a.set_title('',color=c))
+
+        ht.append(ax.set_title('', color=c))
         #fg.colorbar(hi[-1],ax=a).set_label('14-bit data numbers')
         if themisplot is not None:
-            themisplot.overlayrowcol(a, rows, cols)
-
-    T = max([t[0,0] for t in times])
-    Tmax = min([t[-1,0] for t in times])
+            themisplot.overlayrowcol(ax, rows, cols)
 
     fg.tight_layout(h_pad=1.08) #get rid of big white space in between figures
 #%% loop
-    print('generating video until',datetime.fromtimestamp(Tmax))
-    with writer.saving(fg, str(ofn), DPI):
-        while T<=Tmax:
-            for I,Hi,Ht,t in zip(img,hi,ht,times):
-                ft = interp1d(t[:,0],arange(len(t)),kind='nearest')
-                ind = ft(T).astype(int)
-                #print(ind,end=' ')
-                Hi.set_data(I[ind])
-                try:
-                    Ht.set_text(datetime.fromtimestamp(t[ind,0],tz=UTC))
-                except OSError: #file had corrupted time
-                    Ht.set_text('')
+    print('generating video until', time[-1])
+    t = time[0]
+    dt = timedelta(seconds=cadence)
+    while t <= time[-1]:
+        for w,Hi,Ht in zip(wavelength,hi,ht):
+            I = imgs[w].sel(time=t, method='nearest')
+            Hi.set_data(I)
+            try:
+                Ht.set_text(str(I.time.values))
+            except OSError: #file had corrupted time
+                Ht.set_text('')
+
+        draw(), pause(0.05) # the pause avoids random crashes
+        t += dt
 
-            draw(), pause(0.05) # the pause avoids random crashes
-            T += cadence
-            if write:
-                print('generating AVI @ ',datetime.fromtimestamp(T),'from',datetime.fromtimestamp(t[ind][0]), end='\r')
-                writer.grab_frame(facecolor='k')
+        if odir:
+            ofn = odir / (str(t)+'.png')
+            print('saving', ofn,end='\r')
+            fg.savefig(ofn, bbox_inches='tight',facecolor='k')