Coverage for PyFHD/plotting/gridding.py: 0%

1import numpy as np 

2from numpy.typing import NDArray 

3from pathlib import Path 

4from PyFHD.plotting.image import quick_image 

5from logging import Logger 

6 

7 

8def plot_gridding( 

9 obs: dict, 

10 image_uv: NDArray[np.complex128], 

11 weights_uv: NDArray[np.complex128], 

12 variance_uv: NDArray[np.float64], 

13 pyfhd_config: dict, 

14 model_uv: NDArray[np.complex64] | None, 

15 logger: Logger, 

16) -> None: 

17 """ 

18 Plot the continuum uv-planes for data, model, and weights. 

19 

20 Parameters 

21 ---------- 

22 obs : dict 

23 Observation metadata dictionary 

24 image_uv : NDArray[np.complex128] 

25 UV-plane of the gridded data 

26 weights_uv : NDArray[np.complex128] 

27 UV-plane of the gridded weights 

28 variance_uv : NDArray[np.float64] 

29 UV-plane of the gridded variances 

30 model_uv : NDArray[np.complex128], optional 

31 UV-plane of the gridded model 

32 pyfhd_config : dict 

33 Run option dictionary 

34 logger : Logger 

35 PyFHD's logger for displaying errors and info to the log files 

36 """ 

37 

38 # Check if image, weights, and variance contain any non-zero elements 

39 if ( 

40 (not np.any(image_uv != 0)) 

41 or (not np.any(weights_uv != 0)) 

42 or (not np.any(variance_uv != 0)) 

43 ): 

44 logger.warning( 

45 "Gridded image, weights, or variance are all zeros. Plotting skipped." 

46 ) 

47 return 

48 

49 # Plotting paths for apparent gridded image (weighted gridded data), variance, 

50 # and apparent model (weighted gridded model) 

51 obs_id = pyfhd_config["obs_id"] 

52 grid_plot_dir = Path(pyfhd_config["output_dir"], "plots", "gridding") 

53 grid_plot_dir.mkdir(parents=True, exist_ok=True) 

54 save_path_roots = [ 

55 Path(grid_plot_dir, f"{obs_id}_grid_apparent_image"), 

56 Path(grid_plot_dir, f"{obs_id}_grid_variance"), 

57 Path(grid_plot_dir, f"{obs_id}_grid_apparent_model"), 

58 ] 

59 

60 # Calculate the apparent image, setting inds with zero weights to zero 

61 apparent_image = np.abs(np.divide(image_uv, weights_uv, where=weights_uv != 0)) 

62 

63 # Calculate the apparent model, setting inds with zero weights to zero 

64 if model_uv is not None: 

65 apparent_model = np.abs(np.divide(model_uv, weights_uv, where=weights_uv != 0)) 

66 

67 # Calculate the x- and y-axis values and labels 

68 xvals = (np.arange(apparent_image.shape[1]) - apparent_image.shape[1] / 2) * obs[ 

69 "kpix" 

70 ] 

71 xtitle = "u (wavelengths)" 

72 yvals = (np.arange(apparent_image.shape[2]) - apparent_image.shape[2] / 2) * obs[ 

73 "kpix" 

74 ] 

75 ytitle = "v (wavelengths)" 

76 

77 # Get the pol_names 

78 pol_names = obs["pol_names"] 

79 

80 for pol_i in range(obs["n_pol"]): 

81 # Plot the apparent image, variance, and the optional model for each polarization 

82 

83 # Add a suffix to each path 

84 save_path_pol = [ 

85 path.with_stem(path.stem + "_" + pol_names[pol_i]) 

86 for path in save_path_roots 

87 ] 

88 

89 quick_image( 

90 apparent_image[pol_i, :, :], 

91 xvals=xvals, 

92 yvals=yvals, 

93 xtitle=xtitle, 

94 ytitle=ytitle, 

95 title="Apparent Gridded Continuum Data " + pol_names[pol_i], 

96 cb_title="Amplitude (Jy/beam)", 

97 savefile=save_path_pol[0], 

98 missing_value=0, 

99 log=True, 

100 png=True, 

101 ) 

102 quick_image( 

103 variance_uv[pol_i, :, :], 

104 xvals=xvals, 

105 yvals=yvals, 

106 xtitle=xtitle, 

107 ytitle=ytitle, 

108 title="Gridded Continuum Variance " + pol_names[pol_i], 

109 cb_title="(Jy/beam)$^2$", 

110 savefile=save_path_pol[1], 

111 missing_value=0, 

112 log=True, 

113 png=True, 

114 ) 

115 

116 if model_uv is not None: 

117 quick_image( 

118 apparent_model[pol_i, :, :], 

119 xvals=xvals, 

120 yvals=yvals, 

121 xtitle=xtitle, 

122 ytitle=ytitle, 

123 title="Apparent Gridded Continuum Model " + pol_names[pol_i], 

124 cb_title="Amplitude (Jy/beam)", 

125 savefile=save_path_pol[2], 

126 missing_value=0, 

127 log=True, 

128 png=True, 

129 )