another optimisation to curl operator to reduce number of tasks in...

another optimisation to curl operator to reduce number of tasks in compute_curl and prevent memory overload

pyicon/pyicon_calc_aw.py

@@ -203,26 +203,16 @@ class daskicon:
             vector p_vn_c remapped to edges
         """
         
-        edge2cell_coeff_cc_t = (self.ds_aux.edge2cell_coeff_cc_t
-                                        .drop_vars({'clat','clon'})                                         # N.B.: These additional coordinates confuse dask !
-                                        .transpose('cart','edge','nc_e').chunk({'nc_e':-1, 'cart':-1}))     # Reorder for better performance
-                         
-        ic = self.ds_IcD.adjacent_cell_of_edge 
-        
-        p_vn_c = p_vn_c.chunk({'cart':-1}).transpose('time','cart','cell')     # Reorder for better performance
-        
-        ptp_vn = xr.apply_ufunc(
-            lambda v, ind, coeff: (v[:,ind] * coeff).sum(axis=(0,2)),
-            p_vn_c,
-            ic, 
-            edge2cell_coeff_cc_t,
-            input_core_dims=[['cart','cell'], ['edge', 'nc_e'], ['cart', 'edge', 'nc_e']],
-            output_core_dims=[['edge']],
-            vectorize=True,
-            dask='parallelized',
-            output_dtypes=[p_vn_c.dtype],
-            dask_gufunc_kwargs={'output_sizes':{'edge': ic.edge.size}}
-        )
+        edge2cell_coeff_cc_t = self.ds_aux.edge2cell_coeff_cc_t.drop_vars({'clat','clon'})  # N.B.: These additional coordinates confuse dask !
+        
+        ic0 = self.ds_IcD.adjacent_cell_of_edge.isel(nc_e=0).compute()
+        ic1 = self.ds_IcD.adjacent_cell_of_edge.isel(nc_e=1).compute()
+        
+        p_vn_c = p_vn_c.drop_vars(p_vn_c.coords)   # The additional coordinates create a lot of overhead....
+        
+        ptp_vn = (p_vn_c.isel(cell=ic0) * edge2cell_coeff_cc_t.isel(nc_e=0) + 
+                  p_vn_c.isel(cell=ic1) * edge2cell_coeff_cc_t.isel(nc_e=1))  .sum(dim='cart')
+        
         
         return ptp_vn.chunk({'edge':-1}) # Ensure still in a single chunk...
 
@@ -373,15 +363,16 @@ class daskicon:
         """
         assert "edge" in vector.dims
 
-        rot_coeff = self.ds_aux.rot_coeff
+        rot_coeff = self.ds_aux.rot_coeff.chunk({'ne_v':-1})
         
-        eov = self.ds_IcD["edges_of_vertex"]
+        eov = self.ds_IcD.edges_of_vertex.compute()
         
         curl_vec = xr.apply_ufunc(
-            lambda v, ind: (v[ind] * rot_coeff).sum(dim='ne_v'),  # NOTE: This gives very large/wrong values when the vertex has _any_ neighbouring undefined cell....
+            lambda v, ind, coeff: (v[ind] * coeff).sum(axis=1),  # NOTE: This gives very large/wrong values when the vertex has _any_ neighbouring undefined cell....
             vector,
             eov,
-            input_core_dims=[['edge'], ['vertex','ne_v']],
+            rot_coeff,
+            input_core_dims=[['edge'], ['vertex','ne_v'], ['vertex','ne_v']],
             output_core_dims=[['vertex']],
             vectorize=True,
             dask='parallelized',