Coverage for wifa_uq / postprocessing / bayesian_calibration / bayesian_calibration.py: 0%

1from wifa_uq.postprocessing.postprocesser import PostProcesser 

2import numpy as np 

3from wifa_uq.postprocessing.calibration import ( 

4 MinBiasCalibrator, 

5 # DefaultParams, 

6 # LocalParameterPredictor, 

7) 

8 

9try: 

10 import umbra 

11except Exception as e: 

12 print("Umbra not installed.") 

13 print(e) 

14 quit() 

15 

16 

17class BayesianCalibration(PostProcesser): 

18 def __init__(self, system_yaml, params, data): 

19 self.system_yaml = system_yaml 

20 self.params = params 

21 self.data = data 

22 

23 self._flow_model = umbra.flow_model.WIFAModel(open(system_yaml)) 

24 self._initialize_bayesian_model() 

25 

26 self._samples_posterior = None 

27 self._samples_posterior_predictive = None 

28 self._inference_data = None 

29 

30 def _initialize_bayesian_model(self): 

31 # Initalize prior 

32 params = [] 

33 dists = [] 

34 for param, range in self.params.items(): 

35 assert (len(range) == 2) and (range[0] < range[1]) 

36 params.append(param) 

37 dists.append(umbra.distributions.Uniform(range[0], range[1])) 

38 prior = umbra.bayesian_model.Prior(params, dists) 

39 # Initalize likelihood 

40 likelihood = umbra.bayesian_model.LikelihoodTurbinePower( 

41 flow_model=self._flow_model, data=self.data 

42 ) 

43 # Initialize bayesian model 

44 self._bayesian_model = umbra.bayesian_model.BayesianModel( 

45 prior, likelihood, name="umbra4wifa" 

46 ) 

47 

48 def fit(self): 

49 sampler = umbra.sampler.TMCMC(self._bayesian_model, parallel=False) 

50 self._inference_data = sampler.sample() 

51 self._samples_posterior = self._inference_data.samples 

52 return self._samples_posterior 

53 

54 def predict(self): 

55 self._samples_posterior_predictive = self._bayesian_model.posterior_predictive( 

56 self.samples_posterior 

57 ) 

58 return self._samples_posterior_predictive 

59 

60 @property 

61 def samples_posterior(self): 

62 return self._samples_posterior 

63 

64 @property 

65 def samples_posterior_predictive(self): 

66 return self._samples_posterior_predictive 

67 

68 

69class BayesianCalibrationWrapper: 

70 """ 

71 Wrapper to make BayesianCalibration compatible with MainPipeline. 

72 

73 This performs Bayesian inference to get a posterior distribution 

74 of parameters, then uses the MAP (Maximum A Posteriori) estimate 

75 as the calibrated parameters. 

76 """ 

77 

78 def __init__( 

79 self, dataset_train, system_yaml: str = None, param_ranges: dict = None 

80 ): 

81 """ 

82 Args: 

83 dataset_train: xarray Dataset with the model bias database 

84 system_yaml: Path to windIO system YAML (required for WIFA model) 

85 param_ranges: Dict of {param_name: [min, max]} for prior bounds 

86 """ 

87 self.dataset_train = dataset_train 

88 self.system_yaml = system_yaml 

89 self.param_ranges = param_ranges or {} 

90 

91 self.swept_params = dataset_train.attrs.get("swept_params", []) 

92 self.best_idx_ = None 

93 self.best_params_ = None 

94 self.posterior_samples_ = None 

95 

96 # Build param_ranges from database if not provided 

97 if not self.param_ranges: 

98 self._infer_param_ranges() 

99 

100 def _infer_param_ranges(self): 

101 """Infer parameter ranges from database coordinates.""" 

102 for param_name in self.swept_params: 

103 if param_name in self.dataset_train.coords: 

104 values = self.dataset_train.coords[param_name].values 

105 self.param_ranges[param_name] = [ 

106 float(values.min()), 

107 float(values.max()), 

108 ] 

109 

110 def fit(self): 

111 """ 

112 Run Bayesian inference and extract MAP estimate. 

113 

114 Note: This requires UMBRA to be installed and a valid system_yaml. 

115 Falls back to MinBiasCalibrator if UMBRA is unavailable. 

116 """ 

117 try: 

118 if self.system_yaml is None: 

119 raise ValueError("system_yaml path required for Bayesian calibration") 

120 

121 # Prepare observation data from dataset 

122 # This is simplified - you may need to adjust based on your data structure 

123 ref_power = self.dataset_train["ref_power_cap"].isel(sample=0).values 

124 

125 # Create BayesianCalibration instance 

126 bc = BayesianCalibration( 

127 system_yaml=self.system_yaml, params=self.param_ranges, data=ref_power 

128 ) 

129 

130 # Run inference 

131 self.posterior_samples_ = bc.fit() 

132 

133 # Extract MAP estimate (mode of posterior) 

134 self.best_params_ = {} 

135 for i, param_name in enumerate(self.swept_params): 

136 samples = self.posterior_samples_[:, i] 

137 # Use median as robust point estimate 

138 self.best_params_[param_name] = float(np.median(samples)) 

139 

140 # Find closest sample index to MAP estimate 

141 self.best_idx_ = self._find_closest_sample_idx() 

142 

143 except ImportError: 

144 print("WARNING: UMBRA not installed. Falling back to MinBiasCalibrator.") 

145 fallback = MinBiasCalibrator(self.dataset_train) 

146 fallback.fit() 

147 self.best_idx_ = fallback.best_idx_ 

148 self.best_params_ = fallback.best_params_ 

149 

150 except Exception as e: 

151 print( 

152 f"WARNING: Bayesian calibration failed ({e}). Falling back to MinBiasCalibrator." 

153 ) 

154 fallback = MinBiasCalibrator(self.dataset_train) 

155 fallback.fit() 

156 self.best_idx_ = fallback.best_idx_ 

157 self.best_params_ = fallback.best_params_ 

158 

159 return self 

160 

161 def _find_closest_sample_idx(self): 

162 """Find sample index closest to MAP estimate.""" 

163 n_samples = len(self.dataset_train.sample) 

164 distances = np.zeros(n_samples) 

165 

166 for param_name, map_value in self.best_params_.items(): 

167 if param_name in self.dataset_train.coords: 

168 sample_values = self.dataset_train.coords[param_name].values 

169 distances += (sample_values - map_value) ** 2 

170 

171 return int(np.argmin(distances)) 

172 

173 def get_posterior_samples(self): 

174 """Return posterior samples for uncertainty quantification.""" 

175 return self.posterior_samples_