Coverage for src/edelweiss/custom_clfs.py: 100%
51 statements
« prev ^ index » next coverage.py v7.10.1, created at 2025-07-31 10:21 +0000
« prev ^ index » next coverage.py v7.10.1, created at 2025-07-31 10:21 +0000
1# Copyright (C) 2023 ETH Zurich
2# Institute for Particle Physics and Astrophysics
3# Author: Silvan Fischbacher
4# created: Mon Nov 06 2023
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
13from edelweiss.tf_utils import EpochProgressCallback
16class NeuralNetworkClassifier(BaseEstimator, ClassifierMixin):
17 """
18 Neural network classifier based on Keras Sequential model
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 """
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
58 def fit(self, X, y, sample_weight=None, early_stopping_patience=10):
59 """
60 Fit the neural network model
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 """
73 # Encode labels
74 self.label_encoder = LabelEncoder()
75 y_encoded = self.label_encoder.fit_transform(y)
76 self.classes_ = self.label_encoder.classes_
78 # Determine if it's binary or multiclass
79 self.n_classes_ = len(self.classes_)
80 self.is_binary_ = self.n_classes_ == 2
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
96 # Build the neural network model
97 self._build_model(X.shape[1])
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 )
106 # Add early stopping
107 early_stopping = EarlyStopping(
108 monitor="val_loss",
109 patience=early_stopping_patience,
110 restore_best_weights=True,
111 )
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 )
125 def _build_model(self, input_dim):
126 """
127 Build the neural network model
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()
149 def predict(self, X):
150 """
151 Predict the class labels for the provided data
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))
161 def predict_proba(self, X):
162 """
163 Predict the class probabilities for the provided data
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)
172 # for backwards compatibility
173 if not hasattr(self, "is_binary_"): # pragma: no cover
174 self.is_binary_ = True
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