Add description for mean of censored

MeteoSwiss · Sep 6, 2024 · b6a1e48 · b6a1e48
1 parent 949aa69
commit b6a1e48
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/mlpp_lib/probabilistic_layers.py b/mlpp_lib/probabilistic_layers.py
@@ -362,6 +362,15 @@ def _sample_n(self, n, seed=None):
                     return chosen_samples
 
                 def _mean(self):
+                    """
+                    Original: X ~ N(mu, sigma)
+                    Censored: Y = X if 0 <= X <= 1 else 0 if X < 0 else 1
+
+                    Law of total expectations: E[Y] = E[Y | X > 1] * P(X > 1) + E[Y | X < 0] * P(X < 0) + E[Y | 0 <= X <= 1] * P(0 <= X <= 1)
+                                                    = P(X > 1) * 1 + P(X < 0) * 0 + E[X | 0 <= X <= 1] * P(0 <= X <= 1)
+                                                    = 1 - Phi((1 - mu) / sigma) + E[Z ~ TruncNormal(mu, sigma, 0, 1)] * (Phi((1 - mu) / sigma) - Phi(-mu / sigma))
+                    Ref for TruncatedNormal mean: https://en.wikipedia.org/wiki/Truncated_normal_distribution
+                    """
                     original_mean = self.normal.mean()
                     low_bound_standard = (0 - original_mean) / self.normal.stddev()
                     high_bound_standard = (1 - original_mean) / self.normal.stddev()