Add probability distributions: PDF, CDF, Quantile

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 ### Added
 
 - **More statistics:** `Skewness`, `Kurtosis` (sample-adjusted, matching Excel's `SKEW`/`KURT` and `scipy.stats.skew`/`kurtosis` with `bias=False`), `LinearRegression` (least-squares fit, returns `["Array", slope, intercept]`), `PolynomialFit` (least-squares fit of a given degree via the normal equations, no numpy dependency, returns `["Array", c0, c1, ..., cd]` - lowest degree first, the opposite order from `numpy.polyfit`). `PolynomialFit`'s degree is capped at 50.
+- **Probability distributions:** `NormalDistribution`, `BinomialDistribution`, `PoissonDistribution`, `UniformDistribution`, `ExponentialDistribution`, queried via `PDF`/`CDF`/`Quantile` (e.g. `["CDF", ["NormalDistribution", 0, 1], 1.5]`), matching CortexJS's calling convention. `PDF` needs no optional dependency for any of the five, including the two discrete distributions (uses a log-space reformulation - the direct formula overflows a float well within realistic `n`). `Binomial`/`PoissonDistribution`'s `CDF`/`Quantile` use the already-shipped `GammaRegularized`/`BetaRegularized` identities and so require the `special-functions` extra; the other three distributions' `CDF`/`Quantile` don't need it either.
 - **Number theory**, ~45 new functions:
   - stdlib-based: `PowerMod`, `ModularInverse`, `IntegerSqrt`
   - Factorization: `FactorInteger`, `PrimeFactors`, `PrimeNu`, `PrimeOmega`, `Radical`, `IsSquareFree`

README.md

@@ -34,7 +34,9 @@ pip install mathjson-solver
 **Optional:**
 - `pip install mathjson-solver[integration]` (installs numpy; only required for `TrapezoidalIntegrate`)
 - `pip install mathjson-solver[regex]` (installs [google-re2](https://pypi.org/project/google-re2/); only required for `RegExp`/`IsMatch`/`StringMatch`/`StringMatchAll`)
-- `pip install mathjson-solver[special-functions]` (installs [scipy](https://scipy.org/); only required for `BesselJ`/`BesselY`/`BesselI`/`BesselK`, `AiryAi`/`AiryBi`/`AiryAiPrime`/`AiryBiPrime`, `Zeta`, `GammaRegularized`, `BetaRegularized`)
+- `pip install mathjson-solver[special-functions]` (installs [scipy](https://scipy.org/); required for `BesselJ`/`BesselY`/`BesselI`/`BesselK`, `AiryAi`/`AiryBi`/`AiryAiPrime`/`AiryBiPrime`, `Zeta`, `GammaRegularized`, `BetaRegularized`, and `CDF`/`Quantile` on `BinomialDistribution`/`PoissonDistribution`)
+
+Combine any of these by separating them with commas in one `pip install`, e.g. `pip install mathjson-solver[integration,regex,special-functions]` for all three at once.
 
 ## Quick Start
 
@@ -121,6 +123,7 @@ The library supports a comprehensive set of mathematical operations:
 * **Comparison:** Equal, StrictEqual, IdenticallyEqual, NotEqual, Greater, GreaterEqual, Less, LessEqual, Congruent
 * **Logic & Sets:** Any, All, Not, And, Or, Xor, Nand, Nor, Implies, Equivalent, In/Element, NotIn/NotElement, ContainsAnyOf, ContainsAllOf, ContainsNoneOf, bare `True`/`False` literals, Union, Intersection, SetMinus, SymmetricDifference (over arrays - no dedicated Set type)
 * **Statistics:** Average/Mean, Max, Min (both list and variadic forms), Median, Mode, Variance, StandardDeviation, PopulationVariance, PopulationStandardDeviation, Quartiles, InterquartileRange, Covariance, Correlation, Skewness, Kurtosis, LinearRegression, PolynomialFit, Length/Count
+* **Probability Distributions:** NormalDistribution, BinomialDistribution, PoissonDistribution, UniformDistribution, ExponentialDistribution, queried via PDF, CDF, Quantile (e.g. `["CDF", ["NormalDistribution", 0, 1], 1.5]`). No optional dependency needed for PDF on any distribution, or for CDF/Quantile on Normal/Uniform/Exponential — only Binomial/Poisson's CDF and Quantile require the `special-functions` extra
 * **Functional Programming:** Map/StrictMap, Reduce, Filter, Product (all also accept CortexJS calling conventions, including `Function` lambdas)
 * **Arrays:** Array/List creation, GenerateRange, Range, AtIndex, At, Slice, Appended/Append, First, Second, Third, Last, Rest, Most, Reverse, Sort, Unique, Dedup, Join, Zip, IsEmpty, CumulativeSum, CumulativeProduct, Take, Drop, TakeWhile, DropWhile, Contains, IndexOf, IndexWhere, Find, CountIf, Position, RotateLeft, RotateRight, MaxBy, MinBy, ArgMax, ArgMin, Ordering, FlatMap, Scan, Differences, Fold, Insert, DeleteAt, ReplaceAt, Partition, Chunk, GroupBy, ChunkBy, Tally
 * **Control Flow:** If statements (Python pair form and CortexJS flat form), Switch/StrictSwitch (value-equality case dispatch), Which (CortexJS flat condition/value chain), Constants definition

docs/README.md

@@ -992,6 +992,21 @@ Sample-adjusted estimators - `Skewness` matches Excel's `SKEW` and `scipy.stats.
 ["PolynomialFit", ["Array", 0, 1, 2, 3, 4], ["Array", 3, 6, 11, 18, 27], 2]      # ≈ ["Array", 3.0, 2.0, 1.0] (y = x² + 2x + 3)
 ```
 
+#### Probability Distributions
+`NormalDistribution(mean, standardDeviation)`, `BinomialDistribution(n, p)`, `PoissonDistribution(lambda)`, `UniformDistribution(a, b)`, `ExponentialDistribution(lambda)` are "markers" - like `Function`, they return their own unevaluated expression rather than a computed value. `PDF`, `CDF`, and `Quantile` take one of these (inline, or via a solver parameter that holds one) as their first argument and do the actual computation, matching CortexJS's calling convention exactly.
+
+```python
+["PDF", ["BinomialDistribution", 4, ["Rational", 1, 2]], 2]   # 0.375
+["Round", ["CDF", ["NormalDistribution", 0, 1], 1], 5]        # 0.84134  (Φ(1))
+["Quantile", ["PoissonDistribution", 9], 0.95]                # 14
+
+# A distribution can also be a solver parameter, not just inline:
+solver = create_solver({"risk_model": ["NormalDistribution", 0, 1]})
+solver(["CDF", "risk_model", 1.5])
+```
+
+No optional dependency is needed for `PDF` on any of the five distributions, or for `CDF`/`Quantile` on `NormalDistribution`/`UniformDistribution`/`ExponentialDistribution`. Only `CDF`/`Quantile` on `BinomialDistribution`/`PoissonDistribution` require the `special-functions` extra (`pip install mathjson-solver[special-functions]`) - they're built on the already-shipped `GammaRegularized`/`BetaRegularized` (via the identities `P(X≤k) = Q(k+1,λ)` for Poisson and `P(X≤k) = I_{1-p}(n-k,k+1)` for Binomial), which is both far more numerically robust and much faster than direct summation for large `n`/`λ` - the direct formula (`math.comb(n,k) * p**k * (1-p)**(n-k)`) actually overflows a float at `n` as low as 10,000, well within a realistic input range, since `math.comb` alone produces an enormous intermediate integer there. `PDF` on these two avoids that same overflow via a log-space reformulation instead (no scipy needed).
+
 #### Length (alias: Count)
 Returns the number of elements in an array, including non-numeric elements like `None`. `Count` is the CortexJS name for the same function.
 
@@ -2175,6 +2190,7 @@ Evaluating a `["Function", ...]` expression outside of such a context (i.e. not
 - [Covariance, Correlation](#covariance-and-correlation) - Two-array statistics
 - [Skewness, Kurtosis](#skewness-and-kurtosis) - Distribution shape statistics
 - [LinearRegression, PolynomialFit](#linearregression-and-polynomialfit) - Least-squares curve fitting
+- [NormalDistribution, BinomialDistribution, PoissonDistribution, UniformDistribution, ExponentialDistribution, PDF, CDF, Quantile](#probability-distributions) - Probability distributions
 - [Length / Count](#length-alias-count) - Array length
 - [First, Second, Third, Last, Rest, Most](#first-second-third-last-rest-most) - Access or trim array ends
 - [Reverse, Sort](#reverse-and-sort) - Reverse or sort an array

pyproject.toml

@@ -29,7 +29,10 @@ regex = ["google-re2"]
 # Bessel/Airy/Zeta/GammaRegularized/BetaRegularized need scipy.special -
 # real numerical algorithms (stable series/asymptotic-expansion
 # switching depending on argument regime), not something to reimplement
-# from scratch. Optional, same as `integration`/`regex` above.
+# from scratch. Also needed for CDF/Quantile on BinomialDistribution/
+# PoissonDistribution (built on GammaRegularized/BetaRegularized) - PDF
+# on those two, and everything on the other three distributions, needs
+# no optional dependency. Optional, same as `integration`/`regex` above.
 special-functions = ["scipy"]
 
 [project.urls]

src/mathjson_solver/__main__.py

@@ -269,6 +269,14 @@ def translate_v1_mathjson(expr: MathJSONExpression) -> MathJSONExpression:
     Kurtosis,
     LinearRegression,
     PolynomialFit,
+    NormalDistribution,
+    BinomialDistribution,
+    PoissonDistribution,
+    UniformDistribution,
+    ExponentialDistribution,
+    PDF,
+    CDF,
+    Quantile,
 )
 from ._string_constructs import (
     Str,
@@ -676,6 +684,14 @@ def If(f, c, solver_parameters, s):
     "Kurtosis": Kurtosis,
     "LinearRegression": LinearRegression,
     "PolynomialFit": PolynomialFit,
+    "NormalDistribution": NormalDistribution,
+    "BinomialDistribution": BinomialDistribution,
+    "PoissonDistribution": PoissonDistribution,
+    "UniformDistribution": UniformDistribution,
+    "ExponentialDistribution": ExponentialDistribution,
+    "PDF": PDF,
+    "CDF": CDF,
+    "Quantile": Quantile,
     # --- Collections ---
     "First": First,
     "Second": Second,
@@ -1116,6 +1132,14 @@ def extract_variables(
         "Kurtosis",
         "LinearRegression",
         "PolynomialFit",
+        "NormalDistribution",
+        "BinomialDistribution",
+        "PoissonDistribution",
+        "UniformDistribution",
+        "ExponentialDistribution",
+        "PDF",
+        "CDF",
+        "Quantile",
         "First",
         "Second",
         "Third",

src/mathjson_solver/_statistics_constructs.py

@@ -1,7 +1,9 @@
 """
-Descriptive-statistics and regression constructs.
+Descriptive-statistics and regression constructs, plus probability
+distributions (NormalDistribution, ..., PDF, CDF, Quantile).
 """
 
+import math
 from statistics import (
     correlation,
     covariance,
@@ -13,6 +15,26 @@
     variance,
 )
 from ._construct_shared import _arr_vals, _arr_vals_of
+from ._special_functions import _erf_inv
+
+SCIPY_AVAILABLE = False
+try:
+    from scipy import special as _scipy_special
+
+    SCIPY_AVAILABLE = True
+except ImportError:
+    pass
+
+
+def _require_scipy():
+    if not SCIPY_AVAILABLE:
+        raise ImportError(
+            "'CDF'/'Quantile' on a BinomialDistribution or "
+            "PoissonDistribution require 'scipy' (PDF does not). "
+            "Install with 'pip install scipy' (or 'pip install "
+            "mathjson-solver[special-functions]')."
+        )
+
 
 # Hard cap on PolynomialFit's degree - the normal-equations system it
 # builds is (degree+1) x (degree+1), solved by O(degree^3) Gaussian
@@ -213,3 +235,239 @@ def PolynomialFit(f, c, solver_parameters, s):
     matrix = [[power_sums[i + j] for j in range(degree + 1)] for i in range(degree + 1)]
     vector = [sum((xi**k) * yi for xi, yi in zip(x, y)) for k in range(degree + 1)]
     return ["Array"] + _solve_linear_system(f, c, solver_parameters, matrix, vector)
+
+
+# --- Probability distributions ---
+#
+# NormalDistribution/BinomialDistribution/PoissonDistribution/
+# UniformDistribution/ExponentialDistribution are "markers", the same
+# self-quoting pattern already used for Function/Hold: they just
+# return their own unevaluated expression. PDF/CDF/Quantile resolve
+# and dispatch on that expression themselves. This means a
+# distribution's parameters (e.g. ["NormalDistribution", "mu",
+# "sigma"]) are still treated as ordinary sub-expressions by
+# extract_variables - unlike Head/Tail/Hold, they're meant to be
+# evaluated later by whichever query function unpacks them, not held
+# opaque forever, so no special-casing is needed there.
+
+
+def NormalDistribution(f, c, solver_parameters, s):
+    return s
+
+
+def BinomialDistribution(f, c, solver_parameters, s):
+    return s
+
+
+def PoissonDistribution(f, c, solver_parameters, s):
+    return s
+
+
+def UniformDistribution(f, c, solver_parameters, s):
+    return s
+
+
+def ExponentialDistribution(f, c, solver_parameters, s):
+    return s
+
+
+def _normal_pdf(mu, sigma, x):
+    if sigma <= 0:
+        raise ValueError("'NormalDistribution' requires standard deviation > 0.")
+    return math.exp(-0.5 * ((x - mu) / sigma) ** 2) / (sigma * math.sqrt(2 * math.pi))
+
+
+def _normal_cdf(mu, sigma, x):
+    if sigma <= 0:
+        raise ValueError("'NormalDistribution' requires standard deviation > 0.")
+    return 0.5 * (1 + math.erf((x - mu) / (sigma * math.sqrt(2))))
+
+
+def _normal_quantile(mu, sigma, p):
+    if sigma <= 0:
+        raise ValueError("'NormalDistribution' requires standard deviation > 0.")
+    return mu + sigma * math.sqrt(2) * _erf_inv(2 * p - 1)
+
+
+def _uniform_pdf(a, b, x):
+    if not a < b:
+        raise ValueError("'UniformDistribution' requires a < b.")
+    return 1 / (b - a) if a <= x <= b else 0.0
+
+
+def _uniform_cdf(a, b, x):
+    if not a < b:
+        raise ValueError("'UniformDistribution' requires a < b.")
+    if x < a:
+        return 0.0
+    if x > b:
+        return 1.0
+    return (x - a) / (b - a)
+
+
+def _uniform_quantile(a, b, p):
+    if not a < b:
+        raise ValueError("'UniformDistribution' requires a < b.")
+    return a + p * (b - a)
+
+
+def _exponential_pdf(lam, x):
+    if lam <= 0:
+        raise ValueError("'ExponentialDistribution' requires lambda > 0.")
+    return lam * math.exp(-lam * x) if x >= 0 else 0.0
+
+
+def _exponential_cdf(lam, x):
+    if lam <= 0:
+        raise ValueError("'ExponentialDistribution' requires lambda > 0.")
+    return 1 - math.exp(-lam * x) if x >= 0 else 0.0
+
+
+def _exponential_quantile(lam, p):
+    if lam <= 0:
+        raise ValueError("'ExponentialDistribution' requires lambda > 0.")
+    return -math.log(1 - p) / lam
+
+
+def _binomial_pdf(n, p, k):
+    """
+    P(X=k) for Binomial(n,p), via a log-space reformulation - the
+    direct formula (math.comb(n,k) * p**k * (1-p)**(n-k)) overflows a
+    float for even moderately large n (verified: n=10,000 already
+    raises OverflowError converting math.comb's huge intermediate
+    integer to float), well within a realistic input range.
+    """
+    n = int(n)
+    if not (0 <= p <= 1):
+        raise ValueError("'BinomialDistribution' requires 0 <= p <= 1.")
+    k = int(k)
+    if k < 0 or k > n:
+        return 0.0
+    log_pdf = (
+        math.lgamma(n + 1)
+        - math.lgamma(k + 1)
+        - math.lgamma(n - k + 1)
+        + k * math.log(p)
+        + (n - k) * math.log(1 - p)
+    )
+    return math.exp(log_pdf)
+
+
+def _poisson_pdf(lam, k):
+    """Same log-space reasoning as `_binomial_pdf` - math.factorial(k) overflows for large k."""
+    if lam <= 0:
+        raise ValueError("'PoissonDistribution' requires lambda > 0.")
+    k = int(k)
+    if k < 0:
+        return 0.0
+    return math.exp(-lam + k * math.log(lam) - math.lgamma(k + 1))
+
+
+def _binomial_cdf(n, p, k):
+    """
+    P(X<=k) for Binomial(n,p), via the regularized incomplete beta
+    identity I_(1-p)(n-k, k+1) - verified against direct summation to
+    machine precision, and far more numerically robust for large n
+    (direct summation both overflows and is measurably slower well
+    before this identity's cost grows at all).
+    """
+    _require_scipy()
+    n = int(n)
+    if not (0 <= p <= 1):
+        raise ValueError("'BinomialDistribution' requires 0 <= p <= 1.")
+    k = int(math.floor(k))
+    if k < 0:
+        return 0.0
+    if k >= n:
+        return 1.0
+    return _scipy_special.betainc(n - k, k + 1, 1 - p)
+
+
+def _poisson_cdf(lam, k):
+    """P(X<=k) for Poisson(lambda), via Q(k+1, lambda) - see `_binomial_cdf`."""
+    _require_scipy()
+    if lam <= 0:
+        raise ValueError("'PoissonDistribution' requires lambda > 0.")
+    k = int(math.floor(k))
+    if k < 0:
+        return 0.0
+    return _scipy_special.gammaincc(k + 1, lam)
+
+
+def _discrete_quantile(cdf_fn, p):
+    """
+    Smallest non-negative integer k such that cdf_fn(k) >= p, via
+    doubling then bisecting (same pattern as `_is_figurate` in
+    _number_theory.py) - robust for arbitrarily large true quantiles
+    without assuming where they'll land, unlike a linear scan from 0.
+    """
+    if cdf_fn(0) >= p:
+        return 0
+    lo, hi = 0, 1
+    while cdf_fn(hi) < p:
+        hi *= 2
+    while lo < hi:
+        mid = (lo + hi) // 2
+        if cdf_fn(mid) < p:
+            lo = mid + 1
+        else:
+            hi = mid
+    return lo
+
+
+def PDF(f, c, solver_parameters, s):
+    """["PDF", distribution, x] - probability density (or mass) function."""
+    dist = f(s[1], c)
+    x = f(s[2], c)
+    name = dist[0]
+    if name == "NormalDistribution":
+        return _normal_pdf(f(dist[1], c), f(dist[2], c), x)
+    if name == "UniformDistribution":
+        return _uniform_pdf(f(dist[1], c), f(dist[2], c), x)
+    if name == "ExponentialDistribution":
+        return _exponential_pdf(f(dist[1], c), x)
+    if name == "BinomialDistribution":
+        return _binomial_pdf(f(dist[1], c), f(dist[2], c), x)
+    if name == "PoissonDistribution":
+        return _poisson_pdf(f(dist[1], c), x)
+    raise ValueError(f"'PDF' does not support distribution '{name}'.")
+
+
+def CDF(f, c, solver_parameters, s):
+    """["CDF", distribution, x] - cumulative distribution function."""
+    dist = f(s[1], c)
+    x = f(s[2], c)
+    name = dist[0]
+    if name == "NormalDistribution":
+        return _normal_cdf(f(dist[1], c), f(dist[2], c), x)
+    if name == "UniformDistribution":
+        return _uniform_cdf(f(dist[1], c), f(dist[2], c), x)
+    if name == "ExponentialDistribution":
+        return _exponential_cdf(f(dist[1], c), x)
+    if name == "BinomialDistribution":
+        return _binomial_cdf(f(dist[1], c), f(dist[2], c), x)
+    if name == "PoissonDistribution":
+        return _poisson_cdf(f(dist[1], c), x)
+    raise ValueError(f"'CDF' does not support distribution '{name}'.")
+
+
+def Quantile(f, c, solver_parameters, s):
+    """["Quantile", distribution, p] - inverse CDF (0 <= p <= 1)."""
+    dist = f(s[1], c)
+    p = f(s[2], c)
+    if not (0 <= p <= 1):
+        raise ValueError("'Quantile' requires 0 <= p <= 1.")
+    name = dist[0]
+    if name == "NormalDistribution":
+        return _normal_quantile(f(dist[1], c), f(dist[2], c), p)
+    if name == "UniformDistribution":
+        return _uniform_quantile(f(dist[1], c), f(dist[2], c), p)
+    if name == "ExponentialDistribution":
+        return _exponential_quantile(f(dist[1], c), p)
+    if name == "BinomialDistribution":
+        n, prob = f(dist[1], c), f(dist[2], c)
+        return _discrete_quantile(lambda k: _binomial_cdf(n, prob, k), p)
+    if name == "PoissonDistribution":
+        lam = f(dist[1], c)
+        return _discrete_quantile(lambda k: _poisson_cdf(lam, k), p)
+    raise ValueError(f"'Quantile' does not support distribution '{name}'.")