Coverage for src/susi/atlantes/compare.py: 25%

1#!/usr/bin/env python3 

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

3""" 

4Module provides methods to compare a 1D spectrum with an atlas spectrum 

5 

6@author: hoelken 

7""" 

8 

9import numpy as np 

10import matplotlib.pyplot as plt 

11from scipy.interpolate import interp1d 

12from scipy.optimize import minimize 

13 

14from .atlas import Atlas 

15from .hamburg import FTSAtlas 

16from .sss import SSSAtlas 

17from ..base import Logging 

18from scipy import signal 

19 

20 

21log = Logging.get_logger() 

22 

23# Conversion factors 

24ANGSTROM2NM = 1/10 

25NM2ANGSTROM = 10 

26 

27 

28def overplot(spectrum: np.array, central_wl: float, linear_dispersion: float = 0.0015, offset: float = 0, 

29 label: str = 'SUSI', atlas: str = 'fts', normalize: bool = True) -> None: 

30 """ 

31 Method to overplot a given SUSI 1D Spectrum with a reference from an atlas. 

32 

33 Params: 

34 ------ 

35 

36 spectrum: The SUSI 1D spectrum as numpy array 

37 

38 central_wl: The (assumed) central WL of the spectrum in [nm] 

39 

40 offset: The (assumed) offset in [nm] to the reference spectrum 

41 

42 title: An optional label for the plot (Default: `'SUSI'`) 

43 

44 atlas: Select the atlas to use. Default: `'fts'`, Supported: 'fts', 'sss' 

45 

46 normalize: Flag to determine if the spectra should be normalized by their mean (Default: `True`) 

47 """ 

48 wl_range = len(spectrum) * linear_dispersion / 2 

49 lim_low = central_wl - wl_range - 1 

50 lim_high = central_wl + wl_range + 1 

51 

52 if atlas == 'hamburg' or atlas == 'fts': 

53 log.info('Loading Hamburg Atlas from %.2f to %.2f nm', lim_low, lim_high) 

54 ref = FTSAtlas(lim_low * NM2ANGSTROM, lim_high * NM2ANGSTROM).load() 

55 ref.convert_wl(ANGSTROM2NM) 

56 elif atlas == 'sss': 

57 log.info('Loading SSS Atlas from %.2f to %.2f nm', lim_low, lim_high) 

58 ref = SSSAtlas(lim_low * NM2ANGSTROM, lim_high * NM2ANGSTROM).load() 

59 ref.convert_wl(ANGSTROM2NM) 

60 else: 

61 log.critical('Atlas %s not supported. Abort comparison.', atlas) 

62 return 

63 

64 coeff = find_deflection(spectrum, ref, central_wl, offset=offset, linear_dispersion=linear_dispersion) 

65 string_poly = f'{coeff[1]:.3e}x + {coeff[0]:.2e}' 

66 log.info("Dispersion polynomial: %s nm/px", string_poly) 

67 xes = _gen_xes(central_wl, len(spectrum), np.polynomial.polynomial.Polynomial(coeff)) 

68 

69 intensity = ref.intensity 

70 ylabel = 'Intensity' 

71 if normalize: 

72 ylabel = 'Normalized intensity' 

73 intensity = intensity / intensity.mean() 

74 spectrum = spectrum / np.mean(spectrum) 

75 

76 p = np.polyfit(xes, spectrum, 6) 

77 plt.plot(xes, spectrum, label=f'{label} (scaled)') 

78 plt.plot(xes, np.polyval(p, xes), label=f'fitted continuum') 

79 plt.plot(ref.wl, intensity, linestyle='-.', label=f"{atlas.upper()} atlas") 

80 plt.axvline(x=central_wl, color='black', linestyle='-.', linewidth=0.8) 

81 plt.title(f'1D Spectrum comparison around {central_wl} nm\nDispersion {string_poly} nm/px') 

82 plt.grid(axis='x', which='major', linestyle='--') 

83 plt.grid(axis='x', which='minor', linestyle=':') 

84 plt.ylabel(ylabel) 

85 plt.xlabel(r'$\lambda [nm]$') 

86 plt.xlim(xes[0], xes[-1]) 

87 plt.legend() 

88 plt.minorticks_on() 

89 plt.tight_layout() 

90 plt.show() 

91 

92 

93def _gen_xes(center: float, size: int, poly: np.polynomial.polynomial.Polynomial) -> list: 

94 return np.array(poly(range(-size//2, size//2))) + center 

95 

96 

97def find_deflection(spec: np.array, ref: Atlas, central_wl: float, 

98 linear_dispersion: float = 0.0015, offset: float = 0) -> tuple: 

99 par0 = np.array([offset, linear_dispersion]) 

100 log.info("Optimizing linear dispersion regarding to atlas. Assume %s", par0) 

101 bounds = [(-0.1, 0.1), (1e-4, 2e-2)] 

102 res = minimize(_max_delta, par0, (spec / spec.mean(), ref, central_wl), bounds=bounds) 

103 log.debug(res.message) 

104 return res.x 

105 

106 

107def _max_delta(params: np.array, spec: np.array, ref: Atlas, central_wl: float) -> float: 

108 """ 

109 params[2] = [float] quadratic_dispersion the x^2 coefficient 

110 params[1] = [float] linear_dispersion the x^1 coefficient (spectral resolution at center) 

111 params[0] = [float] constant x^0 coefficient (offset against atlas) 

112 """ 

113 xes = _gen_xes(central_wl, len(spec), np.polynomial.polynomial.Polynomial(params)) 

114 stretched = interp1d(xes, spec, kind='cubic') 

115 intersection = [i for i in range(len(ref.wl)) if xes[0] < ref.wl[i] < xes[-1]] 

116 if len(intersection) < len(spec)//4: 

117 return np.infty 

118 diff = stretched(ref.wl[intersection]) - ref.normalized_intensity()[intersection] 

119 return np.abs(np.std(diff)/np.mean(diff))