Coverage for dynasor/post_processing/average_runs.py: 100%

1from copy import deepcopy 

2from typing import Optional 

3import numpy as np 

4from dynasor.sample import Sample 

5 

6 

7def get_sample_averaged_over_independent_runs( 

8 samples: list[Sample], 

9 live_dangerously: Optional[bool] = False, 

10) -> Sample: 

11 """ 

12 Compute an averaged sample from multiple samples obtained from identical independent runs. 

13 

14 Note all the metadata and dimensions in all samples must be the same. 

15 Otherwise a `ValueError` is raised (unless `live_dangerously` is set to `True`). 

16 

17 Parameters 

18 ---------- 

19 samples 

20 List of all sample objects to be averaged over. 

21 live_dangerously 

22 Setting to `True` allows for averaging over samples 

23 which metadata information is not identical. 

24 """ 

25 

26 if len(samples) < 2: 

27 raise ValueError('Averaging requires at least two samples.') 

28 

29 # get metadata and dimensions from first sample 

30 sample_ref = samples[0] 

31 data_dict = dict() 

32 simulation_data = deepcopy(sample_ref.simulation_data) 

33 

34 # test that all samples have identical dimensions and correlation functions 

35 for m, sample in enumerate(samples): 

36 if sorted(sample.dimensions) != sorted(sample_ref.dimensions): 

37 raise ValueError(f'Sample dimensions do not match for sample #{m}.') 

38 for dim in sample_ref.dimensions: 

39 if not np.allclose(sample[dim], sample_ref[dim]): 

40 raise ValueError(f'Sample dimensions do not match for sample #{m}.') 

41 if sample.available_correlation_functions != sample_ref.available_correlation_functions: 

42 raise ValueError(f'Sample correlation functions do not match for sample #{m}.') 

43 

44 for dim in sample_ref.dimensions: 

45 data_dict[dim] = sample_ref[dim] 

46 

47 # test that all samples have identical metadata 

48 if not live_dangerously: 

49 for m, sample in enumerate(samples): 

50 for key, val in simulation_data.items(): 

51 if key not in sample.simulation_data: 

52 raise ValueError( 

53 f'Sample #{m} is missing "{key}" in the simulation_data field.') 

54 match = True 

55 if isinstance(val, dict): 

56 for k, v in val.items(): 

57 match &= sample.simulation_data[key].get(k, None) == val[k] 

58 elif isinstance(val, np.ndarray): 

59 match &= np.allclose(sample.simulation_data[key], val) 

60 elif isinstance(val, float): 

61 match &= np.isclose(sample.simulation_data[key], val) 

62 else: 

63 match &= sample.simulation_data[key] == val 

64 if not match: 

65 raise ValueError(f'Field "{key}" of sample #{m} does not match.') 

66 

67 # average all correlation functions 

68 for key in sample_ref.available_correlation_functions: 

69 data = [] 

70 for sample in samples: 

71 data.append(sample[key]) 

72 data_average = np.nanmean(data, axis=0) 

73 data_dict[key] = data_average 

74 

75 # keep history of original samples 

76 previous_history = [] 

77 for m, s in enumerate(samples): 

78 for h in s.history: 

79 rec = h.copy() 

80 rec['func'] += f'_sample{m}' 

81 previous_history.append(rec) 

82 

83 # compose new sample object 

84 new_sample = sample_ref.__class__( 

85 data_dict, 

86 simulation_data=simulation_data, 

87 history=previous_history) 

88 new_sample._append_history( 

89 'get_sample_averaged_over_independent_runs', 

90 dict( 

91 live_dangerously=live_dangerously, 

92 n_samples=len(samples), 

93 )) 

94 

95 return new_sample