[Numpy-discussion] Optimize speed of for loop using numpy

[Anne Archibald]

On 25/02/2008, Trond Kristiansen <trond@unc.edu> wrote:

>  I have attached the function that the FOR loop is part of as a python file.
>  What I am trying to do is to create a set of functions that will read the
>  output files (NetCDF) from running the ROMS model (ocean model). The output
>  file is organized in xi (x-direction), eta (y-direction), and s
>  (z-direction) where the s-values are vertical layers and not depth. This
>  particular function (z_slice) will find the closest upper and lower s-layer
>  for a given depth in meters (e.g. -100). Then values from the two selcted
>  layers will be interpolated to create a new layer at the selected depth
>  (-100). The problem is that the s-layers follow the bathymetry and a
>  particular s-layer will therefore sometimes be above and sometimes below the
>  selected depth that we want to interpolate to. That's why I need a quick
>  script that searches all of the layers and find the upper and lower layers
>  for a given depth value (which is negative). The z_r is a 3D array
>  (s,eta,xi) that is created using the netcdf module.

If I understand what you're doing correctly, you have a 3D array of
depths, indexed by two unproblematic coordinates (eta and xi), and by
a third coordinate whose values are peculiar.   For each pair (eta,
xi), you want to find the value of the third coordinate that occurs at
a depth closest to some given depth.

Are you looking to do this for many values of depth? It seems like
what you're trying to do is (approximately) invert a function - you
have z as a function of s, and you want s as a function of z. Is the
mapping between depths and s values monotonic?

First of all, numpy.searchsorted() is the right tool for finding where
depth occurs in a list of z values. It gives you the position of the
next lower value; it's pretty straightforward to write down a linear
interpolation (you should be able to do it without any for loop at
all) and more sophisticated interpolation schemes may be
straightforward as well.

If you want a smoother interpolation scheme, you may want to look at
scipy.interpolate. It is not always as vectorial as you might wish,
but it implements splines (optionally with smoothing) quite
efficiently. I believe scipy.ndimage may well have some suitable
interpolation code as well.

In my opinion, the value of numpy/scipy comes from the tools that
allow you to express major parts of your algorithm as a line or two of
code. But it is often difficult to take advantage of this by "trying
to accelerate for loops". I find it really pays to step back and look
at my algorithm and think how to express it as uniform operations on
arrays. (Of course it helps to have a rough idea of what operations
are available; the numpy documentation is sometimes sadly lacking in
terms of letting you know what tools are there.)

You might find useful the new
http://www.scipy.org/Numpy_Functions_by_Category

Good luck,
Anne