Pytorch: How to generate random vectors with length in a certain range?

user_7472132 0
Viewed 75

I want a k by 3 by n tensor representing k batches of n random 3d vectors, each vector has a magnitude (Euclidean norm) between a and b. Other than rescaling the entries of a random kx3xn tensor to n random lengths in a for loop, is there a better/more idiomatic way to do this?

pytorch random numpy python
1 Answers

Assuming a < b, you now have a constraint on the 3rd random number due to the norm. i.e sqrt(a^2 - x^2 - y^2) < z < sqrt(b^2 - x^2 - y^2)

Now a^2 - x^2 - y^2 > 0 which implies that x^2 + y^2 < a^2

We need two sets of generate numbers such that x^2 + y^2 < a^2

import numpy as np

def rand_generator(a,b,n,k):

    req_array = np.zeros((n,k,3))
    # first generate random numbers for x i.e 0<x<a
    
    req_array[:,:,0] = np.random.rand(n,k)*a
    
    # now generate random numbers for y such that 0 < y < a^-x2

    req_array[:,:,1] = np.random.rand( n,k) * np.sqrt(a**2 - req_array[:,:,0]**2)
    
    norm_temp = np.linalg.norm(req_array,axis=2)

    a1 = np.sqrt(a**2 - norm_temp**2) 

    b1 = np.sqrt(b**2 - norm_temp**2)
    
    # generate numbers for z such that they are inbetween a1 and b1

    req_array[:,:,2] = a1 + np.random.rand(n,k)*(b1-a1)

    return req_array


ll = rand_generator(2,5,10,12)
lp = np.linalg.norm(ll,axis=2)

print(np.all(lp>2) and np.all(lp<5))

##output: True

You can also use spherical coordinates for this(which is exactly same as above) x = rsin(theta)cos(phi), y = rsin(theta)sin(phi), z = rcos(theta) with a< r <b 0<theta<pi/2 and 0<phi<pi/2

import numpy as np

def rand_generator(a,b,n,k):

    req_array = np.zeros((n,k,3))
    # first generate random numbers for r in [a,b)
    
    r = a + np.random.rand(n,k)*(b-a)
    
    # now generate random numbers for theta in [0,pi/2)

    theta = np.random.rand( n,k) * np.pi/2
    
    # now generate random numbers for phi in [0,pi/2)

    phi = np.random.rand( n,k) * np.pi/2

    req_array[:,:,0] = r*np.sin(theta)*np.cos(phi)
 
    req_array[:,:,1] = r*np.sin(theta)*np.sin(phi)
    
    req_array[:,:,2] = r*np.cos(theta)
    
    return req_array


ll = rand_generator(2,5,10,12)
lp = np.linalg.norm(ll,axis=2)

print(np.all(lp>2) and np.all(lp<5))

##output: True
user_102890250
Related