Fitting Many Datasets with MetalogFitMethod.Feasible¶
Ordinary least squares can return metalog coefficients whose density goes negative somewhere, especially with many terms. MetalogFitMethod.Feasible returns the best least-squares fit among valid metalogs (the coefficient vector a* of Baucells, Chrisman, Keelin and Xu, 2025), so every fit is a proper distribution.
This notebook fits a batch of datasets with Feasible:
Fixed number of terms: one call to
fit_grid_datasets, which vmapsfitover the datasets.Different numbers of terms: one
fit_grid_datasetscall per term count, then pick the best term count for each dataset by KS distance.
fit_grid cannot be used here: it varies the number of terms inside vmap, while the feasible solver needs it fixed at trace time, so it raises a ValueError for Feasible. fit_grid_datasets compiles once per term count instead.
[1]:
import jax.numpy as jnp
import numpy as np
from scipy.stats import gamma, lognorm, lomax, weibull_min
from metalog_jax.base import (
MetalogBoundedness,
MetalogFitMethod,
MetalogInputData,
MetalogParameters,
)
from metalog_jax.grid_search import extract_metalog, fit_grid, fit_grid_datasets
from metalog_jax.utils import DEFAULT_Y
Data¶
Five positive datasets of 200 samples each, including a heavy tail (Lomax) and a two-peaked mixture. Each is summarized by its empirical quantiles at DEFAULT_Y and stacked into (n_datasets, n_quantiles) arrays. All fits use a lower bound of 0.
[2]:
_rng = np.random.default_rng(5)
_bimodal = np.where(
_rng.random(200) < 0.6, _rng.normal(3.0, 0.5, 200), _rng.normal(7.0, 1.0, 200)
)
_samples = {
"Lognormal": lognorm(s=0.9).rvs(size=200, random_state=1),
"Weibull": weibull_min(c=1.5, scale=2).rvs(size=200, random_state=2),
"Gamma": gamma(a=2).rvs(size=200, random_state=3),
"Lomax": lomax(c=2.5).rvs(size=200, random_state=4),
"Bimodal": _bimodal,
}
dataset_names = list(_samples)
_data = [
MetalogInputData.from_values(jnp.asarray(v), DEFAULT_Y, False)
for v in _samples.values()
]
batched_x = jnp.stack([d.x for d in _data])
batched_y = jnp.stack([d.y for d in _data])
print(f"Batched x shape: {batched_x.shape}")
print(f"Batched y shape: {batched_y.shape}")
Batched x shape: (5, 105)
Batched y shape: (5, 105)
Case 1: Fixed Number of Terms¶
First, fit all datasets with OLS. fit_grid returns every fit, and its feasible mask says which are valid distributions. The mask comes from check_feasibility, which is exact: it finds every local minimum of the density and tests both tails, so it also catches negative density outside the probabilities that fit checks.
[3]:
FIXED_TERMS = 10
params_ols = MetalogParameters(
boundedness=MetalogBoundedness.STRICTLY_LOWER_BOUND,
lower_bound=0.0,
upper_bound=0.0,
method=MetalogFitMethod.OLS,
num_terms=FIXED_TERMS,
)
# fit_grid returns every OLS fit, and marks which are valid distributions
result_ols = fit_grid(batched_x, batched_y, params_ols)
ols_feasible = result_ols.feasible
print(
f"OLS fits that are valid distributions: {int(ols_feasible.sum())} of {len(ols_feasible)}"
)
OLS fits that are valid distributions: 2 of 5
Now fit the same batch with MetalogFitMethod.Feasible. fit_grid_datasets vmaps fit over the datasets, so this is one compiled call for all of them. Where the OLS fit is already feasible, Feasible returns the same coefficients.
[4]:
params_feasible = MetalogParameters(
boundedness=MetalogBoundedness.STRICTLY_LOWER_BOUND,
lower_bound=0.0,
upper_bound=0.0,
method=MetalogFitMethod.Feasible,
num_terms=FIXED_TERMS,
)
result_feasible = fit_grid_datasets(batched_x, batched_y, params_feasible)
feasible_ok = result_feasible.feasible
print(f"Coefficients shape: {result_feasible.metalog.a.shape}")
print(f"KS Distances shape: {result_feasible.ks_dist.shape}")
print()
print(
f"{'Dataset':<10} {'OLS valid':>10} {'OLS KS':>8} {'Feasible valid':>15} {'Feasible KS':>12}"
)
for _i, _name in enumerate(dataset_names):
print(
f"{_name:<10} {str(bool(ols_feasible[_i])):>10} {float(result_ols.ks_dist[_i]):>8.4f}"
f" {str(bool(feasible_ok[_i])):>15} {float(result_feasible.ks_dist[_i]):>12.4f}"
)
Coefficients shape: (5, 10)
KS Distances shape: (5,)
Dataset OLS valid OLS KS Feasible valid Feasible KS
Lognormal False 0.0286 True 0.0286
Weibull True 0.0190 True 0.0190
Gamma True 0.0286 True 0.0286
Lomax False 0.0381 True 0.0381
Bimodal False 0.0571 True 0.0667
An infeasible OLS fit can have a lower KS distance than the feasible fit, because the KS distance only compares quantiles at the data points. It is still not a valid distribution: its density is negative somewhere, so its quantile function is not monotone.
Case 2: Different Numbers of Terms¶
The feasible solver needs the number of terms fixed when it compiles, so loop over the term counts in Python and batch the datasets inside each call. Each term count compiles once. Stacking the KS distances gives a (n_datasets, n_terms) grid.
[5]:
num_terms_list = [4, 6, 8, 10, 12]
# One compiled call per term count; each batches all datasets
results_by_terms = {
_k: fit_grid_datasets(
batched_x, batched_y, params_feasible.replace(num_terms=_k)
)
for _k in num_terms_list
}
ks_grid = jnp.stack([results_by_terms[_k].ks_dist for _k in num_terms_list], axis=1)
print(f"KS grid shape (n_datasets, n_terms): {ks_grid.shape}")
print()
print(f"{'Dataset':<10}" + "".join(f"{_k:>8} terms" for _k in num_terms_list))
for _i, _name in enumerate(dataset_names):
print(f"{_name:<10}" + "".join(f"{float(_v):>14.4f}" for _v in ks_grid[_i]))
KS grid shape (n_datasets, n_terms): (5, 5)
Dataset 4 terms 6 terms 8 terms 10 terms 12 terms
Lognormal 0.0381 0.0286 0.0286 0.0286 0.0190
Weibull 0.0476 0.0286 0.0286 0.0190 0.0190
Gamma 0.0571 0.0381 0.0381 0.0286 0.0190
Lomax 0.0476 0.0286 0.0286 0.0381 0.0381
Bimodal 0.1333 0.0857 0.0667 0.0667 0.0571
Pick the term count with the lowest KS distance for each dataset (ties go to the fewest terms), then use extract_metalog to pull out a ready-to-use Metalog.
[6]:
best_term_idx = jnp.argmin(ks_grid, axis=1)
best_metalogs = {}
print(
f"{'Dataset':<10} {'Terms':>6} {'KS':>8} {'Median':>8} {'P90':>8} {'Valid':>6}"
)
for _i, _name in enumerate(dataset_names):
_k = num_terms_list[int(best_term_idx[_i])]
_m = extract_metalog(results_by_terms[_k], _i)
best_metalogs[_name] = _m
_median, _p90 = _m.ppf(jnp.array([0.5, 0.9]))
_valid = bool(results_by_terms[_k].feasible[_i])
print(
f"{_name:<10} {_k:>6} {float(ks_grid[_i, int(best_term_idx[_i])]):>8.4f}"
f" {float(_median):>8.3f} {float(_p90):>8.3f} {str(_valid):>6}"
)
Dataset Terms KS Median P90 Valid
Lognormal 12 0.0190 1.110 2.944 True
Weibull 12 0.0190 1.451 3.156 True
Gamma 12 0.0190 1.717 4.349 True
Lomax 6 0.0286 0.368 1.785 True
Bimodal 12 0.0571 3.435 7.393 True
extract_metalog returns an ordinary Metalog, so each best fit prints its fitted equation. The fits are bounded below by 0, so Q(y) = 0 + exp(M(y)).
[7]:
for _name, _m in best_metalogs.items():
print(f"--- {_name}")
_m.print()
--- Lognormal
Metalog: 12 terms, Feasible, lower bound 0
Q(y) = 0 + exp(M(y))
M(y) = 0.104676
+ 10.289 ln(y/(1-y))
+ 11.8326 (y - 0.5) ln(y/(1-y))
- 39.3163 (y - 0.5)
- 46.3022 (y - 0.5)^2
- 74.8519 (y - 0.5)^2 ln(y/(1-y))
+ 251.515 (y - 0.5)^3
- 65.0183 (y - 0.5)^3 ln(y/(1-y))
+ 174.852 (y - 0.5)^4
+ 135.934 (y - 0.5)^4 ln(y/(1-y))
- 322.459 (y - 0.5)^5
+ 73.0509 (y - 0.5)^5 ln(y/(1-y))
--- Weibull
Metalog: 12 terms, Feasible, lower bound 0
Q(y) = 0 + exp(M(y))
M(y) = 0.372291
- 1.09448 ln(y/(1-y))
- 1.66134 (y - 0.5) ln(y/(1-y))
+ 5.98058 (y - 0.5)
+ 5.56858 (y - 0.5)^2
+ 3.69382 (y - 0.5)^2 ln(y/(1-y))
- 5.0922 (y - 0.5)^3
+ 10.8938 (y - 0.5)^3 ln(y/(1-y))
- 31.7413 (y - 0.5)^4
+ 3.06623 (y - 0.5)^4 ln(y/(1-y))
- 10.6074 (y - 0.5)^5
- 17.6535 (y - 0.5)^5 ln(y/(1-y))
--- Gamma
Metalog: 12 terms, Feasible, lower bound 0
Q(y) = 0 + exp(M(y))
M(y) = 0.540348
- 11.5595 ln(y/(1-y))
+ 0.307369 (y - 0.5) ln(y/(1-y))
+ 48.0526 (y - 0.5)
- 4.34273 (y - 0.5)^2
+ 69.5203 (y - 0.5)^2 ln(y/(1-y))
- 212.118 (y - 0.5)^3
+ 3.92031 (y - 0.5)^3 ln(y/(1-y))
+ 11.9898 (y - 0.5)^4
- 90.9757 (y - 0.5)^4 ln(y/(1-y))
+ 145.198 (y - 0.5)^5
- 24.9057 (y - 0.5)^5 ln(y/(1-y))
--- Lomax
Metalog: 6 terms, Feasible, lower bound 0
Q(y) = 0 + exp(M(y))
M(y) = -1.00038
- 0.30998 ln(y/(1-y))
- 0.517191 (y - 0.5) ln(y/(1-y))
+ 4.14349 (y - 0.5)
+ 1.36477 (y - 0.5)^2
+ 2.38814 (y - 0.5)^2 ln(y/(1-y))
--- Bimodal
Metalog: 12 terms, Feasible, lower bound 0
Q(y) = 0 + exp(M(y))
M(y) = 1.23413
+ 11.131 ln(y/(1-y))
+ 15.6573 (y - 0.5) ln(y/(1-y))
- 42.9342 (y - 0.5)
- 55.1576 (y - 0.5)^2
- 77.7504 (y - 0.5)^2 ln(y/(1-y))
+ 255.59 (y - 0.5)^3
- 91.0345 (y - 0.5)^3 ln(y/(1-y))
+ 214.656 (y - 0.5)^4
+ 133.956 (y - 0.5)^4 ln(y/(1-y))
- 323.198 (y - 0.5)^5
+ 115.723 (y - 0.5)^5 ln(y/(1-y))
Summary¶
MetalogFitMethod.Feasibleguarantees every fit is a valid distribution, and returns the OLS coefficients unchanged when OLS is already valid.Use
fit_grid_datasetsto fit a batch of datasets withFeasiblein one call.To compare term counts, loop over them in Python with one
fit_grid_datasetscall each.fit_gridraises aValueErrorforFeasible.Every grid result marks its valid fits in
feasible, fromcheck_feasibility, an exact validity check that you can also run on any coefficient vector.