Python algorithm for identification of low-level jets (meteorology)

A low-level jet is a local wind speed maximum near the surface. Baas (2009) defined a low-level jet as follows:


the lowest maximum of the wind speed profile in the lowest 500 m of the atmosphere that is at least both 2 m/s and 25% faster than the next minimum above (...) A minimum is neglected if the wind speed above that minimum increases less than 1 m/s before decreasing again to values lower than the wind speed of that minimum. When no minumum is present, the lowest value of the wind speed profile above the jet is taken as a minimum.

I have a 2D numpy array that contains wind speed values at several measurement heights and times. I've written an algorithm to identify for each time step whether the wind profile obeys this criterion. However, I think my code can be improved. Right now I'm using nested loops and a lot of if statements. Also I need to catch an index error because I need to check whether the 'next' value in the array is lower or higher.

I was hoping to get some feedback from more experienced programmers. Here's my code:

```
def findLLJ(wspd_matrix, time, hgt):
'''Find low-level jets according to Peter Baas' algorithm
input: wind speed matrix with time as first axis and height as second axis
time and height arrays
output: jet index array, jet height array, jet speed array
use: jetidx,jethgt,jetspd = findLLJ(wsp_matrix,time,height)'''
#Allocate output arrays:
jet = np.zeros(time.shape,dtype=bool)
jetspd = np.zeros(time.shape)
jethgt = np.zeros(time.shape)
#Shape of input array:
ntime,nlev = wspd_matrix.shape
#Consistency check:
if not (ntime == len(time) and nlev == len(hgt)):
raise Exception("Shapes of input arrays are not consistent")
#Loop over all wind profiles at given times
for t,prof in enumerate(wspd_matrix):
maxflg = False
minflg = False
maxspd = 0.
#Loop over all vertical points in the profile
for i,x in enumerate(prof):
#Find out whether vertical point is

Naming the variables

• wspd_matrix would be much easier to read if it were called wind_speed_matrix

• jetspd would be much easier to read if it were called jet_speed

• jethgt would be much easier to read if it were called jet_height

Remove the numeric values in the code

0.8 and 2.5 appear quite often in your code. It is hard for someone without a domain knowledge to know what they stand for. And if you have to change them for some reason, this is error-prone. Why not just naming them (max_multiplier = 0.8 per example) at the beginning of the function ?

Redundant condition

if x<0.8*maxspd and x<maxspd-2.5

appears twice in your code. It represents a clear mathematical function :

if x<min(0.8*maxspd , maxspd-2.5)

which is much easier to read (and picture mentally)

Redundant code

Something needs to be done about this three prints :

print 'Maximum found at time',t,'and level',i,'with speed',x,'m/s'
print 'Minimum found at time',t,'and level',i,'with speed',x,'m/s' 
print 'Minimum found at time',t,'and level',i,'with speed',x,'m/s'

Maybe a

print_local_optimum_found(type_optimum, at_time, at_level, at_speed)

Check the input

if maxflg and not minflg:
    print 'probably something wrong with NaNs'

Why waiting until the end of the program to realize something wrong happened ? If you know NAs cause trouble, just check the presence of NAs before

The algorithm itself

It is hard to suggest improvements now with every piece of code in the same function. If you fix all this, such improvements will become much more readable!

StackExchange Code Review Q#121901, answer score: 3