Can not manage memory properly with norm_of_vectorfield

[simon3122]

Hi all,
My problem is the following : I want to process simple statistics using the norm of a vector (u/v fields).
When using norm_of_vectorfield, memory gets busier than when writing manually a norm (as if grids were collocated). Fatal outcome is that norm_of_vectorfield may lead to overfill memory, whereas memory is well bounded with a xarray manual writing of the norm.

Here is a simple script showing (watching the monitor) that memory gets higher with norm_of_vectorfield.

import sys, os
import xarray as xr
from xarray import ufuncs as uf
from contextlib import contextmanager
import time
from oocgcm.oceanmodels.nemo import grids
from oocgcm.core import grids as gridsc

#-----------------------------------------------------------------------------------------------------------
@contextmanager
def timeit_context(name):
    startTime = time.time()
    yield
    elapsedTime = time.time() - startTime
    print('{} takes {} s'.format(name, int(elapsedTime)))

#-----------------------------------------------------------------------------------------------------------

chunks = (1727, 2711)
xr_chunks = {'x': chunks[-1], 'y': chunks[-2]}
xr_chunks_u = {'x': chunks[-1], 'y': chunks[-2], 'time_counter':1, 'depthu': 1}
xr_chunks_v = {'x': chunks[-1], 'y': chunks[-2], 'time_counter':1, 'depthv': 1}

# - Parameter
natl60_path = '/home7/pharos/othr/NATL60/'
coordfile = natl60_path + 'NATL60-I/NATL60_coordinates_v4.nc'
maskfile = natl60_path + 'NATL60-I/NATL60_v4.1_cdf_byte_mask.nc'

grd = grids.nemo_2d_grid(nemo_coordinate_file=coordfile, nemo_byte_mask_file=maskfile, chunks=xr_chunks)


#30 times, 2 fields, 300 vertical levels
filenameu = natl60_path+'NATL60-MJM155-S/5d/2008/NATL60-MJM155_y2008m0*d*.5d_gridU.nc'
filenamev = natl60_path+'NATL60-MJM155-S/5d/2008/NATL60-MJM155_y2008m0*d*.5d_gridV.nc'


# open u and v components over 10 levels
vi_xr = xr.open_mfdataset(filenameu,chunks=xr_chunks_u)['vozocrtx'].isel(depthu=slice(0,10))
vj_xr = xr.open_mfdataset(filenamev,chunks=xr_chunks_v)['vomecrty'].isel(depthv=slice(0,10))

# vector object
vec_xr = gridsc.VectorField2d(vi_xr, 
                              vj_xr,
                              x_component_grid_location='u',
                              y_component_grid_location='v') 

#1 Manual vector norm
with timeit_context('     #1'):
    print '# euclidian norm'   
    vnorm_xr=uf.sqrt(uf.square(vec_xr[0])+uf.square(vec_xr[1]))
    print vnorm_xr.mean(dim='time_counter').isel(depthv=9,depthu=9).values

#2 oocgcm vector norm
with timeit_context('     #2'):     
    print '# euclidian norm 2'
    vnorm2_xr=grd.norm_of_vectorfield(vec_xr)
    print vnorm2_xr.mean(dim='time_counter').isel(depthv=9,depthu=9).values


print 'end'

Thank you very much for any advice,
Simon