Coverage for polars_analysis / analysis / cross_talk_analysis.py: 100%

1import logging 

2 

3import polars as pl 

4 

5from polars_analysis.analysis import constants 

6from polars_analysis.analysis import pulse_analysis as pulse 

7 

8# Instantiate logger 

9log = logging.getLogger(__name__) 

10 

11""" 

12Functions to calculate derived values for cross-talk runs. 

13""" 

14 

15 

16def expr_awg_amp_to_amp(col1: str = "awg_amp", col2: str = "att_val") -> pl.Expr: 

17 return pulse.expr_awg_amp_to_amp(col1, col2) 

18 

19 

20def expr_max_pulse_amp(col: str = "samples") -> pl.Expr: 

21 return pulse.expr_max_pulse_amp(col) 

22 

23 

24def expr_max_phase_indices(mean_interleaved_pulse: str = "mean_interleaved_pulse", phase_shift: int = 0) -> pl.Expr: 

25 """ 

26 Get indices of samples around the maximum pulse amplitude, shifted by a phase. 

27 For cross-talk measurements, these are calculated from the reference pulse only. 

28 

29 Args: 

30 mean_interleaved_pulse: Column name of mean interleaved pulse. 

31 phase_shift: Phase shift in samples. 

32 

33 Returns: 

34 Expression to calculate indices of samples around the maximum pulse amplitude. 

35 

36 """ 

37 samples_around_max = [-2, -1, 0, 1, 2] 

38 return pl.concat_list( 

39 [ 

40 pl.col(mean_interleaved_pulse).sort_by(pl.col(mean_interleaved_pulse).list.max()).last().list.arg_max() 

41 + phase_shift 

42 + (constants.N_PHASES * i) 

43 for i in samples_around_max 

44 ] 

45 ) 

46 

47 

48def expr_is_reference_pulse(meas_chan: str = "meas_chan", channel: str = "channel"): 

49 return pl.col(meas_chan) == pl.col(channel) 

50 

51 

52def pipe_samples_interleaved( 

53 df: pl.DataFrame, samples: str = "samples", is_reference_pulse: str = "is_reference_pulse" 

54) -> pl.DataFrame: 

55 return pulse.pipe_samples_interleaved(df, samples, is_reference_pulse) 

56 

57 

58def pipe_OFCs( 

59 df: pl.DataFrame, 

60 mean_interleaved_pulse: str = "mean_interleaved_pulse", 

61 max_phase_indices: str = "max_phase_indices", 

62 amp: str = "amp", 

63 quantile: float = 1.0, 

64) -> pl.DataFrame: 

65 return pulse.pipe_OFCs(df, mean_interleaved_pulse, max_phase_indices, amp, is_crosstalk=True, quantile=quantile) 

66 

67 

68def pipe_apply_OFCs( 

69 df: pl.DataFrame, 

70 samples_interleaved: str = "samples_interleaved", 

71 max_phase_indices: str = "max_phase_indices", 

72 OFCs_a: str = "OFCs_a", 

73 OFCs_b: str = "OFCs_b", 

74 all_phases=False, 

75) -> pl.DataFrame: 

76 return pulse.pipe_apply_OFCs(df, samples_interleaved, max_phase_indices, OFCs_a, OFCs_b, all_phases) 

77 

78 

79def pipe_rise_time( 

80 df: pl.DataFrame, 

81 samples_interleaved: str = "samples_interleaved", 

82 mean_interleaved_pulse: str = "mean_interleaved_pulse", 

83) -> pl.DataFrame: 

84 return pulse.pipe_rise_time(df, samples_interleaved, mean_interleaved_pulse)