Commit 13f33895 authored by teresa.bister's avatar teresa.bister
Browse files

[coord] enable setting a random seed for shuffling etc

parent a2f2a654
Pipeline #416008 failed with stages
...@@ -467,7 +467,7 @@ def exposure_equatorial(dec, a0=-35.25, zmax=60): ...@@ -467,7 +467,7 @@ def exposure_equatorial(dec, a0=-35.25, zmax=60):
return cov / np.pi # normalize the maximum possible value to 1 return cov / np.pi # normalize the maximum possible value to 1
def rand_fisher(kappa, n=1): def rand_fisher(kappa, n=1, seed=None):
""" """
Random angles to the center of a Fisher distribution with concentration parameter kappa. Random angles to the center of a Fisher distribution with concentration parameter kappa.
...@@ -476,54 +476,64 @@ def rand_fisher(kappa, n=1): ...@@ -476,54 +476,64 @@ def rand_fisher(kappa, n=1):
:param n: number of vectors drawn from fisher distribution :param n: number of vectors drawn from fisher distribution
:return: theta values (angle towards the mean direction) :return: theta values (angle towards the mean direction)
""" """
if seed is not None:
np.random.seed(seed)
if np.ndim(kappa) > 0: if np.ndim(kappa) > 0:
n = kappa.shape n = kappa.shape
return np.arccos(1 + np.log(1 - np.random.random(n) * (1 - np.exp(-2 * kappa))) / kappa) return np.arccos(1 + np.log(1 - np.random.random(n) * (1 - np.exp(-2 * kappa))) / kappa)
def rand_phi(n=1): def rand_phi(n=1, seed=None):
""" """
Random uniform phi (-pi, pi). Random uniform phi (-pi, pi).
:param n: number of samples that are drawn :param n: number of samples that are drawn
:return: random numbers in range (-pi, pi) :return: random numbers in range (-pi, pi)
""" """
if seed is not None:
np.random.seed(seed)
return (np.random.random(n) * 2 - 1) * np.pi return (np.random.random(n) * 2 - 1) * np.pi
def rand_theta(n=1, theta_min=-np.pi/2, theta_max=np.pi/2): def rand_theta(n=1, theta_min=-np.pi/2, theta_max=np.pi/2, seed=None):
""" """
Random theta (pi/2, -pi/2) from uniform cos(theta) distribution. Random theta (pi/2, -pi/2) from uniform cos(theta) distribution.
:param n: number of samples that are drawn :param n: number of samples that are drawn
:return: random theta in range (-pi/2, pi/2) from cos(theta) distribution :return: random theta in range (-pi/2, pi/2) from cos(theta) distribution
""" """
if seed is not None:
np.random.seed(seed)
assert theta_max > theta_min assert theta_max > theta_min
u = np.sin(theta_min) + (np.sin(theta_max) - np.sin(theta_min)) * np.random.random(n) u = np.sin(theta_min) + (np.sin(theta_max) - np.sin(theta_min)) * np.random.random(n)
return np.pi / 2 - np.arccos(u) return np.pi / 2 - np.arccos(u)
def rand_theta_plane(n=1): def rand_theta_plane(n=1, seed=None):
""" """
Random theta (pi/2, 0) on a planar surface from uniform cos(theta)^2 distribution. Random theta (pi/2, 0) on a planar surface from uniform cos(theta)^2 distribution.
:param n: number of samples that are drawn :param n: number of samples that are drawn
:return: random theta on plane in range (pi/2, 0) :return: random theta on plane in range (pi/2, 0)
""" """
if seed is not None:
np.random.seed(seed)
return np.pi / 2 - np.arccos(np.sqrt(np.random.random(n))) return np.pi / 2 - np.arccos(np.sqrt(np.random.random(n)))
def rand_vec(n=1): def rand_vec(n=1, seed=None):
""" """
Random spherical unit vectors. Random spherical unit vectors.
:param n: number of vectors that are drawn :param n: number of vectors that are drawn
:return: random unit vectors of shape (3, n) :return: random unit vectors of shape (3, n)
""" """
return ang2vec(rand_phi(n), rand_theta(n)) if seed is not None:
np.random.seed(seed)
return ang2vec(rand_phi(n, seed=seed), rand_theta(n, seed=seed))
def rand_vec_on_surface(x0, n=1): def rand_vec_on_surface(x0, n=1, seed=None):
""" """
Given unit normal vectors x0 orthogonal on a surface, samples one isotropic Given unit normal vectors x0 orthogonal on a surface, samples one isotropic
direction for each given vector x0 from a cos(theta)*sin(theta) distribution direction for each given vector x0 from a cos(theta)*sin(theta) distribution
...@@ -531,13 +541,15 @@ def rand_vec_on_surface(x0, n=1): ...@@ -531,13 +541,15 @@ def rand_vec_on_surface(x0, n=1):
:param x0: ortogonal unit vector on the surface, shape: (3, N) :param x0: ortogonal unit vector on the surface, shape: (3, N)
:return: isotropic directions for the respective normal vectors x0 :return: isotropic directions for the respective normal vectors x0
""" """
if seed is not None:
np.random.seed(seed)
if np.ndim(np.squeeze(x0)) > 1: if np.ndim(np.squeeze(x0)) > 1:
n = x0.shape[1] n = x0.shape[1]
v = ang2vec(rand_phi(n), rand_theta_plane(n)) # produce random vecs on plane through z-axis v = ang2vec(rand_phi(n, seed=seed), rand_theta_plane(n, seed=seed)) # produce random vecs on plane through z-axis
return rotate_zaxis_to_x(v, x0) # rotation to respective surface vector x0 return rotate_zaxis_to_x(v, x0) # rotation to respective surface vector x0
def rand_vec_on_sphere(n, r=1): def rand_vec_on_sphere(n, r=1, seed=None):
""" """
Calculates n random positions x and correpsonding directions on the surface of the Calculates n random positions x and correpsonding directions on the surface of the
sphere as they would arise from an isotropic distribution of cosmic rays in the universe. sphere as they would arise from an isotropic distribution of cosmic rays in the universe.
...@@ -546,12 +558,14 @@ def rand_vec_on_sphere(n, r=1): ...@@ -546,12 +558,14 @@ def rand_vec_on_sphere(n, r=1):
:param r: size of the sphere, which scales the position vectors, default: 1 :param r: size of the sphere, which scales the position vectors, default: 1
:return: position vectors, directions :return: position vectors, directions
""" """
x = rand_vec(n) if seed is not None:
v = rand_vec_on_surface(x) np.random.seed(seed)
x = rand_vec(n, seed=seed)
v = rand_vec_on_surface(x, seed=seed)
return r*x, v return r*x, v
def rand_exposure_vec(a0=-35.25, zmax=60, n=1, coord_system='gal', res_theta=5000): def rand_exposure_vec(a0=-35.25, zmax=60, n=1, coord_system='gal', res_theta=5000, seed=None):
""" """
Random vecs following the exposure of an experiment (equatorial coordinates). Random vecs following the exposure of an experiment (equatorial coordinates).
If you need other coordinates, change 'coord_system' keyword to eg. 'eq' or 'sgal'. This If you need other coordinates, change 'coord_system' keyword to eg. 'eq' or 'sgal'. This
...@@ -565,11 +579,13 @@ def rand_exposure_vec(a0=-35.25, zmax=60, n=1, coord_system='gal', res_theta=500 ...@@ -565,11 +579,13 @@ def rand_exposure_vec(a0=-35.25, zmax=60, n=1, coord_system='gal', res_theta=500
:param res_theta: resolution of theta, number of bins in (-pi/2, pi/2) :param res_theta: resolution of theta, number of bins in (-pi/2, pi/2)
:return: random unit vectors from exposure of shape (3, n), equatorial coordinates :return: random unit vectors from exposure of shape (3, n), equatorial coordinates
""" """
if seed is not None:
np.random.seed(seed)
eps = 1. / res_theta eps = 1. / res_theta
theta_bin = np.linspace(-np.pi/2.+eps, np.pi/2.-eps, num=res_theta) theta_bin = np.linspace(-np.pi/2.+eps, np.pi/2.-eps, num=res_theta)
p = np.cos(theta_bin) * exposure_equatorial(theta_bin, a0, zmax) p = np.cos(theta_bin) * exposure_equatorial(theta_bin, a0, zmax)
thetas = np.random.choice(theta_bin, n, p=p/p.sum()) + np.random.uniform(-eps, eps, n) thetas = np.random.choice(theta_bin, n, p=p/p.sum()) + np.random.uniform(-eps, eps, n)
v = ang2vec(rand_phi(n), thetas) v = ang2vec(rand_phi(n, seed=seed), thetas)
if coord_system != 'eq': if coord_system != 'eq':
v = globals()['eq2%s' % coord_system](v) v = globals()['eq2%s' % coord_system](v)
...@@ -577,7 +593,7 @@ def rand_exposure_vec(a0=-35.25, zmax=60, n=1, coord_system='gal', res_theta=500 ...@@ -577,7 +593,7 @@ def rand_exposure_vec(a0=-35.25, zmax=60, n=1, coord_system='gal', res_theta=500
return v return v
def rand_fisher_vec(vmean, kappa, n=1): def rand_fisher_vec(vmean, kappa, n=1, seed=None):
""" """
Random Fisher distributed vectors with mean direction vmean and concentration parameter kappa. Random Fisher distributed vectors with mean direction vmean and concentration parameter kappa.
...@@ -586,19 +602,21 @@ def rand_fisher_vec(vmean, kappa, n=1): ...@@ -586,19 +602,21 @@ def rand_fisher_vec(vmean, kappa, n=1):
:param n: number of vectors drawn from fisher distribution, becomes m if vmean has shape (3, m) :param n: number of vectors drawn from fisher distribution, becomes m if vmean has shape (3, m)
:return: vectors from fisher distribution of shape (3, n) :return: vectors from fisher distribution of shape (3, n)
""" """
if seed is not None:
np.random.seed(seed)
vmean = atleast_kd(vmean, k=np.ndim(kappa)+1) vmean = atleast_kd(vmean, k=np.ndim(kappa)+1)
if np.ndim(kappa) > 0: if np.ndim(kappa) > 0:
n = kappa.shape n = kappa.shape
elif np.ndim(vmean) > 1: elif np.ndim(vmean) > 1:
n = vmean.shape[1:] n = vmean.shape[1:]
# create random fisher distributed directions around z-axis (0, 0, 1) # create random fisher distributed directions around z-axis (0, 0, 1)
theta = np.pi / 2 - rand_fisher(kappa, n) theta = np.pi / 2 - rand_fisher(kappa, n, seed=seed)
phi = rand_phi(theta.shape) phi = rand_phi(theta.shape, seed=seed)
# rotate these to reference vector vmean # rotate these to reference vector vmean
return rotate_zaxis_to_x(ang2vec(phi, theta), vmean) return rotate_zaxis_to_x(ang2vec(phi, theta), vmean)
def equatorial_scrambling(v, n=1, coord_system='gal'): def equatorial_scrambling(v, n=1, coord_system='gal', seed=None):
""" """
Performs a scrambling of the input vectors v around the equatorial z-axis. In this sense it keeps Performs a scrambling of the input vectors v around the equatorial z-axis. In this sense it keeps
the declination while assigning new uniform azimuth angles. the declination while assigning new uniform azimuth angles.
...@@ -608,8 +626,10 @@ def equatorial_scrambling(v, n=1, coord_system='gal'): ...@@ -608,8 +626,10 @@ def equatorial_scrambling(v, n=1, coord_system='gal'):
:param coord_system: coordinate system for input vectors :param coord_system: coordinate system for input vectors
:return: scrambled vectors in shape (3, n, ncrs) :return: scrambled vectors in shape (3, n, ncrs)
""" """
if seed is not None:
np.random.seed(seed)
decs = get_declination(v, coord_system) decs = get_declination(v, coord_system)
v_scrambled = ang2vec(rand_phi(v.shape[1] * n), np.tile(decs, n)) v_scrambled = ang2vec(rand_phi(v.shape[1] * n, seed=seed), np.tile(decs, n))
if coord_system != 'eq': if coord_system != 'eq':
v_scrambled = globals()['eq2%s' % coord_system](v_scrambled) v_scrambled = globals()['eq2%s' % coord_system](v_scrambled)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment