pyic_view: Slightly improved chunking that is switched off when zarr files are...

pyic_view: Slightly improved chunking that is switched off when zarr files are opened since this is a large performance loss.

scripts/pyic_view.py

@@ -106,7 +106,7 @@ def str_to_array(string):
 #  return fpath_ckdtree
 
 
-def get_data(ds, var_name, it, iz, res, lon_reg, lat_reg):
+def get_data(ds, var_name, it, iz, res, lon_reg, lat_reg, do_chunking=True):
     isel_dict = dict(time=it)
     if ds[var_name].ndim==3:
       depth_name = pyic.identify_depth_name(ds[var_name])
@@ -121,7 +121,8 @@ def get_data(ds, var_name, it, iz, res, lon_reg, lat_reg):
         chunks = dict(time=1)
     if depth_name!='none':
         chunks[depth_name] = 1
-    da = da.chunk(**chunks)
+    if do_chunking:
+      da = da.chunk(**chunks)
     dai = pyic.interp_to_rectgrid_xr(
         da.isel(**isel_dict), res=res,
         lon_reg=lon_reg, lat_reg=lat_reg,
@@ -142,6 +143,8 @@ class view(object):
         self.flist = flist
         self.path_grid = path_grid
         self.fig_size_fac = fig_size_fac
+
+        self.do_chunking = True
         self.colormaps = [
             "inferno", "viridis", "plasma", 
             "RdYlBu_r", "RdBu_r", "Blues_r", 
@@ -182,6 +185,7 @@ class view(object):
             self.transformer = "None"
 
         # Opean data set
+        self.message("Opening data set")
         self.load_data()
 
         # Default selections
@@ -198,6 +202,7 @@ class view(object):
         self.rect = None
 
         # Create figure and axis
+        self.message("Generating axes")
         self.fig, self.ax, self.cax = generate_axes(asp=0.5)
 
         #print('------')
@@ -205,6 +210,7 @@ class view(object):
         #print(self.ax.get_position())
 
         # TK canvas
+        self.message("Setting up TK")
         frame_plot = tk.Frame(root)
         frame_plot.pack(fill="both", expand=True)
         self.canvas = FigureCanvasTkAgg(self.fig, master=frame_plot)
@@ -234,7 +240,7 @@ class view(object):
         spacer = tk.Frame(frame1, width=30)  # 10 pixels wide
         spacer.pack(side="left")
         
-        # depth slide
+        # depth slider
         btn_dec_d = tk.Button(frame1, text="-", command=lambda: self.decrease_slider(self.slider_d))
         btn_dec_d.pack(side="left")
 
@@ -402,6 +408,9 @@ class view(object):
             self.canvas.mpl_disconnect(self.cid_motion)
 
     def load_data(self):
+        if self.flist[0].endswith('zarr'):
+            self.message('Detected zarr file and switching off chunking.')
+            self.do_chunking = False
         self.message('opening dataset')
         mfdset_kwargs = dict(
             combine='nested', concat_dim='time',
@@ -416,6 +425,7 @@ class view(object):
         self.ds = xr.open_mfdataset(
             self.flist, **mfdset_kwargs, 
         )
+        self.message('Done opening data files')
         delvars = [
             "clon_bnds", "clat_bnds", "elon_bnds", "elat_bnds",
             "vlon_bnds", "vlat_bnds",
@@ -464,7 +474,8 @@ class view(object):
         #self.fpath_ckdtree = get_fpath_ckdtree(self.ds, self.res, self.path_grid)
         self.dai = get_data(
             self.ds, self.var_name, self.it, self.iz, 
-            self.res, self.lon_reg, self.lat_reg
+            self.res, self.lon_reg, self.lat_reg,
+            do_chunking=self.do_chunking,
         )
         self.Lon, self.Lat = np.meshgrid(self.dai.lon.data, self.dai.lat.data)
         if self.proj=="None":
@@ -562,8 +573,6 @@ class view(object):
         self.lon_lat_reg_tk.set(f"{self.lon_reg[0]:.3g},{self.lon_reg[1]:.3g},{self.lat_reg[0]:.3g},{self.lat_reg[1]:.3g}")
         self.xlim, self.ylim = get_xlim_ylim(
             self.lon_reg, lat_reg_axlim, self.proj, self.transformer)
-        print(f'set_default_lon_lat_reg:')
-        print(self.lon_reg, self.lat_reg, self.xlim, self.ylim)
         
     def update_projection(self, *args):
         self.proj = self.proj_dict[self.selected_proj.get()]
@@ -601,11 +610,9 @@ class view(object):
 
     def increase_slider(self, slider):
         slider.set(slider.get() + 1)
-        self.update_data()
     
     def decrease_slider(self, slider):
         slider.set(slider.get() - 1)
-        self.update_data()
     
     # Function to update plot
     def update_data(self, *args):
@@ -628,7 +635,8 @@ class view(object):
         self.update_lon_lat_reg()
         self.dai = get_data(
             self.ds, self.var_name, self.it, self.iz, 
-            self.res, self.lon_reg, self.lat_reg
+            self.res, self.lon_reg, self.lat_reg,
+            do_chunking=self.do_chunking,
         )
         self.Lon, self.Lat = np.meshgrid(self.dai.lon, self.dai.lat)
         self.hm[0].set_array(self.dai.data.flatten())