Coverage for src/galsbi/ucat/galaxy_sampling_util.py: 100%

299 statements  

« prev     ^ index     » next       coverage.py v7.17.2, created at 2026-11-09 18:05 +0000

1# Copyright (C) 2018 ETH Zurich, Institute for Particle Physics and Astrophysics 

2 

3""" 

4Created on Mar 6, 2018 

5author: Joerg Herbel 

6""" 

7 

8import math 

9import operator 

10import warnings 

11 

12import numpy as np 

13import scipy.interpolate 

14import scipy.optimize 

15import scipy.special 

16 

17warnings.filterwarnings("once") 

18EPS = 10e-15 

19 

20 

21class UCatNumGalError(ValueError): 

22 """ 

23 Raised when more galaxies than allowed by the input parameters are sampled 

24 """ 

25 

26 

27class UCatMZInterpError(ValueError): 

28 """ 

29 Raised when lum_fct_m_max is too low for the given redshift range 

30 """ 

31 

32 

33class Catalog: 

34 pass 

35 

36 

37def apply_pycosmo_distfun(dist_fun, z): 

38 z_unique, inv_ind = np.unique(z, return_inverse=True) 

39 dist = dist_fun(a=1 / (1 + z_unique))[inv_ind] # Unit: Mpc 

40 return dist 

41 

42 

43def intp_z_m_cut(cosmo, mag_calc, par): 

44 """ 

45 This function returns an interpolator which predicts the maximum absolute magnitude 

46 for a given redshift such that a galaxy will still fall below the limiting apparent 

47 magnitude in the reference band (par.gals_mag_max). To do this, we do a check for a 

48 number of grid points in redshift. The check consists of finding the template with 

49 the smallest ratio of flux in the ref-band at that redshift (apparent flux) and flux 

50 in the luminosity function band at redshift zero (absolute flux). We then compute 

51 the absolute magnitude a galaxy would need to have to still pass the cut assuming 

52 that its spectrum is given by only that one template which we found. This works 

53 because the mixing-in of other templates will only make the object brighter in the 

54 r-band at this redshift. See also docs/jupyter_notebooks/z_m_cut.ipynb. 

55 """ 

56 

57 def find_max_template_ind(z, n_templates): 

58 coeffs = np.eye(n_templates) 

59 

60 m_lf = mag_calc( 

61 redshifts=np.zeros(n_templates), 

62 excess_b_v=np.zeros(n_templates), 

63 coeffs=coeffs, 

64 filter_names=[par.lum_fct_filter_band], 

65 )[par.lum_fct_filter_band] 

66 

67 m_ref = mag_calc( 

68 redshifts=np.full(n_templates, z), 

69 excess_b_v=np.zeros(n_templates), 

70 coeffs=coeffs, 

71 filter_names=[par.reference_band], 

72 )[par.reference_band] 

73 

74 ind = np.argmin(m_ref - m_lf) 

75 

76 return ind 

77 

78 def app_mag_ref(z, temp_i, m_abs, n_templates): 

79 coeff = np.zeros((1, n_templates)) 

80 coeff[0, temp_i] = 1 

81 coeff[0, temp_i] = 10 ** ( 

82 0.4 

83 * ( 

84 mag_calc( 

85 redshifts=np.zeros(1), 

86 excess_b_v=np.zeros(1), 

87 coeffs=coeff, 

88 filter_names=[par.lum_fct_filter_band], 

89 )[par.lum_fct_filter_band][0] 

90 - m_abs 

91 ) 

92 ) 

93 coeff[0, temp_i] *= (1e-5 / cosmo.background.dist_lum_a(a=1 / (1 + z))) ** 2 / ( 

94 1 + z 

95 ) 

96 m = mag_calc( 

97 redshifts=np.atleast_1d(z), 

98 excess_b_v=np.zeros(1), 

99 coeffs=coeff, 

100 filter_names=[par.reference_band], 

101 )[par.reference_band] 

102 

103 return m[0] 

104 

105 n_templates = mag_calc.n_templates 

106 z_intp = [par.lum_fct_z_res] 

107 temp_ind = find_max_template_ind(z_intp[0], n_templates) 

108 abs_mag = [ 

109 scipy.optimize.brentq( 

110 f=lambda m: app_mag_ref(z_intp[0], temp_ind, m, n_templates) 

111 - par.gals_mag_max, 

112 a=-100, 

113 b=par.gals_mag_max, 

114 ) 

115 ] 

116 

117 i_loop = 0 

118 

119 def magnitude_difference(m, z, template_ind): 

120 """Return the apparent-magnitude residual for root finding.""" 

121 return app_mag_ref(z, template_ind, m, n_templates) - par.gals_mag_max 

122 

123 while True: 

124 i_loop += 1 

125 try: 

126 z_ = z_intp[-1] + 0.02 

127 

128 if z_ > par.lum_fct_z_max: 

129 break 

130 

131 temp_ind = find_max_template_ind(z_, n_templates) 

132 abs_mag_ = scipy.optimize.brentq( 

133 f=magnitude_difference, 

134 a=-100, 

135 b=abs_mag[-1], 

136 args=(z_, temp_ind), 

137 ) 

138 except ValueError: 

139 try: 

140 abs_mag_ = scipy.optimize.brentq( 

141 f=magnitude_difference, 

142 a=-100, 

143 b=par.gals_mag_max, 

144 args=(z_, temp_ind), 

145 ) 

146 except ValueError: 

147 break 

148 z_intp.append(z_) 

149 abs_mag.append(abs_mag_) 

150 

151 intp = scipy.interpolate.interp1d( 

152 x=z_intp, 

153 y=abs_mag, 

154 kind="cubic", 

155 bounds_error=False, 

156 fill_value=(abs_mag[0], abs_mag[-1]), 

157 ) 

158 

159 if np.any(intp(intp.x) > par.lum_fct_m_max): 

160 msg = ( 

161 "par.lum_fct_m_max is too low according to z-m-interpolation," 

162 " some galaxies may be missing\n" 

163 f"gals_mag_max={par.gals_mag_max:2.2f}" 

164 f" lum_fct_z_max={par.lum_fct_z_max:2.2f}" 

165 f" m_max_sample={np.max(intp(intp.x)):2.2f}" 

166 f" lum_fct_m_max={par.lum_fct_m_max:2.2f}" 

167 ) 

168 if par.raise_z_m_interp_error: 

169 raise UCatMZInterpError(msg) 

170 else: 

171 warnings.warn(msg, stacklevel=1) 

172 

173 return intp 

174 

175 

176class TruncatedMVN: 

177 """Create a normal distribution :math:`X \sim N ({\mu}, {\Sigma})` subject to 

178 linear inequality constraints :math:`lb < X < ub` and sample from it using minimax 

179 tilting. Based on the MATLAB implemention by the authors (reference below). 

180 

181 :param np.ndarray mu: (size D) 

182 mean of the normal distribution :math:`\mathbf {\mu}`. 

183 :param np.ndarray cov: (size D x D) 

184 covariance of the normal distribution :math:`\mathbf {\Sigma}`. 

185 :param np.ndarray lb: (size D) 

186 lower bound constrain of the multivariate normal distribution :math:`\mathbf lb` 

187 :param np.ndarray ub: (size D) 

188 upper bound constrain of the multivariate normal distribution :math:`\mathbf ub` 

189 :param Union[int, None] seed: a random seed. 

190 

191 Note that the algorithm may not work if 'cov' is close to being rank deficient. 

192 

193 Reference: 

194 Botev, Z. I., (2016), The normal law under linear restrictions: simulation and 

195 estimation via minimax tilting, 

196 Journal of the Royal Statistical Society Series B, 79, issue 1, p. 125-148, 

197 

198 Reimplementation by Paul Brunzema 

199 """ 

200 

201 def __init__(self, mu, cov, lb, ub, seed=None): 

202 self.dim = len(mu) 

203 if not cov.shape[0] == cov.shape[1]: 

204 raise RuntimeError("Covariance matrix must be of shape DxD!") 

205 if not ( 

206 self.dim == cov.shape[0] and self.dim == len(lb) and self.dim == len(ub) 

207 ): 

208 raise RuntimeError( 

209 "Dimensions D of mean (mu), covariance matric (cov), lower bound (lb) " 

210 "and upper bound (ub) must be the same!" 

211 ) 

212 

213 # The sampling algorithm permutes these arrays internally. Store copies so 

214 # that configuration arrays supplied by the caller remain unchanged. 

215 self.cov = np.array(cov, copy=True) 

216 self.orig_mu = np.array(mu, copy=True) 

217 self.orig_lb = np.array(lb, copy=True) 

218 self.orig_ub = np.array(ub, copy=True) 

219 

220 # permutated 

221 self.lb = self.orig_lb - self.orig_mu # move distr./bounds to have zero mean 

222 self.ub = self.orig_ub - self.orig_mu # move distr./bounds to have zero mean 

223 if np.any(self.ub <= self.lb): 

224 raise RuntimeError( 

225 "Upper bound (ub) must be strictly greater than lower bound (lb) for " 

226 "all D dimensions!" 

227 ) 

228 

229 # scaled Cholesky with zero diagonal, permutated 

230 self.L = np.empty_like(cov) 

231 self.unscaled_L = np.empty_like(cov) 

232 

233 # placeholder for optimization 

234 self.perm = None 

235 self.x = None 

236 self.mu = None 

237 self.psistar = None 

238 

239 # for numerics 

240 self.eps = EPS 

241 

242 # a random state 

243 self.random_state = np.random.RandomState(seed) 

244 

245 def sample(self, n): 

246 """Create n samples from the truncated normal distribution. 

247 

248 :param int n: Number of samples to create. 

249 :return: D x n array with the samples. 

250 :rtype: np.ndarray 

251 """ 

252 # operator.index accepts Python and numpy integers but, like numpy's 

253 # size arguments, rejects floats such as 2.0 to avoid silent rounding. 

254 try: 

255 n = operator.index(n) 

256 except TypeError as err: 

257 raise RuntimeError("Number of samples must be an integer!") from err 

258 

259 # factors (Cholesky, etc.) only need to be computed once! 

260 if self.psistar is None: 

261 self.compute_factors() 

262 

263 # start acceptance rejection sampling 

264 rv = np.array([], dtype=np.float64).reshape(self.dim, 0) 

265 accept, iteration = 0, 0 

266 while accept < n: 

267 logpr, Z = self.mvnrnd(n, self.mu) # simulate n proposals 

268 idx = -np.log(self.random_state.rand(n)) > ( 

269 self.psistar - logpr 

270 ) # acceptance tests 

271 rv = np.concatenate((rv, Z[:, idx]), axis=1) # accumulate accepted 

272 accept = rv.shape[1] # keep track of # of accepted 

273 iteration += 1 

274 if iteration == 10**3: 

275 print("Warning: Acceptance prob. smaller than 0.001.") 

276 elif iteration > 10**4: 

277 accept = n 

278 rv = np.concatenate((rv, Z), axis=1) 

279 print("Warning: Sample is only approximately distributed.") 

280 

281 # finish sampling and postprocess the samples! 

282 order = self.perm.argsort(axis=0) 

283 rv = rv[:, :n] 

284 rv = self.unscaled_L @ rv 

285 rv = rv[order, :] 

286 

287 # retransfer to original mean 

288 rv += np.tile( 

289 self.orig_mu.reshape(self.dim, 1), (1, rv.shape[-1]) 

290 ) # Z = X + mu 

291 return rv 

292 

293 def compute_factors(self): 

294 # compute permutated Cholesky factor and solve optimization 

295 

296 # Cholesky decomposition of matrix with permuation 

297 self.unscaled_L, self.perm = self.colperm() 

298 D = np.diag(self.unscaled_L) 

299 if np.any(self.eps > D): 

300 print("Warning: Method might fail as covariance matrix is singular!") 

301 

302 # rescale 

303 scaled_L = self.unscaled_L / np.tile(D.reshape(self.dim, 1), (1, self.dim)) 

304 self.lb = self.lb / D 

305 self.ub = self.ub / D 

306 

307 # remove diagonal 

308 self.L = scaled_L - np.eye(self.dim) 

309 

310 # get gradient/Jacobian function 

311 gradpsi = self.get_gradient_function() 

312 x0 = np.zeros(2 * (self.dim - 1)) 

313 

314 # find optimal tilting parameter non-linear equation solver 

315 sol = scipy.optimize.root( 

316 gradpsi, x0, args=(self.L, self.lb, self.ub), method="hybr", jac=True 

317 ) 

318 if not sol.success: 

319 print("Warning: Method may fail as covariance matrix is close to singular!") 

320 self.x = sol.x[: self.dim - 1] 

321 self.mu = sol.x[self.dim - 1 :] 

322 

323 # compute psi star 

324 self.psistar = self.psy(self.x, self.mu) 

325 

326 def mvnrnd(self, n, mu): 

327 # generates the proposals from the exponentially tilted sequential importance 

328 # sampling pdf 

329 # output: logpr, log-likelihood of sample 

330 # Z, random sample 

331 mu = np.append(mu, [0.0]) 

332 Z = np.zeros((self.dim, n)) 

333 logpr = 0 

334 for k in range(self.dim): 

335 # compute matrix multiplication L @ Z 

336 col = self.L[k, :k] @ Z[:k, :] 

337 # compute limits of truncation 

338 tl = self.lb[k] - mu[k] - col 

339 tu = self.ub[k] - mu[k] - col 

340 # simulate N(mu,1) conditional on [tl,tu] 

341 Z[k, :] = mu[k] + self.trandn(tl, tu) 

342 # update likelihood ratio 

343 logpr += lnNormalProb(tl, tu) + 0.5 * mu[k] ** 2 - mu[k] * Z[k, :] 

344 return logpr, Z 

345 

346 def trandn(self, lower_bounds, upper_bounds): 

347 """Sample generator for the truncated standard multivariate normal 

348 distribution :math:`X \sim N(0,I)` s.t. :math:`lb<X<ub`. 

349 

350 If you wish to simulate a random variable 'Z' from the non-standard 

351 Gaussian :math:`N(m,s^2)` conditional on :math:`lb<Z<ub`, then first 

352 simulate x=TruncatedMVNSampler.trandn((l-m)/s,(u-m)/s) and set Z=m+s*x. 

353 Infinite values for 'ub' and 'lb' are accepted. 

354 

355 :param np.ndarray lower_bounds: (size D) lower bound constrain of the 

356 normal distribution :math:`\mathbf lb`. 

357 :param np.ndarray upper_bounds: (size D) upper bound constrain of the 

358 normal distribution :math:`\mathbf lb`. 

359 :return: D samples if the truncated normal distribition x ~ N(0, I) 

360 subject to lb < x < ub. 

361 :rtype: np.ndarray 

362 """ 

363 if not len(lower_bounds) == len(upper_bounds): 

364 raise RuntimeError( 

365 "Lower bound (lb) and upper bound (ub) must be of the same " "length!" 

366 ) 

367 

368 x = np.empty_like(lower_bounds) 

369 a = 0.66 # threshold used in MATLAB implementation 

370 # three cases to consider 

371 # case 1: a<lb<ub 

372 case1_mask = lower_bounds > a 

373 if np.any(case1_mask): 

374 tl = lower_bounds[case1_mask] 

375 tu = upper_bounds[case1_mask] 

376 x[case1_mask] = self.ntail(tl, tu) 

377 # case 2: lb<ub<-a 

378 case2_mask = upper_bounds < -a 

379 if np.any(case2_mask): 

380 tl = -upper_bounds[case2_mask] 

381 tu = -lower_bounds[case2_mask] 

382 x[case2_mask] = -self.ntail(tl, tu) 

383 # case 3: otherwise use inverse transform or accept-reject 

384 case3_mask = ~(case1_mask | case2_mask) 

385 if np.any(case3_mask): 

386 tl = lower_bounds[case3_mask] 

387 tu = upper_bounds[case3_mask] 

388 x[case3_mask] = self.tn(tl, tu) 

389 return x 

390 

391 def tn(self, lower_bounds, upper_bounds, tol=2): 

392 # samples a column vector of length=len(lb)=len(ub) from the standard 

393 # multivariate normal distribution truncated over the region [lb,ub], 

394 # where -a<lb<ub<a for some 'a' and lb and ub are column vectors 

395 # uses acceptance rejection and inverse-transform method 

396 

397 # controls switch between methods, 

398 # threshold can be tuned for maximum speed for each platform 

399 sw = tol 

400 x = np.empty_like(lower_bounds) 

401 # case 1: abs(ub-lb)>tol, uses accept-reject from randn 

402 large_interval_mask = abs(upper_bounds - lower_bounds) > sw 

403 if np.any(large_interval_mask): 

404 tl = lower_bounds[large_interval_mask] 

405 tu = upper_bounds[large_interval_mask] 

406 x[large_interval_mask] = self.trnd(tl, tu) 

407 

408 # case 2: abs(u-l)<tol, uses inverse-transform 

409 small_interval_mask = ~large_interval_mask 

410 if np.any(small_interval_mask): 

411 tl = lower_bounds[small_interval_mask] 

412 tu = upper_bounds[small_interval_mask] 

413 pl = scipy.special.erfc(tl / np.sqrt(2)) / 2 

414 pu = scipy.special.erfc(tu / np.sqrt(2)) / 2 

415 x[small_interval_mask] = np.sqrt(2) * scipy.special.erfcinv( 

416 2 * (pl - (pl - pu) * self.random_state.rand(len(tl))) 

417 ) 

418 return x 

419 

420 def trnd(self, lower_bounds, upper_bounds): 

421 # uses acceptance rejection to simulate from truncated normal 

422 x = self.random_state.randn(len(lower_bounds)) # sample normal 

423 test = (x < lower_bounds) | (x > upper_bounds) 

424 rejected_indices = np.where(test)[0] 

425 num_rejected = len(rejected_indices) 

426 while num_rejected > 0: # while there are rejections 

427 ly = lower_bounds[rejected_indices] 

428 uy = upper_bounds[rejected_indices] 

429 y = self.random_state.randn(len(uy)) # resample 

430 accepted_mask = (y > ly) & (y < uy) # accepted 

431 x[rejected_indices[accepted_mask]] = y[accepted_mask] 

432 rejected_indices = rejected_indices[~accepted_mask] 

433 num_rejected = len(rejected_indices) 

434 return x 

435 

436 def ntail(self, lower_bounds, upper_bounds): 

437 # samples a column vector of length=len(lb)=len(ub) from the standard 

438 # multivariate normal distribution truncated over the region [lb,ub], 

439 # where lb>0 and lb and ub are column vectors 

440 # uses acceptance-rejection from Rayleigh distr. similar to Marsaglia (1964) 

441 if not len(lower_bounds) == len(upper_bounds): 

442 raise RuntimeError( 

443 "Lower bound (lb) and upper bound (ub) must be of the same " "length!" 

444 ) 

445 c = (lower_bounds**2) / 2 

446 n = len(lower_bounds) 

447 f = np.expm1(c - upper_bounds**2 / 2) 

448 x = c - np.log(1 + self.random_state.rand(n) * f) # sample using Rayleigh 

449 # keep list of rejected 

450 rejected_indices = np.where(self.random_state.rand(n) ** 2 * x > c)[0] 

451 num_rejected = len(rejected_indices) 

452 while num_rejected > 0: # while there are rejections 

453 cy = c[rejected_indices] 

454 y = cy - np.log( 

455 1 + self.random_state.rand(num_rejected) * f[rejected_indices] 

456 ) 

457 accepted_mask = ( 

458 self.random_state.rand(num_rejected) ** 2 * y 

459 ) < cy # accepted 

460 x[rejected_indices[accepted_mask]] = y[accepted_mask] # store the accepted 

461 rejected_indices = rejected_indices[ 

462 ~accepted_mask 

463 ] # remove accepted from the list 

464 num_rejected = len(rejected_indices) 

465 return np.sqrt(2 * x) # this Rayleigh transform can be delayed till the end 

466 

467 def psy(self, x, mu): 

468 # implements psi(x,mu); assumes scaled 'L' without diagonal 

469 x = np.append(x, [0.0]) 

470 mu = np.append(mu, [0.0]) 

471 c = self.L @ x 

472 lt = self.lb - mu - c 

473 ut = self.ub - mu - c 

474 p = np.sum(lnNormalProb(lt, ut) + 0.5 * mu**2 - x * mu) 

475 return p 

476 

477 def get_gradient_function(self): 

478 # wrapper to avoid dependancy on self 

479 

480 def gradpsi(y, chol_matrix, lower_bounds, upper_bounds): 

481 # implements gradient of psi(x) to find optimal exponential twisting, 

482 # returns also the Jacobian 

483 # NOTE: assumes scaled 'L' with zero diagonal 

484 d = len(upper_bounds) 

485 c = np.zeros(d) 

486 mu, x = c.copy(), c.copy() 

487 x[0 : d - 1] = y[0 : d - 1] 

488 mu[0 : d - 1] = y[d - 1 :] 

489 

490 # compute now ~l and ~u 

491 c[1:d] = chol_matrix[1:d, :] @ x 

492 lt = lower_bounds - mu - c 

493 ut = upper_bounds - mu - c 

494 

495 # compute gradients avoiding catastrophic cancellation 

496 w = lnNormalProb(lt, ut) 

497 pl = np.exp(-0.5 * lt**2 - w) / np.sqrt(2 * math.pi) 

498 pu = np.exp(-0.5 * ut**2 - w) / np.sqrt(2 * math.pi) 

499 P = pl - pu 

500 

501 # output the gradient 

502 dfdx = -mu[0 : d - 1] + (P.T @ chol_matrix[:, 0 : d - 1]).T 

503 dfdm = mu - x + P 

504 grad = np.concatenate((dfdx, dfdm[:-1]), axis=0) 

505 

506 # construct jacobian 

507 lt[np.isinf(lt)] = 0 

508 ut[np.isinf(ut)] = 0 

509 

510 dP = -(P**2) + lt * pl - ut * pu 

511 DL = np.tile(dP.reshape(d, 1), (1, d)) * chol_matrix 

512 mx = DL - np.eye(d) 

513 xx = chol_matrix.T @ DL 

514 mx = mx[:-1, :-1] 

515 xx = xx[:-1, :-1] 

516 jacobian = np.block([[xx, mx.T], [mx, np.diag(1 + dP[:-1])]]) 

517 return (grad, jacobian) 

518 

519 return gradpsi 

520 

521 def colperm(self): 

522 perm = np.arange(self.dim) 

523 chol_matrix = np.zeros_like(self.cov) 

524 z = np.zeros_like(self.orig_mu) 

525 

526 for j in perm.copy(): 

527 pr = np.ones_like(z) * np.inf # compute marginal prob. 

528 remaining_dims = np.arange(j, self.dim) # search remaining dimensions 

529 D = np.diag(self.cov) 

530 s = D[remaining_dims] - np.sum( 

531 chol_matrix[remaining_dims, 0:j] ** 2, axis=1 

532 ) 

533 s[s < 0] = self.eps 

534 s = np.sqrt(s) 

535 tl = ( 

536 self.lb[remaining_dims] - chol_matrix[remaining_dims, 0:j] @ z[0:j] 

537 ) / s 

538 tu = ( 

539 self.ub[remaining_dims] - chol_matrix[remaining_dims, 0:j] @ z[0:j] 

540 ) / s 

541 pr[remaining_dims] = lnNormalProb(tl, tu) 

542 # find smallest marginal dimension 

543 k = np.argmin(pr) 

544 

545 # flip dimensions k-->j 

546 jk = [j, k] 

547 kj = [k, j] 

548 self.cov[jk, :] = self.cov[kj, :] # update rows of cov 

549 self.cov[:, jk] = self.cov[:, kj] # update cols of cov 

550 chol_matrix[jk, :] = chol_matrix[kj, :] # update only rows of L 

551 self.lb[jk] = self.lb[kj] # update integration limits 

552 self.ub[jk] = self.ub[kj] # update integration limits 

553 perm[jk] = perm[kj] # keep track of permutation 

554 

555 # construct L sequentially via Cholesky computation 

556 s = self.cov[j, j] - np.sum(chol_matrix[j, 0:j] ** 2, axis=0) 

557 if s < -0.01: 

558 raise RuntimeError("Sigma is not positive semi-definite") 

559 elif s < 0: 

560 s = self.eps 

561 chol_matrix[j, j] = np.sqrt(s) 

562 new_chol = ( 

563 self.cov[j + 1 : self.dim, j] 

564 - chol_matrix[j + 1 : self.dim, 0:j] @ chol_matrix[j, 0:j].T 

565 ) 

566 chol_matrix[j + 1 : self.dim, j] = new_chol / chol_matrix[j, j] 

567 

568 # find mean value, z(j), of truncated normal 

569 tl = (self.lb[j] - chol_matrix[j, 0 : j - 1] @ z[0 : j - 1]) / chol_matrix[ 

570 j, j 

571 ] 

572 tu = (self.ub[j] - chol_matrix[j, 0 : j - 1] @ z[0 : j - 1]) / chol_matrix[ 

573 j, j 

574 ] 

575 w = lnNormalProb( 

576 tl, tu 

577 ) # aids in computing expected value of trunc. normal 

578 z[j] = (np.exp(-0.5 * tl**2 - w) - np.exp(-0.5 * tu**2 - w)) / np.sqrt( 

579 2 * math.pi 

580 ) 

581 return chol_matrix, perm 

582 

583 

584def lnNormalProb(a, b): 

585 # computes ln(P(a<Z<b)) where Z~N(0,1) very accurately for any 'a', 'b' 

586 p = np.zeros_like(a) 

587 # case b>a>0 

588 positive_mask = a > 0 

589 if np.any(positive_mask): 

590 pa = lnPhi(a[positive_mask]) 

591 pb = lnPhi(b[positive_mask]) 

592 p[positive_mask] = pa + np.log1p(-np.exp(pb - pa)) 

593 # case a<b<0 

594 negative_mask = b < 0 

595 if np.any(negative_mask): 

596 pa = lnPhi(-a[negative_mask]) # log of lower tail 

597 pb = lnPhi(-b[negative_mask]) 

598 p[negative_mask] = pb + np.log1p(-np.exp(pa - pb)) 

599 # case a < 0 < b 

600 mixed_mask = (~positive_mask) & (~negative_mask) 

601 if np.any(mixed_mask): 

602 pa = scipy.special.erfc(-a[mixed_mask] / np.sqrt(2)) / 2 # lower tail 

603 pb = scipy.special.erfc(b[mixed_mask] / np.sqrt(2)) / 2 # upper tail 

604 p[mixed_mask] = np.log1p(-pa - pb) 

605 return p 

606 

607 

608def lnPhi(x): 

609 # computes logarithm of tail of Z~N(0,1) mitigating numerical roundoff errors 

610 out = ( 

611 -0.5 * x**2 - np.log(2) + np.log(scipy.special.erfcx(x / np.sqrt(2)) + EPS) 

612 ) # divide by zeros error -> add eps 

613 return out