Metalog JAX

https://tjefferies.github.io/metalog_jax/
Admin
Metalog JAX is a flexible probability distribution library providing JAX-based implementations with...
Tokens:88,317
Snippets:502
Trust Score:9.7
Update:6 months ago
Show doc for...
Tokens:
Raw
### Install metalog-jax using pip or uv

Source: https://tjefferies.github.io/metalog_jax/_sources/getting-started.rst.txt

Instructions for installing the metalog-jax library using package managers pip and uv. This is the first step to using the library.

```bash
pip install metalog-jax
```

```bash
uv pip install metalog-jax
```

--------------------------------

### Fitting a Metalog Distribution with JAX

Source: https://tjefferies.github.io/metalog_jax/_sources/index.rst.txt

This example demonstrates how to initialize Metalog input data from raw values, configure distribution parameters for an unbounded OLS fit, and generate quantiles. It highlights the use of the fit function and accessing distribution properties like mean and median.

```python
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
)

# 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)

# Generate quantiles
quantiles = m.ppf(jnp.array([0.25, 0.5, 0.75]))

# Access distribution properties
print(f"Mean: {m.mean}, Median: {m.median}, Std: {m.std}")
```

--------------------------------

### Fit Metalog Distribution using OLS in JAX

Source: https://tjefferies.github.io/metalog_jax/index.html

This example demonstrates how to fit a Metalog distribution using the Ordinary Least Squares (OLS) method with JAX. It covers creating input data, configuring Metalog parameters, fitting the distribution, generating quantiles, and accessing distribution properties like mean, median, and standard deviation. Dependencies include jax.numpy and specific modules from metalog_jax.

```python
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)

# Generate quantiles
quantiles = m.ppf(jnp.array([0.25, 0.5, 0.75]))

# Access distribution properties
print(f"Mean: {m.mean}, Median: {m.median}, Std: {m.std}")

```

--------------------------------

### Regression Methods: OLS and LASSO

Source: https://tjefferies.github.io/metalog_jax/getting-started.html

Examples for configuring the fitting process using Ordinary Least Squares (OLS) or LASSO regularization.

```python
# OLS Example
from metalog_jax.base import MetalogParameters, MetalogBoundedness, MetalogFitMethod
from metalog_jax.metalog import fit

metalog_params = MetalogParameters(
    boundedness=MetalogBoundedness.UNBOUNDED,
    method=MetalogFitMethod.OLS,
    lower_bound=0.0,
    upper_bound=0.0,
    num_terms=5
)
m = fit(data, metalog_params)

# LASSO Import Example
from metalog_jax.regression.lasso import LassoParameters
```

--------------------------------

### Load and Print Metalog Documentation (Python)

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

This snippet demonstrates how to load and print the documentation string for a Metalog function using Python. It assumes the Metalog library is installed and accessible within the Python environment.

```python
print(Metalog.load.__doc__)
```

--------------------------------

### Initialize JAX Uniform PRNG Parameters

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Initializes a PRNG parameter object for a JAX uniform distribution. This is a standard JAX PRNG setup used for generating samples from a distribution.

```python
# to generate samples from distribution, we first init a PRNG parameter object
# below is standard jax PRNG
uniform_prng_params = JaxUniformDistributionParameters(seed=42)
uniform_prng_params
```

--------------------------------

### Configure Shared Parameters for Batched Fit (Python)

Source: https://tjefferies.github.io/metalog_jax/fitting_grids.html

Configures shared Metalog parameters for fitting multiple batched datasets. This example sets up parameters for a strictly lower-bounded distribution using OLS fitting with a specified number of terms. Dependencies include MetalogParameters, MetalogBoundedness, MetalogFitMethod.

```python
# Configure shared parameters
params_batch = MetalogParameters(
    boundedness=MetalogBoundedness.STRICTLY_LOWER_BOUND,
    lower_bound=0,
    upper_bound=0,
    method=MetalogFitMethod.OLS,
    num_terms=7,
)
```

--------------------------------

### Initialize Metalog Random Variable Parameters

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Initializes Metalog Random Variable parameters, combining PRNG parameters with the desired size for drawing random variables. This setup is used for drawing samples from a fitted metalog distribution.

```python
# init metalog RV parameters
metalog_rv_params = MetalogRandomVariableParameters(
    prng_params=uniform_prng_params, size=10
)
```

--------------------------------

### GET /metalog/median

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Retrieves the median value of the metalog distribution.

```APIDOC
## GET /metalog/median

### Description
Computes the median of the metalog distribution, which is the 50th percentile, computed exactly via ppf(0.5).

### Method
GET

### Endpoint
/metalog/median

### Response
#### Success Response (200)
- **median** (chex.Scalar) - The median of the distribution.

#### Response Example
{
  "median": 50.077335
}
```

--------------------------------

### MetalogParameters Initialization and Configuration

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Demonstrates how to initialize MetalogParameters with specific configurations for boundedness, fitting method, and number of terms.

```APIDOC
## MetalogParameters Initialization

### Description
Initializes MetalogParameters with specified configuration for metalog distribution fitting.

### Method
Python Class Instantiation

### Endpoint
N/A (Local Class Instantiation)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **boundedness** (MetalogBoundedness) - Required - Type of domain constraint for the distribution.
- **lower_bound** (float) - Optional - Lower boundary value (used if boundedness is not UNBOUNDED).
- **upper_bound** (float) - Optional - Upper boundary value (used if boundedness is not UNBOUNDED).
- **method** (MetalogFitMethod) - Required - Regression method used for fitting.
- **num_terms** (int) - Required - Number of terms in the metalog series expansion.

### Request Example
```python
from metalog_jax.base import MetalogBoundedness, MetalogFitMethod

metalog_params = MetalogParameters(
    boundedness=MetalogBoundedness.UNBOUNDED,
    lower_bound=0,
    upper_bound=1,
    method=MetalogFitMethod.OLS,
    num_terms=5,
)
```

### Response
#### Success Response (200)
N/A (Object Instantiation)

#### Response Example
N/A
```

--------------------------------

### GET /metalog/std

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Computes the standard deviation of the metalog distribution.

```APIDOC
## GET /metalog/std

### Description
Computes the standard deviation of the metalog distribution using Monte Carlo sampling with 20,000 draws.

### Method
GET

### Endpoint
/metalog/std

### Response
#### Success Response (200)
- **std** (float) - The estimated standard deviation of the distribution.

#### Response Example
{
  "std": 5.67292401
}
```

--------------------------------

### GET /metalog/mode

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Retrieves the mode (most likely value) of the metalog distribution.

```APIDOC
## GET /metalog/mode

### Description
Computes the mode of the fitted metalog distribution by locating the maximum of the PDF over a fine grid of probability values.

### Method
GET

### Endpoint
/metalog/mode

### Response
#### Success Response (200)
- **mode** (chex.Scalar) - The mode of the distribution.

#### Response Example
{
  "mode": 50.66968091
}
```

--------------------------------

### Find Best Configuration for Each Dataset

Source: https://tjefferies.github.io/metalog_jax/fitting_grids.html

Demonstrates how to iterate through dataset results from a 3D grid search, identify the best configuration (L1 penalty and number of terms) for each dataset using `find_best_config`, and print the optimal parameters and corresponding KS distance.

```APIDOC
## Find best configuration for each dataset

### Description
This code snippet iterates through the results of a 3D grid search performed by `fit_grid`. For each dataset, it uses `find_best_config` to determine the optimal combination of L1 penalty and number of terms that minimizes the Kolmogorov-Smirnov (KS) distance. It then prints the best found parameters for each distribution.

### Method
Iterative processing of grid search results.

### Endpoint
N/A (This is a code example for a library function)

### Parameters
None directly applicable as this is a code example.

### Request Example
```python
# Assuming result_3d, dist_names, l1_3d, num_terms_3d are pre-defined from a fit_grid operation
print("Best configuration per dataset:")
for _i, _name in enumerate(dist_names):
    _best_idx, _best_ks = find_best_config(result_3d.ks_dist[_i])
    _best_l1 = l1_3d[_best_idx[0]]
    _best_nt = num_terms_3d[_best_idx[1]]
    print(
        f"  {_name}: L1={float(_best_l1):.3f}, {_best_nt} terms (KS = {float(_best_ks):.4f})"
    )

print() # Newline for separation
print("Full 3D grid for first dataset (Lognormal):")
print("         ", end="")
for _nt in num_terms_3d:
    print(f"{_nt:>8} terms", end="")
print()
for _j, _l1 in enumerate(l1_3d):
    print(f"L1={float(_l1):5.3f}", end="")
    for _k in range(len(num_terms_3d)):
        print(f"     {float(result_3d.ks_dist[0, _j, _k]):.4f}", end="")
    print()
```

### Response
Prints the best configuration (L1 penalty, number of terms, and KS distance) for each dataset, followed by a detailed grid of KS distances for the first dataset.

#### Success Response (200)
Output is printed to the console.

#### Response Example
```
Best configuration per dataset:
  Lognormal: L1=0.000, 7 terms (KS = 0.0381)
  Weibull: L1=0.000, 5 terms (KS = 0.0381)
  Gamma: L1=0.000, 5 terms (KS = 0.0286)

Full 3D grid for first dataset (Lognormal):
                5 terms       7 terms       9 terms
L1=0.000     0.0476     0.0381     0.0381
L1=0.010     0.0476     0.0381     0.0381
L1=0.100     0.0762     0.0762     0.0762
```
```

--------------------------------

### GET /metalog/mean

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

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

```APIDOC
## GET /metalog/mean

### Description
Computes the mean of the metalog distribution using Monte Carlo sampling with 20,000 draws.

### Method
GET

### Endpoint
/metalog/mean

### Response
#### Success Response (200)
- **mean** (float) - The estimated mean of the distribution.

#### Response Example
{
  "mean": 50.03155654
}
```

--------------------------------

### GET /spt_metalog/num_terms

Source: https://tjefferies.github.io/metalog_jax/api/modules.html

Retrieves the number of terms used in the SPT metalog expansion.

```APIDOC
## GET /spt_metalog/num_terms

### Description
Returns the number of terms in the SPT metalog expansion. For SPTMetalog, this is always 3.

### Method
GET

### Endpoint
/spt_metalog/num_terms

### Response
#### Success Response (200)
- **num_terms** (int) - The number of terms (always 3).

#### Response Example
{
  "num_terms": 3
}
```

--------------------------------

### Fit SPT Metlog Distribution with JAX

Source: https://tjefferies.github.io/metalog_jax/api/modules.html

Demonstrates fitting an SPT metalog distribution using sample data and custom parameters. It shows how to configure SPT parameters like alpha and boundedness, and asserts the expected number of coefficients in the fitted distribution.

```python
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
```

--------------------------------

### GET /metalog/variance

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Retrieves the estimated variance of the metalog distribution using Monte Carlo sampling.

```APIDOC
## GET /metalog/variance

### Description
Computes the variance of the metalog distribution. This is estimated via Monte Carlo sampling with 20,000 draws.

### Method
GET

### Endpoint
/metalog/variance

### Response
#### Success Response (200)
- **variance** (chex.Scalar) - The estimated variance of the distribution.

#### Response Example
{
  "variance": 32.18206683
}
```

--------------------------------

### Initialize Metalog-JAX Environment

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Imports the necessary classes, functions, and utilities required to perform metalog distribution fitting.

```python
import tempfile
from pathlib import Path
import jax.numpy as jnp
import numpy as np
from scipy.stats import t
from metalog_jax.base import (
    MetalogBoundedness,
    MetalogFitMethod,
    MetalogInputData,
    MetalogParameters,
    MetalogPlotOptions,
    MetalogRandomVariableParameters,
)
from metalog_jax.metalog import (
    Metalog,
    fit,
)
from metalog_jax.utils import (
    DEFAULT_Y,
    HDRPRNGParameters,
    JaxUniformDistributionParameters,
)
```

--------------------------------

### GET /metalog/isf

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Computes the inverse survival function (inverse complementary CDF) for the Metalog distribution.

```APIDOC
## GET /metalog/isf

### Description
Returns the value x such that P(X > x) = q. This is equivalent to the Percent Point Function (ppf) of (1 - q).

### Method
GET

### Endpoint
/metalog/isf

### Parameters
#### Query Parameters
- **q** (array) - Required - Probability values in (0, 1).

### Response
#### Success Response (200)
- **quantiles** (array) - Quantile values corresponding to the provided survival probabilities.

#### Response Example
{
  "quantiles": [73.95, 54.62, 29.95]
}
```

--------------------------------

### Prepare Batched Datasets (Python)

Source: https://tjefferies.github.io/metalog_jax/fitting_grids.html

Generates multiple sample distributions (lognormal, Weibull, gamma) and prepares them as batched input data for Metalog. This demonstrates how to stack data from different distributions into a single array for parallel processing, which is useful when fitting many similar distributions. Dependencies include lognorm, weibull_min, gamma from scipy.stats, MetalogInputData, and DEFAULT_Y.

```python
# Generate multiple distributions (all lower-bounded at 0)
dist_lognorm = lognorm(s=0.5, loc=0, scale=1).rvs(size=200, random_state=42)
dist_weibull = weibull_min(c=2, scale=2).rvs(size=200, random_state=43)
dist_gamma = gamma(a=4, scale=1).rvs(size=200, random_state=44)

# Create input data for each
data1 = MetalogInputData.from_values(dist_lognorm, DEFAULT_Y, False)
data2 = MetalogInputData.from_values(dist_weibull, DEFAULT_Y, False)
data3 = MetalogInputData.from_values(dist_gamma, DEFAULT_Y, False)

# Stack into batched arrays
batched_x = jnp.stack([data1.x, data2.x, data3.x])
batched_y = jnp.stack([data1.y, data2.y, data3.y])

print(f"Batched x shape: {batched_x.shape}")
print(f"Batched y shape: {batched_y.shape}")
```

--------------------------------

### Importing Metalog JAX Dependencies

Source: https://tjefferies.github.io/metalog_jax/_sources/fitting_grids.ipynb.txt

Initializes the environment by importing necessary JAX, SciPy, and metalog_jax modules required for grid search operations.

```python
import jax.numpy as jnp
from scipy.stats import beta, gamma, lognorm, norm, weibull_min
from metalog_jax.base import (
    MetalogBoundedness,
    MetalogFitMethod,
    MetalogInputData,
    MetalogParameters,
)
from metalog_jax.grid_search import find_best_config, fit_grid
from metalog_jax.utils import DEFAULT_Y
```

--------------------------------

### GET /extract_best_from_grid

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/grid_search.html

Extracts the optimal Metalog configuration for a specific dataset from a grid result object.

```APIDOC
## GET /extract_best_from_grid

### Description
Finds the configuration with the minimum KS distance for a specific dataset index and returns the corresponding GridResult.

### Method
GET

### Endpoint
/extract_best_from_grid

### Parameters
#### Request Body
- **grid_results** (GridResult) - Required - The result object from a grid search.
- **dataset_idx** (int) - Required - The index of the dataset to extract.

### Response
#### Success Response (200)
- **metalog** (Metalog) - The best metalog for the specified dataset.
- **ks_dist** (float) - The corresponding KS distance.
```

--------------------------------

### Find Best Configuration for Each Dataset (Python)

Source: https://tjefferies.github.io/metalog_jax/fitting_grids.html

Iterates through datasets to find the best Metalog configuration (L1 penalty and number of terms) based on KS distance. It prints the optimal parameters for each dataset and displays the full 3D grid for the first dataset.

```python
print("Best configuration per dataset:")
for _i, _name in enumerate(dist_names):
    _best_idx, _best_ks = find_best_config(result_3d.ks_dist[_i])
    _best_l1 = l1_3d[_best_idx[0]]
    _best_nt = num_terms_3d[_best_idx[1]]
    print(
        f"  {_name}: L1={float(_best_l1):.3f}, {_best_nt} terms (KS = {float(_best_ks):.4f})"
    )

print()
print("Full 3D grid for first dataset (Lognormal):")
print("         ", end="")
for _nt in num_terms_3d:
    print(f"{_nt:>8} terms", end="")
print()
for _j, _l1 in enumerate(l1_3d):
    print(f"L1={float(_l1):5.3f}", end="")
    for _k in range(len(num_terms_3d)):
        print(f"     {float(result_3d.ks_dist[0, _j, _k]):.4f}", end="")
    print()
```

--------------------------------

### GET /metalog/logpdf

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Computes the natural logarithm of the probability density function (log-PDF) for the metalog distribution.

```APIDOC
## GET /metalog/logpdf

### Description
Computes the natural logarithm of the probability density function (log-PDF) at specified probability values.

### Method
GET

### Endpoint
/metalog/logpdf

### Parameters
#### Query Parameters
- **x** (array/scalar) - Required - Probability values in (0, 1) at which to evaluate log(PDF).

### Response
#### Success Response (200)
- **log_density** (chex.Numeric) - Log-density values with the same shape as input.

#### Response Example
[-7.90282201, -6.80756844, ...]
```

--------------------------------

### Loading a Previously Saved Metalog Instance

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Demonstrates how to load a Metalog distribution instance from a previously saved JSON file.

```APIDOC
## Loading a Previously Saved Metalog Instance

### Description
Loads a fitted metalog distribution from a JSON file. Returns a new instance of the metalog class with loaded coefficients and parameters.

### Method
```python
loaded_metalog = Metalog.load(path)
```

### Parameters
#### Arguments
* `path` (Path): File path to the JSON file containing the serialized distribution.

### Returns
* `T_MetalogBase`: A new instance of the metalog class with loaded coefficients and parameters.

### Request Example
```json
{
  "path": "/path/to/your/metalog_model.json"
}
```
```

--------------------------------

### GET /find_best_config

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/grid_search.html

Finds the configuration with the minimum Kolmogorov-Smirnov (KS) distance within a multi-dimensional grid of results.

```APIDOC
## GET /find_best_config

### Description
Identifies the indices and value of the minimum KS distance in a grid of any dimensionality.

### Method
GET

### Endpoint
/find_best_config

### Parameters
#### Query Parameters
- **ks_dists** (jnp.ndarray) - Required - Array of KS distances with shape (d1, d2, ..., dn).

### Response
#### Success Response (200)
- **best_indices** (jnp.ndarray) - Indices of the minimum value.
- **best_ks** (float) - The minimum KS distance value.

### Response Example
{
  "best_indices": [0, 2],
  "best_ks": 0.012
}
```

--------------------------------

### GET /metalog/pdf

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Computes the probability density function (PDF) of the metalog distribution at specified probability values.

```APIDOC
## GET /metalog/pdf

### Description
Computes the probability density function (PDF) of the metalog distribution. The PDF is calculated as the reciprocal of the derivative of the quantile function (PPF).

### Method
GET

### Endpoint
/metalog/pdf

### Parameters
#### Query Parameters
- **x** (array/scalar) - Required - Probability values in (0, 1) at which to evaluate the PDF.

### Response
#### Success Response (200)
- **density** (chex.Numeric) - Probability density values with the same shape as input.

#### Response Example
[0.0003697, 0.00110538, ...]
```

--------------------------------

### Fit Using Precomputed Quantiles with Metalog (Python)

Source: https://tjefferies.github.io/metalog_jax/fitting_grids.html

Shows how to use the fit_grid function when quantiles and their corresponding probabilities are already known, simulating expert elicitation. It sets precomputed_quantiles=True and prints the KS distance and coefficients.

```python
# Expert-elicited quantiles for a 0-100 bounded variable
quantiles = jnp.array([10.0, 25.0, 40.0, 50.0, 60.0, 75.0, 90.0])
probabilities = jnp.array([0.05, 0.20, 0.40, 0.50, 0.60, 0.80, 0.95])

params_quantiles = MetalogParameters(
    boundedness=MetalogBoundedness.BOUNDED,
    lower_bound=0,
    upper_bound=100,
    method=MetalogFitMethod.OLS,
    num_terms=5,
)

result_quantiles = fit_grid(
    quantiles, probabilities, params_quantiles, precomputed_quantiles=True
)

print(
    f"Precomputed quantiles fit KS distance: {float(result_quantiles.ks_dist):.4f}"
)
print(f"Coefficients: {result_quantiles.metalog.a}")
```

--------------------------------

### Calculating Quantiles with Metalog

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Demonstrates how to compute quantiles from a fitted Metalog distribution using the q method, which is designed to be compatible with the rmetalog API.

```python
# get quantiles from fitted metalog, compatible with rmetalog API
metalog.q(DEFAULT_Y)
```

--------------------------------

### GET /metalog/sf

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Computes the survival function (complement of CDF) for the Metalog distribution, representing P(X > x).

```APIDOC
## GET /metalog/sf

### Description
Computes the survival function S(x) = 1 - F(x) = P(X > x) for a given set of quantile values.

### Method
GET

### Endpoint
/metalog/sf

### Parameters
#### Query Parameters
- **x** (array) - Required - Quantile values at which to evaluate the survival function.

### Response
#### Success Response (200)
- **probabilities** (array) - Survival probabilities in (0, 1) with the same shape as input.

#### Response Example
{
  "probabilities": [0.95, 0.50, 0.05]
}
```

--------------------------------

### GET /metalog/properties

Source: https://tjefferies.github.io/metalog_jax/api/modules.html

Access properties of a fitted Metalog distribution instance, such as the number of terms or the fitting method used.

```APIDOC
## GET /metalog/properties

### Description
Retrieves configuration properties from an existing Metalog distribution instance.

### Method
GET

### Endpoint
/metalog/properties

### Parameters
#### Query Parameters
- **instance_id** (string) - Required - The identifier of the fitted metalog instance.

### Response
#### Success Response (200)
- **num_terms** (int) - Number of basis functions used.
- **method** (string) - The regression method used for fitting (OLS or LASSO).

#### Response Example
{
  "num_terms": 9,
  "method": "OLS"
}
```

--------------------------------

### Initialize MetalogParameters

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Initializes MetalogParameters with specified configurations for boundedness, bounds, fitting method, and number of terms. This object is used for fitting metalog distributions via regression.

```python
metalog_params = MetalogParameters(
    boundedness=MetalogBoundedness.UNBOUNDED,
    lower_bound=0,
    upper_bound=1,
    method=MetalogFitMethod.OLS,
    num_terms=5,
)
```

--------------------------------

### Fit and Predict with OLS Regression in JAX

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/regression/ols.html

Demonstrates how to define an OLS model, fit it to feature data using JAX, and perform predictions. The implementation uses numerical SVD for stability.

```python
import jax.numpy as jnp
from metalog_jax.regression.ols import fit_ordinary_least_squares, predict_ordinary_least_squares

# Prepare data
X = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
y = jnp.array([1.0, 2.0, 3.0])

# Fit model
model = fit_ordinary_least_squares(X, y)

# Predict
X_test = jnp.array([[2.0, 3.0], [4.0, 5.0]])
predictions = predict_ordinary_least_squares(X_test, model)
```

--------------------------------

### GET /metalog/fit-function

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/metalog.html

Retrieves the appropriate regression implementation function based on the specified fit method and optional hyperparameters.

```APIDOC
## GET /metalog/fit-function

### Description
Returns a callable function used to fit a metalog distribution. It uses a dispatch table to map methods like OLS or Lasso to their respective implementations.

### Method
GET

### Endpoint
/metalog/fit-function

### Parameters
#### Query Parameters
- **method** (enum) - Required - The regression method to use (e.g., OLS, Lasso).
- **hyperparams** (object) - Optional - Regularization parameters for Lasso regression.

### Request Example
{
  "method": "Lasso",
  "hyperparams": {"lambda": 0.1}
}

### Response
#### Success Response (200)
- **function** (Callable) - The regression fit function, potentially with hyperparameters bound via partial.

#### Response Example
{
  "status": "success",
  "function_name": "fit_lasso"
}
```

--------------------------------

### Get Fit Function for Method

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/grid_search.html

Retrieves the appropriate fitting function based on the specified MetalogFitMethod (OLS or Lasso).

```APIDOC
## POST /_get_fit_function_for_method

### Description
Returns the appropriate fitting function based on the provided `MetalogFitMethod`. It handles the selection and potential initialization of regression parameters for methods like Lasso.

### Method
POST

### Endpoint
`/_get_fit_function_for_method`

### Parameters
#### Request Body
- **method** (MetalogFitMethod) - Required - The enum value specifying the fitting method (e.g., 'OLS', 'Lasso').
- **regression_params** (Any) - Optional - Parameters specific to the regression method, such as `LassoParameters` for Lasso.

### Request Example
```json
{
  "method": "Lasso",
  "regression_params": {
    "lam": 0.1,
    "learning_rate": 0.01,
    "num_iters": 1000,
    "tol": 1e-05,
    "momentum": 0.9
  }
}
```

### Response
#### Success Response (200)
- **fit_function** (function) - A callable function that takes `(target, quantiles)` and returns a regression model.

#### Response Example
```json
{
  "fit_function": "<function fit_lasso at 0x...>"
}
```
```

--------------------------------

### Fit Metalog Distribution

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Demonstrates how to fit a metalog distribution to input data using specified parameters.

```python
# fit metalog, return distribution object
metalog = fit(data_1, metalog_params)
```

--------------------------------

### GET /grid_result

Source: https://tjefferies.github.io/metalog_jax/api/modules.html

Represents the outcome of a metalog grid search, providing the fitted distribution and the Kolmogorov-Smirnov distance for goodness-of-fit evaluation.

```APIDOC
## GET /grid_result

### Description
Encapsulates the results of a metalog fitting process, typically used during hyperparameter optimization to compare multiple configurations via the KS distance.

### Method
GET

### Endpoint
/grid_result

### Response
#### Success Response (200)
- **metalog** (Object) - The fitted distribution (Standard or SPT).
- **ks_dist** (float) - Kolmogorov-Smirnov distance (0 to 1). Lower values indicate a better fit.
```

--------------------------------

### Calculate Log-PPF and Survival Function

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Demonstrates how to compute the natural logarithm of the percent point function (logppf) and the survival function (sf) for a Metalog distribution object.

```python
metalog.logppf(DEFAULT_Y)
metalog.sf(ppf_1)
```

--------------------------------

### Initialize MetalogInputData with Factory Method

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Demonstrates how to create a validated MetalogInputData instance using the from_values classmethod. This ensures input data is correctly processed for either raw sample data or precomputed quantiles.

```python
data_1 = MetalogInputData.from_values(
    x=rvs, y=DEFAULT_Y, precomputed_quantiles=False
)
```

--------------------------------

### Get Quantiles from Fitted Metalog in Python

Source: https://tjefferies.github.io/metalog_jax/_sources/basic_usage.ipynb.txt

Retrieves quantiles from a fitted metalog distribution object. This is useful for understanding the distribution's behavior across different probability levels.

```python
# get quantiles from fitted metalog
metalog.ppf(DEFAULT_Y)
```

--------------------------------

### Metalog PDF Documentation

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Provides documentation for the `metalog.pdf` function, explaining its purpose, arguments, return values, and potential exceptions. It clarifies that the function computes the PDF by inverting the PPF and handling bounded distributions.

```python
print(Metalog.pdf.__doc__)
```

--------------------------------

### Ordinary Least Squares (OLS) Regression Method

Source: https://tjefferies.github.io/metalog_jax/_sources/getting-started.rst.txt

Example of configuring metalog-jax to use the Ordinary Least Squares (OLS) regression method for fitting. OLS is the default method and does not involve regularization.

```python
from metalog_jax.base import MetalogInputData, MetalogParameters
from metalog_jax.base import MetalogBoundedness, MetalogFitMethod
from metalog_jax.metalog import fit

metalog_params = MetalogParameters(
    boundedness=MetalogBoundedness.UNBOUNDED,

```

--------------------------------

### Prepare Metalog Input Data

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Demonstrates how to create a MetalogInputData object from either precomputed quantiles or raw sample data. The from_values class method is required to ensure data validity.

```python
# Setup sample data
dist = t(df=6, loc=50, scale=5)
rvs = jnp.array(dist.rvs(size=1000, random_state=42))
ppf = jnp.array(dist.ppf(DEFAULT_Y))

# Init from precomputed quantiles
data = MetalogInputData.from_values(x=ppf, y=DEFAULT_Y, precomputed_quantiles=True)

# Init from raw data
data_1 = MetalogInputData.from_values(x=rvs, y=DEFAULT_Y, precomputed_quantiles=False)
```

--------------------------------

### Get Fit Function for Regression Method

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/grid_search.html

Selects and returns the appropriate fitting function based on the specified MetalogFitMethod. Supports OLS and Lasso, returning a lambda function for Lasso to include regression parameters.

```python
def _get_fit_function_for_method(
    method: MetalogFitMethod, regression_params: Any = None
):
    """Get the appropriate fit function for a regression method.

    Args:
        method: The MetalogFitMethod enum value.
        regression_params: Optional regression parameters (for Lasso).

    Returns:
        A function that takes (target, quantiles) and returns a regression model.
    """
    if method == MetalogFitMethod.OLS:
        return fit_ordinary_least_squares
    elif method == MetalogFitMethod.Lasso:
        from metalog_jax.regression.lasso import DEFAULT_LASSO_PARAMETERS

        params = regression_params or DEFAULT_LASSO_PARAMETERS
        return lambda target, quantiles: fit_lasso(target, quantiles, params)
    else:
        raise ValueError(f"Unsupported method: {method}")
```

--------------------------------

### Fit Metalog Distribution with OLS

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/metalog.html

Demonstrates how to prepare input data using MetalogInputData and fit a metalog distribution using Ordinary Least Squares (OLS) regression in JAX.

```python
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
)

# 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)
```

--------------------------------

### Find Best Configuration (Python)

Source: https://tjefferies.github.io/metalog_jax/fitting_grids.html

Finds the index and value of the best configuration from a list of KS distances. This is typically used after a grid search to identify the optimal parameters. It assumes the input is a 1D array of KS distances.

```python
best_idx_nt, best_ks_nt = find_best_config(result_num_terms.ks_dist)
print(
    f"\nBest: {num_terms_grid[int(best_idx_nt)]} terms (KS = {float(best_ks_nt):.4f})"
)
```

--------------------------------

### Build Single Fit Configuration (Jax)

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/grid_search.html

Constructs a fitting function for a single configuration of Metalog parameters. This function is designed to be traceable by `vmap`, utilizing Jax-based helper functions for quantile transformation and design matrix construction.

```python
def _build_fit_single_config(
    params: MetalogParameters,
    max_terms_output: int,
    use_lasso: bool,
    learning_rate: float,
    num_iters: int,
    tol: float,
    momentum: float,
):
    """Build a fit function for a single configuration.

    Returns a function that fits a single (x, y, l1, n_terms) configuration.
    This is extracted to reduce cyclomatic complexity of fit_grid.
    """
    lower_bound = params.lower_bound
    upper_bound = params.upper_bound
    boundedness_value = params.boundedness.value

    def fit_single_config(
        x_i: jnp.ndarray,
        y_i: jnp.ndarray,
        l1_penalty: jnp.ndarray,
        n_terms: jnp.ndarray,
    ) -> GridResult:
        """Fit a single configuration (traceable for vmap)."""
        quantiles = _get_quantiles_vmap(
            x_i, boundedness_value, lower_bound, upper_bound
        )
        target = _get_target_vmap(y=y_i, num_terms=n_terms)

        if use_lasso:
            lasso_p = LassoParameters(
                lam=l1_penalty,
                learning_rate=learning_rate,
                num_iters=num_iters,
                tol=tol,
                momentum=momentum,
            )
            model = fit_lasso(target, quantiles, lasso_p)
        else:
            model = fit_ordinary_least_squares(target, quantiles)

        coeff_indices = jnp.arange(MAX_TERMS)
        coeff_mask = (coeff_indices < n_terms).astype(model.weights.dtype)
        masked_weights = model.weights * coeff_mask
        padded_weights = masked_weights[:max_terms_output]

        delta_term = y_i - 0.5
        log_term = jnp.log(y_i / (1 - y_i))

        a_padded = jnp.zeros(MAX_TERMS)
        a_padded = a_padded.at[: masked_weights.shape[0]].set(masked_weights)

        ppf = a_padded[0] + (a_padded[1] * log_term)
        ppf = ppf + jnp.where(n_terms > 2, a_padded[2] * delta_term * log_term, 0.0)
        ppf = ppf + jnp.where(n_terms > 3, a_padded[3] * delta_term, 0.0)

        terms = jnp.arange(5, MAX_TERMS + 1)
        active_mask = terms <= n_terms
        odd_mask = terms % 2 == 1

```

--------------------------------

### Metalog JAX Grid Search Module Imports

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/grid_search.html

Initializes the grid search module by importing necessary JAX primitives, base Metalog classes, and regression fitting methods. This setup provides the foundation for vectorized grid search operations.

```python
import jax
import jax.numpy as jnp
from metalog_jax.base import MetalogBaseData, MetalogParameters
from metalog_jax.base.enums import MetalogFitMethod
from metalog_jax.metalog import GridResult, Metalog, SPTMetalog, fit
from metalog_jax.regression import LassoParameters, RegularizedParameters
from metalog_jax.regression.lasso import fit_lasso
from metalog_jax.regression.ols import fit_ordinary_least_squares
from metalog_jax.utils import DEFAULT_Y, ks_distance
```

--------------------------------

### Fit Strictly Lower-Bounded Metalog Coefficients (JAX)

Source: https://tjefferies.github.io/metalog_jax/_modules/metalog_jax/metalog.html

Calculates the three metalog coefficients [a1, a2, a3] for strictly lower-bounded distributions. It transforms quantiles to gamma values, performs feasibility checks, and then computes the coefficients using JAX for JIT compilation.

```python
import chex
import jax.numpy as jnp

def _strictly_lower_bound_fit(q: chex.Numeric, y: chex.Numeric, lower_bound: float) -> chex.Numeric:
    """Fit SPT metalog coefficients for strictly lower-bounded distributions.

    Computes the three metalog coefficients [a1, a2, a3] for a distribution
    with only a lower bound using the analytical SPT formulas. The distribution
    has support on (lower_bound, infinity).

    The method first transforms the quantiles to gamma values by subtracting
    the lower bound: gamma_i = q_i - lower_bound.
    Then, it computes coefficients in log-space:
    - a1 = log(gamma_median)
    - a2 = 0.5 * log(gamma_complement / gamma_alpha) / log((1-alpha)/alpha)
    - a3 = log((gamma_complement * gamma_alpha) / gamma_median^2) /
           ((1 - 2*alpha) * log((1-alpha)/alpha))

    Args:
        q: Array of length 3 containing [q_alpha, median, q_complement],
            the empirical quantiles at probabilities [alpha, 0.5, 1-alpha].
            All values must be strictly greater than lower_bound.
        y: Array of length 1 containing the alpha parameter.
        lower_bound: The lower bound of the distribution.

    Returns:
        Metalog coefficients [a1, a2, a3] in log-space.
    """
    # ... (feasibility check code omitted for brevity)
    # ... (_log_term definition omitted for brevity)
    # ... (_second_term definition omitted for brevity)
    # ... (_third_term definition omitted for brevity)

    alpha = y[0]

    q_alpha = q[0]
    median = q[1]
    q_compliment = q[2]

    gamma_alpha = q_alpha - lower_bound
    gamma_median = median - lower_bound
    gamma_compliment = q_compliment - lower_bound

    gamma = jnp.array([gamma_alpha, gamma_median, gamma_compliment])
    _strictly_lower_bound_fit_feasibility_check(q, gamma, alpha, lower_bound)
    return jnp.array(
        [
            jnp.log(gamma_median),
            _second_term(gamma, alpha),
            _third_term(gamma, alpha),
        ]
    )
```

--------------------------------

### Get Documentation for Metalog logppf Function

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

This snippet retrieves and prints the docstring for the Metalog.logppf function. The docstring explains that the function computes the natural logarithm of the percent point function, detailing its arguments (probability values) and return value (log-quantile values).

```python
print(Metalog.logppf.__doc__)
```

--------------------------------

### Get Documentation for Metalog sf Function

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

This snippet retrieves and prints the docstring for the Metalog.sf function. The docstring explains that the function computes the survival function (1 - CDF), detailing its arguments (quantile values) and return value (survival probabilities in the range (0, 1)).

```python
print(Metalog.sf.__doc__)
```

--------------------------------

### POST /metalog/prng

Source: https://tjefferies.github.io/metalog_jax/basic_usage.html

Initializes the configuration parameters for a JAX uniform distribution PRNG.

```APIDOC
## POST /metalog/prng

### Description
Initializes the PRNG parameters required for generating samples from the distribution.

### Method
POST

### Endpoint
/metalog/prng

### Parameters
#### Request Body
- **seed** (int) - Required - Seed for the pseudorandom number generator.

### Response
#### Success Response (200)
- **params** (object) - The initialized JaxUniformDistributionParameters object.
```
Tokens:
Raw