Coverage for /usr/lib/python3/dist-packages/scipy/stats/_sensitivity_analysis.py: 28%

128 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-14 15:55 +0200

1from __future__ import annotations 

2 

3import inspect 

4from dataclasses import dataclass 

5from typing import ( 

6 Callable, Literal, Protocol, TYPE_CHECKING 

7) 

8 

9import numpy as np 

10 

11from scipy.stats._common import ConfidenceInterval 

12from scipy.stats._qmc import check_random_state 

13from scipy.stats._resampling import BootstrapResult 

14from scipy.stats import qmc, bootstrap 

15 

16 

17if TYPE_CHECKING: 

18 import numpy.typing as npt 

19 from scipy._lib._util import DecimalNumber, IntNumber, SeedType 

20 

21 

22__all__ = [ 

23 'sobol_indices' 

24] 

25 

26 

27def f_ishigami(x: npt.ArrayLike) -> np.ndarray: 

28 r"""Ishigami function. 

29 

30 .. math:: 

31 

32 Y(\mathbf{x}) = \sin x_1 + 7 \sin^2 x_2 + 0.1 x_3^4 \sin x_1 

33 

34 with :math:`\mathbf{x} \in [-\pi, \pi]^3`. 

35 

36 Parameters 

37 ---------- 

38 x : array_like ([x1, x2, x3], n) 

39 

40 Returns 

41 ------- 

42 f : array_like (n,) 

43 Function evaluation. 

44 

45 References 

46 ---------- 

47 .. [1] Ishigami, T. and T. Homma. "An importance quantification technique 

48 in uncertainty analysis for computer models." IEEE, 

49 :doi:`10.1109/ISUMA.1990.151285`, 1990. 

50 """ 

51 x = np.atleast_2d(x) 

52 f_eval = ( 

53 np.sin(x[0]) 

54 + 7 * np.sin(x[1])**2 

55 + 0.1 * (x[2]**4) * np.sin(x[0]) 

56 ) 

57 return f_eval 

58 

59 

60def sample_A_B( 

61 n: IntNumber, 

62 dists: list[PPFDist], 

63 random_state: SeedType = None 

64) -> np.ndarray: 

65 """Sample two matrices A and B. 

66 

67 Uses a Sobol' sequence with 2`d` columns to have 2 uncorrelated matrices. 

68 This is more efficient than using 2 random draw of Sobol'. 

69 See sec. 5 from [1]_. 

70 

71 Output shape is (d, n). 

72 

73 References 

74 ---------- 

75 .. [1] Saltelli, A., P. Annoni, I. Azzini, F. Campolongo, M. Ratto, and 

76 S. Tarantola. "Variance based sensitivity analysis of model 

77 output. Design and estimator for the total sensitivity index." 

78 Computer Physics Communications, 181(2):259-270, 

79 :doi:`10.1016/j.cpc.2009.09.018`, 2010. 

80 """ 

81 d = len(dists) 

82 A_B = qmc.Sobol(d=2*d, seed=random_state, bits=64).random(n).T 

83 A_B = A_B.reshape(2, d, -1) 

84 try: 

85 for d_, dist in enumerate(dists): 

86 A_B[:, d_] = dist.ppf(A_B[:, d_]) 

87 except AttributeError as exc: 

88 message = "Each distribution in `dists` must have method `ppf`." 

89 raise ValueError(message) from exc 

90 return A_B 

91 

92 

93def sample_AB(A: np.ndarray, B: np.ndarray) -> np.ndarray: 

94 """AB matrix. 

95 

96 AB: rows of B into A. Shape (d, d, n). 

97 - Copy A into d "pages" 

98 - In the first page, replace 1st rows of A with 1st row of B. 

99 ... 

100 - In the dth page, replace dth row of A with dth row of B. 

101 - return the stack of pages 

102 """ 

103 d, n = A.shape 

104 AB = np.tile(A, (d, 1, 1)) 

105 i = np.arange(d) 

106 AB[i, i] = B[i] 

107 return AB 

108 

109 

110def saltelli_2010( 

111 f_A: np.ndarray, f_B: np.ndarray, f_AB: np.ndarray 

112) -> tuple[np.ndarray, np.ndarray]: 

113 r"""Saltelli2010 formulation. 

114 

115 .. math:: 

116 

117 S_i = \frac{1}{N} \sum_{j=1}^N 

118 f(\mathbf{B})_j (f(\mathbf{AB}^{(i)})_j - f(\mathbf{A})_j) 

119 

120 .. math:: 

121 

122 S_{T_i} = \frac{1}{N} \sum_{j=1}^N 

123 (f(\mathbf{A})_j - f(\mathbf{AB}^{(i)})_j)^2 

124 

125 Parameters 

126 ---------- 

127 f_A, f_B : array_like (s, n) 

128 Function values at A and B, respectively 

129 f_AB : array_like (d, s, n) 

130 Function values at each of the AB pages 

131 

132 Returns 

133 ------- 

134 s, st : array_like (s, d) 

135 First order and total order Sobol' indices. 

136 

137 References 

138 ---------- 

139 .. [1] Saltelli, A., P. Annoni, I. Azzini, F. Campolongo, M. Ratto, and 

140 S. Tarantola. "Variance based sensitivity analysis of model 

141 output. Design and estimator for the total sensitivity index." 

142 Computer Physics Communications, 181(2):259-270, 

143 :doi:`10.1016/j.cpc.2009.09.018`, 2010. 

144 """ 

145 # Empirical variance calculated using output from A and B which are 

146 # independent. Output of AB is not independent and cannot be used 

147 var = np.var([f_A, f_B], axis=(0, -1)) 

148 

149 # We divide by the variance to have a ratio of variance 

150 # this leads to eq. 2 

151 s = np.mean(f_B * (f_AB - f_A), axis=-1) / var # Table 2 (b) 

152 st = 0.5 * np.mean((f_A - f_AB) ** 2, axis=-1) / var # Table 2 (f) 

153 

154 return s.T, st.T 

155 

156 

157@dataclass 

158class BootstrapSobolResult: 

159 first_order: BootstrapResult 

160 total_order: BootstrapResult 

161 

162 

163@dataclass 

164class SobolResult: 

165 first_order: np.ndarray 

166 total_order: np.ndarray 

167 _indices_method: Callable 

168 _f_A: np.ndarray 

169 _f_B: np.ndarray 

170 _f_AB: np.ndarray 

171 _A: np.ndarray | None = None 

172 _B: np.ndarray | None = None 

173 _AB: np.ndarray | None = None 

174 _bootstrap_result: BootstrapResult | None = None 

175 

176 def bootstrap( 

177 self, 

178 confidence_level: DecimalNumber = 0.95, 

179 n_resamples: IntNumber = 999 

180 ) -> BootstrapSobolResult: 

181 """Bootstrap Sobol' indices to provide confidence intervals. 

182 

183 Parameters 

184 ---------- 

185 confidence_level : float, default: ``0.95`` 

186 The confidence level of the confidence intervals. 

187 n_resamples : int, default: ``999`` 

188 The number of resamples performed to form the bootstrap 

189 distribution of the indices. 

190 

191 Returns 

192 ------- 

193 res : BootstrapSobolResult 

194 Bootstrap result containing the confidence intervals and the 

195 bootstrap distribution of the indices. 

196 

197 An object with attributes: 

198 

199 first_order : BootstrapResult 

200 Bootstrap result of the first order indices. 

201 total_order : BootstrapResult 

202 Bootstrap result of the total order indices. 

203 See `BootstrapResult` for more details. 

204 

205 """ 

206 def statistic(idx): 

207 f_A_ = self._f_A[:, idx] 

208 f_B_ = self._f_B[:, idx] 

209 f_AB_ = self._f_AB[..., idx] 

210 return self._indices_method(f_A_, f_B_, f_AB_) 

211 

212 n = self._f_A.shape[1] 

213 

214 res = bootstrap( 

215 [np.arange(n)], statistic=statistic, method="BCa", 

216 n_resamples=n_resamples, 

217 confidence_level=confidence_level, 

218 bootstrap_result=self._bootstrap_result 

219 ) 

220 self._bootstrap_result = res 

221 

222 first_order = BootstrapResult( 

223 confidence_interval=ConfidenceInterval( 

224 res.confidence_interval.low[0], res.confidence_interval.high[0] 

225 ), 

226 bootstrap_distribution=res.bootstrap_distribution[0], 

227 standard_error=res.standard_error[0], 

228 ) 

229 total_order = BootstrapResult( 

230 confidence_interval=ConfidenceInterval( 

231 res.confidence_interval.low[1], res.confidence_interval.high[1] 

232 ), 

233 bootstrap_distribution=res.bootstrap_distribution[1], 

234 standard_error=res.standard_error[1], 

235 ) 

236 

237 return BootstrapSobolResult( 

238 first_order=first_order, total_order=total_order 

239 ) 

240 

241 

242class PPFDist(Protocol): 

243 @property 

244 def ppf(self) -> Callable[..., float]: 

245 ... 

246 

247 

248def sobol_indices( 

249 *, 

250 func: Callable[[np.ndarray], npt.ArrayLike] | 

251 dict[Literal['f_A', 'f_B', 'f_AB'], np.ndarray], # noqa 

252 n: IntNumber, 

253 dists: list[PPFDist] | None = None, 

254 method: Callable | Literal['saltelli_2010'] = 'saltelli_2010', 

255 random_state: SeedType = None 

256) -> SobolResult: 

257 r"""Global sensitivity indices of Sobol'. 

258 

259 Parameters 

260 ---------- 

261 func : callable or dict(str, array_like) 

262 If `func` is a callable, function to compute the Sobol' indices from. 

263 Its signature must be:: 

264 

265 func(x: ArrayLike) -> ArrayLike 

266 

267 with ``x`` of shape ``(d, n)`` and output of shape ``(s, n)`` where: 

268 

269 - ``d`` is the input dimensionality of `func` 

270 (number of input variables), 

271 - ``s`` is the output dimensionality of `func` 

272 (number of output variables), and 

273 - ``n`` is the number of samples (see `n` below). 

274 

275 Function evaluation values must be finite. 

276 

277 If `func` is a dictionary, contains the function evaluations from three 

278 different arrays. Keys must be: ``f_A``, ``f_B`` and ``f_AB``. 

279 ``f_A`` and ``f_B`` should have a shape ``(s, n)`` and ``f_AB`` 

280 should have a shape ``(d, s, n)``. 

281 This is an advanced feature and misuse can lead to wrong analysis. 

282 n : int 

283 Number of samples used to generate the matrices ``A`` and ``B``. 

284 Must be a power of 2. The total number of points at which `func` is 

285 evaluated will be ``n*(d+2)``. 

286 dists : list(distributions), optional 

287 List of each parameter's distribution. The distribution of parameters 

288 depends on the application and should be carefully chosen. 

289 Parameters are assumed to be independently distributed, meaning there 

290 is no constraint nor relationship between their values. 

291 

292 Distributions must be an instance of a class with a ``ppf`` 

293 method. 

294 

295 Must be specified if `func` is a callable, and ignored otherwise. 

296 method : Callable or str, default: 'saltelli_2010' 

297 Method used to compute the first and total Sobol' indices. 

298 

299 If a callable, its signature must be:: 

300 

301 func(f_A: np.ndarray, f_B: np.ndarray, f_AB: np.ndarray) 

302 -> Tuple[np.ndarray, np.ndarray] 

303 

304 with ``f_A, f_B`` of shape ``(s, n)`` and ``f_AB`` of shape 

305 ``(d, s, n)``. 

306 These arrays contain the function evaluations from three different sets 

307 of samples. 

308 The output is a tuple of the first and total indices with 

309 shape ``(s, d)``. 

310 This is an advanced feature and misuse can lead to wrong analysis. 

311 random_state : {None, int, `numpy.random.Generator`}, optional 

312 If `random_state` is an int or None, a new `numpy.random.Generator` is 

313 created using ``np.random.default_rng(random_state)``. 

314 If `random_state` is already a ``Generator`` instance, then the 

315 provided instance is used. 

316 

317 Returns 

318 ------- 

319 res : SobolResult 

320 An object with attributes: 

321 

322 first_order : ndarray of shape (s, d) 

323 First order Sobol' indices. 

324 total_order : ndarray of shape (s, d) 

325 Total order Sobol' indices. 

326 

327 And method: 

328 

329 bootstrap(confidence_level: float, n_resamples: int) 

330 -> BootstrapSobolResult 

331 

332 A method providing confidence intervals on the indices. 

333 See `scipy.stats.bootstrap` for more details. 

334 

335 The bootstrapping is done on both first and total order indices, 

336 and they are available in `BootstrapSobolResult` as attributes 

337 ``first_order`` and ``total_order``. 

338 

339 Notes 

340 ----- 

341 The Sobol' method [1]_, [2]_ is a variance-based Sensitivity Analysis which 

342 obtains the contribution of each parameter to the variance of the 

343 quantities of interest (QoIs; i.e., the outputs of `func`). 

344 Respective contributions can be used to rank the parameters and 

345 also gauge the complexity of the model by computing the 

346 model's effective (or mean) dimension. 

347 

348 .. note:: 

349 

350 Parameters are assumed to be independently distributed. Each 

351 parameter can still follow any distribution. In fact, the distribution 

352 is very important and should match the real distribution of the 

353 parameters. 

354 

355 It uses a functional decomposition of the variance of the function to 

356 explore 

357 

358 .. math:: 

359 

360 \mathbb{V}(Y) = \sum_{i}^{d} \mathbb{V}_i (Y) + \sum_{i<j}^{d} 

361 \mathbb{V}_{ij}(Y) + ... + \mathbb{V}_{1,2,...,d}(Y), 

362 

363 introducing conditional variances: 

364 

365 .. math:: 

366 

367 \mathbb{V}_i(Y) = \mathbb{\mathbb{V}}[\mathbb{E}(Y|x_i)] 

368 \qquad 

369 \mathbb{V}_{ij}(Y) = \mathbb{\mathbb{V}}[\mathbb{E}(Y|x_i x_j)] 

370 - \mathbb{V}_i(Y) - \mathbb{V}_j(Y), 

371 

372 Sobol' indices are expressed as 

373 

374 .. math:: 

375 

376 S_i = \frac{\mathbb{V}_i(Y)}{\mathbb{V}[Y]} 

377 \qquad 

378 S_{ij} =\frac{\mathbb{V}_{ij}(Y)}{\mathbb{V}[Y]}. 

379 

380 :math:`S_{i}` corresponds to the first-order term which apprises the 

381 contribution of the i-th parameter, while :math:`S_{ij}` corresponds to the 

382 second-order term which informs about the contribution of interactions 

383 between the i-th and the j-th parameters. These equations can be 

384 generalized to compute higher order terms; however, they are expensive to 

385 compute and their interpretation is complex. 

386 This is why only first order indices are provided. 

387 

388 Total order indices represent the global contribution of the parameters 

389 to the variance of the QoI and are defined as: 

390 

391 .. math:: 

392 

393 S_{T_i} = S_i + \sum_j S_{ij} + \sum_{j,k} S_{ijk} + ... 

394 = 1 - \frac{\mathbb{V}[\mathbb{E}(Y|x_{\sim i})]}{\mathbb{V}[Y]}. 

395 

396 First order indices sum to at most 1, while total order indices sum to at 

397 least 1. If there are no interactions, then first and total order indices 

398 are equal, and both first and total order indices sum to 1. 

399 

400 .. warning:: 

401 

402 Negative Sobol' values are due to numerical errors. Increasing the 

403 number of points `n` should help. 

404 

405 The number of sample required to have a good analysis increases with 

406 the dimensionality of the problem. e.g. for a 3 dimension problem, 

407 consider at minima ``n >= 2**12``. The more complex the model is, 

408 the more samples will be needed. 

409 

410 Even for a purely addiditive model, the indices may not sum to 1 due 

411 to numerical noise. 

412 

413 References 

414 ---------- 

415 .. [1] Sobol, I. M.. "Sensitivity analysis for nonlinear mathematical 

416 models." Mathematical Modeling and Computational Experiment, 1:407-414, 

417 1993. 

418 .. [2] Sobol, I. M. (2001). "Global sensitivity indices for nonlinear 

419 mathematical models and their Monte Carlo estimates." Mathematics 

420 and Computers in Simulation, 55(1-3):271-280, 

421 :doi:`10.1016/S0378-4754(00)00270-6`, 2001. 

422 .. [3] Saltelli, A. "Making best use of model evaluations to 

423 compute sensitivity indices." Computer Physics Communications, 

424 145(2):280-297, :doi:`10.1016/S0010-4655(02)00280-1`, 2002. 

425 .. [4] Saltelli, A., M. Ratto, T. Andres, F. Campolongo, J. Cariboni, 

426 D. Gatelli, M. Saisana, and S. Tarantola. "Global Sensitivity Analysis. 

427 The Primer." 2007. 

428 .. [5] Saltelli, A., P. Annoni, I. Azzini, F. Campolongo, M. Ratto, and 

429 S. Tarantola. "Variance based sensitivity analysis of model 

430 output. Design and estimator for the total sensitivity index." 

431 Computer Physics Communications, 181(2):259-270, 

432 :doi:`10.1016/j.cpc.2009.09.018`, 2010. 

433 .. [6] Ishigami, T. and T. Homma. "An importance quantification technique 

434 in uncertainty analysis for computer models." IEEE, 

435 :doi:`10.1109/ISUMA.1990.151285`, 1990. 

436 

437 Examples 

438 -------- 

439 The following is an example with the Ishigami function [6]_ 

440 

441 .. math:: 

442 

443 Y(\mathbf{x}) = \sin x_1 + 7 \sin^2 x_2 + 0.1 x_3^4 \sin x_1, 

444 

445 with :math:`\mathbf{x} \in [-\pi, \pi]^3`. This function exhibits strong 

446 non-linearity and non-monotonicity. 

447 

448 Remember, Sobol' indices assumes that samples are independently 

449 distributed. In this case we use a uniform distribution on each marginals. 

450 

451 >>> import numpy as np 

452 >>> from scipy.stats import sobol_indices, uniform 

453 >>> rng = np.random.default_rng() 

454 >>> def f_ishigami(x): 

455 ... f_eval = ( 

456 ... np.sin(x[0]) 

457 ... + 7 * np.sin(x[1])**2 

458 ... + 0.1 * (x[2]**4) * np.sin(x[0]) 

459 ... ) 

460 ... return f_eval 

461 >>> indices = sobol_indices( 

462 ... func=f_ishigami, n=1024, 

463 ... dists=[ 

464 ... uniform(loc=-np.pi, scale=2*np.pi), 

465 ... uniform(loc=-np.pi, scale=2*np.pi), 

466 ... uniform(loc=-np.pi, scale=2*np.pi) 

467 ... ], 

468 ... random_state=rng 

469 ... ) 

470 >>> indices.first_order 

471 array([0.31637954, 0.43781162, 0.00318825]) 

472 >>> indices.total_order 

473 array([0.56122127, 0.44287857, 0.24229595]) 

474 

475 Confidence interval can be obtained using bootstrapping. 

476 

477 >>> boot = indices.bootstrap() 

478 

479 Then, this information can be easily visualized. 

480 

481 >>> import matplotlib.pyplot as plt 

482 >>> fig, axs = plt.subplots(1, 2, figsize=(9, 4)) 

483 >>> _ = axs[0].errorbar( 

484 ... [1, 2, 3], indices.first_order, fmt='o', 

485 ... yerr=[ 

486 ... indices.first_order - boot.first_order.confidence_interval.low, 

487 ... boot.first_order.confidence_interval.high - indices.first_order 

488 ... ], 

489 ... ) 

490 >>> axs[0].set_ylabel("First order Sobol' indices") 

491 >>> axs[0].set_xlabel('Input parameters') 

492 >>> axs[0].set_xticks([1, 2, 3]) 

493 >>> _ = axs[1].errorbar( 

494 ... [1, 2, 3], indices.total_order, fmt='o', 

495 ... yerr=[ 

496 ... indices.total_order - boot.total_order.confidence_interval.low, 

497 ... boot.total_order.confidence_interval.high - indices.total_order 

498 ... ], 

499 ... ) 

500 >>> axs[1].set_ylabel("Total order Sobol' indices") 

501 >>> axs[1].set_xlabel('Input parameters') 

502 >>> axs[1].set_xticks([1, 2, 3]) 

503 >>> plt.tight_layout() 

504 >>> plt.show() 

505 

506 .. note:: 

507 

508 By default, `scipy.stats.uniform` has support ``[0, 1]``. 

509 Using the parameters ``loc`` and ``scale``, one obtains the uniform 

510 distribution on ``[loc, loc + scale]``. 

511 

512 This result is particularly interesting because the first order index 

513 :math:`S_{x_3} = 0` whereas its total order is :math:`S_{T_{x_3}} = 0.244`. 

514 This means that higher order interactions with :math:`x_3` are responsible 

515 for the difference. Almost 25% of the observed variance 

516 on the QoI is due to the correlations between :math:`x_3` and :math:`x_1`, 

517 although :math:`x_3` by itself has no impact on the QoI. 

518 

519 The following gives a visual explanation of Sobol' indices on this 

520 function. Let's generate 1024 samples in :math:`[-\pi, \pi]^3` and 

521 calculate the value of the output. 

522 

523 >>> from scipy.stats import qmc 

524 >>> n_dim = 3 

525 >>> p_labels = ['$x_1$', '$x_2$', '$x_3$'] 

526 >>> sample = qmc.Sobol(d=n_dim, seed=rng).random(1024) 

527 >>> sample = qmc.scale( 

528 ... sample=sample, 

529 ... l_bounds=[-np.pi, -np.pi, -np.pi], 

530 ... u_bounds=[np.pi, np.pi, np.pi] 

531 ... ) 

532 >>> output = f_ishigami(sample.T) 

533 

534 Now we can do scatter plots of the output with respect to each parameter. 

535 This gives a visual way to understand how each parameter impacts the 

536 output of the function. 

537 

538 >>> fig, ax = plt.subplots(1, n_dim, figsize=(12, 4)) 

539 >>> for i in range(n_dim): 

540 ... xi = sample[:, i] 

541 ... ax[i].scatter(xi, output, marker='+') 

542 ... ax[i].set_xlabel(p_labels[i]) 

543 >>> ax[0].set_ylabel('Y') 

544 >>> plt.tight_layout() 

545 >>> plt.show() 

546 

547 Now Sobol' goes a step further: 

548 by conditioning the output value by given values of the parameter 

549 (black lines), the conditional output mean is computed. It corresponds to 

550 the term :math:`\mathbb{E}(Y|x_i)`. Taking the variance of this term gives 

551 the numerator of the Sobol' indices. 

552 

553 >>> mini = np.min(output) 

554 >>> maxi = np.max(output) 

555 >>> n_bins = 10 

556 >>> bins = np.linspace(-np.pi, np.pi, num=n_bins, endpoint=False) 

557 >>> dx = bins[1] - bins[0] 

558 >>> fig, ax = plt.subplots(1, n_dim, figsize=(12, 4)) 

559 >>> for i in range(n_dim): 

560 ... xi = sample[:, i] 

561 ... ax[i].scatter(xi, output, marker='+') 

562 ... ax[i].set_xlabel(p_labels[i]) 

563 ... for bin_ in bins: 

564 ... idx = np.where((bin_ <= xi) & (xi <= bin_ + dx)) 

565 ... xi_ = xi[idx] 

566 ... y_ = output[idx] 

567 ... ave_y_ = np.mean(y_) 

568 ... ax[i].plot([bin_ + dx/2] * 2, [mini, maxi], c='k') 

569 ... ax[i].scatter(bin_ + dx/2, ave_y_, c='r') 

570 >>> ax[0].set_ylabel('Y') 

571 >>> plt.tight_layout() 

572 >>> plt.show() 

573 

574 Looking at :math:`x_3`, the variance 

575 of the mean is zero leading to :math:`S_{x_3} = 0`. But we can further 

576 observe that the variance of the output is not constant along the parameter 

577 values of :math:`x_3`. This heteroscedasticity is explained by higher order 

578 interactions. Moreover, an heteroscedasticity is also noticeable on 

579 :math:`x_1` leading to an interaction between :math:`x_3` and :math:`x_1`. 

580 On :math:`x_2`, the variance seems to be constant and thus null interaction 

581 with this parameter can be supposed. 

582 

583 This case is fairly simple to analyse visually---although it is only a 

584 qualitative analysis. Nevertheless, when the number of input parameters 

585 increases such analysis becomes unrealistic as it would be difficult to 

586 conclude on high-order terms. Hence the benefit of using Sobol' indices. 

587 

588 """ 

589 random_state = check_random_state(random_state) 

590 

591 n_ = int(n) 

592 if not (n_ & (n_ - 1) == 0) or n != n_: 

593 raise ValueError( 

594 "The balance properties of Sobol' points require 'n' " 

595 "to be a power of 2." 

596 ) 

597 n = n_ 

598 

599 if not callable(method): 

600 indices_methods: dict[str, Callable] = { 

601 "saltelli_2010": saltelli_2010, 

602 } 

603 try: 

604 method = method.lower() # type: ignore[assignment] 

605 indices_method_ = indices_methods[method] 

606 except KeyError as exc: 

607 message = ( 

608 f"{method!r} is not a valid 'method'. It must be one of" 

609 f" {set(indices_methods)!r} or a callable." 

610 ) 

611 raise ValueError(message) from exc 

612 else: 

613 indices_method_ = method 

614 sig = inspect.signature(indices_method_) 

615 

616 if set(sig.parameters) != {'f_A', 'f_B', 'f_AB'}: 

617 message = ( 

618 "If 'method' is a callable, it must have the following" 

619 f" signature: {inspect.signature(saltelli_2010)}" 

620 ) 

621 raise ValueError(message) 

622 

623 def indices_method(f_A, f_B, f_AB): 

624 """Wrap indices method to ensure proper output dimension. 

625 

626 1D when single output, 2D otherwise. 

627 """ 

628 return np.squeeze(indices_method_(f_A=f_A, f_B=f_B, f_AB=f_AB)) 

629 

630 if callable(func): 

631 if dists is None: 

632 raise ValueError( 

633 "'dists' must be defined when 'func' is a callable." 

634 ) 

635 

636 def wrapped_func(x): 

637 return np.atleast_2d(func(x)) 

638 

639 A, B = sample_A_B(n=n, dists=dists, random_state=random_state) 

640 AB = sample_AB(A=A, B=B) 

641 

642 f_A = wrapped_func(A) 

643 

644 if f_A.shape[1] != n: 

645 raise ValueError( 

646 "'func' output should have a shape ``(s, -1)`` with ``s`` " 

647 "the number of output." 

648 ) 

649 

650 def funcAB(AB): 

651 d, d, n = AB.shape 

652 AB = np.moveaxis(AB, 0, -1).reshape(d, n*d) 

653 f_AB = wrapped_func(AB) 

654 return np.moveaxis(f_AB.reshape((-1, n, d)), -1, 0) 

655 

656 f_B = wrapped_func(B) 

657 f_AB = funcAB(AB) 

658 else: 

659 message = ( 

660 "When 'func' is a dictionary, it must contain the following " 

661 "keys: 'f_A', 'f_B' and 'f_AB'." 

662 "'f_A' and 'f_B' should have a shape ``(s, n)`` and 'f_AB' " 

663 "should have a shape ``(d, s, n)``." 

664 ) 

665 try: 

666 f_A, f_B, f_AB = np.atleast_2d( 

667 func['f_A'], func['f_B'], func['f_AB'] 

668 ) 

669 except KeyError as exc: 

670 raise ValueError(message) from exc 

671 

672 if f_A.shape[1] != n or f_A.shape != f_B.shape or \ 

673 f_AB.shape == f_A.shape or f_AB.shape[-1] % n != 0: 

674 raise ValueError(message) 

675 

676 # Normalization by mean 

677 # Sobol', I. and Levitan, Y. L. (1999). On the use of variance reducing 

678 # multipliers in monte carlo computations of a global sensitivity index. 

679 # Computer Physics Communications, 117(1) :52-61. 

680 mean = np.mean([f_A, f_B], axis=(0, -1)).reshape(-1, 1) 

681 f_A -= mean 

682 f_B -= mean 

683 f_AB -= mean 

684 

685 # Compute indices 

686 # Filter warnings for constant output as var = 0 

687 with np.errstate(divide='ignore', invalid='ignore'): 

688 first_order, total_order = indices_method(f_A=f_A, f_B=f_B, f_AB=f_AB) 

689 

690 # null variance means null indices 

691 first_order[~np.isfinite(first_order)] = 0 

692 total_order[~np.isfinite(total_order)] = 0 

693 

694 res = dict( 

695 first_order=first_order, 

696 total_order=total_order, 

697 _indices_method=indices_method, 

698 _f_A=f_A, 

699 _f_B=f_B, 

700 _f_AB=f_AB 

701 ) 

702 

703 if callable(func): 

704 res.update( 

705 dict( 

706 _A=A, 

707 _B=B, 

708 _AB=AB, 

709 ) 

710 ) 

711 

712 return SobolResult(**res)