Sampling from k circles with different radii

Question

Description: I would like to sample n_points points from k circles whose radii and centers are given as lists.

Note: This is my first question :) And k is not a very large number, usually k<7.

My code: I guess the parts where I compute the number of points for each circle, and the for loop can be improved. Perhaps with a simple trick, the for loop can be removed, but I am not really sure. Any help to beautify the code is appreciated.

# given parameters
radii = [3,1,5]
n_points_total = 250
centers = [[3,0],[7,0],[13,0]]

# my code starts here
import numpy as np

r_sum = sum(radii) # proportional to total arc length
probabilities = [r/r_sum for r in radii] # probability of a point being in the i'th circle.
positions = np.random.choice(len(radii), n_points_total, p=probabilities) # each point is randomly assigned to a circle
n_points = np.unique(positions, return_counts = True)[1] # counts how many points each circle will contain


data = []

for n, center, r in zip(n_points, centers, radii):
    t = np.random.random(n)*2*np.pi
    x = r*np.cos(t) + center[0]
    y = r*np.sin(t) + center[1]
    data.append(np.vstack([x,y]).T)
    
data = np.concatenate(data)

Answer 1

Let's assume that your zip-shortest behaviour is intended.

radii and centers should immediately be made np.ndarrays. If they're fed from some external source like a file or stdin, there are easy ways to get them into arrays.

import at the very top, not after your givens.

Everything here can be vectorised. Avoid like the plague for loops, lists and iterative vstack calls.

First, recognise that your centres and radii are going to be repeated by the respective counts represented by the elements of n_points. This is exactly what np.repeat does.

Rather than indexing into the result of np.unique(), prefer tuple-unpacking.

Do not use np.random.random; that's deprecated. Prefer the new Generator stuff and call into uniform instead of post-multiplying.

Add unit tests with a predictable random seed.

Suggested

import numpy as np
from numpy.random import default_rng
from numpy.random._generator import Generator


def new(
    centers: np.array,
    radii: np.array,
    rand: Generator,
    n_points_total: int,
) -> np.ndarray:
    probabilities = radii / radii.sum()  # probability of a point being in the i'th circle.
    positions = rand.choice(a=len(radii), size=n_points_total, p=probabilities)  # each point is randomly assigned to a circle
    _, n_points = np.unique(positions, return_counts=True)  # counts how many points each circle will contain

    centers_spread = centers[:len(n_points)].repeat(repeats=n_points, axis=0)
    radii_spread = radii[:len(n_points)].repeat(repeats=n_points)

    t = rand.uniform(low=0, high=2*np.pi, size=n_points_total)
    cossin = np.empty_like(centers_spread)
    np.cos(t, out=cossin[:, 0])
    np.sin(t, out=cossin[:, 1])

    data = centers_spread + radii_spread[:, np.newaxis] * cossin
    return data


def old(
    centers: np.array,
    radii: np.array,
    rand: Generator,
    n_points_total: int,
) -> np.ndarray:
    r_sum = sum(radii)  # proportional to total arc length
    probabilities = [r / r_sum for r in radii]  # probability of a point being in the i'th circle.
    positions = rand.choice(len(radii), n_points_total,
                            p=probabilities)  # each point is randomly assigned to a circle
    n_points = np.unique(positions, return_counts=True)[1]  # counts how many points each circle will contain

    data = []

    for n, center, r in zip(n_points, centers, radii):
        t = rand.random(n) * 2 * np.pi
        x = r * np.cos(t) + center[0]
        y = r * np.sin(t) + center[1]
        data.append(np.vstack([x, y]).T)

    data = np.concatenate(data)
    return data


def test() -> None:
    centers = np.array((
        (3, 0), (7, 0), (13, 0),
    ), dtype=np.float64)
    radii = np.array((3, 1, 5), dtype=np.float64)

    for method in (old, new):
        rand = default_rng(seed=0)
        assert np.allclose(
            method(centers, radii, rand, n_points_total=10),
            np.array((
                (4.20619233, -2.74683455),
                (5.99955592,  0.05161697),
                (4.87431947, -2.34241895),
                (7.97781686,  0.20946168),
                (6.87251924, -0.99184104),
                (7.45031626,  0.89286912),
                (7.65268587, -0.75762864),
                (6.03374088, -0.25757196),
                (6.69270516,  0.95161435),
                (6.11568453,  0.46688988),
            )),
        )

        data = method(centers, radii, rand, n_points_total=250)
        assert data.shape == (250, 2)
        assert np.isclose(-4.999847197562797, data.min())
        assert np.isclose(17.99763106087105, data.max())



if __name__ == '__main__':
    test()

Answer 2

I don't have much to add to the answer of @Reinderien, maybe a slightly different solution.

import numpy as np

def generate_points_on_circles(radii, centers, num_pts):
    """
    usage:

    generate_points_on_circles(
        radii = np.array([3,1,1]), 
        centers = np.array([[3, 7, 13],
                            [0, 0, 0]]),
        num_pts = 10
    )
    """
    assert(radii.shape[0] == centers.shape[1] and centers.shape[0] == 2 and num_pts >= 0)

    circle_indices = np.random.choice(len(radii), num_pts, p=radii / radii.sum())
    
    angles = 2*np.pi*np.random.random(num_pts)

    return centers[:, circle_indices] + radii[circle_indices] * np.vstack([np.cos(angles), 
                                                                           np.sin(angles)])

This does assume row vectors, but you can always transpose with a .T. To keep the points grouped by the circle they are from, add a np.sort call to the line defining the circle indices.