如何使用python对numpy数组（来自netCDF文件）执行比较？

Question

I'm writing a python script that reads two different netCDF files and, after performing an evaluation on the variables makes some calculation, the idea is (an example of the idea of the code the real one is too long):

import netCDF4
import numpy as np
#other modules loaded...

#Values
a = 2
b = 4
c = 1

def srf(r, h):
    if r[:] == 2:
        if h[:] > 0:
            surf = 1 + b
        else:
            surf = a + b 
    else:
        surf = a - c

return surf

path_file : /home/file.nc
fhp = Dataset(path_file, r+)
ra = fhp.variables['VAR'][:]
path_file2 : /home/file2.nc
fhp2 = Dataset(path_file2, r+)
hu = fhp2.variables['VAR2'][:]   

#Call the Function
srf(ra, hu)       

The netCDF files each have 3 dimensions, if I try to run this code I get this error

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I need the function to be perform in the las two dimensions that are the ones containing the information to check in a domain [TSTEP, ROW, COL]. So I would need to loop through the dimensions, but I'm not sure how to do it and I don't know if maybe there is an easiest way. Thanks

Answer 1

This question is unrelated to netCDF, but is rather about the manipulation numpy arrays, obtained from a netCDF file.

The problem is that given a numpy vector r , r[:] == 2 (or r[:] > 0 for that matter) returns a boolean array and not a single boolean. It should therefore not be used in an if construction. Depending on what you are trying to accomplish you can use following approaches,

• use (r[:] == 2).any() or (r[:] == 2).all() instead

• use more complex indexing, for instance,

   import numpy as np def srf(r, h): mask_r = (r[:] == 2) mask_h = (h[:] > 0) surf = np.ones(r.shape)*(ac) surf[mask_r&mask_h] = 1 + b surf[mask_r&(~mask_h)] = a + b return surf 

please refer to the numpy documentation about advanced indexing for more details. This approach is significantly more efficient then loops over indices in python and should be used instead when possible.

Answer 2

Also, if you want to apply a function along a particular axis, you can use numpy.apply_along_axis.

For example, if you wanted to apply function srf along the time axis, you would do

import numpy as np srf_arr = np.apply_along_axis(srf, axisnumber, arrayname)