Coverage for src/galsbi/ucat/galaxy_population_models/galaxy_position.py: 96%

1# Copyright (C) 2018 ETH Zurich, Institute for Particle Physics and Astrophysics 

2 

3""" 

4Created 2021 

5author: Tomasz Kacprzak 

6""" 

7 

8import healpy as hp 

9import numpy as np 

10from cosmic_toolbox import logger 

11from ufig import coordinate_util 

12 

13########################################################## 

14# 

15# Sampling position 

16# 

17########################################################## 

18 

19LOGGER = logger.get_logger(__file__) 

20 

21# Consistent with PINOCCHIO/SHAM pipeline 

22GALAXY_TYPE_FLAGS = { 

23 "blue": 1, 

24 "red": 0, 

25} 

26 

27 

28def sample_position_uniform(numobj, w, pixel_index, nside): 

29 """ 

30 Sample a Healpix pixel uniformly 

31 

32 :param numobj: Number of uniform samples 

33 :param w: wcs-object containing all the relevant wcs-information 

34 :param pixel_index: Index of the Healpix pixels sampled 

35 :param nside: NSIDE of the Healpix map 

36 :return: Uniformly drawn x-coordinate (in pixels) 

37 :return: Uniformly drawn y-coordinate (in pixels) 

38 """ 

39 

40 corners = hp.boundaries(nside, pixel_index, 1) 

41 theta, phi = hp.vec2ang(np.transpose(corners)) 

42 

43 # removing this as it will cause sampling uniform to sample over too big range and 

44 # the while loop to take forever the subsequent code should handle this 

45 phi[phi > np.pi] -= 2 * np.pi # To deal with periodic boundaries 

46 

47 cos_theta_sample = np.random.uniform( 

48 low=np.min(np.cos(theta)), high=np.max(np.cos(theta)), size=numobj 

49 ) 

50 phi_sample = np.random.uniform(low=np.min(phi), high=np.max(phi), size=numobj) 

51 

52 n_reps = 0 

53 

54 while True: 

55 mask = hp.ang2pix(nside, np.arccos(cos_theta_sample), phi_sample) != pixel_index 

56 

57 if np.sum(mask) == 0: 

58 break 

59 

60 cos_theta_sample[mask] = np.random.uniform( 

61 low=np.min(np.cos(theta)), high=np.max(np.cos(theta)), size=np.sum(mask) 

62 ) 

63 

64 phi_sample[mask] = np.random.uniform( 

65 low=np.min(phi), high=np.max(phi), size=np.sum(mask) 

66 ) 

67 

68 n_reps += 1 

69 

70 if n_reps > 10_000: 

71 raise RuntimeError("more than 10_000 many iterations") 

72 if w is None: 

73 # If no wcs is provided, ra and dec are returned 

74 return coordinate_util.thetaphi2radec(np.arccos(cos_theta_sample), phi_sample) 

75 else: 

76 # If a wcs is provided, x and y are returned 

77 return coordinate_util.thetaphi2xy(w, np.arccos(cos_theta_sample), phi_sample)