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 a maximum
      if np.any(x&gt;(prof[i:]+2.5)) and np.any(x&gt;1.25*prof[i:]) and x&gt;maxspd:
        maxflg = True
        maxind = i
        maxspd = x
        print 'Maximum found at time',t,'and level',i,'with speed',x,'m/s'
      #Find out whether vertical point is a minimum
      try:
        if x&lt;0.8*maxspd and x&lt;maxspd-2.5 and x&lt;prof[i+1] and not prof[i+2]&lt;x:
          minflg = True
          minind = i
          minspd = x
          print 'Minimum found at time',t,'and level',i,'with speed',x,'m/s'
      except IndexError:
        if i == nlev-2 and x&lt;0.8*maxspd and x&lt;maxspd-2.5 and x&lt;(prof[nlev-1]-1):
          minflg = True
          minind = i
          minspd = x
          print 'Minimum found at time',t,'and level',i,'with speed',x,'m/s'   
        else:
          minflg = True
          minind = nlev-1
          minspd = prof[nlev-1] # this should be changed to minumum of the profile, but that might return NaN, need to find an update for that.
          print 'No minumum found, using lowest value from profile'
      #Determine wheter a LLJ occured
      if maxflg and minflg:
        print 'Low-level jet found at time',t,'and height',hgt[maxind],'m with speed',maxspd,'m/s and minimum at',hgt[minind],'m with speed',minspd,'m/s'
        jet[t] = True
        jetspd[t] = maxspd
        jethgt[t] = hgt[maxind]
        break
      else:
        continue
  if maxflg and not minflg:
    print 'probably something wrong with NaNs'
  #End of loop, next profile...
  return jet,jethgt,jetspd
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&lt;0.8*maxspd and x&lt;maxspd-2.5


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

if x&lt;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!