CODEWITHSUNDEEP
pythonarraysnumpy
James Briggs
James Briggs

Reputation: 914

More Pythonic way to build random clusters in Python

I would like to create a function which will create a uniform distribution random cluster centered around a set of co-ordinates and with a specified radius, I have done this with the below:

import numpy as np

# create cluster builder
def cluster(center, radius=10, n=50):
    xx = np.random.uniform(center[0]-radius, center[0]+radius, size=n)
    yy = np.random.uniform(center[1]-radius, center[1]+radius, size=n)
    zz = np.random.uniform(center[2]-radius, center[2]+radius, size=n)
    return xx, yy, zz

# create random cluster
xx1, yy1, zz1 = cluster((25, 15, 5))

This works as expected, but I just feel that they're must be a more Pythonic way to build the cluster function. Does anyone have any suggestions?

Upvotes: 2

Views: 123

Answers (2)

Nils Werner
Nils Werner

Reputation: 36849

You can define cluster() so that it vectorizes all operations

def cluster(center, radius=10, n=50):
    center = np.atleast_1d(np.asarray(center))[:, None, None]
    radius = np.atleast_1d(np.asarray(radius))[None, :, None]
    shape = (center.shape[0], radius.shape[1], n)
    return np.random.uniform(center - radius, center + radius, size=shape)

and get all values in one single call:

cluster(25, 10).shape  # (1, 1, 50)
cluster((25, 15, 5), 10).shape  # (3, 1, 50)
cluster((25, 15, 5), (10, 5)).shape  # (3, 2, 50)
cluster((25, 15, 5), (10, 5), n=100).shape  # (3, 2, 100)

of course you can still separate the result into xx, yy, zz:

xx, yy, zz = cluster((25, 15, 5), (10, 5), n=100)
xx.shape  # (2, 100)
yy.shape  # (2, 100)
zz.shape  # (2, 100)

Upvotes: 2

Divakar
Divakar

Reputation: 221754

np.random.uniform also accepts low and high as arrays/lists. Hence, we can simply do -

c = np.asarray(center)
xx,yy,zz = np.random.uniform(c-radius, c+radius, size=(n,3)).T

If any older version only supported scalar low and high, we can use some scaling -

xx,yy,zz = np.random.uniform(size=(3,n))*radius*2 + c[:,None] - radius

Upvotes: 4

Related Questions