Coverage for src/edelweiss/custom_clfs.py: 100%

1# Copyright (C) 2023 ETH Zurich 

2# Institute for Particle Physics and Astrophysics 

3# Author: Silvan Fischbacher 

4# created: Mon Nov 06 2023 

5 

6import numpy as np 

7import tensorflow as tf 

8from sklearn.base import BaseEstimator, ClassifierMixin 

9from sklearn.preprocessing import LabelEncoder 

10from tensorflow.keras.callbacks import EarlyStopping 

11from tensorflow.keras.layers import Dropout 

12 

13from edelweiss.tf_utils import EpochProgressCallback 

14 

15 

16class NeuralNetworkClassifier(BaseEstimator, ClassifierMixin): 

17 """ 

18 Neural network classifier based on Keras Sequential model 

19 

20 :param hidden_units: tuple/list, optional (default=(64, 32)) 

21 The number of units per hidden layer 

22 :param learning_rate: float, optional (default=0.001) 

23 The learning rate for the Adam optimizer 

24 :param epochs: int, optional (default=10) 

25 The number of epochs to train the model 

26 :param batch_size: int, optional (default=32) 

27 The batch size for training the model 

28 :param loss: str, optional (default="auto") 

29 The loss function to use, defaults to binary_crossentropy if binary and 

30 sparse_categorical_crossentropy if multiclass 

31 :param activation: str, optional (default="relu") 

32 The activation function to use for the hidden layers 

33 :param activation_output: str, optional (default="auto") 

34 The activation function to use for the output layer, defaults to sigmoid for 

35 single class and softmax for multiclass 

36 :param sample_weight_col: int, optional (default=None) 

37 """ 

38 

39 def __init__( 

40 self, 

41 hidden_units=(64, 32), 

42 learning_rate=0.001, 

43 epochs=10, 

44 batch_size=32, 

45 loss="auto", 

46 activation="relu", 

47 activation_output="auto", 

48 ): 

49 self.hidden_units = hidden_units 

50 self.learning_rate = learning_rate 

51 self.epochs = epochs 

52 self.batch_size = batch_size 

53 self.loss = loss 

54 self.activation = activation 

55 self.activation_output = activation_output 

56 self.model = None 

57 

58 def fit(self, X, y, sample_weight=None, early_stopping_patience=10): 

59 """ 

60 Fit the neural network model 

61 

62 :param X: array-like, shape (n_samples, n_features) 

63 The training input samples 

64 :param y: array-like, shape (n_samples,) 

65 The target values 

66 :param sample_weight: array-like, shape (n_samples,), optional (default=None) 

67 Sample weights 

68 :param early_stopping_patience: int, optional (default=10) 

69 The number of epochs with no improvement after which training will be 

70 stopped 

71 """ 

72 

73 # Encode labels 

74 self.label_encoder = LabelEncoder() 

75 y_encoded = self.label_encoder.fit_transform(y) 

76 self.classes_ = self.label_encoder.classes_ 

77 

78 # Determine if it's binary or multiclass 

79 self.n_classes_ = len(self.classes_) 

80 self.is_binary_ = self.n_classes_ == 2 

81 

82 # Adjust loss and activation_output based on problem type 

83 if self.loss == "auto": 

84 self.loss_ = ( 

85 "binary_crossentropy" 

86 if self.is_binary_ 

87 else "sparse_categorical_crossentropy" 

88 ) 

89 else: 

90 self.loss_ = self.loss 

91 if self.activation_output == "auto": 

92 self.activation_output_ = "sigmoid" if self.is_binary_ else "softmax" 

93 else: 

94 self.activation_output_ = self.activation_output 

95 

96 # Build the neural network model 

97 self._build_model(X.shape[1]) 

98 

99 # Compile the model 

100 self.model.compile( 

101 optimizer=tf.keras.optimizers.Adam(learning_rate=self.learning_rate), 

102 loss=self.loss_, 

103 metrics=["accuracy"], 

104 ) 

105 

106 # Add early stopping 

107 early_stopping = EarlyStopping( 

108 monitor="val_loss", 

109 patience=early_stopping_patience, 

110 restore_best_weights=True, 

111 ) 

112 

113 # Fit the model 

114 self.model.fit( 

115 X, 

116 y_encoded, 

117 sample_weight=sample_weight, 

118 epochs=self.epochs, 

119 batch_size=self.batch_size, 

120 validation_split=0.2, # use 20% of the training data as validation data 

121 callbacks=[early_stopping, EpochProgressCallback(total_epochs=self.epochs)], 

122 verbose=0, 

123 ) 

124 

125 def _build_model(self, input_dim): 

126 """ 

127 Build the neural network model 

128 

129 :param input_dim: int 

130 The number of input features 

131 """ 

132 self.model = tf.keras.Sequential() 

133 try: 

134 self.model.add(tf.keras.layers.InputLayer(shape=(input_dim,))) 

135 except Exception: # pragma: no cover 

136 # backwards compatibility for tf<2.16 

137 self.model.add(tf.keras.layers.InputLayer(input_shape=(input_dim,))) 

138 for units in self.hidden_units: 

139 self.model.add(tf.keras.layers.Dense(units, activation=self.activation)) 

140 self.model.add(Dropout(0.2)) 

141 self.model.add( 

142 tf.keras.layers.Dense( 

143 1 if self.is_binary_ else self.n_classes_, 

144 activation=self.activation_output_, 

145 ) 

146 ) 

147 self.model.summary() 

148 

149 def predict(self, X): 

150 """ 

151 Predict the class labels for the provided data 

152 

153 :param X: array-like, shape (n_samples, n_features) 

154 The input samples 

155 :return: array-like, shape (n_samples,) 

156 The predicted class labels 

157 """ 

158 y_pred = self.model.predict(X, verbose=0) 

159 return self.label_encoder.inverse_transform(np.argmax(y_pred, axis=1)) 

160 

161 def predict_proba(self, X): 

162 """ 

163 Predict the class probabilities for the provided data 

164 

165 :param X: array-like, shape (n_samples, n_features) 

166 The input samples 

167 :return: array-like, shape (n_samples, n_classes) 

168 The predicted class probabilities 

169 """ 

170 y_prob = self.model.predict(X, verbose=0) 

171 

172 # for backwards compatibility 

173 if not hasattr(self, "is_binary_"): # pragma: no cover 

174 self.is_binary_ = True 

175 

176 if self.is_binary_: 

177 y_prob = y_prob.flatten() 

178 return np.column_stack((1 - y_prob, y_prob)) 

179 else: 

180 return y_prob