add option to pass in absolute tolerance

modules/deeplearning/quantile_regression.py

@@ -38,8 +38,7 @@ def make_data(num_points=1000):
     return X_train, X_test, Y_train, Y_test, X, Y
 
 
-# Function to create a quantile regression model
-def build_quantile_model(q):
+def build_model(loss=tf.keras.losses.MeanSquaredError()):
     model = tf.keras.models.Sequential([
         tf.keras.layers.InputLayer(shape=(1,)),
         tf.keras.layers.Dense(64, activation='relu'),
@@ -47,42 +46,9 @@ def build_quantile_model(q):
         tf.keras.layers.Dense(64, activation='relu'),
         tf.keras.layers.Dense(1)  # Output layer
     ])
-    model.compile(optimizer='adam', loss=quantile_loss(q))
+    model.compile(optimizer='adam', loss=loss)
     return model
 
-def build_bulk_quantile_model():
-    model = tf.keras.models.Sequential([
-        tf.keras.layers.InputLayer(shape=(1,)),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(1)  # Output layer
-    ])
-    model.compile(optimizer='adam', loss=bulk_quantile_loss())
-    return model
-
-def build_mae_model():
-    model = tf.keras.models.Sequential([
-        tf.keras.layers.InputLayer(shape=(1,)),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(1)  # Output layer
-    ])
-    # model.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError())
-    model.compile(optimizer='adam', loss=tf.keras.losses.MeanAbsoluteError())
-    return model
-
-def build_mse_model():
-    model = tf.keras.models.Sequential([
-        tf.keras.layers.InputLayer(shape=(1,)),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(64, activation='relu'),
-        tf.keras.layers.Dense(1)  # Output layer
-    ])
-    model.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError())
-    return model
 
 def run(num_points=1000, num_plot_pts=200):
     # Define quantiles
@@ -95,24 +61,24 @@ def run(num_points=1000, num_plot_pts=200):
     # Train a model for each quantile
     for q in quantiles:
         print(f"Training quantile {q} model...")
-        models[q] = build_quantile_model(q)
+        models[q] = build_model(loss=quantile_loss(q))
         models[q].fit(X_train, Y_train, epochs=100, batch_size=32, verbose=0)
 
     # Generate test data predictions
     X_range = np.linspace(X.min(), X.max(), num_plot_pts).reshape(-1, 1)
     predictions = {q: models[q].predict(X_range) for q in quantiles}
 
-    model = build_mae_model()
+    model = build_model(loss=tf.keras.losses.MeanAbsoluteError())
     print(f"Training MAE model...")
     model.fit(X_train, Y_train, epochs=100, batch_size=32, verbose=0)
     mae_predictions = model.predict(X_range)
 
-    model = build_mse_model()
+    model = build_model()
     print(f"Training MSE model...")
     model.fit(X_train, Y_train, epochs=100, batch_size=32, verbose=0)
     mse_predictions = model.predict(X_range)
 
-    model = build_bulk_quantile_model()
+    model = build_model(loss=bulk_quantile_loss())
     print(f"Training bulk quantile model...")
     model.fit(X_train, Y_train, epochs=100, batch_size=32, verbose=0)
     bulk_predictions = model.predict(X_range)