Tips for golfing with numpy, scipy, or pylab

Question

There's already a comprehensive list of tips for python here, so what I'm asking for are tips that specifically apply to using the numpy, scipy or pylab libraries.

These can be either ways to shorten code already using numpy, or ways to shorten common python operations by using these libraries.

One tip per answer, please.

Answer 1

Make use of Numpy's broadcasting

Broadcasting means replicating a multidimensional array along some of its singleton dimensions to match the size of another array. This happens automatically for Numpy arrays when arithmetic operators are applied to them.

For example, to generate a 10×10 multiplication table you can use

import numpy
t=numpy.arange(1,11)
print(t*t[:,None]) # Or replace t[:,None] by [*zip(t)]

Try it online!

Here t is created as the Numpy array [1, 2, ..., 10]. This has shape (10,), which is equivalent to (1,10). The other operand array, t[:,None], has size (10,1). Multiplying the two arrays implicitly replicates them, so they behave as if they both had shape (10,10). The result, which also has shape (10,10), contains the products for all pairs of entries in the original arrays.

Answer 2

Use r_[...] instead of range(...)

Numpy provides matlab like syntax for array creation using r_[...]. Any slice notation in between the brackets is interpreted as an array with the range indicated. So, for instance

r_[:30:4]

is equivalent to

arange(0,30,4)

and for most uses can replace

range(0,30 4)

It can also handle more complex expressions. For example to get indices from 0 up to 10 and back down again,

r_[:10,10:-1:-1]

Answer 3

The shorter than shortest infinite for comprehension

If numpy is *-imported then using this "feature" of numpy.r_/numpy.c_ an infinte counter is as cheap as:

for[i]in r_:print(i)

This is significantly shorter than this pure Python trick and has the added benefit of actually creating a proper counter.

This can also be used for a cheaper enumerate:

S="abcde"
for x,[i]in zip(S,r_):print(i,x)

If you can live with a "boxed" counter, i.e. [0],[1],[2],... rather than 0,1,2,... then you can save the brackets around i.

The boxes can even be useful: For example, if we have set up a loop anyway, then we can as a byproduct create a reversed "range" of the same length as the looped over sequence S:

C=[]
for e,C[:0]in zip(S,r_):f(e)

If for some reason you need double boxes [[0]],[[1]],[[2]],... use c_ instead of r_.

Answer 4

Don't forget the matrix class

It may be aesthetically challenged and semi-deprecated since all but forever but it can be useful for golfing.

Here is an example How to solve the LCM in 50 bytes of Python where on a non numpy specific challenge it allows us to outgolf @dingledooper's extremely clever pure Python approach by ~10%. And that's counting the numpy import and without requiring us to be particularly smart:

from numpy import*
print(lcm(*mat(input()).A1))

Try it online!

This shows off three useful features:

1. Short constructor name mat although not as short as r_ or c_

2. Hilariously graceful input parser. For example, it will accept

mat("""[1 2]\t[3,4];5  6 7        
8""")

without complaint.

3. a handful of very short-named attributes, viz. A,A1,T,H,I which are the underlying array object,the same flattened,the transpose,the conjugate transpose and the pseudo inverse.

Finally, there is one more useful feature that we didn't use in the example:

4. Overloaded * operator. That in itself doesn't gain us much over the @ operator for standard numpy arrays but for example taking the matrix power can be done using ** much shorter than linalg.matrix_power.