COSIMA / initial_conditions_access-om2

WOA13 temperature and salinity initial conditions for ACCESS-OM2
1 stars 1 forks source link

numba not supported for numpy masked array #3

Open minghangli-uni opened 2 months ago

minghangli-uni commented 2 months ago

I am trying to generate iniitial conditions of temperature and salinity for 0.25deg with 75 vertical layers, related to issue136. However, I received an error related to numba,

Unsupported array type: numpy.ma.MaskedArray.

The error prevents the generation of initial fields.

minghangli-uni commented 2 months ago

Following steps in README , an error was encountered as follows,

File "/g/data/hh5/public/apps/miniconda3/envs/analysis3-23.10/lib/python3.10/site-packages/numba/core/typing/typeof.py", line 244, in _typeof_ndarray raise errors.NumbaTypeError(msg) numba.core.errors.NumbaTypeError: ^[[1mUnsupported array type: numpy.ma.MaskedArray.

The error arises from the numba library, specifically within the apply_weights function used for ESMF regridding. This library does not support numpy masked arrays, and this support has not been added yet. I am unsure how we previously generated these initial conditions with this script. After removing numba, I managed to generate the expected field, but with a noticeable decrease in speed.

A sanity check was conducted between the existing data and newly generated data without using numba, and they were found to be identical.

minghangli-uni commented 2 months ago

For anyone interested in reproducing the error, consider the following simple snippet:

import numpy as np
import numba
data = np.ma.masked_array([1,10,15], mask=[0,1,0])

@numba.jit
def masked_array_mean(arr):
    return np.mean(arr)
print(masked_array_mean(data))

It returns the same error, NumbaTypeError: Unsupported array type: numpy.ma.MaskedArray.