Coverage for polars_analysis / linearity.py: 45%

1import logging 

2import os 

3from pathlib import Path 

4from typing import List 

5 

6import polars as pl 

7from polars.exceptions import ComputeError 

8 

9from polars_analysis import pulse 

10from polars_analysis.analysis import linearity_analysis 

11from polars_analysis.data_sources import DataSource 

12from polars_analysis.plotting import linearity_plotting 

13 

14log = logging.getLogger(__name__) 

15 

16 

17def calc_derived( 

18 df: pl.DataFrame, 

19 cross_ofc_df: pl.DataFrame, 

20 plot_dir: Path, 

21 all_phases: bool, 

22 reference_att: float = 35, 

23) -> pl.DataFrame: 

24 return ( 

25 df.pipe(linearity_analysis.calc_derived, cross_ofc_df, all_phases) 

26 .pipe(linearity_analysis.calc_linearity_slopes) 

27 .pipe(linearity_analysis.calibrate_scale_factors, reference_att=reference_att) 

28 .pipe(linearity_plotting.awg_noise_fit_plot, plot_dir, att_val_name="calibrated_att_val") 

29 ) 

30 

31 

32def plot_all(df: pl.DataFrame, plot_dir: Path): 

33 for plot_log_scale in [False, True]: 

34 linearity_plotting.linearity_plot( 

35 df, 

36 plot_dir, 

37 full_range=True, 

38 plot_log_scale=plot_log_scale, 

39 ) 

40 linearity_plotting.linearity_plot( 

41 df, 

42 plot_dir, 

43 full_range=False, 

44 plot_log_scale=plot_log_scale, 

45 ) 

46 for awg_subtracted in [False, True]: 

47 linearity_plotting.energy_resolution(df, plot_dir, awg_subtracted) 

48 

49 linearity_plotting.energy_range_resolution(df, plot_dir) 

50 

51 

52def calc_plot_all( 

53 loader: DataSource, 

54 run_numbers: List[int], 

55 plot_dir: Path, 

56 all_phases: bool, 

57): 

58 if not plot_dir.exists(): 

59 plot_dir.mkdir(parents=True, exist_ok=True) 

60 os.chmod(plot_dir, 0o775) 

61 

62 # pick out the middle attenuation for which to use OFCs 

63 cross_ofc_run_number = loader.get_middle_attenuation_run(run_numbers) 

64 

65 log.info(f"Loading middle attenuation run {cross_ofc_run_number}") 

66 try: 

67 cross_ofc_df = loader.load_derived_data(cross_ofc_run_number, "pulse") 

68 except ComputeError: 

69 log.info(f"Could not find derived data for middle attenuation run {cross_ofc_run_number}, calculating...") 

70 cross_ofc_raw_data = loader.load_raw_data(cross_ofc_run_number, require_unsaturated=True, require_nonempty=True) 

71 cross_ofc_df = pulse.calc_derived(cross_ofc_raw_data, all_phases=True) 

72 loader.save_derived_data(cross_ofc_df, cross_ofc_run_number, "pulse") 

73 

74 dfs: List[pl.DataFrame] = [] 

75 for run_number in run_numbers: 

76 try: 

77 derived_df = loader.load_derived_data(run_number=run_number, meas_type="pulse") 

78 except ComputeError: 

79 log.info(f"Failed to load derived data for run {run_number}, calculating...") 

80 raw_df = loader.load_raw_data(run_number, require_unsaturated=True, require_nonempty=True) 

81 derived_df = pulse.calc_derived(raw_df, all_phases=True) 

82 loader.save_derived_data(derived_df, run_number, "pulse") 

83 

84 dfs.append(derived_df) 

85 

86 df = pl.concat(dfs, how="diagonal_relaxed") 

87 

88 reference_att: float = cross_ofc_df["att_val"][0] 

89 df = calc_derived(df, cross_ofc_df, plot_dir, all_phases, reference_att=reference_att) 

90 

91 plot_all(df, plot_dir)