Coverage for src/susi/io/fits_cube.py: 24%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4module provides calsses to deal with susi fits cubes 

5 

6@author: iglesias 

7""" 

8from astropy.io import fits 

9import numpy as np 

10from ..base import Logging 

11from ..base.header_keys import * 

12from astropy.io import fits 

13 

14logger = Logging.get_logger() 

15 

16 

17class SusiCube: 

18 """ 

19 Class to read/write susi fits cubes. Supported formats are: 

20 

21 - 4dView FITS cube: [stks,scan,slit,wl] 

22 - 3dExt FITS cube: [stks,slit,wl] and nscans one in each FITS extension to preserve the headers 

23 

24 Returns allways data array in format [stks,scan,slit,wl] and all avilable headers 

25 """ 

26 

27 def __init__(self, path): 

28 self.path = path 

29 self.hdul = None 

30 # data has allways [stks,scan,slit,wl] dimensions 

31 self.data = None 

32 #: header can be a list of headers in 3dExt format 

33 self.header = None 

34 #: cube format that was read and that will be written 

35 self.format = None 

36 

37 def read(self, convert_to_4dView=False): 

38 # detect cube format and read the data accordingly 

39 # self.data has always [stks,scan,slit,wl] dimensions 

40 # self.headers can be one (4dView) or a list of headers (3dExt) 

41 self.hdul = fits.open(self.path) 

42 if len(self.hdul) == 1: 

43 logger.info("Reading 4dView cube") 

44 self.data = self.hdul[0].data 

45 self.header = self.hdul[0].header 

46 self.format = '4dView' 

47 else: 

48 if convert_to_4dView: 

49 logger.info("Reading 3dExt cube and converting to 4dView format") 

50 susi_cube = SusiCube.read_3dExt_as_4dView(self.hdul) 

51 self.data = susi_cube.data 

52 self.header = susi_cube.header 

53 self.format = '4dView' 

54 else: 

55 logger.info("Reading 3dExt cube") 

56 self.data, self.header = SusiCube.read_3dExt(self.hdul) 

57 self.format = '3dExt' 

58 return self.data, self.header 

59 

60 def write_to_disk(self, overwrite=False): 

61 # updates hdul with the data and headers and writes to disk 

62 if self.format == '4dView': 

63 fits.writeto(self.path, self.data, self.header, overwrite=True) 

64 elif self.format == '3dExt': 

65 mhdu = fits.HDUList() 

66 mhdu.append(fits.PrimaryHDU(self.data[:, 0, :, :], header=self.header[0])) 

67 for i in range(1, len(self.header)): 

68 mhdu.append(fits.ImageHDU(self.data[:, i, :, :], header=self.header[i])) 

69 mhdu.writeto(self.path, overwrite=overwrite) 

70 else: 

71 raise ValueError("Unknown cube format") 

72 

73 @staticmethod 

74 def read_3dExt(hdul): 

75 # Reads susi 3dExt cube with dimensions [stks,slit,wl] and nscans extensions 

76 # Returns data with dimensions [stks,scan,slit,wl] and all headers 

77 data = [] 

78 headers = [] 

79 for i in range(1, len(hdul)): 

80 data.append(hdul[i].data) 

81 headers.append(hdul[i].header) 

82 data = np.array(data).transpose(1, 0, 2, 3) 

83 return data, headers 

84 

85 @staticmethod 

86 def read_3dExt_as_4dView(hdul): 

87 # Reads susi 3dExt cube with dimensions [stks,slit,wl] and nscans extensions 

88 # Returns a susi 4dView compatible cube with dimensions [stks,scan,slit,wl] 

89 # Modifies the first header to be 4dView compatible 

90 data = [] 

91 for i in range(1, len(hdul)): 

92 data.append(hdul[i].data) 

93 data = np.array(data).transpose(1, 0, 2, 3) 

94 header = hdul[0].header 

95 header['NAXIS'] = 4 

96 header['NAXIS4'] = data.shape[3] 

97 header['NAXIS3'] = data.shape[2] 

98 header['NAXIS2'] = data.shape[1] 

99 header['NAXIS1'] = data.shape[0] 

100 return fits.PrimaryHDU(data, header)