Convert scalar arguments to arrays, and check that lengths are equal

Question

I've got a Python function that takes two arguments, lat and lon. These arguments can be either scalar values (52.3) or any sort of iterable (e.g. a list or a NumPy array).

At the beginning of my function I need to check what I've been given, and convert both arguments to arrays (if needed) and check that both arguments are the same length.

# Deal with scalar values
try:
    lat_count = len(lat)
except:
    lat = [lat]
    lat_count = 1

try:
    lon_count = len(lon)
except:
    lon = [lon]
    lon_count = 1

if lat_count != lon_count:
    raise ValueError("lan and lon arrays must be the same length")

lat = np.array(lat)
lon = np.array(lon)

It feels horribly messy, but I can't work out how to make it better. Any ideas?

Answer 1

np.atleast_1d handles the conversion in one step. It's doc reads:

Scalar inputs are converted to 1-dimensional arrays, whilst higher-dimensional inputs are preserved.

In [217]: np.atleast_1d(1)
Out[217]: array([1])
In [218]: np.atleast_1d([1,2,3])
Out[218]: array([1, 2, 3])

So your code simplifies to:

lat = np.atleast_1d(lat)
lon = np.atleast_1d(lon)
if lat.shape != lon.shape:
    <raise error>

Look at the code for atleast_1d; it's instructive. Bascially it does:

lat = np.array(lat)
if len(lat.shape)==0:
   lat = lat.reshape(1)

Check for yourself what np.array(1) produces.

I could simplify it further with just one call to atleast_1d, since it takes multiple inputs:

lat, lon = np.atleast_1d(lat, lon)

Answer 2

Create a function to factorize some code:

def input_to_array(values):
    """Deal with scalar values as well as array-like objects.

    Convert them into numpy array.
    """
    try:
        single_value = float(values)
    except TypeError: # Trying to float() a list
        return np.array(values)
    except ValueError: # String does not contain a float
        # Return an array of strings, you can remove it if not needed.
        return np.array((values,))
    else:
        return np.array((single_value,))

lat = input_to_array(lat)
lon = input_to_array(lon)

if len(lat) != len(lon):
    raise ValueError("lat and lon arrays must be the same length")

Also note that you should provide the type of the exception you are expecting in an except clause or you would potentially be masking more serious problems.

And last, there might exist something more suited to numpy than len().

Answer 3

Mathias Ettinger already showed you some ways to clean up the code, but I have an important note. Never use a bare except. The rare occasion that you want to catch all exceptions you should still be passing in except Exception because otherwise exit() calls or KeyboardInterrupts will be excepted and those are usually triggered manually. You should always try to except a specific error, one you're expecting could happen if the data is not of the intended form. Since len() failing raises a TypeError, you should use that.

try:
    lat_count = len(lat)
except TypeError:
    lat = [lat]
    lat_count = 1