fix numerical issues for lookahead_levelset_at_xstar (#370)

Summary:
Pull Request resolved: #370

The `lookahead_levelset_at_xstar` in `lookahead_utils,py` has several divisions where the denominator could be 0.

For example, P1 = Z_qs / py1: py1=normal cdf(a star) could potentially be 0 when a_star is very negative number, causing division by 0 and P1 equals to nans.

Reviewed By: crasanders

Differential Revision: D59839473

fbshipit-source-id: c2c795449214c4fa1c5345f8cb25c02a9da47ba7

aepsych/acquisition/lookahead_utils.py

@@ -57,6 +57,7 @@ def lookahead_levelset_at_xstar(
     Xstar: Tensor,
     Xq: Tensor,
     posterior_transform: Optional[PosteriorTransform] = None,
+    eps: float = 1e-8,
     **kwargs: Dict[str, Any],
 ):
     """
@@ -86,18 +87,18 @@ def lookahead_levelset_at_xstar(
     # Compute look-ahead components
     Norm = torch.distributions.Normal(0, 1)
     Sigma_q = torch.sqrt(Sigma2_q)
-    b_q = (gamma - Mu_q) / Sigma_q
+    b_q = (gamma - Mu_q) / (Sigma_q + eps)
     Phi_bq = Norm.cdf(b_q)
     denom = torch.sqrt(1 + Sigma2_s)
     a_s = Mu_s / denom
     Phi_as = Norm.cdf(a_s)
-    Z_rho = -Sigma_sq / (Sigma_q * denom)
+    Z_rho = -Sigma_sq / (Sigma_q * denom + eps)
     Z_qs = bvn_cdf(a_s, b_q, Z_rho)
 
     Px = Phi_bq
     py1 = Phi_as
-    P1 = Z_qs / py1
-    P0 = (Phi_bq - Z_qs) / (1 - py1)
+    P1 = Z_qs / (py1 + eps)
+    P0 = (Phi_bq - Z_qs) / (1 - py1 + eps)
     return Px, P1, P0, py1