Coverage for backend/ahuora-builder/src/ahuora_builder/ml_wizard.py: 94%

1import idaes 

2from pydantic import BaseModel 

3from ahuora_builder_types.payloads.ml_request_schema import MLTrainRequestPayload, MLTrainingCompletionPayload 

4import contextlib 

5import json 

6from dataclasses import dataclass 

7from io import StringIO 

8from typing import Any 

9 

10 

11import pandas as pd 

12import numpy as np 

13from sklearn.model_selection import train_test_split 

14from sklearn.metrics import mean_squared_error, r2_score 

15from idaes.core.surrogate.pysmo_surrogate import PysmoRBFTrainer, PysmoPolyTrainer, PysmoSurrogate 

16 

17from ahuora_builder_types.payloads.ml_request_schema import ( 

18 MLTrainingChartPayload as MLChart, 

19 MLTrainingRegressionMetricPayload as RegressionMetric, 

20 MLTrainingResultPayload as MLResult, 

21) 

22 

23 

24MAX_QQ_CHART_POINTS = 250 

25 

26@dataclass 

27class TrainingOutput: 

28 """Intermediate ML training artefacts before object-storage upload.""" 

29 

30 result_payload: MLResult 

31 test_inputs_df: pd.DataFrame 

32 test_outputs_df: pd.DataFrame 

33 

34 

35def _json_indexed_records(frame: pd.DataFrame) -> dict[str, dict[str, float]]: 

36 """Convert a dataframe to a JSON-safe indexed mapping with string row keys.""" 

37 return { 

38 str(index): { 

39 str(column): float(value) 

40 for column, value in row.items() 

41 } 

42 for index, row in frame.to_dict(orient="index").items() 

43 } 

44 

45 

46def train_dataframe( 

47 df: pd.DataFrame, 

48 *, 

49 input_labels: list[str], 

50 output_labels: list[str], 

51 model_type: str, 

52) -> TrainingOutput: 

53 """Train a PySMO surrogate model from a dataframe and return upload-ready artefacts.""" 

54 train_df, test_df = train_test_split(df, test_size=0.2, random_state=42) 

55 

56 if model_type == "polynomial_regression": 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true

57 trainer = PysmoPolyTrainer( 

58 input_labels=input_labels, output_labels=output_labels, training_dataframe=train_df) 

59 elif model_type == "rbf_regression": 59 ↛ 65line 59 didn't jump to line 65 because the condition on line 59 was always true

60 trainer = PysmoRBFTrainer( 

61 input_labels=input_labels, output_labels=output_labels, training_dataframe=train_df) 

62 trainer.config.basis_function = 'gaussian' 

63 

64 # Train surrogate (calls PySMO through IDAES Python wrapper) 

65 stream = StringIO() 

66 with contextlib.redirect_stdout(stream): 

67 train = trainer.train_surrogate() 

68 

69 # create callable surrogate model 

70 surr = PysmoSurrogate(train , input_labels, 

71 output_labels, input_bounds=None) 

72 f = StringIO() 

73 surr.save(f) 

74 content = f.getvalue() 

75 json_data = json.loads(content) 

76 

77 df_evaluate = surr.evaluate_surrogate(test_df) 

78 

79 metrics: list[RegressionMetric] = [] 

80 charts: list[MLChart] = [] 

81 

82 for output_label in output_labels: 

83 charts.append(compute_chart( 

84 test_df[output_label], df_evaluate[output_label], output_label)) 

85 metrics.append( 

86 RegressionMetric( 

87 mean_squared_error=round( 

88 mean_squared_error(df_evaluate[output_label], test_df[output_label]), 

89 4, 

90 ), 

91 r2_score=round( 

92 r2_score(df_evaluate[output_label], test_df[output_label]), 

93 4, 

94 ), 

95 ) 

96 ) 

97 

98 return TrainingOutput( 

99 result_payload=MLResult( 

100 surrogate_model=json_data, 

101 charts=charts, 

102 metrics=metrics, 

103 test_results_bucket="", 

104 test_results_key="", 

105 timing={}, 

106 ), 

107 test_inputs_df=test_df[input_labels].copy(), 

108 test_outputs_df=df_evaluate[output_labels].copy(), 

109 ) 

110 

111 

112def compute_chart( 

113 test_data_df: pd.Series, 

114 eval_data_df: pd.Series, 

115 output_label: str, 

116) -> MLChart: 

117 """Build the chart payload for one output label.""" 

118 minn = round(np.min([np.min(test_data_df), np.min(eval_data_df)]), 4) 

119 maxx = round(np.max([np.max(test_data_df), np.max(eval_data_df)]), 4) 

120 qq_plot_data = compute_qq_coordinates(test_data_df, eval_data_df) 

121 

122 return MLChart( 

123 min=minn, 

124 max=maxx, 

125 qq_plot_data=qq_plot_data, 

126 output_label=output_label, 

127 ) 

128 

129 

130def compute_qq_coordinates(test_data: pd.Series, eval_data: pd.Series) -> str: 

131 """Compute QQ plot coordinates for two datasets as a frontend-ready JSON string.""" 

132 test_values = test_data.to_numpy().flatten() 

133 eval_values = eval_data.to_numpy().flatten() 

134 

135 # Sort values 

136 test_values.sort() 

137 eval_values.sort() 

138 

139 # Generate QQ plot data points 

140 n_points = min(len(test_values), MAX_QQ_CHART_POINTS) 

141 quantiles = np.linspace(0, 1, n_points) 

142 test_quantiles = np.quantile(test_values, quantiles) 

143 eval_quantiles = np.quantile(eval_values, quantiles) 

144 

145 # Prepare JSON response for frontend 

146 qq_data = [{"x": round(float(t), 4), "y": round(float(e), 4)} 

147 for t, e in zip(test_quantiles, eval_quantiles)] 

148 return json.dumps(qq_data)