API Reference

This section contains the complete API reference for metalog-jax.

metalog_jax

Main Module

Metalog JAX library.

Metalog Distribution

Metalog distribution implementations and fitting functions.

This module provides the main metalog distribution classes and fitting functions:

Classes:

Metalog: Standard metalog distribution fitted via regression methods. SPTMetalog: Symmetric Percentile Triplet metalog with closed-form coefficients. GridResult: Results container for grid search and hyperparameter optimization.

Functions:

fit: Fit a standard metalog distribution using configurable regression methods. fit_spt_metalog: Fit an SPT metalog distribution using closed-form formulas.

The fit function uses a dispatch table pattern to route to the appropriate regression implementation based on the MetalogFitMethod specified in MetalogParameters:

  • MetalogFitMethod.OLS -> metalog_jax.regression.ols.fit_ordinary_least_squares

  • MetalogFitMethod.Lasso -> metalog_jax.regression.lasso.fit_lasso

  • MetalogFitMethod.Feasible -> metalog_jax.feasibility.a_star.fit_feasible

See also

metalog_jax.base: Base classes and parameter configurations. metalog_jax.regression: Regression implementations for fitting.

class metalog_jax.metalog.SPTMetalog(metalog_params, a)[source]

Bases: MetalogBase

Symmetric Percentile Triplet (SPT) metalog distribution.

SPTMetalog is a specialized metalog implementation that uses exactly three terms fitted from three quantiles: the alpha-quantile, median (0.5), and (1-alpha)-quantile, where alpha < 0.5. This method provides a computationally efficient way to fit metalog distributions with closed-form coefficient solutions, avoiding the need for optimization or linear regression.

The SPT approach is particularly useful when you need: - A quick approximation using only three data points - Guaranteed feasibility through explicit feasibility checks - Analytical coefficient formulas for different boundedness types - A distribution that matches specific symmetric or asymmetric quantile triplets

Unlike the standard Metalog class which uses OLS or LASSO regression to fit arbitrary numbers of terms, SPTMetalog always has exactly 3 terms and computes coefficients directly from the quantile triplet. This makes it faster for small datasets but less flexible than the full metalog approach.

Variables:

  • metalog_params (SPTMetalogParameters) – Configuration parameters specific to SPT fitting, including boundedness type, boundary values, and the alpha parameter that determines which quantiles are used (alpha, 0.5, 1-alpha).

  • a (Union[Array, ndarray, bool, number, float, int]) – Coefficient vector of shape (3,) representing the three fitted metalog distribution parameters. These are computed using closed-form formulas specific to the boundedness type. A list or tuple is converted to a float JAX array on construction.

Examples

Fit an unbounded SPT metalog from data:

>>> import jax.numpy as jnp
>>> from metalog_jax.base import MetalogBoundedness, SPTMetalogParameters
>>> from metalog_jax.metalog import SPTMetalog, fit_spt_metalog
>>>
>>> # Generate sample data
>>> data = jnp.array([1.2, 2.3, 2.8, 3.5, 4.1, 5.6, 6.2, 7.8, 9.1])
>>>
>>> # Configure SPT parameters
>>> params = SPTMetalogParameters(
...     boundedness=MetalogBoundedness.UNBOUNDED,
...     alpha=0.1,  # Use 10th, 50th, and 90th percentiles
...     lower_bound=0.0,
...     upper_bound=0.0,
... )
>>>
>>> # Fit the SPT metalog
>>> spt = fit_spt_metalog(data, params)
>>>
>>> # The fit passes through the data's 10th, 50th and 90th percentiles
>>> spt.ppf(jnp.array([0.1, 0.5, 0.9])).round(2).tolist()
[2.08, 4.1, 8.06]
>>>
>>> # Summary statistics are properties
>>> mean, std = spt.mean, spt.std
>>>
>>> # Serialize and load it back (save and load do the same with a file)
>>> SPTMetalog.loads(spt.dumps()) == spt
True

Show the fitted quantile function. print writes one line (wrapped here), repr the same line at full precision, and the print method one term per line. With a lower bound of 0, Q(y) = 0 + exp(M(y)):

>>> lower = fit_spt_metalog(
...     data, params.replace(boundedness=MetalogBoundedness.STRICTLY_LOWER_BOUND)
... )
>>> print(lower)
SPTMetalog(3 terms, alpha = 0.1, lower bound 0): Q(y) = 0 + exp(M(y)),
M(y) = 1.41099 + 0.30824 ln(y/(1-y)) - 0.00153176 (y - 0.5) ln(y/(1-y))
>>> lower.print()
SPTMetalog: 3 terms, alpha = 0.1, lower bound 0
  Q(y) = 0 + exp(M(y))
  M(y) = 1.41099
         + 0.30824 ln(y/(1-y))
         - 0.00153176 (y - 0.5) ln(y/(1-y))
>>> lower.print(m=False, verbose=True)
SPTMetalog: 3 terms, alpha = 0.1, lower bound 0
  Q(y) = 0 + exp(1.41099 + 0.30824 ln(y/(1-y))
         - 0.00153176 (y - 0.5) ln(y/(1-y)))

Note

SPT metalog distributions may fail feasibility checks if the data does not satisfy specific constraints on the symmetry ratio r. The feasibility constraints depend on the boundedness type and alpha parameter. When feasibility fails, fit_spt_metalog will raise an AssertionError.

The SPT method is described in Keelin (2016) as a special case of the metalog family that enables rapid approximation with minimal data.

See also

Metalog: Standard metalog implementation using regression fitting. MetalogBase: Base class defining the common interface. fit_spt_metalog: Function to fit SPT metalog distributions from data.

References

Keelin, T. W. (2016). The Metalog Distributions. Decision Analysis, 13(4), 243-277. https://doi.org/10.1287/deca.2016.0338

Parameters:
metalog_params: SPTMetalogParameters
a: Union[Array, ndarray, bool, number, float, int]
__eq__(other)[source]

Check equality between two SPTMetalog distribution instances.

Delegates to the parent MetalogBase.__eq__ method.

Parameters:

other (object) – The object to compare against.

Returns:

True if the two instances are equal, False otherwise.

Return type:

bool

__repr__()[source]

Show the SPTMetalog and its quantile function on one line, at full precision.

Delegates to MetalogBase.__repr__, which flax.struct.dataclass would otherwise replace with a list of the dataclass fields.

Returns:

The one-line equation, or a summary if the model is traced or

batched.

Return type:

str

property num_terms: int

Get the number of terms in the metalog expansion.

Returns:

3 - an SPT Metalog always has three terms.

property method: NoReturn

An SPT metalog has no fit method.

fit_spt_metalog computes the coefficients in closed form (Keelin 2016, Propositions 1-4) rather than by regression, so SPTMetalogParameters has no method.

Raises:

AttributeError – Always, so hasattr(spt, "method") is False.

__init__(metalog_params, a)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.metalog.Metalog(metalog_params, a)[source]

Bases: MetalogBase

Standard metalog distribution fitted via regression.

Metalog is the full-featured metalog implementation that extends MetalogBase and uses ordinary least squares (OLS) or LASSO regression to fit metalog distributions with arbitrary numbers of terms (typically 2-30). This approach provides maximum flexibility and accuracy for modeling complex distributions from larger datasets.

Unlike SPTMetalog which uses exactly 3 terms with closed-form coefficient formulas, the standard Metalog class: - Supports any number of terms (configurable via num_terms parameter) - Uses regression-based fitting (OLS or LASSO methods) - Can model more complex distribution shapes with higher-order terms - Requires more data for reliable fitting (recommended minimum: 3 x num_terms) - May encounter numerical issues with very high term counts

This class inherits all distribution methods from MetalogBase including ppf (quantile function), cdf, pdf, and statistical moments (mean, variance, std, skewness, kurtosis). It also provides serialization capabilities for saving and loading fitted distributions.

Variables:

  • metalog_params (MetalogParameters) – Configuration parameters for the standard metalog fit, including boundedness type (UNBOUNDED, STRICTLY_LOWER_BOUND, STRICTLY_UPPER_BOUND, or BOUNDED), boundary values, fitting method (OLS or LASSO), and number of terms in the expansion.

  • a (Union[Array, ndarray, bool, number, float, int]) – Coefficient vector of shape (num_terms,) representing the fitted metalog distribution parameters. Each coefficient corresponds to one basis function in the metalog expansion. These are computed via regression from the input data. A list or tuple is converted to a float JAX array on construction.

Examples

Fit an unbounded metalog with 9 terms:

>>> import jax.numpy as jnp
>>> from metalog_jax.base import MetalogBoundedness, MetalogFitMethod
>>> from metalog_jax.base import MetalogInputData, MetalogParameters
>>> from metalog_jax.metalog import fit
>>> from metalog_jax.utils import DEFAULT_Y
>>>
>>> # Sample data
>>> samples = jnp.array([1.2, 2.3, 2.8, 3.5, 4.1, 5.6, 6.2, 7.8, 9.1, 10.5])
>>> data = MetalogInputData.from_values(
...     samples, DEFAULT_Y, precomputed_quantiles=False
... )
>>>
>>> # Configure metalog parameters
>>> params = MetalogParameters(
...     boundedness=MetalogBoundedness.UNBOUNDED,
...     method=MetalogFitMethod.Feasible,
...     num_terms=9,
...     lower_bound=0.0,
...     upper_bound=0.0,
... )
>>>
>>> # Fit the metalog distribution
>>> m = fit(data, params)
>>>
>>> # Evaluate the quantile function at probability 0.75
>>> round(float(m.ppf(jnp.array([0.75]))[0]), 2)
7.4
>>>
>>> # Summary statistics are properties
>>> mean, variance, skewness = m.mean, m.var, m.skewness

Show the fitted quantile function of a 5-term OLS fit bounded on (0, 15). print writes one line (wrapped here), repr the same line at full precision, and the print method one term per line, with options:

>>> bounded = fit(
...     data,
...     MetalogParameters(
...         boundedness=MetalogBoundedness.BOUNDED,
...         method=MetalogFitMethod.OLS,
...         num_terms=5,
...         lower_bound=0.0,
...         upper_bound=15.0,
...     ),
... )
>>> print(bounded)
Metalog(5 terms, OLS, bounded on (0, 15)): Q(y) = (0 + 15 exp(M(y))) /
(1 + exp(M(y))), M(y) = -0.754464 + 0.0719273 ln(y/(1-y)) - 0.110759
(y - 0.5) ln(y/(1-y)) + 2.45794 (y - 0.5) + 1.1187 (y - 0.5)^2
>>> bounded.print()
Metalog: 5 terms, OLS, bounded on (0, 15)
  Q(y) = (0 + 15 exp(M(y))) / (1 + exp(M(y)))
  M(y) = -0.754464
         + 0.0719273 ln(y/(1-y))
         - 0.110759 (y - 0.5) ln(y/(1-y))
         + 2.45794 (y - 0.5)
         + 1.1187 (y - 0.5)^2
>>> bounded.print(m=False, verbose=True)
Metalog: 5 terms, OLS, bounded on (0, 15)
  Q(y) = (0 + 15 exp(-0.754464 + 0.0719273 ln(y/(1-y))
         - 0.110759 (y - 0.5) ln(y/(1-y)) + 2.45794 (y - 0.5)
         + 1.1187 (y - 0.5)^2))
         / (1 + exp(-0.754464 + 0.0719273 ln(y/(1-y))
         - 0.110759 (y - 0.5) ln(y/(1-y)) + 2.45794 (y - 0.5)
         + 1.1187 (y - 0.5)^2))

Note

The standard Metalog class is the recommended choice for most use cases involving sufficient data. Use SPTMetalog only when you need a quick three-term approximation or when you have exactly three representative quantiles. For datasets with at least 10-30 observations, the standard metalog typically provides superior fit quality.

Higher term counts generally improve accuracy but require more data and may cause overfitting or numerical instability. As a rule of thumb, ensure your sample size is at least 3 times the number of terms.

See also

SPTMetalog: Specialized three-term metalog using closed-form coefficients. MetalogBase: Base class defining the common distribution interface. fit_metalog: Function to fit standard metalog distributions from data.

References

Keelin, T. W. (2016). The Metalog Distributions. Decision Analysis, 13(4), 243-277. https://doi.org/10.1287/deca.2016.0338

Parameters:
metalog_params: MetalogParameters
a: Union[Array, ndarray, bool, number, float, int]
__eq__(other)[source]

Check equality between two Metalog distribution instances.

Delegates to the parent MetalogBase.__eq__ method.

Parameters:

other (object) – The object to compare against.

Returns:

True if the two instances are equal, False otherwise.

Return type:

bool

__repr__()[source]

Show the Metalog and its quantile function on one line, at full precision.

Delegates to MetalogBase.__repr__, which flax.struct.dataclass would otherwise replace with a list of the dataclass fields.

Returns:

The one-line equation, or a summary if the model is traced or

batched.

Return type:

str

property num_terms: int

Get the number of terms in the metalog expansion.

Returns:

Integer representing the number of basis functions used in the metalog approximation. Higher values generally provide better accuracy but require more data.

property method: MetalogFitMethod

Get the method used to fit the metalog distribution.

Returns:

MetalogFitMethod enum of regression methods for fitting metalog distribution.

__init__(metalog_params, a)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

metalog_jax.metalog.fit(data, metalog_params, regression_hyperparams=None)[source]

Fit a metalog distribution to data using the specified configuration.

Estimates the metalog distribution coefficients by solving a linear regression problem that maps the basis functions (target matrix) to the transformed quantiles. The regression method specified in metalog_params determines the fitting approach, and optional regression_hyperparams allow fine-tuning of regularization settings.

This function processes a MetalogInputData instance containing validated input data (quantiles and probability levels) and fits a metalog distribution according to the specified parameters. The data must be created using MetalogInputData.from_values() which performs comprehensive validation.

Parameters:

  • data (MetalogInputData) – Validated input data created via MetalogInputData.from_values(). Contains: - x: Quantile values (precomputed or computed from raw samples) - y: Probability levels in (0, 1) - precomputed_quantiles: Flag indicating data type Direct instantiation of MetalogInputData is prevented by __post_init__ validation to ensure data integrity.

  • metalog_params (MetalogParameters) – Configuration parameters specifying the distribution characteristics (boundedness, bounds, number of terms) and the regression method to use for fitting. The method field determines which regression approach is used: - MetalogFitMethod.OLS: Ordinary Least Squares (no regularization) - MetalogFitMethod.Lasso: L1 regularization - MetalogFitMethod.Feasible: Best least-squares fit among valid metalogs

  • regression_hyperparams (RegularizedParameters) – Optional regularization hyperparameters for controlling the fitting process when using LASSO regression method. Must be None for OLS and Feasible. If None, default hyperparameters are used. Must be an instance of LassoParameters for MetalogFitMethod.Lasso.

Return type:

Metalog

Returns:

Metalog containing the fitted coefficient vector and the configuration parameters used to fit the metalog distribution. The returned object is validated via assert_fitted() to ensure the fit produces a feasible distribution with strictly positive PDF values.

Raises:

  • TypeError – If metalog_params.method is not a valid MetalogFitMethod instance, or if regression_hyperparams is given for OLS or Feasible.

  • checkify.JaxRuntimeError – If the fitted distribution produces non-positive PDF values, indicating an infeasible fit. This validation is performed by assert_fitted() before returning.

  • checkify.JaxRuntimeError – For OLS and Feasible, if the basis matrix is rank deficient so the fit is not unique (e.g. n = k symmetric probabilities with 7, 11, 15, … terms; Baucells et al. 2025, Proposition 2), or if Feasible cannot certify its result.

Example

Basic usage with OLS (no regularization):

>>> import jax.numpy as jnp
>>> from metalog_jax.base import MetalogInputData, MetalogParameters
>>> from metalog_jax.base import MetalogBoundedness, MetalogFitMethod
>>> from metalog_jax.metalog import fit
>>>
>>> # Create validated input data using from_values()
>>> raw_data = jnp.array([1.2, 2.3, 2.8, 3.5, 4.1, 5.6])
>>> data = MetalogInputData.from_values(
...     x=raw_data,
...     y=jnp.array([0.1, 0.25, 0.5, 0.75, 0.9, 0.95]),
...     precomputed_quantiles=False  # Raw samples, not precomputed quantiles
... )
>>>
>>> # Configure metalog parameters with OLS
>>> metalog_params = MetalogParameters(
...     boundedness=MetalogBoundedness.UNBOUNDED,
...     method=MetalogFitMethod.OLS,
...     lower_bound=0.0,
...     upper_bound=0.0,
...     num_terms=5
... )
>>>
>>> # Fit the metalog distribution
>>> m = fit(data, metalog_params)

Note

  • The regression_hyperparams parameter is only applicable for LASSO method. Passing it with OLS or Feasible raises a TypeError.

  • If regression_hyperparams is None, sensible defaults are used for each method.

  • For production use, consider tuning regularization hyperparameters via cross-validation to prevent overfitting while maintaining good fit quality.

See also

fit_spt_metalog: Alternative fitting method using Symmetric Percentile Triplet. MetalogInputData.from_values: Required method for creating validated input data. metalog_jax.regression.lasso.LassoParameters: Hyperparameters for LASSO regression. metalog_jax.feasibility.best_feasible_fit: Best feasible fit (Baucells et al.).

References

Keelin, T. W. (2016). The Metalog Distributions. Decision Analysis, 13(4), 243-277. https://doi.org/10.1287/deca.2016.0338

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

metalog_jax.metalog.fit_spt_metalog(array, spt_metalog_params, *, precomputed_quantiles=False)[source]

Fit a Symmetric Percentile Triplet (SPT) metalog to data or assessed quantiles.

This function fits a 3-term metalog distribution using the Symmetric Percentile Triplet (SPT) parameterization method described in Keelin (2016). Unlike the standard metalog fit which uses many quantiles, the SPT method uses exactly three symmetric quantiles: alpha, 0.5 (median), and (1 - alpha), where 0 < alpha < 0.5.

The SPT approach provides a quick, analytical solution for fitting a metalog distribution when: - You have limited data or only three quantiles available - You want a simple, closed-form solution without regression - You need a fast approximation with 3 terms only

The method computes the 3 metalog coefficients directly from the three quantile values without requiring least squares regression. The formulas vary based on the boundedness type (unbounded, semi-bounded, or bounded).

Parameters:

  • array (Union[Array, ndarray, bool, number, float, int]) – Input data array with at least 3 elements from which to compute empirical quantiles. The array will be used to compute the alpha-th, 50th, and (1-alpha)-th percentiles. If precomputed_quantiles is True, it is instead the triplet [q_alpha, q_0.5, q_1-alpha] itself, for example an expert’s assessed 10th, 50th and 90th percentiles.

  • spt_metalog_params (SPTMetalogParameters) –

    Configuration parameters containing:

    • alpha: Lower percentile parameter in (0, 0.5). Common values are 0.1 (10-50-90 percentiles) or 0.25 (25-50-75 percentiles/IQR).

    • boundedness: Domain constraint type (UNBOUNDED, STRICTLY_LOWER_BOUND, STRICTLY_UPPER_BOUND or BOUNDED).

    • lower_bound: Lower boundary value (used for STRICTLY_LOWER_BOUND and BOUNDED).

    • upper_bound: Upper boundary value (used for STRICTLY_UPPER_BOUND and BOUNDED).

  • precomputed_quantiles (bool) – Whether array holds the three quantiles themselves rather than data. Defaults to False.

Returns:

Fitted SPT metalog distribution containing:

  • metalog_params: The input configuration parameters

  • a: Coefficient vector of length 3 containing the fitted metalog parameters [a1, a2, a3] that define the quantile function

Return type:

SPTMetalog

Raises:

  • AssertionError – If array has fewer than 3 elements, is not rank 1, or contains non-numeric values.

  • AssertionError – If precomputed_quantiles is True and array doesn’t have shape (3,).

  • AssertionError – If alpha is not positive or alpha >= 0.5.

  • AssertionError – If the computed quantiles violate feasibility constraints (q_alpha < median < q_complement).

  • AssertionError – If a quantile lies outside the support: at or below lower_bound (STRICTLY_LOWER_BOUND and BOUNDED), or at or above upper_bound (STRICTLY_UPPER_BOUND and BOUNDED).

  • AssertionError – If the triplet is infeasible for the boundedness type (Keelin 2016, Propositions 2-4, and the mirror image of Proposition 3 for STRICTLY_UPPER_BOUND).

  • AssertionError – If the coefficients are not finite, for example because the data or bounds contain NaN.

  • checkify.JaxRuntimeError – Under jax.vmap, for the value checks above (array shape and type checks still raise AssertionError).

Example

>>> import jax.numpy as jnp
>>> from metalog_jax.metalog import fit_spt_metalog, SPTMetalogParameters
>>> from metalog_jax.metalog import MetalogBoundedness
>>> # Generate sample data
>>> data = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0])
>>> # Configure SPT parameters using 10-50-90 percentiles
>>> params = SPTMetalogParameters(
...     alpha=0.1,
...     boundedness=MetalogBoundedness.UNBOUNDED,
...     lower_bound=0.0,
...     upper_bound=1.0
... )
>>> # Fit the SPT metalog
>>> spt_metalog = fit_spt_metalog(data, params)
>>> # The result contains 3 coefficients
>>> assert len(spt_metalog.a) == 3
>>>
>>> # Fit several datasets at once
>>> import jax
>>> batch = jnp.stack([data, 2 * data])
>>> jax.vmap(lambda d: fit_spt_metalog(d, params).a)(batch).shape
(2, 3)
>>>
>>> # Fit an expert's 10th, 50th and 90th percentile assessments directly
>>> assessed = jnp.array([2.08, 4.1, 8.06])
>>> spt = fit_spt_metalog(assessed, params, precomputed_quantiles=True)
>>> spt.ppf(jnp.array([0.1, 0.5, 0.9])).round(2).tolist()
[2.08, 4.1, 8.06]

Note

  • The SPT method always produces exactly 3 metalog terms, regardless of data size. For more flexible fits with more terms, use the standard fit() function.

  • The method performs multiple feasibility checks to ensure the quantiles and parameters produce a valid probability distribution.

  • Pass precomputed_quantiles=True for quantiles you already have. By default a 3-element array is treated as data, and jnp.quantile interpolates between its three values, so the fit would not pass through them.

  • For unbounded distributions, the median position ratio r must satisfy specific constraints related to k_alpha to ensure feasibility, where r = (median - q_alpha) / (q_complement - q_alpha).

  • Works under jax.vmap and jax.jit over the data, alpha and bounds, and the parameters can be passed as a jit argument (boundedness is a static field, so each boundedness type compiles once). Inside transforms the value checks use checkify, so wrap a jitted call with checkify.checkify and call throw() on the returned error.

References

Keelin, T. W. (2016). The Metalog Distributions. Decision Analysis, 13(4).

class metalog_jax.metalog.GridResult(metalog, ks_dist, feasible=None)[source]

Bases: object

Results from metalog grid search or hyperparameter optimization.

GridResult encapsulates the outcome of fitting a metalog distribution along with a goodness-of-fit metric. This dataclass is typically used when performing grid searches over hyperparameters (e.g., number of terms, regularization penalties) or when comparing multiple metalog configurations to select the best-fitting distribution.

The primary use case is hyperparameter tuning via grid search with JAX’s vmap to fit multiple metalog configurations in parallel, then selecting the best configuration based on the Kolmogorov-Smirnov distance.

The class uses Flax’s struct.dataclass decorator to ensure immutability and compatibility with JAX transformations (jit, vmap, grad), making it suitable for use in vectorized grid search operations.

Variables:

  • metalog (Metalog | SPTMetalog) – The fitted metalog distribution. Can be either a standard Metalog (fitted via regression with configurable terms) or SPTMetalog (fitted using Symmetric Percentile Triplet with exactly 3 terms). Contains: - metalog_params: Configuration parameters (boundedness, method, etc.) - a: Coefficient vector representing the fitted distribution

  • ks_dist (Union[float, int]) – Scalar Kolmogorov-Smirnov distance measuring the goodness-of-fit between the fitted metalog distribution and the input random variable. This is the maximum absolute difference between the empirical CDF of the input data and the fitted metalog’s CDF: - 0: Perfect fit (empirical CDFs are identical) - 1: Worst possible fit (completely non-overlapping distributions) - Lower values indicate better fit quality Used to select the best metalog configuration from a grid search.

  • feasible (Union[Array, ndarray, bool, number, None]) – Whether each fit is a valid distribution, with the same shape as ks_dist: True when M'(y) >= 0 on all of (0, 1), as metalog_jax.feasibility.check_feasibility decides exactly (Baucells et al. 2025, Proposition 5 with Algorithm 1). An infeasible fit’s quantile function isn’t increasing, so its ks_dist isn’t a distance between two distributions, and find_best_config and extract_best_from_grid skip it. Every grid function fills this in; it’s None for a GridResult built without it.

Note

  • This class is immutable due to the struct.dataclass decorator

  • The ks_dist attribute provides a single goodness-of-fit metric for comparing different metalog configurations

  • When used with jax.vmap, the metalog and ks_dist fields will be batched arrays/PyTrees, allowing vectorized comparison

  • Lower KS distances indicate better fit quality, but extremely low values may indicate overfitting (especially with high term counts)

  • For production use, consider cross-validation instead of or in addition to KS distance for hyperparameter selection

See also

ks_distance: Function to compute the Kolmogorov-Smirnov distance. fit: Standard metalog fitting function. fit_spt_metalog: SPT metalog fitting function. Metalog: Standard metalog distribution class. SPTMetalog: Symmetric Percentile Triplet metalog class.

References

Kolmogorov, A. N. (1933). “Sulla determinazione empirica di una legge di distribuzione”. Giornale dell’Istituto Italiano degli Attuari, 4: 83-91.

Keelin, T. W. (2016). The Metalog Distributions. Decision Analysis, 13(4), 243-277. https://doi.org/10.1287/deca.2016.0338

Parameters:
metalog: Metalog | SPTMetalog
ks_dist: Union[float, int]
feasible: Union[Array, ndarray, bool, number, None] = None
__init__(metalog, ks_dist, feasible=None)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

Base Classes and Configuration

Core

Core MetalogBase class for metalog distributions.

This module contains the MetalogBase class, which is the abstract base class for all metalog distribution implementations. It provides the common interface and core functionality shared by Metalog and SPTMetalog classes defined in metalog_jax.metalog.

The MetalogBase class includes:

  • Quantile function (PPF) computation

  • Probability density function (PDF) computation

  • Cumulative distribution function (CDF) computation

  • Statistical moments (mean, median, variance, std, mode)

  • Random variable sampling (rvs)

  • Serialization/deserialization (save, load, dumps, loads)

  • Visualization (plot)

  • Display of the fitted equation (str, repr, _repr_latex_, print)

See also

metalog_jax.base.data: Input data container classes. metalog_jax.base.display: Text and LaTeX formatting of the fitted equation. metalog_jax.base.enums: Enumeration types for boundedness and fit methods. metalog_jax.base.parameters: Parameter configuration classes. metalog_jax.metalog: Metalog and SPTMetalog distribution classes.

class metalog_jax.base.core.MetalogBase(metalog_params, a)[source]

Bases: object

Base class for metalog distribution implementations.

This abstract base class provides the common interface and core functionality shared by all metalog distribution types (Metalog and SPTMetalog). It implements the fundamental metalog distribution operations including quantile function (PPF), probability density function (PDF), cumulative distribution function (CDF), and statistical moments.

Subclasses:
Metalog: Standard metalog fitted via regression (OLS, Ridge, ElasticNet, Lasso).

Defined in metalog_jax.metalog.

SPTMetalog: Symmetric Percentile Triplet metalog with closed-form coefficients.

Defined in metalog_jax.metalog.

Variables:

  • metalog_params (SPTMetalogParameters | MetalogParameters) – Configuration parameters defining the distribution type. Can be MetalogParameters (for Metalog) or SPTMetalogParameters (for SPTMetalog). Defined in metalog_jax.base.parameters.

  • a (Union[Array, ndarray, bool, number, float, int]) – Coefficient vector defining the metalog quantile function. Shape is (num_terms,) where num_terms depends on the configuration. A list or tuple is converted to a float JAX array on construction.

Example

MetalogBase is not instantiated directly. Use Metalog or SPTMetalog:

>>> import jax.numpy as jnp
>>> from metalog_jax.base import MetalogBoundedness, MetalogFitMethod
>>> from metalog_jax.base import MetalogInputData, MetalogParameters
>>> from metalog_jax.metalog import fit
>>> from metalog_jax.utils import DEFAULT_Y
>>>
>>> samples = jnp.array([1.2, 2.3, 2.8, 3.5, 4.1, 5.6, 6.2, 7.8, 9.1, 10.5])
>>> data = MetalogInputData.from_values(
...     samples, DEFAULT_Y, precomputed_quantiles=False
... )
>>> params = MetalogParameters(
...     boundedness=MetalogBoundedness.UNBOUNDED,
...     method=MetalogFitMethod.Feasible,
...     num_terms=5, lower_bound=0.0, upper_bound=0.0
... )
>>> m = fit(data, params)  # Returns a Metalog
>>> isinstance(m, MetalogBase)
True

See also

metalog_jax.metalog.Metalog: Standard metalog distribution class. metalog_jax.metalog.SPTMetalog: SPT metalog distribution class. metalog_jax.base.parameters.MetalogParameters: Configuration for Metalog. metalog_jax.base.parameters.SPTMetalogParameters: Configuration for SPTMetalog.

Parameters:
__post_init__()[source]

Coerce a list or tuple of coefficients to a float JAX array.

Every method assumes a is an array (jnp.result_type, .shape, .tolist(), @), so a plain sequence such as hand-copied coefficients would otherwise fail deep inside JAX (issue #9). Only lists and tuples are converted: pytree unflattening under jit/vmap/tree.map re-runs __init__ with tracers, ShapeDtypeStruct or None as a, and those must pass through untouched.

Return type:

None

__eq__(other)[source]

Check equality between two metalog distribution instances.

Compares the metalog parameters and coefficient arrays to determine equality. For JAX arrays, uses jnp.allclose for numerical comparison.

Parameters:

other (object) – The object to compare against.

Returns:

True if the two instances are equal, False otherwise.

Return type:

bool

__str__()[source]

Show the model and its quantile function on one line.

The coefficients and bounds have 6 significant digits. A traced or batched model shows a one-line summary instead.

Returns:

The one-line equation, such as

Metalog(3 terms, OLS, unbounded): Q(y) = 1.5 - 0.25 ln(y/(1-y)) ..., or a one-line summary if the model is traced or batched.

Return type:

str

__repr__()[source]

Show the model and its quantile function on one line, at full precision.

The same line as str, with each number written as Python’s shortest repr(float), so the coefficients parse back to a exactly. Each subclass defines __repr__ to call this one, because flax.struct.dataclass would otherwise generate its own.

Returns:

The one-line equation, or a one-line summary if the model is traced

or batched.

Return type:

str

print(*, header=True, q=True, m=True, verbose=False, precision=6, file=None)[source]

Print the quantile function, one term per line.

For a bounded or semi-bounded model, Q(y) maps M(y) back to the support, and the options choose which lines appear. An unbounded model’s Q(y) is M(y), so its terms print once, labelled Q(y) (or M(y) when q is False). A traced or batched model prints a one-line summary instead. The builtin print(m) still prints str(m).

Parameters:

  • header (bool) – Print a first line describing the model, such as Metalog: 5 terms, OLS, bounded on (0, 15), and indent the equations under it. Without it, the equations start at the left margin.

  • q (bool) – Print the quantile function Q(y).

  • m (bool) – Print M(y), one term per line. For an unbounded model, m=False keeps the terms on Q(y)’s line.

  • verbose (bool) – Write Q(y) with M(y)’s terms in place of M(y), so it reads on its own. A bounded Q(y) uses M(y) twice, so the terms appear twice.

  • precision (int) – Significant digits of the coefficients and bounds.

  • file (TextIO | None) – Where to write, such as an open file or io.StringIO; standard output by default.

Raises:

ValueError – If q and m are both False, verbose is True while q is False, or precision isn’t a positive integer.

Return type:

None

property boundedness: MetalogBoundedness

Get the boundedness type of the metalog distribution.

Returns:

The boundedness type (UNBOUNDED, STRICTLY_LOWER_BOUND,

STRICTLY_UPPER_BOUND, or BOUNDED).

Return type:

MetalogBoundedness

property lower_bound: Union[float, int]

Get the lower bound of the metalog distribution.

Returns:

The lower bound value. Only meaningful when boundedness

is STRICTLY_LOWER_BOUND or BOUNDED.

Return type:

chex.Scalar

property upper_bound: Union[float, int]

Get the upper bound of the metalog distribution.

Returns:

The upper bound value. Only meaningful when boundedness

is STRICTLY_UPPER_BOUND or BOUNDED.

Return type:

chex.Scalar

property num_terms: int

Get the number of terms in the metalog expansion.

Returns:

The number of terms (coefficients) in the metalog quantile function.

Return type:

int

property method: MetalogFitMethod

Get the fit method used for the metalog distribution.

Returns:

The regression method (OLS or Lasso).

Return type:

MetalogFitMethod

assert_fitted()[source]

Validate that the metalog distribution has been fitted and is feasible.

Performs two validation checks:

  1. Verifies that coefficients exist (distribution has been fit).

  2. Computes the PDF at standard quantile points to verify feasibility, which internally calls assert_fit_valid to ensure all PDF values are strictly positive.

This method is vmap-compatible through the use of jax.experimental.checkify in the underlying PDF validation.

Raises:

  • NotFittedError – If the distribution has not been fitted (no coefficients).

  • checkify.JaxRuntimeError – If the fitted distribution produces non-positive PDF values, indicating an infeasible fit.

Return type:

None

static assert_fit_valid(m)[source]

Validate that PDF values are strictly positive.

Checks that all probability density values are positive, which is a necessary condition for a valid probability distribution. A metalog fit that produces non-positive PDF values is theoretically infeasible.

This method is vmap-compatible by using jax.experimental.checkify for runtime error checking inside JAX transforms.

Parameters:

m (Union[Array, ndarray, bool, number, float, int]) – Array of PDF values to validate. All elements must be > 0.

Raises:

checkify.JaxRuntimeError – If any PDF value is less than or equal to zero, indicating an infeasible metalog fit.

Return type:

None

static strictly_lower_bound_quantile_transform(x, lower_bound)[source]

Apply logarithmic transformation for strictly lower-bounded quantiles.

Transforms quantiles from (lower_bound, inf) to (-inf, inf) using log(x - L).

Parameters:
Returns:

Transformed quantile values on the unbounded scale.

Return type:

Union[Array, ndarray, bool, number, float, int]

static strictly_upper_bound_quantile_transform(x, upper_bound)[source]

Apply negative logarithmic transformation for strictly upper-bounded quantiles.

Transforms quantiles from (-inf, upper_bound) to (-inf, inf) using -log(U - x).

Parameters:
Returns:

Transformed quantile values on the unbounded scale.

Return type:

Union[Array, ndarray, bool, number, float, int]

static bounded_quantile_transform(x, lower_bound, upper_bound)[source]

Apply logit-like transformation for fully bounded quantiles.

Transforms quantiles from (lower_bound, upper_bound) to (-inf, inf) using log((x - L) / (U - x)).

Parameters:
Returns:

Transformed quantile values on the unbounded scale.

Return type:

Union[Array, ndarray, bool, number, float, int]

static quantile_transform(x, boundedness, lower_bound, upper_bound)[source]

Transform quantiles to the scale on which the metalog is unbounded.

Applies the transform that boundedness selects: none for UNBOUNDED, then strictly_lower_bound_quantile_transform, strictly_upper_bound_quantile_transform or bounded_quantile_transform. boundedness is a static parameter field, so this dispatches in Python and works inside jax.jit and jax.vmap.

Parameters:
Returns:

Transformed quantile values on the unbounded scale.

Return type:

Union[Array, ndarray, bool, number, float, int]

static inverse_quantile_transform(m, boundedness, lower_bound, upper_bound)[source]

Map values of the unbounded metalog back to the distribution’s support.

The inverse of quantile_transform (Keelin 2016, Section 4): L + e^M for STRICTLY_LOWER_BOUND, U - e^(-M) for STRICTLY_UPPER_BOUND and (L + U e^M) / (1 + e^M) for BOUNDED.

Parameters:
Returns:

Quantile values in the distribution’s support.

Return type:

Union[Array, ndarray, bool, number, float, int]

static delta_term(x)[source]

Compute the delta term (x - 0.5) used in metalog basis functions.

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1).

Returns:

The delta term values (x - 0.5).

Return type:

Union[Array, ndarray, bool, number, float, int]

static log_term(x)[source]

Compute the logit transformation log(x / (1 - x)).

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1).

Returns:

The logit-transformed values.

Return type:

Union[Array, ndarray, bool, number, float, int]

static product_term(x)[source]

Compute the product term x * (1 - x).

This term appears in the derivative of the logit function and is used in PDF computation.

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1).

Returns:

The product term values x * (1 - x).

Return type:

Union[Array, ndarray, bool, number, float, int]

static get_quantiles(data, metalog_params)[source]

Extract and transform quantiles from validated input data.

Applies the appropriate boundedness transformation to convert quantiles from their natural domain to the unbounded (-inf, inf) scale used internally by the metalog.

Parameters:

  • data (MetalogInputData) – Validated input data containing quantile values (x) and probability levels (y).

  • metalog_params (MetalogParameters) – Configuration parameters specifying boundedness type and bound values.

Returns:

Transformed quantile values on the unbounded scale.

Return type:

Union[Array, ndarray, bool, number, float, int]

Raises:

chex.AssertionError – If the number of data points is less than num_terms.

static get_target(data, metalog_params)[source]

Construct the design matrix from input data and parameters.

Convenience wrapper around _get_target that extracts the probability levels and num_terms from the input objects.

Parameters:

  • data (MetalogInputData) – Validated input data containing probability levels (y).

  • metalog_params (MetalogParameters) – Configuration parameters specifying num_terms.

Returns:

Design matrix of shape (m, num_terms).

Return type:

Union[Array, ndarray, bool, number, float, int]

ppf(x)[source]

Compute the percent point function (quantile function) of the distribution.

The PPF is the inverse of the CDF: for a probability p, ppf(p) returns the value x such that P(X <= x) = p.

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1) at which to evaluate the quantile function. Can be a scalar or array.

Returns:

Quantile values corresponding to the input probabilities.

Same shape as input.

Return type:

Union[Array, ndarray, bool, number, float, int]

Raises:
pdf(x)[source]

Compute the probability density function of the metalog distribution.

The PDF is computed as the reciprocal of the derivative of the quantile function (PPF), with appropriate transformations for bounded distributions.

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1) at which to evaluate the PDF. Can be a scalar or array.

Returns:

Probability density values. Same shape as input.

Return type:

Union[Array, ndarray, bool, number, float, int]

Raises:

  • ValueError – If x contains values outside (0, 1).

  • NotFittedError – If the distribution has not been fitted.

  • checkify.JaxRuntimeError – If the PDF contains non-positive values, indicating an infeasible metalog fit.

cdf(q)[source]

Compute the cumulative distribution function.

Returns P(X <= q) for the fitted metalog distribution. Uses numerical inversion of the quantile function via nearest-neighbor lookup.

Parameters:

q (Union[Array, ndarray, bool, number, float, int]) – Quantile values at which to evaluate the CDF. Can be a scalar or 1D array.

Returns:

Probability values in (0, 1). Same shape as input.

Return type:

Union[Array, ndarray, bool, number, float, int]

rvs(rv_params)[source]

Generate random variates from the fitted metalog distribution.

Samples are generated by drawing uniform random values and applying the quantile function (inverse transform sampling). With HDRPRNGParameters, draw i uses HDR trial trial + i, so the default trial=1 draws trials 1 to size, and trial=size + 1 continues the same stream.

Parameters:

rv_params (MetalogRandomVariableParameters) – Parameters controlling random variate generation, including PRNG configuration (HDR or JAX-based) and sample size.

Returns:

Random samples from the distribution, shape (size,).

Return type:

Union[Array, ndarray, bool, number, float, int]

Raises:

  • chex.AssertionError – If size is not positive or exceeds max_draws.

  • checkify.JaxRuntimeError – If the PDF is infeasible (non-positive values).

logpdf(x)[source]

Compute the natural logarithm of the probability density function.

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1) at which to evaluate log(PDF).

Returns:

Log-density values. Same shape as input.

Return type:

Union[Array, ndarray, bool, number, float, int]

logppf(x)[source]

Compute the natural logarithm of the percent point function.

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1) at which to evaluate log(PPF).

Returns:

Log-quantile values. Same shape as input.

Return type:

Union[Array, ndarray, bool, number, float, int]

sf(x)[source]

Compute the survival function (complement of CDF).

The survival function is S(x) = 1 - F(x) = P(X > x).

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Quantile values at which to evaluate the survival function.

Returns:

Survival probabilities in (0, 1). Same shape as input.

Return type:

Union[Array, ndarray, bool, number, float, int]

logsf(x)[source]

Compute the natural logarithm of the survival function.

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Quantile values at which to evaluate log(SF).

Returns:

Log-survival probabilities. Same shape as input.

Return type:

Union[Array, ndarray, bool, number, float, int]

isf(q)[source]

Compute the inverse survival function (inverse complementary CDF).

Returns the value x such that P(X > x) = q, equivalent to ppf(1 - q).

Parameters:

q (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1).

Returns:

Quantile values corresponding to survival probabilities.

Return type:

Union[Array, ndarray, bool, number, float, int]

property mean: Union[float, int]

Compute the mean (expected value) of the metalog distribution.

Exact for unbounded metalogs, via the closed-form integrals of Baucells, Chrisman, Keelin and Xu (2025) (Lemma 1, Proposition 3). Bounded and semi-bounded metalogs have no closed form after the log/logit back-transform, so their mean is estimated by Monte Carlo sampling with 20,000 draws.

Returns:

The mean of the distribution.

Return type:

chex.Scalar

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

property median: Union[float, int]

Compute the median of the metalog distribution.

The median is the 50th percentile, computed exactly via ppf(0.5).

Returns:

The median of the distribution.

Return type:

chex.Scalar

property var: Union[float, int]

Compute the variance of the metalog distribution.

Exact for unbounded metalogs (see mean); estimated by Monte Carlo sampling with 20,000 draws for bounded and semi-bounded metalogs.

Returns:

The variance of the distribution.

Return type:

chex.Scalar

property std: Union[float, int]

Compute the standard deviation of the metalog distribution.

Exact for unbounded metalogs (see mean); estimated by Monte Carlo sampling with 20,000 draws for bounded and semi-bounded metalogs.

Returns:

The standard deviation of the distribution.

Return type:

chex.Scalar

property skewness: Union[float, int]

Compute the skewness of the metalog distribution.

Skewness is the third standardized central moment, E[(X - mean)^3] / std^3. Exact for unbounded metalogs (see mean); estimated from 20,000 draws for bounded and semi-bounded metalogs.

Returns:

The skewness of the distribution.

Return type:

chex.Scalar

property kurtosis: Union[float, int]

Compute the kurtosis of the metalog distribution.

Pearson kurtosis, E[(X - mean)^4] / var^2: 3 for a normal distribution and 4.2 for the logistic. Subtract 3 for excess kurtosis (scipy’s default convention). Exact for unbounded metalogs (see mean); estimated from 20,000 draws for bounded and semi-bounded metalogs.

Returns:

The kurtosis of the distribution.

Return type:

chex.Scalar

property mode: Union[float, int]

Compute the mode (most likely value) of the fitted metalog distribution.

For unbounded metalogs the mode is exact: the highest-density root of M'', found by Algorithm 1 of Baucells, Chrisman, Keelin and Xu (2025). This also finds modes that a grid misses, such as the binding modes of a Feasible fit, which can sit at y < 0.001. Bounded and semi-bounded metalogs (whose back-transform moves the modes) and metalogs without an interior mode use the maximum of the PDF over a fine grid of probabilities.

Returns:

The mode of the distribution.

Return type:

chex.Scalar

plot(plot_option)[source]

Visualize the fitted metalog distribution using interactive Plotly plots.

Creates an interactive Plotly figure showing the selected distribution function (PDF, CDF, or survival function).

Parameters:

plot_option (MetalogPlotOptions) – Which function to plot. One of MetalogPlotOptions.PDF, MetalogPlotOptions.CDF, or MetalogPlotOptions.SF.

Returns:

The Plotly figure object.

Return type:

Figure

Raises:

NotFittedError – If the distribution has not been fitted.

save(path)[source]

Save the fitted metalog distribution to a JSON file.

Writes the JSON that dumps returns. Read it back with load on the same class: Metalog.load or SPTMetalog.load.

Parameters:

path (Path) – File path where the JSON file will be written.

Return type:

None

dumps()[source]

Serialize the fitted metalog distribution to a JSON string.

The JSON holds the coefficients a and the parameters’ to_dict(): boundedness, fit method, bounds and number of terms for a Metalog; boundedness, bounds and alpha for an SPTMetalog.

Returns:

JSON string representation of the distribution.

Return type:

str

classmethod load(path)[source]

Load a fitted metalog distribution from a JSON file.

Parameters:

path (Path) – File path to the JSON file containing the serialized distribution.

Returns:

A new instance of the metalog class with loaded

coefficients and parameters.

Return type:

T_MetalogBase

Raises:

  • KeyError – If the JSON has no a or metalog_params.

  • ValueError – If the parameters were saved by the other metalog class (for example, an SPTMetalog’s JSON loaded with Metalog.load), or hold an unknown enum value.

classmethod loads(json_string)[source]

Load a fitted metalog distribution from a JSON string.

Parameters:

json_string (str) – JSON string containing the serialized distribution.

Returns:

A new instance of the metalog class with loaded

coefficients and parameters.

Return type:

T_MetalogBase

Raises:

  • KeyError – If the JSON has no a or metalog_params.

  • ValueError – If the parameters were saved by the other metalog class (for example, an SPTMetalog’s JSON loaded with Metalog.loads), or hold an unknown enum value.

q(x)[source]

Quantile function alias for R metalog package compatibility.

This method provides compatibility with the R metalog package’s qmetalog function naming convention. Equivalent to ppf().

Parameters:

x (Union[Array, ndarray, bool, number, float, int]) – Probability values in (0, 1) at which to evaluate the quantile function.

Returns:

Quantile values corresponding to the input probabilities.

Return type:

Union[Array, ndarray, bool, number, float, int]

Raises:
p(q)[source]

Probability function alias for R metalog package compatibility.

This method provides compatibility with the R metalog package’s pmetalog function naming convention. Equivalent to cdf().

Parameters:

q (Union[Array, ndarray, bool, number, float, int]) – Quantile values at which to evaluate the CDF.

Returns:

Probability values in (0, 1).

Return type:

Union[Array, ndarray, bool, number, float, int]

d(q)[source]

Density function alias for R metalog package compatibility.

This method provides compatibility with the R metalog package’s dmetalog function naming convention. Computes pdf(cdf(q)).

Parameters:

q (Union[Array, ndarray, bool, number, float, int]) – Quantile values at which to evaluate the density.

Returns:

Probability density values.

Return type:

Union[Array, ndarray, bool, number, float, int]

__init__(metalog_params, a)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

Display

Text and LaTeX display of fitted metalog distributions.

These functions write out a model’s quantile function with its coefficients. The model methods call them: str(m) and repr(m) use format_line, Jupyter’s _repr_latex_ uses format_latex, and m.print() uses format_block.

The terms follow the library’s Keelin (2016, Definition 1) order, the one MetalogBase._ppf evaluates: 1, ln(y/(1-y)), (y - 0.5) ln(y/(1-y)), (y - 0.5), then alternating powers of (y - 0.5), with and without the logit. Bounded and semi-bounded models map M(y) back to their support (Keelin 2016, Section 4): L + exp(M(y)), U - exp(-M(y)) and (L + U exp(M(y))) / (1 + exp(M(y))).

Formatting never raises on a model, because JAX, flax and debuggers call repr on traced and placeholder values. A model whose coefficients or bounds aren’t concrete (under jax.jit, jax.vmap or jax.eval_shape), or that holds a batch of models, prints a one-line summary instead of its equation, and format_latex returns None for it, which tells Jupyter to fall back to repr.

See also

metalog_jax.base.core.MetalogBase: The __str__, __repr__,

_repr_latex_ and print methods.

metalog_jax.base.display.format_line(model, precision=6)[source]

Write a model and its equation on one line.

str(m) is format_line(m), and repr(m) is format_line(m, precision=None).

Parameters:

  • model (Any) – A Metalog or SPTMetalog.

  • precision (int | None) – Significant digits of the coefficients and bounds, or None for Python’s shortest round-trip repr(float), which parses back to the exact values.

Returns:

The one-line equation, such as

Metalog(3 terms, OLS, lower bound 0): Q(y) = 0 + exp(M(y)), M(y) = ..., or a one-line summary if the model is traced or batched.

Return type:

str

metalog_jax.base.display.format_latex(model, precision=6)[source]

Write a model’s equation in LaTeX, for Jupyter’s _repr_latex_.

Up to four terms stay on one line, between $ signs. A longer model is an aligned environment between $$ signs, with Q(y) on its own line and M(y)’s terms wrapped after every four, because MathJax doesn’t break long equations. \ell(y) stands for \ln(y / (1 - y)).

Parameters:

  • model (Any) – A Metalog or SPTMetalog.

  • precision (int) – Significant digits of the coefficients and bounds.

Returns:

The LaTeX, or None if the model is traced or batched, which

tells IPython to show repr instead.

Return type:

str | None

metalog_jax.base.display.format_block(model, *, header=True, q=True, m=True, verbose=False, precision=6)[source]

Write a model’s equation one term per line, as m.print() prints it.

Parameters:

  • model (Any) – A Metalog or SPTMetalog.

  • header (bool) – Start with a line describing the model, such as Metalog: 5 terms, OLS, bounded on (0, 15), and indent the equations under it.

  • q (bool) – Include the quantile function Q(y).

  • m (bool) – Include M(y), one term per line. An unbounded model’s Q(y) is M(y), so its terms appear once, labelled Q(y), or M(y) when q is False; with m=False they stay on Q(y)’s line.

  • verbose (bool) – Write Q(y) with M(y)’s terms in place of M(y), so it can be read on its own.

  • precision (int) – Significant digits of the coefficients and bounds.

Returns:

The lines, without a trailing newline, or a one-line summary if the

model is traced or batched.

Return type:

str

Raises:

ValueError – If q and m are both False, verbose is True while q is False, or precision isn’t a positive integer.

Input Data

Data container classes for the Metalog JAX library.

This module provides data container classes for metalog input data:

Classes:
MetalogBaseData: Base container for metalog input data without validation.

Intended for internal use within JAX-traced functions.

MetalogInputData: Validated container for metalog input data, created

exclusively via the from_values() factory method.

See also

metalog_jax.base.core: MetalogBase class. metalog_jax.base.parameters: Parameter configuration classes. metalog_jax.metalog.fit: Function that uses MetalogInputData for fitting.

class metalog_jax.base.data.MetalogBaseData(x, y, precomputed_quantiles)[source]

Bases: object

Base container for metalog input data without validation.

This is a minimal dataclass that holds the core data required for metalog fitting without performing any validation or transformation. It serves as the parent class for MetalogInputData and is primarily used for internal operations where validation has already been performed or needs to be bypassed, such as within JAX-traced functions like vmap.

IMPORTANT: This class is intended for internal library use and advanced users only. Most users should use MetalogInputData.from_values() instead, which provides comprehensive validation and proper error handling.

Use MetalogBaseData directly only when:

  • Working inside JAX-traced functions (vmap, jit, etc.) where validation functions would fail due to tracer boolean conversion issues

  • Data has already been validated externally

  • Building custom pipelines that require manual control over validation

Variables:

  • x (Union[Array, ndarray, bool, number, float, int]) – Quantile values array corresponding to the probability levels in y. This should be a 1D JAX array of numeric values. When precomputed_quantiles=False, these are typically computed via jnp.quantile(samples, y). When precomputed_quantiles=True, these are user-provided quantile values.

  • y (Union[Array, ndarray, bool, number, float, int]) – Probability levels array in the open interval (0, 1), sorted in strictly ascending order. This should be a 1D JAX array with the same length as x, representing the cumulative probabilities at which the quantiles in x are evaluated.

  • precomputed_quantiles (bool) – Boolean flag indicating whether x contains pre-computed quantiles (True) or was originally raw sample data (False). This is marked as pytree_node=False to exclude it from JAX transformations, as it’s metadata rather than data to be transformed.

Note

This class uses Flax’s struct.dataclass decorator to make instances immutable and compatible with JAX transformations (jit, vmap, grad, etc.). The precomputed_quantiles field is excluded from the pytree structure via pytree_node=False since it’s metadata.

See also

MetalogInputData: Subclass with validation via from_values() factory method. metalog_jax.base.parameters.MetalogParameters: Configuration for fitting. metalog_jax.metalog.fit: Function that processes input data for fitting.

Parameters:
__init__(x, y, precomputed_quantiles)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.base.data.MetalogInputData(x, y, precomputed_quantiles, _from_factory=False)[source]

Bases: MetalogBaseData

Container for input data used to fit metalog distributions.

This dataclass encapsulates the input data for metalog fitting, supporting two modes: (1) fitting from pre-computed quantiles at specified probability levels, or (2) fitting from raw sample data where quantiles are computed automatically.

The mode is controlled by the precomputed_quantiles flag, which determines how the x array should be interpreted and processed during fitting.

IMPORTANT: Instances MUST be created using the from_values() classmethod. Direct instantiation via __init__() is prevented by a __post_init__ validation that raises TypeError if the instance was not created through from_values().

Variables:

  • x (Union[Array, ndarray, bool, number, float, int]) – Input data array. The interpretation depends on precomputed_quantiles: - If precomputed_quantiles=True: Array of quantile values - If precomputed_quantiles=False: Array of raw sample data

  • y (Union[Array, ndarray, bool, number, float, int]) – Array of probability levels in the open interval (0, 1), sorted in ascending order.

  • precomputed_quantiles (bool) – Flag indicating the interpretation of x.

See also

MetalogBaseData: Parent class without validation (for advanced use). from_values: Required factory method for creating validated instances. metalog_jax.metalog.fit: Function that uses MetalogInputData for fitting.

Parameters:
classmethod from_values(x, y, precomputed_quantiles)[source]

Create a validated MetalogInputData instance with comprehensive input checks.

This is the ONLY way to create valid MetalogInputData instances.

Parameters:
Return type:

T_MetalogInputData

Returns:

A validated MetalogInputData instance ready for use in metalog fitting.

Raises:

AssertionError – If any validation check fails.

__post_init__()[source]

Post-initialization validation enforcing factory pattern.

__init__(x, y, precomputed_quantiles, _from_factory=False)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

Enumerations

Enumeration types for the Metalog JAX library.

This module provides enumeration types for configuring metalog distributions:

Classes:

CustomIntEnum: Base class for integer enumerations with enhanced value lookup. MetalogBoundedness: Boundedness types (UNBOUNDED, STRICTLY_LOWER_BOUND, …). MetalogFitMethod: Fit methods (Feasible, OLS, Lasso). MetalogPlotOptions: Enumeration of plot types (PDF, CDF, SF).

See also

metalog_jax.base.parameters: Parameter configuration classes that use these enums. metalog_jax.regression: OLS and Lasso implementations. metalog_jax.feasibility: Feasible implementation (Baucells et al. 2025).

class metalog_jax.base.enums.CustomIntEnum(new_class_name, /, names, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Base class for custom integer enumerations with enhanced value lookup.

Extends Python’s IntEnum to provide additional utility methods for working with enumeration values. This class serves as a base for all custom integer enumerations in the metalog JAX library, providing consistent behavior for value-to-member conversion and validation.

The class maintains all IntEnum functionality (enum members are instances of int and can be used in arithmetic operations) while adding convenient class methods for value lookup and conversion.

Note

This is an abstract base class. Use specific subclasses like MetalogBoundedness or MetalogFitMethod for actual enum definitions.

classmethod from_value(value)[source]

Convert an integer value to the corresponding enum member.

Searches through all members of the enumeration and returns the member whose value matches the input. This provides a safe way to reconstruct enum members from integer values, such as when loading from JSON files or deserializing data.

Parameters:

value (int) – Integer value to convert to an enum member. Must match the value of an existing enum member exactly.

Returns:

The enum member whose value matches the input value.

Raises:

ValueError – If no enum member with the specified value exists.

Example

>>> boundedness = MetalogBoundedness.from_value(1)
>>> boundedness == MetalogBoundedness.UNBOUNDED
True
>>> MetalogBoundedness.from_value(99)
Traceback (most recent call last):
    ...
ValueError: No enum member with value 99
class metalog_jax.base.enums.MetalogBoundedness(*values)[source]

Bases: CustomIntEnum

Enumeration of boundedness types for metalog distributions.

Defines the domain constraints for metalog probability distributions, specifying whether the distribution has lower bounds, upper bounds, both, or neither.

Variables:

  • UNBOUNDED – Distribution with no bounds (support on entire real line).

  • STRICTLY_LOWER_BOUND – Distribution with only a lower bound.

  • STRICTLY_UPPER_BOUND – Distribution with only an upper bound.

  • BOUNDED – Distribution with both lower and upper bounds.

class metalog_jax.base.enums.MetalogFitMethod(*values)[source]

Bases: CustomIntEnum

Enumeration of regression methods for fitting metalog distributions.

Defines the available regression techniques that can be used to estimate metalog distribution coefficients from data.

Variables:

  • OLS – Ordinary Least Squares regression (no regularization, closed-form solution). Implementation: metalog_jax.regression.ols.fit_ordinary_least_squares

  • Lasso – LASSO regression (L1 regularization via proximal gradient descent). Implementation: metalog_jax.regression.lasso.fit_lasso

  • Feasible – Best feasible least-squares fit a* (Baucells, Chrisman, Keelin and Xu, 2025): the closest coefficients whose density is non-negative everywhere. Recommended default; equals OLS whenever the OLS fit is already valid. Implementation: metalog_jax.feasibility.a_star.fit_feasible

See also

metalog_jax.regression: Module containing all regression implementations. metalog_jax.regression.ols: OLS regression implementation. metalog_jax.regression.lasso: LASSO regression implementation. metalog_jax.feasibility: Best feasible fit implementation.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

class metalog_jax.base.enums.MetalogPlotOptions(*values)[source]

Bases: IntEnum

Enumeration of available plot types for metalog distributions.

Defines the types of statistical plots that can be generated for visualizing fitted metalog distributions. Each option corresponds to a different view of the probability distribution.

Variables:

  • PDF – Probability Density Function plot showing the probability density across the distribution’s support.

  • CDF – Cumulative Distribution Function plot showing the cumulative probability from the lower bound to each point.

  • SF – Survival Function plot showing the survival function from the lower bound to each point.

Parameters

Parameter configuration classes for the Metalog JAX library.

This module provides parameter configuration classes for metalog distributions:

Classes:

MetalogRandomVariableParameters: Configuration for random variable sampling. MetalogParametersBase: Base class for metalog distribution parameters. MetalogParameters: Configuration for standard regression-based metalog fitting. SPTMetalogParameters: Configuration for Symmetric Percentile Triplet metalog fitting.

See also

metalog_jax.base.enums: Enumeration types (MetalogBoundedness, MetalogFitMethod). metalog_jax.base.core: MetalogBase class that uses these parameters. metalog_jax.metalog: Fitting functions that accept these parameter configurations.

class metalog_jax.base.parameters.MetalogRandomVariableParameters(prng_params, size, max_draws=100000)[source]

Bases: object

Configuration parameters for drawing rvs from fitted metalog distribution.

size and max_draws set array shapes, so they are static fields: they belong to the pytree structure rather than its leaves, and the parameters can be passed to jax.jit. The PRNG parameters are leaves, so jax.vmap can run over seeds or HDR trials.

Variables:

  • prng_params (JaxUniformDistributionParameters | HDRPRNGParameters) – Params for the psuedorandom number generator.

  • size (int) – The number of random variables to draw. Must be strictly positive. Static.

  • max_draws (int) – Maximum number of draws for rejection sampling. Static.

Parameters:
__post_init__()[source]

Post init data validation.

Raises:

AssertionError – If size is not a positive scalar.

__init__(prng_params, size, max_draws=100000)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.base.parameters.MetalogParametersBase(boundedness, lower_bound, upper_bound)[source]

Bases: object

Base class for metalog distribution parameter configurations.

This abstract base class defines the common parameters shared by all metalog distribution types, specifically the boundedness constraints and boundary values that determine the support of the distribution.

Fields that select code paths or array shapes (boundedness here, and method and num_terms in MetalogParameters) are static: they belong to the pytree structure rather than its leaves. So a model or parameters object can be passed to jax.jit, which compiles once per combination of static values, and models stacked for jax.vmap must share them. The bounds are ordinary leaves and can be traced.

Variables:

  • boundedness (MetalogBoundedness) – Type of domain constraint for the distribution. Static.

  • lower_bound (Union[float, int]) – Lower boundary value for the distribution’s support.

  • upper_bound (Union[float, int]) – Upper boundary value for the distribution’s support.

Parameters:
__init__(boundedness, lower_bound, upper_bound)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.base.parameters.MetalogParameters(boundedness, lower_bound, upper_bound, method, num_terms)[source]

Bases: MetalogParametersBase

Configuration parameters for standard regression-based metalog distributions.

Extends MetalogParametersBase with regression-specific parameters for fitting metalog distributions via regression methods including OLS and Lasso.

Variables:

  • boundedness (MetalogBoundedness) – Inherited. Type of domain constraint for the distribution. Defined in metalog_jax.base.enums.MetalogBoundedness.

  • lower_bound (Union[float, int]) – Inherited. Lower boundary value.

  • upper_bound (Union[float, int]) – Inherited. Upper boundary value.

  • method – Regression method used for fitting. Static. One of: - MetalogFitMethod.OLS: metalog_jax.regression.ols - MetalogFitMethod.Lasso: metalog_jax.regression.lasso

  • num_terms (int) – Number of terms in the metalog series expansion. Static.

See also

MetalogParametersBase: Parent class with boundedness and bound attributes. metalog_jax.base.enums.MetalogFitMethod: Enumeration of regression methods. metalog_jax.metalog.fit: Function that uses MetalogParameters for fitting. metalog_jax.regression: Module containing regression implementations.

Parameters:
to_dict()[source]

Return the parameters as the JSON-ready dict that Metalog.dumps writes.

Returns:

boundedness and method as integers, the bounds as floats,

and num_terms as an integer.

Return type:

dict

classmethod from_dict(params)[source]

Rebuild parameters from a dict written by to_dict.

Parameters:

params (dict) – The metalog_params dict of a saved Metalog.

Returns:

The parameters, with enum members for boundedness

and method.

Return type:

Self

Raises:

ValueError – If the keys aren’t exactly those that to_dict writes (as for an SPTMetalog’s parameters), or an enum value is unknown.

__init__(boundedness, lower_bound, upper_bound, method, num_terms)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.base.parameters.SPTMetalogParameters(boundedness, lower_bound, upper_bound, alpha)[source]

Bases: MetalogParametersBase

Configuration parameters for Symmetric Percentile Triplet (SPT) metalog distributions.

Extends MetalogParametersBase with the alpha parameter for SPT metalog fitting, which uses exactly 3 terms fitted from three symmetric quantiles.

Variables:

  • boundedness (MetalogBoundedness) – Inherited. Type of domain constraint for the distribution.

  • lower_bound (Union[float, int]) – Inherited. Lower boundary value.

  • upper_bound (Union[float, int]) – Inherited. Upper boundary value.

  • alpha (float) – The lower percentile parameter for SPT fitting. Must be in (0, 0.5).

Parameters:
to_dict()[source]

Return the parameters as the JSON-ready dict that SPTMetalog.dumps writes.

Returns:

boundedness as an integer, and the bounds and alpha as

floats.

Return type:

dict

classmethod from_dict(params)[source]

Rebuild parameters from a dict written by to_dict.

Parameters:

params (dict) – The metalog_params dict of a saved SPTMetalog.

Returns:

The parameters, with an enum member for

boundedness.

Return type:

Self

Raises:

ValueError – If the keys aren’t exactly those that to_dict writes (as for a Metalog’s parameters), or the boundedness value is unknown.

__init__(boundedness, lower_bound, upper_bound, alpha)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

Regression

Base Classes

Base classes for regression models and parameters.

This module provides the abstract base classes for regression models and their parameters used in metalog distribution fitting.

Classes:
RegularizedParameters: Base class for regularized regression hyperparameters.

Subclasses: LassoParameters

RegressionModel: Base class for trained regression model weights.

Subclasses: OLSModel, LassoModel

Functions:

has_full_column_rank: Whether a design matrix has full column rank. assert_full_column_rank: Raise when least squares has no unique solution.

See also

metalog_jax.regression.ols: OLS regression implementation. metalog_jax.regression.lasso: LASSO regression implementation.

class metalog_jax.regression.base.RegularizedParameters[source]

Bases: object

Base class for regularized regression hyperparameters.

This abstract base class serves as the parent for all regularized regression parameter configurations in the metalog JAX library. It provides a common type hierarchy for regression methods that incorporate regularization penalties (L1, L2, or both) to prevent overfitting and improve generalization.

Regularization adds penalty terms to the regression objective function to constrain model complexity. This base class enables polymorphic handling of different regularization strategies while maintaining type safety and consistency across the library.

The class uses Flax’s struct.dataclass decorator to ensure immutability and compatibility with JAX transformations (jit, vmap, grad, etc.).

Subclasses:
LassoParameters: L1 regularization parameters (LASSO regression).

Defined in metalog_jax.regression.lasso. Penalizes the absolute magnitude of coefficients: λ||w||₁

Design Pattern:

This class follows the Template Method pattern, providing a common interface while allowing subclasses to define specific regularization configurations. It enables functions to accept any regularized regression parameters through polymorphic type hints.

Example

Using the base class for polymorphic type hints:

>>> from typing import Union
>>> from metalog_jax.regression.base import RegularizedParameters
>>> from metalog_jax.regression.lasso import LassoParameters
>>>
>>> def validate_regularization(params: RegularizedParameters):
...     '''Accept any regularized regression parameters.'''
...     if isinstance(params, LassoParameters):
...         print(f"Lasso with L1={params.lam}")
>>>
>>> lasso_params = LassoParameters(
...     lam=1.0,
...     learning_rate=0.01,
...     num_iters=500,
...     tol=1e-6,
...     momentum=0.9
... )
>>> validate_regularization(lasso_params)  # Valid
Lasso with L1=1.0

Creating subclass instances:

>>> # Lasso regression (L1 only)
>>> from metalog_jax.regression.lasso import LassoParameters
>>> lasso = LassoParameters(
...     lam=0.5,
...     learning_rate=0.01,
...     num_iters=500,
...     tol=1e-6,
...     momentum=0.9
... )
>>> isinstance(lasso, RegularizedParameters)
True

Note

  • This is an abstract base class with no attributes or methods

  • Direct instantiation is possible but not meaningful (use subclasses)

  • All instances are immutable due to struct.dataclass decorator

  • Subclasses must define their own specific regularization parameters

See also

metalog_jax.regression.lasso.LassoParameters: LASSO (L1) regression parameters. metalog_jax.regression.lasso.fit_lasso: Function that uses LassoParameters.

References

Hastie, T., Tibshirani, R., & Friedman, J. (2009). The Elements of Statistical Learning: Data Mining, Inference, and Prediction (2nd ed.). Springer. Chapter 3: Linear Methods for Regression.

__init__()
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.regression.base.RegressionModel(weights)[source]

Bases: object

Structure for trained Regression model weights.

Variables:

weights (Union[Array, ndarray, bool, number]) – Coefficient vector of shape (n_features,).

Parameters:

weights (Union[Array, ndarray, bool, number])

weights: Union[Array, ndarray, bool, number]
__init__(weights)
Parameters:

weights (Union[Array, ndarray, bool, number])

replace(**updates)

Returns a new object replacing the specified fields with new values.

metalog_jax.regression.base.has_full_column_rank(X)[source]

Whether X of shape (n, k) has rank k, using numpy’s default SVD tolerance.

Parameters:

X (Union[Array, ndarray, bool, number]) – Design matrix of shape (n_samples, n_features).

Return type:

Union[Array, ndarray, bool, number]

Returns:

Boolean scalar array; jit/vmap compatible.

metalog_jax.regression.base.assert_full_column_rank(X)[source]

Raise when X is rank deficient, so least squares has no unique solution.

Uses checkify so it also works under vmap (e.g. fit_grid_datasets).

Parameters:

X (Union[Array, ndarray, bool, number]) – Design matrix of shape (n_samples, n_features).

Raises:

checkify.JaxRuntimeError – If X does not have full column rank.

Return type:

None

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416 (Proposition 2).

Ordinary Least Squares

Ordinary Least Squares (OLS) regression implementation.

This module provides OLS regression for fitting metalog distributions when no regularization is needed.

Classes:

OLSModel: Structure for trained OLS model weights.

Functions:

fit_ordinary_least_squares: Fit an OLS regression model. predict_ordinary_least_squares: Make predictions using a fitted OLS model.

This module is used when MetalogFitMethod.OLS is specified in MetalogParameters.

See also

metalog_jax.regression.lasso: LASSO (L1 regularization). metalog_jax.metalog.fit: High-level fitting function that dispatches to OLS.

class metalog_jax.regression.ols.OLSModel(weights)[source]

Bases: RegressionModel

Structure for trained Ordinary Least Squares model weights.

Variables:
Parameters:

weights (Union[Array, ndarray, bool, number])

__init__(weights)
Parameters:

weights (Union[Array, ndarray, bool, number])

replace(**updates)

Returns a new object replacing the specified fields with new values.

metalog_jax.regression.ols.fit_ordinary_least_squares(X, y)[source]

Fit an Ordinary Least Squares (OLS) regression model.

Computes the closed-form solution to linear regression using the normal equations via least squares. This method finds the optimal weights and bias that minimize the sum of squared residuals with no regularization.

The model solves: min ||y - (Xw + b)||^2

Parameters:
Return type:

OLSModel

Returns:

OLSModel containing the fitted weights and bias that minimize the squared error.

Example

>>> X = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
>>> y = jnp.array([1.0, 2.0, 3.0])
>>> model = fit_ordinary_least_squares(X, y)
metalog_jax.regression.ols.predict_ordinary_least_squares(X, model)[source]

Make predictions using a fitted Regression model.

Computes predictions for new data using the trained weights and bias: y_pred = X @ weights + bias

Parameters:
Return type:

Union[Array, ndarray, bool, number]

Returns:

Predictions of shape (n_samples,).

Example

>>> import jax.numpy as jnp
>>> X = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
>>> model = fit_ordinary_least_squares(X, jnp.array([1.0, 2.0, 3.0]))
>>> X_test = jnp.array([[2.0, 3.0], [4.0, 5.0]])
>>> predict_ordinary_least_squares(X_test, model).round(4).tolist()
[1.5, 2.5]

LASSO

LASSO regression implementation with L1 regularization.

This module provides LASSO (Least Absolute Shrinkage and Selection Operator) regression for fitting metalog distributions with L1 regularization using proximal gradient descent with Nesterov acceleration.

Classes:

LassoParameters: Hyperparameters for LASSO regression. LassoModel: Structure for trained LASSO model weights.

Functions:

fit_lasso: Fit a LASSO regression model via proximal gradient descent. soft_thresholding: Apply the soft-thresholding operator (proximal for L1).

Constants:

DEFAULT_LASSO_LAMBDA: Default L1 regularization strength (0). DEFAULT_LASSO_LEARNING_RATE: Default learning rate (0.01). DEFAULT_LASSO_ITERATIONS: Default maximum iterations (500). DEFAULT_LASSO_TOLERANCE: Default convergence tolerance (1e-6). DEFAULT_LASSO_MOMENTUM: Default Nesterov momentum factor (0.9). DEFAULT_LASSO_PARAMETERS: Default LassoParameters instance.

This module is used when MetalogFitMethod.Lasso is specified in MetalogParameters.

See also

metalog_jax.regression.ols: OLS regression (no regularization). metalog_jax.metalog.fit: High-level fitting function that dispatches to LASSO.

class metalog_jax.regression.lasso.LassoParameters(lam, learning_rate, num_iters, tol, momentum)[source]

Bases: RegularizedParameters

Structure for Lasso parameters.

Variables:

  • lam – L1 regularization strength λ ≥ 0. Controls sparsity of the solution: - λ = 0: No regularization (equivalent to OLS, but solved iteratively) - λ > 0: Promotes sparsity; larger values → more coefficients set to zero - Typical values: 0.001 to 10.0, depending on data scale

  • learning_rate – Learning rate (step size) for gradient descent. Default: 0.01. Controls the size of parameter updates. Should be tuned based on problem: - Too large: May cause divergence or oscillation - Too small: Slow convergence - Typical values: 0.001 to 0.1

  • num_iters – Maximum number of iterations. Default: 500. Training will stop early if convergence is reached before this limit.

  • tol – Convergence tolerance for weight changes. Default: 1e-6. Training stops when ||w_new - w|| < tol, indicating the solution has stabilized. Smaller values require more precise convergence but may take longer.

  • momentum – Nesterov momentum factor, must satisfy 0 < momentum < 1. Default: 0.9. Controls acceleration of convergence: - Higher values (e.g., 0.9, 0.99): Faster convergence but less stable - Lower values (e.g., 0.5, 0.7): More stable but slower convergence - Standard choice: 0.9

Parameters:
lam: float
learning_rate: float
num_iters: int
tol: float
momentum: float
__init__(lam, learning_rate, num_iters, tol, momentum)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.regression.lasso.LassoModel(weights)[source]

Bases: RegressionModel

Structure for trained LASSO Regression model weights.

Variables:
Parameters:

weights (Union[Array, ndarray, bool, number])

__init__(weights)
Parameters:

weights (Union[Array, ndarray, bool, number])

replace(**updates)

Returns a new object replacing the specified fields with new values.

metalog_jax.regression.lasso.soft_thresholding(x, lam)[source]

Apply the soft-thresholding operator (proximal operator for L1 norm).

The soft-thresholding operator is the proximal operator for the L1 norm and is fundamental to LASSO regression and other sparse optimization methods. It shrinks values toward zero and sets values below the threshold to exactly zero, promoting sparsity in the solution.

The operator is defined element-wise as:

soft_threshold(x, λ) = sign(x) * max(|x| - λ, 0)

Equivalently:

  • If x > λ: return x - λ

  • If x < -λ: return x + λ

  • If |x| ≤ λ: return 0

This function is the proximal operator for the L1 penalty:

prox_{λ||·||₁}(x) = argmin_z { (1/2)||z - x||² + λ||z||₁ }

Parameters:

  • x (Union[Array, ndarray, bool, number, float, int]) – Input value(s) to threshold. Can be a scalar or array of any shape.

  • lam (float) – Threshold parameter λ ≥ 0. Controls the amount of shrinkage: - λ = 0: No shrinkage (returns x unchanged) - λ > 0: Shrinks values toward zero, setting small values to 0 - Larger λ produces sparser solutions with more zeros

Return type:

Union[Array, ndarray, bool, number, float, int]

Returns:

Thresholded value(s) with the same shape as x. Values are shrunk toward zero by amount λ, with values |x| ≤ λ set to exactly zero.

Example

Scalar inputs:

>>> import jax.numpy as jnp
>>> float(soft_thresholding(5.0, 2.0))
3.0
>>> float(soft_thresholding(-3.0, 1.0))
-2.0
>>> float(soft_thresholding(1.5, 2.0))  # Below threshold -> zero
0.0

Array inputs:

>>> x = jnp.array([-5.0, -1.0, 0.5, 2.0, 4.0])
>>> soft_thresholding(x, 1.5).tolist()
[-3.5, -0.0, 0.0, 0.5, 2.5]

Note

  • This function is JIT-compiled for efficient execution

  • The operation is element-wise and preserves the input shape

  • Setting lam=0 returns the input unchanged

  • This is also known as the “shrinkage operator” or “soft thresholding”

See also

fit_lasso: LASSO regression that uses this operator in proximal gradient descent.

References

Parikh, N., & Boyd, S. (2014). Proximal Algorithms. Foundations and Trends in Optimization, 1(3), 127-239.

Beck, A., & Teboulle, M. (2009). A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse Problems. SIAM Journal on Imaging Sciences, 2(1), 183-202.

metalog_jax.regression.lasso.fit_lasso(X, y, params=LassoParameters(lam=0.0, learning_rate=0.01, num_iters=500, tol=1e-06, momentum=0.9))[source]

Fit a LASSO regression model using proximal gradient descent with Nesterov acceleration.

Trains a linear regression model with L1 regularization (LASSO - Least Absolute Shrinkage and Selection Operator) using an accelerated proximal gradient method. LASSO promotes sparse solutions by shrinking coefficients toward zero and setting many coefficients to exactly zero, making it useful for feature selection and interpretable models.

This implementation uses:

  • Proximal gradient descent: Combines gradient steps on the smooth loss (MSE) with the proximal operator (soft-thresholding) for the non-smooth L1 penalty

  • Nesterov momentum: Accelerates convergence using momentum-based lookahead steps

  • Early stopping: Terminates when weight changes fall below tolerance threshold

  • JAX compatibility: Fully vectorized and compatible with JAX transformations

The LASSO objective function is:

min_w { (1/2n)||y - Xw||² + λ||w||₁ }

where λ is the L1 regularization strength that controls sparsity.

Algorithm:

  1. Initialize weights w and velocity v to zero

  2. For each iteration (until convergence or max iterations): a. Compute lookahead weights: w_lookahead = w + momentum * v b. Compute gradient at lookahead position c. Update velocity with momentum and gradient: v = momentum * v - learning_rate * grad d. Apply proximal operator: w = soft_threshold(w + v, learning_rate * λ) e. Check convergence: stop if ||w_new - w|| < tol

  3. Return final weights

Parameters:

  • X (Union[Array, ndarray, bool, number, float, int]) – Feature matrix of shape (n_samples, n_features). The design matrix containing the independent variables for regression.

  • y (Union[Array, ndarray, bool, number, float, int]) – Target vector of shape (n_samples,). The dependent variable values to predict.

  • params (LassoParameters) – LassoParameters containing optimization hyperparameters. Defaults to DEFAULT_LASSO_PARAMETERS. The object contains: - lam: L1 regularization strength λ ≥ 0. Default: 0. - learning_rate (float): Step size for gradient descent. Default: 0.01. - num_iters (int): Maximum number of iterations. Default: 500. - tol (float): Convergence tolerance for weight changes. Default: 1e-6. - momentum (float): Nesterov momentum factor (0 < momentum < 1). Default: 0.9.

Return type:

LassoModel

Returns:

LassoModel containing the fitted weights vector of shape (n_features,). Many coefficients may be exactly zero due to L1 regularization, achieving feature selection and model sparsity.

Example

Basic LASSO regression with default parameters:

>>> import jax.numpy as jnp
>>> from metalog_jax.regression.lasso import fit_lasso
>>>
>>> # Create sample data
>>> X = jnp.array([[1.0, 2.0, 0.5],
...                [3.0, 4.0, 1.5],
...                [5.0, 6.0, 2.5]])
>>> y = jnp.array([1.0, 2.0, 3.0])
>>>
>>> # Fit LASSO with defaults
>>> model = fit_lasso(X, y)
>>> model.weights.shape  # Some coefficients may be exactly zero
(3,)

Note

  • Unlike Ridge regression, LASSO does not have a closed-form solution and requires iterative optimization

  • The algorithm may converge before num_iters if tol is reached

  • Feature scaling (standardization) is recommended before fitting LASSO

See also

soft_thresholding: The proximal operator used in each iteration. fit_ordinary_least_squares: No regularization (from regression.ols). LassoModel: Return type containing fitted weights. LassoParameters: Configuration dataclass for LASSO hyperparameters.

References

Tibshirani, R. (1996). Regression Shrinkage and Selection via the Lasso. Journal of the Royal Statistical Society: Series B (Methodological), 58(1), 267-288.

Beck, A., & Teboulle, M. (2009). A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse Problems. SIAM Journal on Imaging Sciences, 2(1), 183-202.

Feasibility (Metalog 2.0)

Feasibility tools of Baucells, Chrisman, Keelin and Xu (2025).

JAX ports of the algorithms in “On the Properties of the Metalog Distribution” and its reference implementation by Zixin (Stephen) Xu (CC BY 4.0; see NOTICE):

  • best_feasible_fit() - the optimal feasible coefficients a*.

  • check_feasibility() - air-tight feasibility test (Algorithm 1 + Prop. 5).

  • inflection_points() - all modes and anti-modes of the density.

  • highest_mode() - the highest-density mode (used by Metalog.mode).

  • summary_stats() / raw_moment() - exact moments (Lemma 1, Prop. 3).

Every function is jit-compatible and can be vmapped over batches. They use the Keelin (2016) term assignment of metalog_jax by default; TermOrder.METALOG_2 selects the paper’s Metalog 2.0 assignment, which differs from the 7th term on.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416 Xu, Z. S. metalog_algorithm (reference implementation), licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/); adapted and ported to JAX, see NOTICE. https://github.com/Stephenxuu/metalog_algorithm

Best Feasible Fit

Best feasible metalog fit a* (semi-infinite QP solved by the exchange method).

Solves min_a ||x - Y a||^2 subject to G(y) = y(1-y) M'(y) >= epsilon for all y in [0, 1], where G(0) = s(0) and G(1) = s(1). The constraint is linear in a, so each violated point becomes one linear cut. The loop alternates between (i) finding the most violated points of the current fit and (ii) re-solving a small QP with all cuts so far, and stops when an air-tight check finds no violation. This mirrors find_a_star in the authors’ reference code, with these differences:

  • The QP is whitened with a QR factorization of Y (Hessian becomes the identity) and solved by a fixed-iteration Mehrotra predictor-corrector interior-point method.

  • Shapes are static (a masked buffer of max_cuts cuts), so the whole fit runs under jit and can be vmap-ed over many datasets.

  • Violations are found on a logit grid (reaching y ~ 1e-16 from either edge) with golden-section refinement, then Algorithm 1 plus the Proposition 5 tail tests as an exact fallback. A point is a violation when G(y) < 0 beyond rounding, the same sign test check_feasibility() uses, and the final certificate is check_feasibility() itself.

  • As in the paper, each iteration adds every flagged point: all confirmed grid minima (at least floor((k + 1) / 2) + 1 are refined, the most G can have) or every inflection point with M' < 0 plus any infeasible tail.

  • epsilon is relative to the spread of x, so results do not depend on the units of the data.

Note: “a*” is the paper’s name for the optimal coefficient vector; this is not the A* graph-search algorithm.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416 Xu, Z. S. metalog_algorithm (reference implementation), licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/); adapted and ported to JAX, see NOTICE. https://github.com/Stephenxuu/metalog_algorithm

class metalog_jax.feasibility.a_star.FeasibleFitResult(a_star, a_ols, rss_star, rss_ols, iterations, converged, feasible, full_rank, cut_points)[source]

Bases: object

Result of best_feasible_fit().

Variables:
Parameters:
__init__(a_star, a_ols, rss_star, rss_ols, iterations, converged, feasible, full_rank, cut_points)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.feasibility.a_star.FeasibleModel(weights)[source]

Bases: RegressionModel

Regression-style model returned by fit_feasible() (weights = a*).

Parameters:

weights (Union[Array, ndarray, bool, number])

__init__(weights)
Parameters:

weights (Union[Array, ndarray, bool, number])

replace(**updates)

Returns a new object replacing the specified fields with new values.

metalog_jax.feasibility.a_star.design_matrix(engine, y)[source]

Metalog basis matrix of shape (n, k) for the engine’s term ordering.

Parameters:
Return type:

Union[Array, ndarray, bool, number]

metalog_jax.feasibility.a_star.project_onto_polyhedron(z0, A, b, mask, iters=40)[source]

argmin 0.5||z - z0||^2 s.t. A z >= b on rows where mask is True.

Mehrotra predictor-corrector interior-point method with a fixed iteration budget. The iterate is frozen once converged or if a step becomes non-finite (the barrier Hessian overflows as active slacks reach zero).

Parameters:

iters (int)

metalog_jax.feasibility.a_star.best_feasible_fit(x, y, num_terms, order=TermOrder.KEELIN_2016, epsilon=1e-06, tol=0.0, max_cuts=64, n_new=8, max_iter=60, n_grid=2049)[source]

Fit the best feasible k-term metalog to quantile/probability pairs.

Parameters:

  • x (Union[Array, ndarray, bool, number]) – Quantile values, shape (n,).

  • y (Union[Array, ndarray, bool, number]) – Cumulative probabilities in (0, 1), shape (n,).

  • num_terms (int) – Number of metalog terms k (2 <= k <= n).

  • order (TermOrder) – Basis ordering (KEELIN_2016 matches metalog_jax; use METALOG_2 to reproduce the paper and the reference code).

  • epsilon (float) – Margin in G(y) >= epsilon, relative to max|x - median(x)|.

  • tol (float) – Extra slack for the violation and certificate sign tests, relative to the magnitude of G’s terms. With 0 (the default) a violation is any G(y) < 0 beyond rounding, as in the paper.

  • max_cuts (int) – Capacity of the cut buffer (static).

  • n_new (int) – Minimum number of grid minima refined per iteration (static); at least floor((k + 1) / 2) + 1 are always refined.

  • max_iter (int) – Maximum number of exchange iterations (static).

  • n_grid (int) – Number of logit grid points used to locate violations (static).

Return type:

FeasibleFitResult

Returns:

FeasibleFitResult. Check feasible (the certificate), converged and full_rank (with n = k symmetric probabilities the Keelin 2016 basis is singular for 7, 11, 15, … terms; use METALOG_2 or more quantiles).

Example

>>> import jax.numpy as jnp
>>> from metalog_jax.feasibility import best_feasible_fit
>>> x = jnp.array([1.0, 2.0, 4.0, 8.0, 12.0])
>>> y = jnp.array([0.1, 0.3, 0.5, 0.7, 0.9])
>>> res = best_feasible_fit(x, y, num_terms=4)
>>> bool(res.feasible)
True

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

metalog_jax.feasibility.a_star.fit_feasible(X, y)[source]

Regression-dispatch entry point: best feasible fit from a metalog design matrix.

X must be a metalog_jax design matrix (Keelin 2016 term ordering), as built by MetalogBase.get_target; y holds the (possibly transformed) quantiles. Used by fit when method=MetalogFitMethod.Feasible.

Parameters:
Return type:

FeasibleModel

Returns:

FeasibleModel whose weights are the best feasible coefficients.

Raises:

checkify.JaxRuntimeError – If the result cannot be certified feasible (the exchange loop hit its iteration or cut limit). vmap compatible.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

Feasibility Checks and Exact Moments

Feasibility, modes and exact moments of an (unbounded) metalog.

All functions are jit/vmap/grad compatible once the engine is fixed. For bounded metalogs these statements apply to the underlying unbounded metalog fitted to the transformed quantiles; feasibility is preserved by the monotone back-transform, but the moment formulas are not.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

class metalog_jax.feasibility.analysis.FeasibilityReport(feasible, interior_feasible, tail_feasible_zero, tail_feasible_one, roots, slopes, is_mode)[source]

Bases: object

Result of check_feasibility().

Variables:
Parameters:
__init__(feasible, interior_feasible, tail_feasible_zero, tail_feasible_one, roots, slopes, is_mode)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

metalog_jax.feasibility.analysis.feasibility_function(engine, a, u)[source]

G(y) = y(1-y) M'(y) at y = sigmoid(u) (u = -/+inf gives s(0), s(1)).

Parameters:
Return type:

Union[Array, ndarray, bool, number, float, int]

metalog_jax.feasibility.analysis.zero_test_width(tol=0.0)[source]

Width of the zero tests, in machine epsilons of the terms’ magnitude.

A value counts as zero when |v| <= width * eps * magnitude, where magnitude is the sum of the absolute values of the terms that produced v (a running bound on its rounding error). tol adds slack relative to the same magnitude; tol = 0 decides every sign exactly up to rounding.

Parameters:

tol (float)

Return type:

Union[Array, ndarray, bool, number, float, int]

metalog_jax.feasibility.analysis.tail_feasibility(engine, a, tol=0.0)[source]

Proposition 5 tail conditions at y -> 0 and y -> 1.

Each of s, s' and mu' at the edge counts as zero only within its own rounding error (see zero_test_width()), so a small but resolved s(0) > 0 is tail feasible whatever the scale of a.

Parameters:

  • engine (Engine) – Engine for the coefficient layout.

  • a (Union[Array, ndarray, bool, number]) – Coefficients, shape (k,).

  • tol (float) – Extra zero-test slack, relative to the terms’ magnitude.

Return type:

tuple[Union[Array, ndarray, bool, number], Union[Array, ndarray, bool, number]]

Returns:

(feasible at y -> 0, feasible at y -> 1).

metalog_jax.feasibility.analysis.check_feasibility(engine, a, tol=0.0)[source]

Air-tight feasibility test (Proposition 5 with Algorithm 1).

Unlike a grid test this cannot miss a negative-density region: M' is checked at every root of M'' (its local extrema) plus both tails. Every sign is decided exactly up to rounding (see zero_test_width()), so the verdict does not depend on the scale of a.

Parameters:

  • engine (Engine) – Engine for the coefficient layout.

  • a (Union[Array, ndarray, bool, number]) – Coefficients, shape (k,).

  • tol (float) – Extra zero-test slack, relative to the terms’ magnitude (0: exact up to rounding).

Return type:

FeasibilityReport

Returns:

FeasibilityReport.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

metalog_jax.feasibility.analysis.is_feasible(engine, a, tol=0.0)[source]

Shorthand for check_feasibility(engine, a, tol).feasible.

Parameters:
Return type:

Union[Array, ndarray, bool, number]

metalog_jax.feasibility.analysis.highest_mode(engine, a)[source]

Highest-density interior mode of a feasible unbounded metalog (Algorithm 1).

The density is 1 / M'(y), so among the modes (roots of M'' where the third derivative is non-negative) the highest has the smallest slope. A binding mode of a best feasible fit, where M'(y) is close to 0, has the highest density.

Parameters:
Returns:

the mode’s quantile M(y), its probability y, and whether an interior mode exists (a uniform metalog, for example, has none).

Return type:

tuple[Union[Array, ndarray, bool, number, float, int], Union[Array, ndarray, bool, number, float, int], Union[Array, ndarray, bool, number]]

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

metalog_jax.feasibility.analysis.moment_table(engine, max_order=4)[source]

Float table I(m, u) large enough for moments up to max_order.

Parameters:
Return type:

Union[Array, ndarray, bool, number]

metalog_jax.feasibility.analysis.raw_moment(engine, a, p, table=None)[source]

Exact E[X^p] = int_0^1 M(y)^p dy for an unbounded metalog (Proposition 3).

Expands (mu + s L)^p = sum_n C(p, n) mu^(p-n) s^n L^n by polynomial convolution and integrates each monomial with the exact I(m, n).

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

Parameters:
Return type:

Union[Array, ndarray, bool, number, float, int]

metalog_jax.feasibility.analysis.mean_and_variance(engine, a)[source]

Exact mean and variance of an unbounded metalog (Lemma 1, Proposition 3).

Parameters:
Return type:

tuple[Union[Array, ndarray, bool, number, float, int], Union[Array, ndarray, bool, number, float, int]]

Returns:

(mean, variance).

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

metalog_jax.feasibility.analysis.summary_stats(engine, a)[source]

Exact mean, variance, standard deviation, skewness and kurtosis.

Central moments are the raw moments of the metalog shifted by its mean, which avoids the cancellation of expanding E[(X - m)^p].

Parameters:
Return type:

dict[str, Union[Array, ndarray, bool, number, float, int]]

Returns:

Dict with keys mean, variance, std, skewness, kurtosis. The kurtosis is Pearson’s (not excess): 3 for a normal, 4.2 for the logistic (the 2-term metalog).

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

Engine

Static-shape engine for the feasibility algorithms of Baucells et al. (Algorithm 1).

A k-term metalog is written as M(y) = mu(x) + s(x) * logit(y) with x = y - 1/2. The engine splits a coefficient vector into the ascending coefficients of mu and s and evaluates the scaled derivative levels G_j(y) = (y(1-y))^j M^(j)(y) in logit coordinates u = logit(y), which keeps the tails (y within 1e-300 of 0 or 1) well conditioned.

Two term orderings are supported:

  • TermOrder.KEELIN_2016 - the ordering used by metalog_jax (Keelin 2016): 1, L, xL, x, x^2, x^2 L, x^3, x^3 L, …

  • TermOrder.METALOG_2 - the Metalog 2.0 ordering of the paper: 1, L, xL, x, x^2, x^2 L, x^3 L, x^3, x^4, …

They coincide for k <= 6 and differ from the 7th term on (paper, footnote 2).

All loops are over static sizes, so every public function here is jit/vmap compatible.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

class metalog_jax.feasibility.engine.TermOrder(*values)[source]

Bases: IntEnum

Ordering of the metalog basis terms.

Variables:

  • KEELIN_2016 – Ordering used by metalog_jax (Keelin 2016): 1, L, xL, x, x^2, x^2 L, x^3, x^3 L, …

  • METALOG_2 – Metalog 2.0 ordering of the paper: 1, L, xL, x, x^2, x^2 L, x^3 L, x^3, x^4, …

metalog_jax.feasibility.engine.term_layout(num_terms, order)[source]

Coefficient indices of the mu- and s-polynomials, sorted by power of (y - 1/2).

Parameters:

  • num_terms (int) – Number of metalog terms k.

  • order (TermOrder) – Basis ordering.

Return type:

tuple[list[int], list[int]]

Returns:

(mu_idx, s_idx) where a[mu_idx[p]] multiplies x^p and a[s_idx[p]] multiplies x^p * logit(y).

class metalog_jax.feasibility.engine.Engine(num_terms, order=TermOrder.KEELIN_2016)[source]

Bases: object

Trace-time tables and level evaluators for one (num_terms, order) pair.

Engines are hashable by identity and cached by get_engine(), so they can be passed as static arguments to jax.jit.

Parameters:
__init__(num_terms, order=TermOrder.KEELIN_2016)[source]

Build the exact tables for a k-term metalog.

Parameters:

  • num_terms (int) – Number of metalog terms (>= 2).

  • order (TermOrder) – Basis ordering.

Raises:

ValueError – If num_terms is less than 2.

split(a)[source]

Ascending coefficients (mu, s) of the location and scale polynomials.

Parameters:

a (Union[Array, ndarray, bool, number])

Return type:

tuple[Union[Array, ndarray, bool, number], Union[Array, ndarray, bool, number]]

level(j, mu, s, u)[source]

Value of G_j = (y(1-y))^j M^(j)(y) at y = sigmoid(u).

G_1(y) = y(1-y) M'(y) is the paper’s feasibility function; its limits at u = -inf, +inf are s(0) and s(1).

Parameters:
Return type:

Union[Array, ndarray, bool, number, float, int]

level_sign(j, mu, s, u, kappa=64.0)[source]

Sign of M^(j) at sigmoid(u) with a running-error zero test.

Both parts are rescaled by their joint log-magnitude before combining, so the sign is exact up to rounding even when (y(1-y))^j underflows. Returns 0 where the value is indistinguishable from zero (e.g. levels that vanish identically).

Parameters:
tail_sign(j, mu, s, at_one)[source]

Lemma 2: sign of lim M^(j) as y -> 0 (at_one=False) or y -> 1.

Parameters:
Return type:

Union[Array, ndarray, bool, number, float, int]

metalog_jax.feasibility.engine.get_engine(num_terms, order=TermOrder.KEELIN_2016)[source]

Cached Engine (stable identity keeps jit caches warm).

Parameters:
Return type:

Engine

metalog_jax.feasibility.engine.inflection_points(engine, a)[source]

All roots of M'' on (0, 1), i.e. the modes and anti-modes (Algorithm 1).

The highest level (y(1-y))^i M^(i) is an exact polynomial of degree i-1 (Lemma 4). Its roots are bracketed by running the same cascade over its own polynomial derivatives, so no eigenvalue solver is needed, and the roots are then carried down one level at a time to M'' using Lemma 2 edge signs.

Parameters:
Return type:

tuple[Union[Array, ndarray, bool, number], Union[Array, ndarray, bool, number]]

Returns:

(u, mask) - roots in logit coordinates (y = sigmoid(u)), sorted and padded with UMAX, and a boolean mask of valid entries. The static capacity is 2 * engine.istar - 3, exactly Proposition 1’s bound on the roots of M'' (2 * floor((k - 1) / 2) - 1), so no root can be dropped.

References

Baucells, M., Chrisman, L., Keelin, T. W., & Xu, Z. S. (2025). On the Properties of the Metalog Distribution. Darden Business School Working Paper No. 5279416. https://doi.org/10.2139/ssrn.5279416

Utilities

Utility functions for the Metalog JAX library.

exception metalog_jax.utils.NotFittedError[source]

Bases: Exception

A custom exception that raises when metalog has not been fit yet.

exception metalog_jax.utils.InvalidPDFError[source]

Bases: Exception

A custom exception that raises when metalog returns an invalid PDF.

class metalog_jax.utils.JaxUniformDistributionParameters(seed)[source]

Bases: object

Configuration parameters for a JAX uniform distribution.

Variables:

seed (int) – Seed for the psuedorandom number generator.

Parameters:

seed (int)

__init__(seed)
Parameters:

seed (int)

replace(**updates)

Returns a new object replacing the specified fields with new values.

class metalog_jax.utils.HDRPRNGParameters(trial=1, variable=0, entity=0, time=0, agent=0)[source]

Bases: object

Structure for Hubbard Decision Research Parameterized PRNG.

See paper below:

Douglas W. Hubbard. 2019. A multi-dimensional, counter-based psuedo random number generator as a standard for Monte Carlo simulations. In Proceedings of the Winter Simulation Conference. IEEE Press, 3064 - 3073.

Variables:

  • trial (int) – First simulation trial. MetalogBase.rvs draws trials trial to trial + size - 1, one per variate. Defaults to 1.

  • variable (int) – Parameter representing the random variable being sampled. Defaults to 0.

  • entity (int) – Parameter representing the entity or object being simulated. Defaults to 0.

  • time (int) – Parameter representing the time step in the simulation. Defaults to 0.

  • agent (int) – Parameter representing the agent or actor in the simulation. Defaults to 0.

Parameters:
__init__(trial=1, variable=0, entity=0, time=0, agent=0)
Parameters:
replace(**updates)

Returns a new object replacing the specified fields with new values.

metalog_jax.utils.assert_numeric_array(array)[source]

Assert that an array contains all float or int values.

Parameters:

array (Union[Array, ndarray, bool, number]) – Input array to validate.

Raises:

AssertionError – If the array does not contain numeric (float or int) values.

Return type:

None

metalog_jax.utils.assert_float_array(array)[source]

Assert that an array contains all float values.

Parameters:

array (Union[Array, ndarray, bool, number]) – Input array to validate.

Raises:

AssertionError – If the array does not contain numeric float values.

Return type:

None

metalog_jax.utils.assert_condition(predicate, message, *format_args, error=<class 'AssertionError'>)[source]

Check a scalar condition, inside or outside JAX transforms.

Outside transforms the predicate is concrete, and a false one raises error with message.format(*format_args). Under jax.vmap or jax.jit it is traced, so the check goes through checkify.check, which formats the message with the runtime values: jax.vmap raises checkify.JaxRuntimeError, and jax.jit needs the caller to functionalize it with checkify.checkify. Comparisons with NaN are false, so NaN fails the check.

Parameters:
Raises:

  • AssertionError – If the predicate is concrete and false, or error if given instead.

  • checkify.JaxRuntimeError – Under jax.vmap, if the predicate is false for any element of the batch.

Return type:

None

Example

>>> assert_condition(jnp.array(1.0) > 0, "x must be positive")  # No error
>>> assert_condition(jnp.array(-1.0) > 0, "x must be positive")
Traceback (most recent call last):
    ...
AssertionError: x must be positive
metalog_jax.utils.assert_probability_range(array)[source]

Assert that all array values are in the open interval (0, 1).

Validates that the minimum value in the array is strictly greater than 0 and the maximum value is strictly less than 1, ensuring all values represent valid probabilities in the open unit interval. Works inside jax.jit and jax.vmap through assert_condition.

Parameters:

array (Union[Array, ndarray, bool, number]) – Input array to validate.

Raises:

  • ValueError – If any value is <= 0 or >= 1, or is NaN.

  • checkify.JaxRuntimeError – Under jax.vmap, for the same values.

Return type:

None

Example

>>> valid_probs = jnp.array([0.1, 0.5, 0.9])
>>> assert_probability_range(valid_probs)  # No error
>>> invalid_probs = jnp.array([0.0, 0.5, 1.0])
>>> assert_probability_range(invalid_probs)
Traceback (most recent call last):
    ...
ValueError: Array minimum 0.0 must be strictly greater than 0
metalog_jax.utils.assert_strictly_ascending(x, tol=1e-12)[source]

Assert that array values are strictly ascending (monotonically increasing).

Validates that each element in the array is strictly greater than the previous element, with differences exceeding a specified tolerance threshold. This ensures the array represents a strictly increasing sequence, which is required for valid probability levels and quantile values in metalog distributions.

The function computes consecutive differences between array elements and verifies that all differences are strictly positive (greater than the tolerance). This is more robust than simple inequality checks as it accounts for numerical precision issues in floating-point comparisons.

Parameters:

  • x (Union[Array, ndarray, bool, number, float, int]) – Input array to validate. Should be a 1D numeric array representing values that must be in strictly ascending order (e.g., probability levels, quantile values, or time steps).

  • tol (float) – Tolerance threshold for considering differences as strictly positive. Must be a small positive value. Default is 1e-12, which is appropriate for double-precision floating-point arithmetic. Values separated by less than this tolerance are considered equal (not strictly ascending).

Raises:

AssertionError – If any consecutive difference is less than or equal to the tolerance threshold, indicating the array is not strictly ascending. This can occur if: - The array contains duplicate values (difference = 0) - The array is not sorted in ascending order (negative differences) - Consecutive values are too close together (difference < tol)

Return type:

None

Note

This function is used to validate inputs for metalog fitting, particularly probability levels (y values) which must be strictly increasing to ensure a valid probability distribution. The tolerance parameter helps handle numerical precision issues while still enforcing strict monotonicity.

Unlike simple sorting checks, this function ensures strict inequality (no duplicate values allowed), which is essential for metalog distributions where each probability level must map to a unique quantile value.

Example

Valid strictly ascending array:

>>> import jax.numpy as jnp
>>> from metalog_jax.utils import assert_strictly_ascending
>>>
>>> x = jnp.array([0.1, 0.2, 0.5, 0.9])
>>> assert_strictly_ascending(x)  # No error
>>> # Differences: [0.1, 0.3, 0.4] all > 1e-12

Invalid array with duplicate values:

>>> x_dup = jnp.array([0.1, 0.2, 0.2, 0.9])
>>> assert_strictly_ascending(x_dup)
Traceback (most recent call last):
    ...
AssertionError: Array values must be strictly ascending, with consecutive differences above 1e-12

Invalid array not in ascending order:

>>> x_unsorted = jnp.array([0.1, 0.5, 0.2, 0.9])
>>> assert_strictly_ascending(x_unsorted)
Traceback (most recent call last):
    ...
AssertionError: Array values must be strictly ascending, with consecutive differences above 1e-12

Using custom tolerance:

>>> # Values closer together than the default tolerance (1e-12)
>>> x_close = jnp.array([0.1, 0.1 + 1e-13, 0.2])
>>> assert_strictly_ascending(x_close)
Traceback (most recent call last):
    ...
AssertionError: Array values must be strictly ascending, with consecutive differences above 1e-12
>>> assert_strictly_ascending(x_close, tol=1e-14)  # Passes

See also

assert_probability_range: Validates array values are in (0, 1). jnp.diff: Computes consecutive differences between array elements.

References

Used extensively in metalog fitting to validate probability level arrays (y values) before computing quantile functions or fitting coefficients.

metalog_jax.utils.hdrprng(hdr_parameters)[source]

Generate a pseudo-random number using the Hubbard Decision Research PRNG algorithm.

Implements a multi-dimensional, counter-based pseudo-random number generator designed for reproducible Monte Carlo simulations. This method validates the input parameters and delegates to the JIT-compiled internal implementation.

This PRNG is particularly useful for large-scale simulations where: - Reproducibility is critical - Multiple independent random streams are needed - Parallel execution requires uncorrelated random sequences

Parameters:

hdr_parameters (HDRPRNGParameters) – HDRPRNGParameters instance containing the five integer parameters (trial, variable, entity, time, agent) that determine the random output.

Return type:

float

Returns:

A pseudo-random float in the open interval (0, 1).

Raises:

AssertionError – If hdr_parameters.trial is not a positive scalar (validated by HDRPRNGParameters.__post_init__).

Example

>>> from metalog_jax.utils import HDRPRNGParameters, hdrprng
>>> params = HDRPRNGParameters(trial=1, variable=0, entity=0, time=0, agent=0)
>>> random_value = hdrprng(params)
>>> # Same parameters always produce the same output
>>> random_value_2 = hdrprng(params)
>>> assert random_value == random_value_2

References

Douglas W. Hubbard. 2019. A multi-dimensional, counter-based pseudo random number generator as a standard for Monte Carlo simulations. In Proceedings of the Winter Simulation Conference. IEEE Press, 3064-3073.

Note

This is the public API method. The actual computation is performed by the internal _hdrprng() function, which is JIT-compiled for performance. The algorithm uses a combination of large prime numbers and modular arithmetic to generate uniformly distributed pseudo-random values.

metalog_jax.utils.find_nearest_index(array, value)[source]

Find the index of the nearest value in an array.

Parameters:
Return type:

Union[Array, ndarray, bool, number, float, int]

Returns:

Index of the nearest value

metalog_jax.utils.ks_distance(x, y)[source]

Compute the two-sample Kolmogorov-Smirnov (KS) distance between two samples.

Calculates the maximum absolute difference between the empirical cumulative distribution functions (ECDFs) of two samples. This is the Kolmogorov-Smirnov statistic for the two-sample test, which measures the degree of dissimilarity between two probability distributions.

The KS distance is defined as D = sup_t |F_x(t) - F_y(t)|, where F_x and F_y are the empirical CDFs of samples x and y, and the supremum is taken over all values t.

This implementation is fully vectorized and JIT-compiled for efficient execution with JAX, making it suitable for use in gradient-based optimization and parallel computation workflows. The function uses jnp.searchsorted for efficient CDF computation, avoiding explicit loops.

Parameters:

  • x (Union[Array, ndarray, bool, number, float, int]) – First sample array of shape (n,). Contains n independent observations from the first distribution. Can be unsorted (will be sorted internally).

  • y (Union[Array, ndarray, bool, number, float, int]) – Second sample array of shape (m,). Contains m independent observations from the second distribution. Can be unsorted (will be sorted internally). Does not need to have the same length as x.

Return type:

Union[float, int]

Returns:

Scalar KS distance in the range [0, 1]. The KS statistic represents the maximum vertical distance between the two empirical CDFs: - 0: The two samples have identical empirical distributions - 1: The two samples are completely non-overlapping - Values closer to 0 indicate more similar distributions

Note

  • Both input samples are sorted internally, so there’s no need to pre-sort

  • The function is JIT-compiled, so the first call may be slower due to compilation overhead, but subsequent calls will be very fast

  • This computes only the KS distance (D statistic), not the p-value for the hypothesis test

  • The function is differentiable and can be used in gradient-based optimization (e.g., for distribution matching)

  • Memory complexity is O(n + m) due to the combined sorted values array

Example

Basic usage with identical distributions:

>>> import jax.numpy as jnp
>>> from metalog_jax.utils import ks_distance
>>>
>>> # Two samples from the same distribution
>>> x = jnp.array([0.1, 0.5, 0.3, 0.7, 0.9])
>>> y = jnp.array([0.2, 0.4, 0.6, 0.8])
>>> round(float(ks_distance(x, y)), 2)  # Small: similar distributions
0.2

Comparing different distributions:

>>> # Sample from uniform(0, 1)
>>> uniform_sample = jnp.array([0.1, 0.3, 0.5, 0.7, 0.9])
>>> # Sample from distribution concentrated near 0
>>> skewed_sample = jnp.array([0.05, 0.1, 0.15, 0.2, 0.25])
>>> round(float(ks_distance(uniform_sample, skewed_sample)), 2)
0.8

Using for metalog goodness-of-fit:

>>> from metalog_jax.base import MetalogBoundedness, MetalogFitMethod
>>> from metalog_jax.base import MetalogInputData, MetalogParameters
>>> from metalog_jax.base import MetalogRandomVariableParameters
>>> from metalog_jax.metalog import fit
>>>
>>> # Fit metalog to data
>>> samples = jnp.array([1.2, 2.3, 2.8, 3.5, 4.1, 5.6, 6.2, 7.8, 9.1, 10.5])
>>> data = MetalogInputData.from_values(samples, DEFAULT_Y, False)
>>> params = MetalogParameters(
...     boundedness=MetalogBoundedness.UNBOUNDED,
...     method=MetalogFitMethod.Feasible,
...     num_terms=5, lower_bound=0.0, upper_bound=0.0
... )
>>> metalog = fit(data, params)
>>>
>>> # Draw samples from the fitted distribution
>>> rv_params = MetalogRandomVariableParameters(
...     prng_params=JaxUniformDistributionParameters(seed=0), size=1000
... )
>>> fitted_samples = metalog.rvs(rv_params)
>>>
>>> # A low KS distance indicates a good fit
>>> bool(ks_distance(samples, fitted_samples) < 0.2)
True

See also

jnp.searchsorted: Efficient binary search used for CDF computation. scipy.stats.ks_2samp: SciPy equivalent (includes p-value calculation).

References

Kolmogorov, A. N. (1933). “Sulla determinazione empirica di una legge di distribuzione”. Giornale dell’Istituto Italiano degli Attuari, 4: 83-91.

Smirnov, N. V. (1948). “Table for estimating the goodness of fit of empirical distributions”. Annals of Mathematical Statistics, 19(2): 279-281.

Implementation Details:

The algorithm works as follows: 1. Sort both input samples (O(n log n + m log m)) 2. Combine all unique values from both samples and sort (O((n+m) log(n+m))) 3. Use binary search to compute empirical CDFs at all combined points (O((n+m) log n)) 4. Find maximum absolute difference between the CDFs (O(n+m)) Overall complexity: O((n+m) log(n+m))