|
| 1 | +# Working with Blobs |
| 2 | + |
| 3 | +This example demonstrates how to use "blobs" in Tempest to store auxiliary quantities computed during likelihood evaluation. Blobs allow you to cache expensive calculations, store diagnostic information, and compute derived quantities without re-evaluating the likelihood. |
| 4 | + |
| 5 | +--- |
| 6 | + |
| 7 | +## What Are Blobs? |
| 8 | + |
| 9 | +Blobs are additional return values from your likelihood function beyond the log-likelihood itself. When your likelihood computes quantities of interest (e.g., chi-squared, model predictions, physical parameters), you can return them as blobs and retrieve them later alongside your posterior samples. |
| 10 | + |
| 11 | +**Key Benefits:** |
| 12 | +- Avoid recomputing expensive quantities during analysis |
| 13 | +- Store diagnostic information for debugging |
| 14 | +- Compute derived quantities that depend on the likelihood evaluation |
| 15 | + |
| 16 | +**Important Limitation:** Blobs are not compatible with `vectorize=True`. Use scalar likelihood evaluation when you need blobs. |
| 17 | + |
| 18 | +--- |
| 19 | + |
| 20 | +## Problem Setup: 2D Gaussian Distribution |
| 21 | + |
| 22 | +We'll demonstrate blobs using a simple 2D Gaussian likelihood where we can compute various auxiliary quantities. |
| 23 | + |
| 24 | +**Mathematical Formulation:** |
| 25 | +- Likelihood: $\mathcal{N}(\mu=[0, 0], \Sigma=I)$ |
| 26 | +- Prior: Uniform $[-10, 10]$ for both dimensions |
| 27 | +- True parameters: $\mu = (0, 0)$ |
| 28 | + |
| 29 | +```python |
| 30 | +import numpy as np |
| 31 | +import tempest as tp |
| 32 | + |
| 33 | +n_dim = 2 |
| 34 | + |
| 35 | +def prior_transform(u): |
| 36 | + """Uniform prior from unit hypercube to physical parameters.""" |
| 37 | + return 20 * u - 10 |
| 38 | +``` |
| 39 | + |
| 40 | +--- |
| 41 | + |
| 42 | +## Case 1: Single Blob |
| 43 | + |
| 44 | +Start with a simple likelihood that returns chi-squared as a single blob value. |
| 45 | + |
| 46 | +```python |
| 47 | +def log_likelihood_single_blob(x): |
| 48 | + """Likelihood returning chi-squared as a single blob.""" |
| 49 | + chi2 = np.sum(x**2) |
| 50 | + logl = -0.5 * chi2 |
| 51 | + return logl, chi2 |
| 52 | + |
| 53 | +# Create sampler with single blob |
| 54 | +sampler = tp.Sampler( |
| 55 | + prior_transform=prior_transform, |
| 56 | + log_likelihood=log_likelihood_single_blob, |
| 57 | + n_dim=n_dim, |
| 58 | + n_effective=512, |
| 59 | + blobs_dtype=float, # Each blob is a scalar float |
| 60 | +) |
| 61 | + |
| 62 | +# Run the sampler |
| 63 | +sampler.run(n_total=2048, progress=False) |
| 64 | + |
| 65 | +# Retrieve posterior samples and blobs |
| 66 | +samples, weights, logl, blobs = sampler.posterior(return_blobs=True) |
| 67 | + |
| 68 | +# Expected shapes: |
| 69 | +# samples.shape = (n_samples, n_dim) = (2048, 2) |
| 70 | +# blobs.shape = (n_samples,) = (2048,) |
| 71 | + |
| 72 | +# Compute weighted statistics from blobs |
| 73 | +mean_chi2 = np.average(blobs, weights=weights) |
| 74 | +std_chi2 = np.sqrt(np.average((blobs - mean_chi2)**2, weights=weights)) |
| 75 | + |
| 76 | +print(f"Mean chi-squared: {mean_chi2:.2f} ± {std_chi2:.2f}") |
| 77 | +``` |
| 78 | + |
| 79 | +--- |
| 80 | + |
| 81 | +## Case 2: Multiple Blobs |
| 82 | + |
| 83 | +Return multiple auxiliary quantities. We'll demonstrate both unnamed and named field approaches. |
| 84 | + |
| 85 | +```python |
| 86 | +def log_likelihood_multiple_blobs(x): |
| 87 | + """Likelihood returning multiple blob quantities.""" |
| 88 | + chi2 = np.sum(x**2) |
| 89 | + radius = np.sqrt(chi2) |
| 90 | + max_abs = np.max(np.abs(x)) |
| 91 | + logl = -0.5 * chi2 |
| 92 | + return logl, chi2, radius, max_abs |
| 93 | + |
| 94 | +# Approach 1: Unnamed fields (tuple unpacking) |
| 95 | +# Blobs are returned as a numeric array |
| 96 | +sampler_unnamed = tp.Sampler( |
| 97 | + prior_transform=prior_transform, |
| 98 | + log_likelihood=log_likelihood_multiple_blobs, |
| 99 | + n_dim=n_dim, |
| 100 | + n_effective=512, |
| 101 | + blobs_dtype=(float, 3), # 3 float values per sample |
| 102 | +) |
| 103 | + |
| 104 | +sampler_unnamed.run(n_total=2048, progress=False) |
| 105 | +samples, weights, logl, blobs = sampler_unnamed.posterior(return_blobs=True) |
| 106 | + |
| 107 | +# Expected: blobs.shape = (n_samples, 3) = (2048, 3) |
| 108 | +chi2_values = blobs[:, 0] # First field |
| 109 | +radius_values = blobs[:, 1] # Second field |
| 110 | +max_abs_values = blobs[:, 2] # Third field |
| 111 | + |
| 112 | +# Approach 2: Named fields (recommended) |
| 113 | +# Use numpy structured arrays for clarity |
| 114 | +sampler_named = tp.Sampler( |
| 115 | + prior_transform=prior_transform, |
| 116 | + log_likelihood=log_likelihood_multiple_blobs, |
| 117 | + n_dim=n_dim, |
| 118 | + n_effective=512, |
| 119 | + blobs_dtype=[ |
| 120 | + ("chi2", float), |
| 121 | + ("radius", float), |
| 122 | + ("max_abs", float) |
| 123 | + ], |
| 124 | +) |
| 125 | + |
| 126 | +sampler_named.run(n_total=2048, progress=False) |
| 127 | +samples, weights, logl, blobs = sampler_named.posterior(return_blobs=True) |
| 128 | + |
| 129 | +# Expected: blobs.shape = (n_samples,) = (2048,) |
| 130 | +# Each element is a structured numpy array |
| 131 | + |
| 132 | +# Access fields by name |
| 133 | +chi2_values = blobs["chi2"] # Shape: (n_samples,) |
| 134 | +radius_values = blobs["radius"] # Shape: (n_samples,) |
| 135 | +max_abs_values = blobs["max_abs"] # Shape: (n_samples,) |
| 136 | + |
| 137 | +# Compute correlations with parameters |
| 138 | +corr_x0_radius = np.corrcoef(samples[:, 0], radius_values)[0, 1] |
| 139 | +print(f"Correlation between x[0] and radius: {corr_x0_radius:.3f}") |
| 140 | +``` |
| 141 | + |
| 142 | +--- |
| 143 | + |
| 144 | +## Case 3: Higher-Dimensional Blob Arrays |
| 145 | + |
| 146 | +Blobs can be multi-dimensional arrays, useful for storing per-parameter or per-data-point quantities. |
| 147 | + |
| 148 | +```python |
| 149 | +def log_likelihood_array_blob(x): |
| 150 | + """Likelihood returning array-valued blobs.""" |
| 151 | + chi2_per_dim = x**2 # chi2 for each dimension |
| 152 | + chi2_total = np.sum(chi2_per_dim) |
| 153 | + logl = -0.5 * chi2_total |
| 154 | + return logl, chi2_per_dim |
| 155 | + |
| 156 | +# Single array blob |
| 157 | +sampler_array = tp.Sampler( |
| 158 | + prior_transform=prior_transform, |
| 159 | + log_likelihood=log_likelihood_array_blob, |
| 160 | + n_dim=n_dim, |
| 161 | + n_effective=512, |
| 162 | + blobs_dtype=(float, n_dim), # Each blob is length-n_dim array |
| 163 | +) |
| 164 | + |
| 165 | +sampler_array.run(n_total=2048, progress=False) |
| 166 | +samples, weights, logl, blobs = sampler_array.posterior(return_blobs=True) |
| 167 | + |
| 168 | +# Expected: blobs.shape = (n_samples, n_dim) = (2048, 2) |
| 169 | +chi2_dim0 = blobs[:, 0] # chi2 from first dimension |
| 170 | +chi2_dim1 = blobs[:, 1] # chi2 from second dimension |
| 171 | + |
| 172 | +# Structured with mixed scalar and array fields |
| 173 | +sampler_mixed = tp.Sampler( |
| 174 | + prior_transform=prior_transform, |
| 175 | + log_likelihood=log_likelihood_array_blob, |
| 176 | + n_dim=n_dim, |
| 177 | + n_effective=512, |
| 178 | + blobs_dtype=[ |
| 179 | + ("chi2_total", float), |
| 180 | + ("chi2_per_dim", float, n_dim) # Array field |
| 181 | + ], |
| 182 | +) |
| 183 | + |
| 184 | +sampler_mixed.run(n_total=2048, progress=False) |
| 185 | +samples, weights, logl, blobs = sampler_mixed.posterior(return_blobs=True) |
| 186 | + |
| 187 | +# Expected: blobs.shape = (n_samples,) = (2048,) |
| 188 | +# blobs["chi2_total"].shape = (n_samples,) = (2048,) |
| 189 | +# blobs["chi2_per_dim"].shape = (n_samples, n_dim) = (2048, 2) |
| 190 | + |
| 191 | +total_chi2 = blobs["chi2_total"] |
| 192 | +chi2_per_dim = blobs["chi2_per_dim"] |
| 193 | +``` |
| 194 | + |
| 195 | +--- |
| 196 | + |
| 197 | +## Case 4: Using Blob Data for Analysis |
| 198 | + |
| 199 | +Practical examples of working with retrieved blobs. |
| 200 | + |
| 201 | +```python |
| 202 | +# Run the sampler with named blobs |
| 203 | +sampler = tp.Sampler( |
| 204 | + prior_transform=prior_transform, |
| 205 | + log_likelihood=log_likelihood_multiple_blobs, |
| 206 | + n_dim=n_dim, |
| 207 | + n_effective=512, |
| 208 | + blobs_dtype=[("chi2", float), ("radius", float), ("max_abs", float)], |
| 209 | +) |
| 210 | + |
| 211 | +sampler.run(n_total=2048, progress=False) |
| 212 | +samples, weights, logl, blobs = sampler.posterior(return_blobs=True) |
| 213 | + |
| 214 | +# Weighted statistics |
| 215 | +mean_radius = np.average(blobs["radius"], weights=weights) |
| 216 | +std_radius = np.sqrt(np.average((blobs["radius"] - mean_radius)**2, weights=weights)) |
| 217 | + |
| 218 | +print(f"Mean radius: {mean_radius:.2f} ± {std_radius:.2f}") |
| 219 | + |
| 220 | +# Conditional analysis: examine parameters for high-radius samples |
| 221 | +high_radius_threshold = np.percentile(blobs["radius"], 90) |
| 222 | +high_radius_mask = blobs["radius"] > high_radius_threshold |
| 223 | + |
| 224 | +samples_high_r = samples[high_radius_mask] |
| 225 | +weights_high_r = weights[high_radius_mask] |
| 226 | + |
| 227 | +mean_x0_high_r = np.average(samples_high_r[:, 0], weights=weights_high_r) |
| 228 | +print(f"Mean x[0] for high-radius samples: {mean_x0_high_r:.3f}") |
| 229 | + |
| 230 | +# Blob-only posterior (ignore parameters) |
| 231 | +logz, logz_err = sampler.evidence() |
| 232 | +chi2_mean, chi2_cov = tp.utils.weighted_avg_and_cov(blobs["chi2"], weights) |
| 233 | +``` |
| 234 | + |
| 235 | +--- |
| 236 | + |
| 237 | +## Important Considerations |
| 238 | + |
| 239 | +### Memory Usage |
| 240 | +Blobs are stored for all active particles at every iteration. Memory scales as: |
| 241 | +``` |
| 242 | +memory ≈ n_iterations × n_active × blob_size × dtype_bytes |
| 243 | +``` |
| 244 | +For large blobs or long runs, consider periodic saving and clearing. |
| 245 | + |
| 246 | +### Common Errors |
| 247 | + |
| 248 | +**Wrong blob return shape:** |
| 249 | +```python |
| 250 | +# Wrong: return logl, [chi2, radius] # List instead of tuple |
| 251 | +# Wrong: return logl, chi2, radius # Missing tuple wrapping if single value |
| 252 | +# Correct: return logl, (chi2, radius) # Always return tuple |
| 253 | +``` |
| 254 | + |
| 255 | +**Forgetting `return_blobs=True`:** |
| 256 | +```python |
| 257 | +# Wrong: samples, weights, logl, blobs = sampler.posterior() |
| 258 | +# Raises: ValueError: not enough values to unpack |
| 259 | + |
| 260 | +# Correct: samples, weights, logl, blobs = sampler.posterior(return_blobs=True) |
| 261 | +``` |
| 262 | + |
| 263 | +**Blob dtype mismatch:** |
| 264 | +```python |
| 265 | +# If likelihood returns 2 values but blobs_dtype expects 3 |
| 266 | +# Result: Runtime error or incorrect array shapes |
| 267 | +``` |
| 268 | + |
| 269 | +### Best Practices |
| 270 | + |
| 271 | +1. **Use structured dtypes**: Named fields (`blobs["chi2"]`) are clearer than indexing (`blobs[:, 0]`) |
| 272 | + |
| 273 | +2. **Specify dtypes explicitly**: Prevents ambiguous array creation |
| 274 | +```python |
| 275 | +# Good: blobs_dtype=[("chi2", float)] |
| 276 | +# Avoid: blobs_dtype=None (infers from first result, may cause issues) |
| 277 | +``` |
| 278 | + |
| 279 | +3. **Return numpy scalars/arrays**: Not Python lists |
| 280 | +```python |
| 281 | +# Good: return logl, np.array([chi2, radius]) |
| 282 | +# Avoid: return logl, [chi2, radius] |
| 283 | +``` |
| 284 | + |
| 285 | +4. **Profile memory usage** for large blob arrays |
| 286 | + |
| 287 | +5. **Consider post-processing trade-offs**: Blobs are ideal for expensive calculations, but simple quantities may be faster to recompute |
| 288 | + |
| 289 | +### Compatibility Notes |
| 290 | + |
| 291 | +- **Clustering**: Blobs work correctly with clustering enabled |
| 292 | +- **MCMC moves**: Blobs are properly handled during mutation steps |
| 293 | +- **Resampling**: Blobs are automatically resampled along with particles |
| 294 | +- **Saving/Loading**: Blobs are preserved in sampler state files |
| 295 | +- **Parallelization**: Works correctly with `pool` parameter |
| 296 | + |
| 297 | +### Performance Trade-offs |
| 298 | + |
| 299 | +**When to use blobs:** |
| 300 | +- Likelihood evaluations are expensive |
| 301 | +- Quantities require significant computation |
| 302 | +- Need to analyze quantities at posterior parameter values |
| 303 | +- Debugging or diagnostic information |
| 304 | + |
| 305 | +**When to compute post-hoc:** |
| 306 | +- Blob arrays would be extremely large |
| 307 | +- Quantities are trivial to compute |
| 308 | +- Only need quantities for subset of samples |
| 309 | +- Memory-constrained environments |
| 310 | + |
| 311 | +--- |
| 312 | + |
| 313 | +## Complete Working Example |
| 314 | + |
| 315 | +Here is a complete, runnable example combining all blob concepts: |
| 316 | + |
| 317 | +```python |
| 318 | +import numpy as np |
| 319 | +import tempest as tp |
| 320 | + |
| 321 | +# Setup |
| 322 | +n_dim = 2 |
| 323 | + |
| 324 | +def prior_transform(u): |
| 325 | + return 20 * u - 10 |
| 326 | + |
| 327 | +def log_likelihood_with_blobs(x): |
| 328 | + """Example likelihood returning structured blobs.""" |
| 329 | + chi2_total = np.sum(x**2) |
| 330 | + chi2_per_dim = x**2 |
| 331 | + radius = np.sqrt(chi2_total) |
| 332 | + max_deviation = np.max(np.abs(x)) |
| 333 | + logl = -0.5 * chi2_total |
| 334 | + |
| 335 | + return logl, (chi2_total, radius, max_deviation, chi2_per_dim) |
| 336 | + |
| 337 | +# Create sampler with multiple blob types |
| 338 | +sampler = tp.Sampler( |
| 339 | + prior_transform=prior_transform, |
| 340 | + log_likelihood=log_likelihood_with_blobs, |
| 341 | + n_dim=n_dim, |
| 342 | + n_effective=512, |
| 343 | + blobs_dtype=[ |
| 344 | + ("chi2_total", float), |
| 345 | + ("radius", float), |
| 346 | + ("max_deviation", float), |
| 347 | + ("chi2_per_dim", float, n_dim) |
| 348 | + ], |
| 349 | + random_state=42, |
| 350 | +) |
| 351 | + |
| 352 | +# Run sampling |
| 353 | +sampler.run(n_total=2048, progress=False) |
| 354 | + |
| 355 | +# Retrieve results with blobs |
| 356 | +samples, weights, logl, blobs = sampler.posterior(return_blobs=True) |
| 357 | + |
| 358 | +# Expected shapes: |
| 359 | +# samples.shape = (2048, 2) |
| 360 | +# blobs.shape = (2048,) |
| 361 | +# blobs["chi2_total"].shape = (2048,) |
| 362 | +# blobs["chi2_per_dim"].shape = (2048, 2) |
| 363 | + |
| 364 | +# Weighted statistics |
| 365 | +print("=== Blob Statistics ===") |
| 366 | +print(f"Mean chi2_total: {np.average(blobs['chi2_total'], weights=weights):.2f}") |
| 367 | +print(f"Mean radius: {np.average(blobs['radius'], weights=weights):.2f}") |
| 368 | +print(f"Mean max_deviation: {np.average(blobs['max_deviation'], weights=weights):.2f}") |
| 369 | + |
| 370 | +# Extract per-dimension chi2 values |
| 371 | +chi2_dim0 = blobs["chi2_per_dim"][:, 0] |
| 372 | +chi2_dim1 = blobs["chi2_per_dim"][:, 1] |
| 373 | + |
| 374 | +print(f"Mean chi2_dim0: {np.average(chi2_dim0, weights=weights):.2f}") |
| 375 | +print(f"Mean chi2_dim1: {np.average(chi2_dim1, weights=weights):.2f}") |
| 376 | + |
| 377 | +# Evidence and effective sample size |
| 378 | +logz, logz_err = sampler.evidence() |
| 379 | +print(f"\nLog-evidence: {logz:.2f} ± {logz_err:.2f}") |
| 380 | +``` |
| 381 | + |
| 382 | +--- |
| 383 | + |
| 384 | +## Summary |
| 385 | + |
| 386 | +Blobs provide a flexible mechanism to store auxiliary data from your likelihood function: |
| 387 | + |
| 388 | +- **Simple blobs**: Single values (`blobs_dtype=float`) |
| 389 | +- **Multiple blobs**: Structured arrays with named fields |
| 390 | +- **Array blobs**: Multi-dimensional data per sample |
| 391 | +- **Automatic handling**: Preserved during resampling and MCMC steps |
| 392 | + |
| 393 | +The key is to define appropriate `blobs_dtype` that matches what your likelihood returns, and remember to use `return_blobs=True` when retrieving results. |
0 commit comments