0
# PyMC Mathematical Utilities
1

2
PyMC provides essential mathematical functions for probabilistic modeling, including link functions, numerical utilities, and specialized operations for working with probability distributions and statistical models.
3

4
## Link Functions
5

6
Link functions transform parameters between different scales and domains, crucial for generalized linear models and constraint handling.
7

8
### Logit and Inverse Logit
9

10
```python { .api }
11
import pymc as pm
12

13
def logit(x):
14
    """
15
    Logit function (log-odds transformation).
16
    
17
    Parameters:
18
    - x (array-like): Input values in [0, 1]
19
    
20
    Returns:
21
    - logit_x: Log-odds values in (-∞, ∞)
22
    
23
    Formula: logit(x) = log(x / (1 - x))
24
    """
25
    
26
def invlogit(x):
27
    """
28
    Inverse logit function (logistic/sigmoid transformation).
29
    
30
    Parameters:
31
    - x (array-like): Input values in (-∞, ∞)
32
    
33
    Returns:
34
    - prob_x: Probability values in [0, 1]
35
    
36
    Formula: invlogit(x) = 1 / (1 + exp(-x))
37
    """
38

39
# Transform probabilities to log-odds scale
40
probability = 0.7
41
log_odds = pm.logit(probability)
42
print(f"P = {probability} -> logit(P) = {log_odds:.3f}")
43

44
# Transform log-odds back to probability scale  
45
back_to_prob = pm.invlogit(log_odds)
46
print(f"logit^(-1)({log_odds:.3f}) = {back_to_prob:.3f}")
47

48
# Vectorized operations
49
probabilities = np.array([0.1, 0.3, 0.5, 0.7, 0.9])
50
log_odds_vec = pm.logit(probabilities)
51
recovered_probs = pm.invlogit(log_odds_vec)
52

53
# Use in models for bounded parameters
54
with pm.Model() as logistic_model:
55
    # Unconstrained parameter
56
    alpha_raw = pm.Normal('alpha_raw', 0, 2)
57
    
58
    # Transform to probability scale
59
    success_prob = pm.Deterministic('success_prob', pm.invlogit(alpha_raw))
60
    
61
    # Binomial likelihood
62
    y = pm.Binomial('y', n=10, p=success_prob, observed=successes)
63
```
64

65
### Probit and Inverse Probit
66

67
```python { .api }
68
def probit(x):
69
    """
70
    Probit function (inverse normal CDF transformation).
71
    
72
    Parameters:
73
    - x (array-like): Input values in [0, 1]
74
    
75
    Returns:
76
    - probit_x: Standard normal quantiles in (-∞, ∞)
77
    """
78
    
79
def invprobit(x):
80
    """
81
    Inverse probit function (normal CDF transformation).
82
    
83
    Parameters:
84
    - x (array-like): Input values in (-∞, ∞)
85
    
86
    Returns:
87
    - prob_x: Probability values in [0, 1]
88
    """
89

90
# Probit transformation (alternative to logit)
91
prob_value = 0.8
92
probit_value = pm.probit(prob_value)
93
recovered = pm.invprobit(probit_value)
94

95
print(f"P = {prob_value} -> probit(P) = {probit_value:.3f}")
96
print(f"probit^(-1)({probit_value:.3f}) = {recovered:.3f}")
97

98
# Probit regression model
99
with pm.Model() as probit_regression:
100
    # Linear predictor
101
    beta = pm.Normal('beta', 0, 1, shape=X.shape[1])
102
    linear_pred = pm.math.dot(X, beta)
103
    
104
    # Probit link
105
    success_prob = pm.Deterministic('success_prob', pm.invprobit(linear_pred))
106
    
107
    # Bernoulli likelihood
108
    y = pm.Bernoulli('y', p=success_prob, observed=binary_outcomes)
109
```
110

111
## Numerical Utilities
112

113
### Logarithmic Operations
114

115
```python { .api }
116
def logsumexp(x, axis=None, keepdims=False):
117
    """
118
    Numerically stable computation of log(sum(exp(x))).
119
    
120
    Parameters:
121
    - x (array-like): Input array
122
    - axis (int, optional): Axis along which to compute
123
    - keepdims (bool): Keep reduced dimensions
124
    
125
    Returns:
126
    - result: log(sum(exp(x))) computed stably
127
    """
128

129
def logaddexp(x1, x2):
130
    """
131
    Numerically stable computation of log(exp(x1) + exp(x2)).
132
    
133
    Parameters:
134
    - x1, x2 (array-like): Input values
135
    
136
    Returns:
137
    - result: log(exp(x1) + exp(x2)) computed stably
138
    """
139

140
# Stable computation of mixture probabilities
141
log_weights = np.array([-2.3, -1.8, -3.1])  # Log mixture weights
142
log_normalizing_constant = pm.logsumexp(log_weights)
143
normalized_log_weights = log_weights - log_normalizing_constant
144

145
print(f"Log weights: {log_weights}")
146
print(f"Log Z: {log_normalizing_constant:.3f}")
147
print(f"Normalized log weights: {normalized_log_weights}")
148

149
# Stable addition in log space
150
log_a = -10.5
151
log_b = -10.2
152
log_sum = pm.logaddexp(log_a, log_b)
153
print(f"log(exp({log_a}) + exp({log_b})) = {log_sum:.6f}")
154

155
# Compare with naive computation (may underflow)
156
naive_sum = np.log(np.exp(log_a) + np.exp(log_b))  # May be -inf
157
print(f"Naive computation: {naive_sum}")
158
```
159

160
### Matrix Operations
161

162
```python { .api }
163
def expand_packed_triangular(n, packed, lower=True, diagonal_only=False):
164
    """
165
    Expand packed triangular matrix to full matrix.
166
    
167
    Parameters:
168
    - n (int): Dimension of output matrix
169
    - packed (array): Packed triangular values
170
    - lower (bool): Expand as lower triangular (default: True)
171
    - diagonal_only (bool): Only diagonal elements
172
    
173
    Returns:
174
    - matrix: Expanded n×n matrix
175
    """
176

177
# Pack lower triangular matrix
178
n_dim = 3
179
n_packed = n_dim * (n_dim + 1) // 2  # Number of packed elements
180

181
# Packed values for lower triangle (column-major order)
182
packed_values = np.array([1.0, 0.5, 2.0, 0.3, 0.8, 1.5])  # [L11, L21, L22, L31, L32, L33]
183

184
# Expand to full matrix
185
L_matrix = pm.expand_packed_triangular(n_dim, packed_values, lower=True)
186
print("Lower triangular matrix:")
187
print(L_matrix.eval())
188

189
# Use in Cholesky decomposition parameterization
190
with pm.Model() as cholesky_model:
191
    # Packed Cholesky factors
192
    L_packed = pm.Normal('L_packed', 0, 1, shape=n_packed)
193
    
194
    # Expand to Cholesky matrix
195
    L = pm.expand_packed_triangular(n_dim, L_packed, lower=True)
196
    
197
    # Covariance matrix
198
    Sigma = pm.math.dot(L, L.T)
199
    
200
    # Multivariate normal with Cholesky parameterization
201
    y = pm.MvNormal('y', mu=np.zeros(n_dim), chol=L, observed=data)
202
```
203

204
## Specialized Mathematical Functions
205

206
### Tensor Operations
207

208
```python { .api }
209
# Common tensor operations available through pm.math
210

211
# Matrix multiplication
212
A = pm.Data('A', np.random.randn(5, 3))
213
B = pm.Data('B', np.random.randn(3, 4))
214
C = pm.math.dot(A, B)  # Matrix product
215

216
# Element-wise operations
217
x = pm.Normal('x', 0, 1, shape=10)
218
squared = pm.math.sqr(x)           # Element-wise square
219
sqrt_x = pm.math.sqrt(pm.math.abs(x))  # Element-wise square root
220
exp_x = pm.math.exp(x)             # Element-wise exponential
221
log_x = pm.math.log(pm.math.abs(x)) # Element-wise logarithm
222

223
# Trigonometric functions
224
sin_x = pm.math.sin(x)             # Sine
225
cos_x = pm.math.cos(x)             # Cosine
226
tan_x = pm.math.tan(x)             # Tangent
227

228
# Hyperbolic functions
229
sinh_x = pm.math.sinh(x)           # Hyperbolic sine
230
cosh_x = pm.math.cosh(x)           # Hyperbolic cosine
231
tanh_x = pm.math.tanh(x)           # Hyperbolic tangent
232
```
233

234
### Statistical Functions
235

236
```python { .api }
237
# Statistical operations
238
data_tensor = pm.Data('data', np.random.randn(100, 5))
239

240
# Summary statistics
241
mean_val = pm.math.mean(data_tensor, axis=0)    # Column means
242
std_val = pm.math.std(data_tensor, axis=0)      # Column standard deviations  
243
var_val = pm.math.var(data_tensor, axis=0)      # Column variances
244

245
# Quantiles and percentiles
246
median_val = pm.math.median(data_tensor, axis=0)
247
max_val = pm.math.max(data_tensor, axis=0)
248
min_val = pm.math.min(data_tensor, axis=0)
249

250
# Cumulative operations
251
cumsum = pm.math.cumsum(data_tensor, axis=0)    # Cumulative sum
252
cumprod = pm.math.cumprod(data_tensor, axis=0)  # Cumulative product
253
```
254

255
### Comparison and Logical Operations
256

257
```python { .api }
258
# Logical and comparison operations
259
x = pm.Normal('x', 0, 1, shape=5)
260
y = pm.Normal('y', 0, 1, shape=5)
261

262
# Comparisons
263
greater = pm.math.gt(x, y)         # x > y
264
less_equal = pm.math.le(x, 0)      # x <= 0
265
equal = pm.math.eq(x, 0)           # x == 0
266

267
# Logical operations
268
and_result = pm.math.and_(pm.math.gt(x, 0), pm.math.lt(x, 1))  # 0 < x < 1
269
or_result = pm.math.or_(pm.math.lt(x, -1), pm.math.gt(x, 1))   # x < -1 OR x > 1
270

271
# Conditional operations (switch/where)
272
positive_x = pm.math.switch(pm.math.gt(x, 0), x, 0)  # max(x, 0)
273
clipped_x = pm.math.clip(x, -2, 2)                   # Clip x to [-2, 2]
274
```
275

276
## Advanced Mathematical Utilities
277

278
### Custom Mathematical Functions
279

280
```python { .api }
281
# Define custom mathematical functions using PyTensor operations
282
def softplus(x):
283
    """Softplus function: log(1 + exp(x))"""
284
    return pm.math.log(1 + pm.math.exp(x))
285

286
def stable_softplus(x):
287
    """Numerically stable softplus"""
288
    return pm.math.switch(pm.math.gt(x, 20), x, pm.math.log(1 + pm.math.exp(x)))
289

290
def swish(x):
291
    """Swish activation function: x * sigmoid(x)"""
292
    return x * pm.invlogit(x)
293

294
def gelu(x):
295
    """Gaussian Error Linear Unit (approximate)"""
296
    return 0.5 * x * (1 + pm.math.tanh(np.sqrt(2/np.pi) * (x + 0.044715 * x**3)))
297

298
# Use custom functions in models
299
with pm.Model() as custom_model:
300
    z = pm.Normal('z', 0, 1, shape=10)
301
    
302
    # Apply custom transformations
303
    softplus_z = pm.Deterministic('softplus_z', stable_softplus(z))
304
    swish_z = pm.Deterministic('swish_z', swish(z))
305
```
306

307
### Numerical Stability Utilities
308

309
```python { .api }
310
def log1p(x):
311
    """Compute log(1 + x) accurately for small x"""
312
    return pm.math.log1p(x)
313

314
def expm1(x):  
315
    """Compute exp(x) - 1 accurately for small x"""
316
    return pm.math.expm1(x)
317

318
def log1mexp(x):
319
    """Compute log(1 - exp(x)) stably for x < 0"""
320
    return pm.math.switch(
321
        pm.math.gt(x, -0.693),  # log(0.5)
322
        pm.math.log(-pm.math.expm1(x)),
323
        pm.math.log1p(-pm.math.exp(x))
324
    )
325

326
# Example: Stable computation of survival function
327
def log_survival_function(log_hazard_cumulative):
328
    """Log survival function: log(1 - CDF) = log(1 - exp(log_CDF))"""
329
    return log1mexp(log_hazard_cumulative)
330

331
with pm.Model() as survival_model:
332
    # Log cumulative hazard
333
    log_cum_hazard = pm.Normal('log_cum_hazard', -2, 1)
334
    
335
    # Stable log survival probability
336
    log_survival = pm.Deterministic('log_survival', 
337
                                   log_survival_function(log_cum_hazard))
338
```
339

340
### Matrix Decompositions and Operations
341

342
```python { .api }
343
# Cholesky decomposition utilities
344
def cholesky_decomposition(matrix):
345
    """Cholesky decomposition of positive definite matrix"""
346
    return pm.math.cholesky(matrix)
347

348
def solve_triangular(A, b, lower=True):
349
    """Solve triangular system Ax = b"""
350
    return pm.math.solve_triangular(A, b, lower=lower)
351

352
# Example: Efficient multivariate normal computation
353
with pm.Model() as efficient_mvn:
354
    # Covariance parameters
355
    sigma = pm.HalfNormal('sigma', 1, shape=n_dim)
356
    L_corr = pm.LKJCholeskyCov('L_corr', n=n_dim, eta=1, sd_dist=sigma)
357
    
358
    # Efficient sampling using Cholesky
359
    z = pm.Normal('z', 0, 1, shape=n_dim)
360
    x = pm.Deterministic('x', pm.math.dot(L_corr, z))
361
    
362
    # Likelihood
363
    y = pm.Normal('y', mu=x, sigma=0.1, observed=data)
364
```
365

366
### Broadcasting and Reshaping
367

368
```python { .api }
369
# Array manipulation utilities
370
x = pm.Normal('x', 0, 1, shape=(5, 3))
371

372
# Reshaping
373
x_flat = pm.math.flatten(x)                    # Flatten to 1D
374
x_reshaped = pm.math.reshape(x, (3, 5))        # Reshape to (3, 5)
375

376
# Broadcasting
377
y = pm.Normal('y', 0, 1, shape=5)
378
broadcasted = x + y[:, None]                   # Broadcast y across columns
379

380
# Dimension manipulation
381
x_expanded = pm.math.expand_dims(y, axis=1)    # Add dimension
382
x_squeezed = pm.math.squeeze(x_expanded)       # Remove dimensions of size 1
383

384
# Stacking and concatenation
385
z = pm.Normal('z', 0, 1, shape=(5, 3))
386
stacked = pm.math.stack([x, z], axis=0)        # Stack along new axis
387
concatenated = pm.math.concatenate([x, z], axis=1)  # Concatenate along axis
388
```
389

390
## Integration with Probability Distributions
391

392
### Custom Distribution Support
393

394
```python { .api }
395
# Using mathematical utilities in custom distributions
396
def custom_logp(value, mu, sigma):
397
    """Custom log-probability function using mathematical utilities"""
398
    
399
    # Standardize
400
    standardized = (value - mu) / sigma
401
    
402
    # Custom transformation
403
    transformed = stable_softplus(standardized)
404
    
405
    # Log-probability computation
406
    logp_val = -0.5 * pm.math.sum(transformed**2)
407
    logp_val -= pm.math.sum(pm.math.log(sigma))  # Jacobian
408
    
409
    return logp_val
410

411
# Use in CustomDist
412
with pm.Model() as custom_dist_model:
413
    mu = pm.Normal('mu', 0, 1)
414
    sigma = pm.HalfNormal('sigma', 1)
415
    
416
    x = pm.CustomDist('x', mu=mu, sigma=sigma, 
417
                      logp=custom_logp, 
418
                      observed=data)
419
```
420

421
### Transformation Utilities
422

423
```python { .api }
424
# Common transformations for constrained parameters
425
def constrain_positive(x):
426
    """Ensure parameter is positive using softplus"""
427
    return stable_softplus(x)
428

429
def constrain_unit_interval(x):
430
    """Constrain parameter to [0, 1] using sigmoid"""
431
    return pm.invlogit(x)
432

433
def constrain_correlation(x):
434
    """Constrain to [-1, 1] using tanh"""
435
    return pm.math.tanh(x)
436

437
def stick_breaking(x):
438
    """Stick-breaking transform for simplex"""
439
    # x should be shape (..., K-1) for K-simplex
440
    stick_segments = pm.invlogit(x)
441
    
442
    # Compute stick-breaking weights
443
    remaining_stick = 1.0
444
    weights = []
445
    
446
    for i in range(x.shape[-1]):
447
        weight = stick_segments[..., i] * remaining_stick
448
        weights.append(weight)
449
        remaining_stick -= weight
450
    
451
    # Last weight gets remaining stick
452
    weights.append(remaining_stick)
453
    
454
    return pm.math.stack(weights, axis=-1)
455

456
# Example usage in Dirichlet-like model
457
with pm.Model() as stick_breaking_model:
458
    # Unconstrained parameters
459
    raw_weights = pm.Normal('raw_weights', 0, 1, shape=K-1)
460
    
461
    # Transform to simplex
462
    simplex_weights = pm.Deterministic('simplex_weights', 
463
                                      stick_breaking(raw_weights))
464
    
465
    # Categorical likelihood
466
    y = pm.Categorical('y', p=simplex_weights, observed=categories)
467
```
468

469
## Advanced Mathematical Functions
470

471
### Log-Probability and CDF Functions
472

473
```python { .api }
474
def logp(dist, value):
475
    """
476
    Compute log-probability density/mass function.
477
    
478
    Parameters:
479
    - dist: PyMC distribution object
480
    - value: Values at which to evaluate log-probability
481
    
482
    Returns:
483
    - log_prob: Log-probability values
484
    """
485

486
def logcdf(dist, value):
487
    """
488
    Compute log cumulative distribution function.
489
    
490
    Parameters:
491
    - dist: PyMC distribution object
492
    - value: Values at which to evaluate log-CDF
493
    
494
    Returns:
495
    - log_cdf: Log-CDF values
496
    """
497

498
def icdf(dist, q):
499
    """
500
    Compute inverse cumulative distribution function (quantile function).
501
    
502
    Parameters:
503
    - dist: PyMC distribution object
504
    - q: Quantile values in [0, 1]
505
    
506
    Returns:
507
    - quantiles: Values corresponding to given quantiles
508
    """
509

510
# Example usage
511
normal_dist = pm.Normal.dist(mu=0, sigma=1)
512

513
# Log-probability at specific values
514
log_prob_values = pm.logp(normal_dist, np.array([-1, 0, 1]))
515

516
# Log-CDF at specific values
517
log_cdf_values = pm.logcdf(normal_dist, np.array([-1, 0, 1]))
518

519
# Quantiles (inverse CDF)
520
quantiles = pm.icdf(normal_dist, np.array([0.025, 0.5, 0.975]))  # 95% CI
521
```
522

523
### Matrix and Triangular Operations
524

525
```python { .api }
526
def expand_packed_triangular(n, packed):
527
    """
528
    Expand packed triangular values into full triangular matrix.
529
    
530
    Parameters:
531
    - n (int): Dimension of the triangular matrix
532
    - packed (array): Packed triangular values (n*(n+1)/2 elements)
533
    
534
    Returns:
535
    - triangular_matrix: Lower triangular matrix of shape (n, n)
536
    """
537

538
# Example: create correlation matrix from packed values
539
n_dims = 3
540
n_packed = n_dims * (n_dims - 1) // 2  # Number of off-diagonal elements
541
packed_corr = np.array([0.5, -0.3, 0.7])  # 3 correlation values
542

543
# Expand to correlation matrix
544
corr_matrix = pm.expand_packed_triangular(n_dims, packed_corr, diag=1.0)
545
```
546

547
### Additional Numerical Utilities
548

549
```python { .api }
550
# Advanced numerical functions for edge cases
551
def logaddexp_m(x, y):
552
    """Numerically stable log(exp(x) + exp(y)) with memory optimization."""
553
    
554
def floatX(x):
555
    """Convert to PyTensor's floating point type."""
556
    
557
def constant(x):
558
    """Create constant tensor."""
559
    
560
def shared(value, name=None):
561
    """Create shared variable."""
562

563
# Usage in models
564
with pm.Model() as advanced_model:
565
    # Create shared parameter for updating
566
    shared_param = pm.math.shared(np.array([1.0, 2.0]), 'shared_param')
567
    
568
    # Convert to correct dtype
569
    precise_value = pm.math.floatX(3.14159265359)
570
    
571
    # Create constant
572
    fixed_value = pm.math.constant(2.718281828)
573
```
574

575
PyMC's mathematical utilities provide the essential building blocks for creating sophisticated probabilistic models with proper numerical stability and efficient computation.