add exponential distribution to negloglik and losses; add causal loss metric (#13)

Author: gmgeorg

pypsps/keras/callbacks.py

@@ -25,15 +25,19 @@ def on_epoch_end(self, epoch, logs=None):
             print(log_str)
 
 
-def recommended_callbacks(monitor="val_loss") -> List[tf.keras.callbacks.Callback]:
+def recommended_callbacks(
+    monitor="val_loss", patience: int = 50, mode="min"
+) -> List[tf.keras.callbacks.Callback]:
     """Return a list of recommended callbacks.
 
     This list is subject to change w/o notice. Do not rely on this in production.
     """
     callbacks = [
-        tf.keras.callbacks.EarlyStopping(monitor=monitor, patience=20, restore_best_weights=True),
-        tf.keras.callbacks.ReduceLROnPlateau(patience=10),
+        tf.keras.callbacks.EarlyStopping(
+            monitor=monitor, patience=patience, restore_best_weights=True, mode=mode
+        ),
+        tf.keras.callbacks.ReduceLROnPlateau(patience=patience // 3),
         tf.keras.callbacks.TerminateOnNaN(),
-        VerboseNEpochs(n=10),
+        VerboseNEpochs(n=20),
     ]
     return callbacks

pypsps/keras/metrics.py

@@ -4,6 +4,7 @@
 import tensorflow as tf
 
 from .. import utils
+from . import losses
 
 
 @tf.keras.utils.register_keras_serializable(package="pypsps")
@@ -96,3 +97,53 @@ def predictive_state_df(y_true, y_pred) -> tf.Tensor:
         return pypress.utils.tr_kernel(weights)
 
     return predictive_state_df
+
+
+def causal_loss_metric_gen(
+    outcome_loss: losses.OutcomeLoss,
+    treatment_loss: losses.TreatmentLoss,
+    alpha: float = 1.0,
+    outcome_loss_weight: float = 1.0,
+):
+    """
+    Function wrapper that returns a metric function computing the causal loss.
+
+    The returned function takes (y_true, y_pred) as inputs and computes:
+
+        causal_loss = outcome_loss_weight * outcome_loss(y_true, y_pred)
+                      + alpha * treatment_loss(y_true, y_pred)
+
+    This metric function can be passed to model.compile(metrics=[...]).
+
+    Parameters
+    ----------
+    outcome_loss : OutcomeLoss
+        Instance of an outcome loss (e.g. Normal log-likelihood loss).
+    treatment_loss : TreatmentLoss
+        Instance of a treatment loss (e.g. binary cross-entropy for treatment prediction).
+    alpha : float, default=1.0
+        Penalty parameter for treatment loss.
+    outcome_loss_weight : float, default=1.0
+        Weight for the outcome loss.
+
+    Returns
+    -------
+    function
+        A function metric that takes (y_true, y_pred) and returns the causal loss as a float value (can be passed as metric).
+    """
+    # Construct an instance of CausalLoss with the given parameters.
+    causal_loss_obj = losses.CausalLoss(
+        outcome_loss=outcome_loss,
+        treatment_loss=treatment_loss,
+        alpha=alpha,
+        outcome_loss_weight=outcome_loss_weight,
+    )
+
+    def causal_loss_metric(y_true, y_pred) -> tf.Tensor:
+        """Metric function computing the causal loss."""
+        # Call the causal loss object to compute the loss per example.
+        # Here we assume causal_loss_obj returns per-example loss.
+        return causal_loss_obj(y_true, y_pred)
+
+    causal_loss_metric.__name__ = "causal_loss"
+    return causal_loss_metric

pypsps/keras/neglogliks.py

@@ -71,3 +71,80 @@ def call(self, y_true, y_pred):
         ):
             return tf.reduce_mean(losses, axis=-1)
         raise NotImplementedError("reduction='%s' is not implemented", self.reduction)
+
+
+def _negloglik_exponential(
+    event_time: tf.Tensor, event_indicator: tf.Tensor, rate: tf.Tensor
+) -> tf.Tensor:
+    """
+    Computes the negative log-likelihood for an exponential distribution with censoring.
+
+    For each observation i:
+      - If an event occurs (event_indicator[i] == 1):
+            log-likelihood = log(rate[i]) - rate[i] * event_time[i]
+      - If censored (event_indicator[i] == 0):
+            log-likelihood = - rate[i] * event_time[i]
+
+    Therefore, the negative log-likelihood for observation i is:
+      loss_i = rate[i] * event_time[i] - event_indicator[i] * log(rate[i])
+
+    Parameters
+    ----------
+    event_time : tf.Tensor, shape (n,)
+        The observed event or censoring times.
+    event_indicator : tf.Tensor, shape (n,)
+        Binary indicator (1 if event occurred, 0 if censored).
+    rate : tf.Tensor, shape (n,)
+        The predicted rate (λ) of the exponential distribution.
+
+    Returns
+    -------
+    tf.Tensor
+        A tensor of shape (n,) containing the negative log-likelihood for each observation.
+    """
+    rate = tf.cast(rate, tf.float32)
+    log_rate = tf.math.log(rate + _EPS)
+    # Ensure inputs are float32
+    event_time = tf.cast(event_time, tf.float32)
+    event_indicator = tf.cast(event_indicator, tf.float32)
+
+    # Compute the negative log likelihood per observation
+    nll = rate * event_time - event_indicator * log_rate
+    return nll
+
+
+class NegloglikExponential(tf.keras.losses.Loss):
+    """Computes the negative log-likelihood of an Exponential survival model with censorship."""
+
+    def __init__(
+        self,
+        reduction=tf.keras.losses.Reduction.AUTO,
+        log_rate: bool = False,
+        name="negloglik_exponential",
+    ):
+        super().__init__(reduction=reduction, name=name)
+        self._log_rate = log_rate
+
+    def call(self, y_true, y_pred):
+        """Implements the loss function call."""
+        event_time = y_true[:, 0]
+        event_indicator = y_true[:, 1]
+
+        if self._log_rate:
+            y_pred = tf.exp(y_pred)
+
+        # y_pred is the rate
+        losses = _negloglik_exponential(
+            tf.squeeze(event_time), tf.squeeze(event_indicator), rate=tf.squeeze(y_pred)
+        )
+
+        if self.reduction == tf.keras.losses.Reduction.NONE:
+            return losses
+        if self.reduction == tf.keras.losses.Reduction.SUM:
+            return tf.reduce_sum(losses, axis=-1)
+        if self.reduction in (
+            tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE,
+            tf.keras.losses.Reduction.AUTO,
+        ):
+            return tf.reduce_mean(losses, axis=-1)
+        raise NotImplementedError(f"reduction='{self.reduction}' is not implemented")

pypsps/tests/test_neglogliks.py

@@ -4,13 +4,27 @@
 
 import numpy as np
 import pytest
+import tensorflow as tf
 import tensorflow_probability as tfp
 
 from ..keras import neglogliks
 
 tfd = tfp.distributions
 
 
+def _create_sample_data_exponential():
+    """
+    Creates a simple test case with two observations:
+      - First observation: event_time=10, event_indicator=1, log_hazard=log(0.1)
+      - Second observation: event_time=10, event_indicator=0, log_hazard=log(0.2)
+    """
+    # y_true has shape (n, 2): columns are event_time and event_indicator.
+    y_true = tf.constant([[10.0, 1.0], [10.0, 0.0]])
+    # y_pred has shape (n, 1): log_hazard predictions.
+    y_pred = tf.constant([[0.1], [0.2]])
+    return y_true, y_pred
+
+
 def _test_data() -> Tuple[np.ndarray, np.ndarray]:
     y_true = np.array([0.0, 1.0, 2.0])
     y_pred = np.array([[0.0, 1.0], [-1, 0.1], [0.1, 0.5]])
@@ -46,3 +60,98 @@ def test_negloglik_loss_class_works(reduction):
     print(loss_normal)
     print(loss_class_normal)
     assert loss_normal.numpy() == pytest.approx(loss_class_normal.numpy(), 0.0001)
+
+
+# --------------------------------------------------------------------
+# Tests for _negloglik_exponential function.
+# --------------------------------------------------------------------
+
+
+def test_negloglik_exponential_event():
+    """
+    Test when an event occurs (event_indicator == 1).
+
+    For an observation with:
+      event_time = 10,
+      event_indicator = 1,
+      log_hazard = log(0.1) (so rate = 0.1),
+    the loss should be: rate*event_time - log_hazard = 0.1*10 - log(0.1).
+    """
+    event_time = tf.constant([10.0])
+    event_indicator = tf.constant([1.0])
+    rate = tf.constant([0.1])
+
+    loss = neglogliks._negloglik_exponential(event_time, event_indicator, rate)
+    expected = 0.1 * 10 - np.log(0.1)
+    np.testing.assert_allclose(loss.numpy(), [expected], atol=1e-5)
+
+
+def test_negloglik_exponential_censored():
+    """
+    Test when an observation is censored (event_indicator == 0).
+
+    For an observation with:
+      event_time = 10,
+      event_indicator = 0,
+      log_hazard = log(0.1) (so rate = 0.1),
+    the loss should be: rate*event_time = 0.1*10.
+    """
+    event_time = tf.constant([10.0])
+    event_indicator = tf.constant([0.0])
+    rate = tf.constant([0.1])
+
+    loss = neglogliks._negloglik_exponential(event_time, event_indicator, rate)
+    expected = 0.1 * 10
+    np.testing.assert_allclose(loss.numpy(), [expected], atol=1e-5)
+
+
+def test_NegloglikExponential_none():
+    """
+    Test NegloglikExponential with reduction NONE.
+
+    Expected losses:
+      Observation 1: 0.1*10 - log(0.1)
+      Observation 2: 0.2*10
+    """
+    loss_obj = neglogliks.NegloglikExponential(reduction=tf.keras.losses.Reduction.NONE)
+    y_true, y_pred = _create_sample_data_exponential()
+    losses = loss_obj(y_true, y_pred)
+
+    expected1 = 0.1 * 10 - np.log(0.1)
+    expected2 = 0.2 * 10
+    expected = np.array([expected1, expected2])
+    np.testing.assert_allclose(losses.numpy(), expected, atol=1e-5)
+
+
+def test_NegloglikExponential_sum():
+    """
+    Test NegloglikExponential with reduction SUM.
+
+    Expected loss: sum over observations.
+    """
+    loss_obj = neglogliks.NegloglikExponential(reduction=tf.keras.losses.Reduction.SUM)
+    y_true, y_pred = _create_sample_data_exponential()
+    loss_value = loss_obj(y_true, y_pred)
+
+    expected1 = 0.1 * 10 - np.log(0.1)
+    expected2 = 0.2 * 10
+    expected = expected1 + expected2
+    np.testing.assert_allclose(loss_value.numpy(), expected, atol=1e-5)
+
+
+def test_NegloglikExponential_sum_over_batch_size():
+    """
+    Test NegloglikExponential with reduction SUM_OVER_BATCH_SIZE.
+
+    Expected loss: average loss over observations.
+    """
+    loss_obj = neglogliks.NegloglikExponential(
+        reduction=tf.keras.losses.Reduction.SUM_OVER_BATCH_SIZE
+    )
+    y_true, y_pred = _create_sample_data_exponential()
+    loss_value = loss_obj(y_true, y_pred)
+
+    expected1 = 0.1 * 10 - np.log(0.1)
+    expected2 = 0.2 * 10
+    expected = (expected1 + expected2) / 2.0
+    np.testing.assert_allclose(loss_value.numpy(), expected, atol=1e-5)