freealg 0.7.16__py3-none-any.whl → 0.7.18__py3-none-any.whl

freealg/distributions/_fuss_catalan.py

@@ -0,0 +1,269 @@
+# SPDX-FileCopyrightText: Copyright 2026, Siavash Ameli <sameli@berkeley.edu>
+# SPDX-License-Identifier: BSD-3-Clause
+# SPDX-FileType: SOURCE
+#
+# This program is free software: you can redistribute it and/or modify it
+# under the terms of the license found in the LICENSE.txt file in the root
+# directory of this source tree.
+
+
+# =======
+# Imports
+# =======
+
+import numpy
+from .._algebraic_form._sheets_util import _pick_physical_root_scalar
+
+__all__ = ['FussCatalan']
+
+
+# ============
+# Fuss-Catalan
+# ============
+
+class FussCatalan(object):
+    """
+    Fuss-Catalan (a.k.a. free Bessel / Raney) distribution.
+
+    This is the law of squared singular values of a product of :math:`p`
+    independent (square) Ginibre matrices in the large-n limit. For :math:`p=1`
+    this reduces to the Marchenko--Pastur law with :math:`c = 1`.
+
+    Parameters
+    ----------
+
+    p : int, default=2
+        Order of the Fuss-Catalan distribution. :math:`p = 1` is MP
+        (:math:`c=1`).
+
+    Notes
+    -----
+
+    Let :math:`p \\geq 1` be an integer. The Stieltjes transform :math:`m(z)`
+    of the :math:`p`-th Fuss-Catalan distribution solves the algebraic equation
+
+    ..math::
+
+        (-1)^p z^p m(z)^{p+1} - z m(z) - 1 = 0,
+
+    where the physical branch is selected by the Herglotz condition
+    Im m(z) > 0 for Im z > 0.
+
+    Equivalently, in terms of w(z) = -m(z):
+
+    .. math::
+
+        z = (1 + w)^{p+1} / w^p.
+
+    The support is a single interval [0, x_max] with
+
+    .. math::
+
+        x_max = (p+1)^{p+1} / p^p.
+
+    (All in the standard normalization where the mean is 1.)
+
+    **Application:**
+
+    This law might be applicable to the Jacobian of neural netowkrs.
+    """
+
+    # ====
+    # init
+    # ====
+
+    def __init__(self, p=2):
+        """
+        Initialization.
+        """
+
+        p = int(p)
+
+        if p < 1:
+            raise ValueError("p must be an integer >= 1.")
+
+        self.p = p
+
+    # ============
+    # roots scalar
+    # ============
+
+    def _roots_scalar(self, z):
+        """
+        Return all algebraic roots of the defining polynomial at scalar z.
+        """
+
+        p = int(self.p)
+        z = numpy.asarray(z, dtype=numpy.complex128).reshape(())
+        coeffs = numpy.zeros(p + 2, dtype=numpy.complex128)
+
+        # Polynomial in m:
+        #   (-1)^p z^p m^{p+1} - z m - 1 = 0
+        coeffs[0] = ((-1.0)**p) * (z**p)       # m^{p+1}
+        coeffs[-2] = -z                        # m^1
+        coeffs[-1] = -1.0                      # m^0
+
+        return numpy.roots(coeffs)
+
+    # =========
+    # stieltjes
+    # =========
+
+    def stieltjes(self, z, max_iter=100, tol=1e-12):
+        """
+        Compute the physical Stieltjes transform m(z) by Newton, with robust
+        fallback to algebraic roots + physical picking.
+        """
+
+        p = int(self.p)
+
+        z = numpy.asarray(z, dtype=numpy.complex128)
+        scalar = (z.ndim == 0)
+        if scalar:
+            z = z.reshape((1,))
+
+        # Initial guess: m ~ -1/z for |z| large
+        m = -1.0 / z
+        active = numpy.isfinite(m)
+
+        # Newton on f(m) = (-1)^p z^p m^{p+1} - z m - 1
+        for _ in range(int(max_iter)):
+            if not numpy.any(active):
+                break
+
+            ma = m[active]
+            za = z[active]
+
+            # f and f'
+            f = ((-1.0)**p) * (za**p) * (ma**(p + 1)) - za * ma - 1.0
+            fp = ((-1.0)**p) * (za**p) * (p + 1) * (ma**p) - za
+
+            step = f / fp
+            mn = ma - step
+            m[active] = mn
+
+            conv = numpy.abs(step) < tol * (1.0 + numpy.abs(mn))
+            idx = numpy.where(active)[0]
+            active[idx[conv]] = False
+
+        # Herglotz sanity: sign(Im z) == sign(Im m)
+        sign = numpy.where(numpy.imag(z) >= 0.0, 1.0, -1.0)
+        bad = (~numpy.isfinite(m)) | (sign * numpy.imag(m) <= 0.0)
+
+        if numpy.any(bad):
+            zb = z.ravel()
+            mb = m.ravel()
+            bad_idx = numpy.flatnonzero(bad)
+            for i in bad_idx:
+                zi = zb[i]
+                r = self._roots_scalar(zi)
+                mb[i] = _pick_physical_root_scalar(zi, r)
+            m = mb.reshape(z.shape)
+
+        if scalar:
+            return m.reshape(())
+        return m
+
+    # =======
+    # density
+    # =======
+
+    def density(self, x, eta=2e-4):
+        """
+        Density rho(x) from Im m(x + i eta) / pi.
+        """
+
+        x = numpy.asarray(x, dtype=numpy.float64)
+        z = x + 1j * float(eta)
+        m = self.stieltjes(z)
+        rho = numpy.imag(m) / numpy.pi
+        return numpy.maximum(rho, 0.0)
+
+    # =====
+    # roots
+    # =====
+
+    def roots(self, z):
+        """
+        Return all algebraic branches at scalar z.
+        """
+
+        z = numpy.asarray(z, dtype=numpy.complex128)
+        if z.ndim != 0:
+            raise ValueError("roots(z) expects scalar z.")
+        return self._roots_scalar(z.reshape(()))
+
+    # =======
+    # support
+    # =======
+
+    def support(self):
+        """
+        Return the single support interval [0, x_max].
+        """
+
+        p = float(self.p)
+        x_max = ((p + 1.0)**(p + 1.0)) / (p**p)
+        return [(0.0, float(x_max))]
+
+    # ======
+    # matrix
+    # ======
+
+    def matrix(self, size, seed=None):
+        """
+        Generate a PSD random matrix whose ESD approximates Fuss-Catalan(p).
+
+        Construction
+        ------------
+        Let G_k be i.i.d. n x n Gaussian (Ginibre) matrices. Define
+
+            P = (1/sqrt(n)) G_1 (1/sqrt(n)) G_2 ... (1/sqrt(n)) G_p,
+            A = P P^T.
+
+        Then the ESD of A converges to Fuss-Catalan(p) as n->infty.
+        """
+
+        p = int(self.p)
+        n = int(size)
+        if n <= 0:
+            raise ValueError("size must be a positive integer.")
+
+        rng = numpy.random.default_rng(seed)
+
+        P = numpy.eye(n, dtype=numpy.float64)
+        scale = 1.0 / numpy.sqrt(float(n))
+        for _ in range(p):
+            G = rng.standard_normal((n, n))
+            P = P @ (scale * G)
+
+        A = P @ P.T
+        return A
+
+    # ====
+    # poly
+    # ====
+
+    def poly(self):
+        """
+        Return coeffs for the exact polynomial P(z,m)=0.
+
+        P(z,m) = (-1)^p z^p m^{p+1} - z m - 1.
+
+        coeffs[i, j] is the coefficient of z^i m^j.
+        Shape is (deg_z+1, deg_m+1) = (p+1, p+2).
+        """
+
+        p = int(self.p)
+        a = numpy.zeros((p + 1, p + 2), dtype=numpy.complex128)
+
+        # constant: -1
+        a[0, 0] = -1.0
+
+        # - z m
+        a[1, 1] = -1.0
+
+        # (-1)^p z^p m^{p+1}
+        a[p, p + 1] = ((-1.0)**p)
+
+        return a

freealg/distributions/_kesten_mckay.py

@@ -12,15 +12,9 @@
 # =======
 
 import numpy
-from scipy.interpolate import interp1d
-from .._free_form._plot_util import plot_density, plot_hilbert, \
-    plot_stieltjes, plot_stieltjes_on_disk, plot_samples
-
-try:
-    from scipy.integrate import cumtrapz
-except ImportError:
-    from scipy.integrate import cumulative_trapezoid as cumtrapz
-from scipy.stats import qmc
+from ..visualization._plot_util import plot_density, plot_hilbert, \
+    plot_stieltjes, plot_stieltjes_on_disk
+from ._base_distribution import BaseDistribution
 
 __all__ = ['KestenMcKay']
 
@@ -29,7 +23,7 @@ __all__ = ['KestenMcKay']
 # Kesten McKay
 # ============
 
-class KestenMcKay(object):
+class KestenMcKay(BaseDistribution):
     """
     Kesten-McKay distribution.
 
@@ -433,126 +427,6 @@ class KestenMcKay(object):
 
         return m1, m2
 
-    # ======
-    # sample
-    # ======
-
-    def sample(self, size, x_min=None, x_max=None, method='qmc', seed=None,
-               plot=False, latex=False, save=False):
-        """
-        Sample from distribution.
-
-        Parameters
-        ----------
-
-        size : int
-            Size of sample.
-
-        x_min : float, default=None
-            Minimum of sample values. If `None`, the left edge of the support
-            is used.
-
-        x_max : float, default=None
-            Maximum of sample values. If `None`, the right edge of the support
-            is used.
-
-        method : {``'mc'``, ``'qmc'``}, default= ``'qmc'``
-            Method of drawing samples from uniform distribution:
-
-            * ``'mc'``: Monte Carlo
-            * ``'qmc'``: Quasi Monte Carlo
-
-        seed : int, default=None,
-            Seed for random number generator.
-
-        plot : bool, default=False
-            If `True`, samples histogram is plotted.
-
-        latex : bool, default=False
-            If `True`, the plot is rendered using LaTeX. This option is
-            relevant only if ``plot=True``.
-
-        save : bool, default=False
-            If not `False`, the plot is saved. If a string is given, it is
-            assumed to the save filename (with the file extension). This option
-            is relevant only if ``plot=True``.
-
-        Returns
-        -------
-
-        s : numpy.ndarray
-            Samples.
-
-        Notes
-        -----
-
-        This method uses inverse transform sampling.
-
-        Examples
-        --------
-
-        .. code-block::python
-
-            >>> from freealg.distributions import KestenMcKay
-            >>> km = KestenMcKay(3)
-            >>> s = km.sample(2000)
-
-        .. image:: ../_static/images/plots/km_samples.png
-            :align: center
-            :class: custom-dark
-        """
-
-        if x_min is None:
-            x_min = self.lam_m
-
-        if x_max is None:
-            x_max = self.lam_p
-
-        # Grid and PDF
-        xs = numpy.linspace(x_min, x_max, size)
-        pdf = self.density(xs)
-
-        # CDF (using cumulative trapezoidal rule)
-        cdf = cumtrapz(pdf, xs, initial=0)
-        cdf /= cdf[-1]  # normalize CDF to 1
-
-        # Inverse CDF interpolator
-        inv_cdf = interp1d(cdf, xs, bounds_error=False,
-                           fill_value=(x_min, x_max))
-
-        # Random generator
-        rng = numpy.random.default_rng(seed)
-
-        # Draw from uniform distribution
-        if method == 'mc':
-            u = rng.random(size)
-
-        elif method == 'qmc':
-            try:
-                engine = qmc.Halton(d=1, scramble=True, rng=rng)
-            except TypeError:
-                # Older scipy versions
-                engine = qmc.Halton(d=1, scramble=True, seed=rng)
-            u = engine.random(size).ravel()
-
-        else:
-            raise NotImplementedError('"method" is invalid.')
-
-        # Draw from distribution by mapping from inverse CDF
-        samples = inv_cdf(u).ravel()
-
-        if plot:
-            radius = 0.5 * (self.lam_p - self.lam_m)
-            center = 0.5 * (self.lam_p + self.lam_m)
-            scale = 1.25
-            x_min = numpy.floor(center - radius * scale)
-            x_max = numpy.ceil(center + radius * scale)
-            x = numpy.linspace(x_min, x_max, 500)
-            rho = self.density(x)
-            plot_samples(x, rho, x_min, x_max, samples, latex=latex, save=save)
-
-        return samples
-
     # ===============
     # haar orthogonal
     # ===============

freealg/distributions/_marchenko_pastur.py

@@ -12,16 +12,10 @@
 # =======
 
 import numpy
-from scipy.interpolate import interp1d
-from .._free_form._plot_util import plot_density, plot_hilbert, \
-    plot_stieltjes, plot_stieltjes_on_disk, plot_samples
+from ..visualization._plot_util import plot_density, plot_hilbert, \
+    plot_stieltjes, plot_stieltjes_on_disk
 from ..visualization import glue_branches
-
-try:
-    from scipy.integrate import cumtrapz
-except ImportError:
-    from scipy.integrate import cumulative_trapezoid as cumtrapz
-from scipy.stats import qmc
+from ._base_distribution import BaseDistribution
 
 __all__ = ['MarchenkoPastur']
 
@@ -30,7 +24,7 @@ __all__ = ['MarchenkoPastur']
 # Marchenko Pastur
 # ================
 
-class MarchenkoPastur(object):
+class MarchenkoPastur(BaseDistribution):
     """
     Marchenko-Pastur distribution.
 
@@ -103,8 +97,6 @@ class MarchenkoPastur(object):
         self.lam = lam
         self.sigma = sigma
 
-        # self.lam_p = (1 + numpy.sqrt(self.lam))**2
-        # self.lam_m = (1 - numpy.sqrt(self.lam))**2
         self.lam_p = sigma**2 * (1.0 + numpy.sqrt(lam))**2
         self.lam_m = sigma**2 * (1.0 - numpy.sqrt(lam))**2
 
@@ -123,7 +115,7 @@ class MarchenkoPastur(object):
 
         x : numpy.array, default=None
             The locations where density is evaluated at. If `None`, an interval
-            slightly larger than the supp interval of the spectral density
+            slightly larger than the support interval of the spectral density
             is used.
 
         rho : numpy.array, default=None
@@ -138,7 +130,7 @@ class MarchenkoPastur(object):
 
         save : bool, default=False
             If not `False`, the plot is saved. If a string is given, it is
-        assumed to the save filename (with the file extension). This option
+            assumed to the save filename (with the file extension). This option
             is relevant only if ``plot=True``.
 
         eig : numpy.array, default=None
@@ -223,7 +215,7 @@ class MarchenkoPastur(object):
 
         x : numpy.array, default=None
             The locations where Hilbert transform is evaluated at. If `None`,
-            an interval slightly larger than the supp interval of the
+            an interval slightly larger than the support interval of the
             spectral density is used.
 
         plot : bool, default=False
@@ -288,46 +280,6 @@ class MarchenkoPastur(object):
 
         return hilb
 
-    # =======================
-    # m mp numeric vectorized
-    # =======================
-
-    # def _m_mp_numeric_vectorized(self, z, alt_branch=False, tol=1e-8):
-    #     """
-    #     Stieltjes transform (principal or secondary branch)
-    #     for Marchenko-Pastur distribution on upper half-plane.
-    #     """
-    #
-    #     sigma = 1.0
-    #     m = numpy.empty_like(z, dtype=complex)
-    #
-    #     # When z is too small, do not use quadratic form.
-    #     mask = numpy.abs(z) < tol
-    #     m[mask] = 1 / (sigma**2 * (1 - self.lam))
-    #
-    #     # Use quadratic form
-    #     not_mask = ~mask
-    #     if numpy.any(not_mask):
-    #
-    #         sign = -1 if alt_branch else 1
-    #         A = self.lam * sigma**2 * z[not_mask]
-    #         B = z[not_mask] - sigma**2 * (1 - self.lam)
-    #         D = B**2 - 4 * A
-    #         sqrtD = numpy.sqrt(D)
-    #         m1 = (-B + sqrtD) / (2 * A)
-    #         m2 = (-B - sqrtD) / (2 * A)
-    #
-    #         # pick correct branch only for non-masked entries
-    #         upper = z[not_mask].imag >= 0
-    #         branch = numpy.empty_like(m1)
-    #         branch[upper] = numpy.where(sign*m1[upper].imag > 0, m1[upper],
-    #                                     m2[upper])
-    #         branch[~upper] = numpy.where(sign*m1[~upper].imag < 0,
-    #                                      m1[~upper], m2[~upper])
-    #         m[not_mask] = branch
-    #
-    #     return m
-
     # =============
     # sqrt pos imag
     # =============
@@ -342,26 +294,6 @@ class MarchenkoPastur(object):
 
         return sq
 
-    # ============
-    # m mp reflect
-    # ============
-
-    # def _m_mp_reflect(self, z, alt_branch=False):
-    #     """
-    #     Analytic continuation using Schwarz reflection.
-    #     """
-    #
-    #     mask_p = z.imag >= 0.0
-    #     mask_n = z.imag < 0.0
-    #
-    #     m = numpy.zeros_like(z)
-    #
-    #     f = self._m_mp_numeric_vectorized
-    #     m[mask_p] = f(z[mask_p], alt_branch=False)
-    #     m[mask_n] = f(z[mask_n], alt_branch=alt_branch)
-    #
-    #     return m
-
     # ================
     # stieltjes branch
     # ================
@@ -429,7 +361,7 @@ class MarchenkoPastur(object):
 
         x : numpy.array, default=None
             The x axis of the grid where the Stieltjes transform is evaluated.
-            If `None`, an interval slightly larger than the supp interval of
+            If `None`, an interval slightly larger than the support interval of
             the spectral density is used.
 
         y : numpy.array, default=None
@@ -557,126 +489,6 @@ class MarchenkoPastur(object):
         else:
             return m1
 
-    # ======
-    # sample
-    # ======
-
-    def sample(self, size, x_min=None, x_max=None, method='qmc', seed=None,
-               plot=False, latex=False, save=False):
-        """
-        Sample from distribution.
-
-        Parameters
-        ----------
-
-        size : int
-            Size of sample.
-
-        x_min : float, default=None
-            Minimum of sample values. If `None`, the left edge of the supp
-            is used.
-
-        x_max : float, default=None
-            Maximum of sample values. If `None`, the right edge of the supp
-            is used.
-
-        method : {``'mc'``, ``'qmc'``}, default= ``'qmc'``
-            Method of drawing samples from uniform distribution:
-
-            * ``'mc'``: Monte Carlo
-            * ``'qmc'``: Quasi Monte Carlo
-
-        seed : int, default=None,
-            Seed for random number generator.
-
-        plot : bool, default=False
-            If `True`, samples histogram is plotted.
-
-        latex : bool, default=False
-            If `True`, the plot is rendered using LaTeX. This option is
-            relevant only if ``plot=True``.
-
-        save : bool, default=False
-            If not `False`, the plot is saved. If a string is given, it is
-            assumed to the save filename (with the file extension). This option
-            is relevant only if ``plot=True``.
-
-        Returns
-        -------
-
-        s : numpy.ndarray
-            Samples.
-
-        Notes
-        -----
-
-        This method uses inverse transform sampling.
-
-        Examples
-        --------
-
-        .. code-block::python
-
-            >>> from freealg.distributions import MarchenkoPastur
-            >>> mp = MarchenkoPastur(1/50)
-            >>> s = mp.sample(2000)
-
-        .. image:: ../_static/images/plots/mp_samples.png
-            :align: center
-            :class: custom-dark
-        """
-
-        if x_min is None:
-            x_min = self.lam_m
-
-        if x_max is None:
-            x_max = self.lam_p
-
-        # Grid and PDF
-        xs = numpy.linspace(x_min, x_max, size)
-        pdf = self.density(xs)
-
-        # CDF (using cumulative trapezoidal rule)
-        cdf = cumtrapz(pdf, xs, initial=0)
-        cdf /= cdf[-1]  # normalize CDF to 1
-
-        # Inverse CDF interpolator
-        inv_cdf = interp1d(cdf, xs, bounds_error=False,
-                           fill_value=(x_min, x_max))
-
-        # Random generator
-        rng = numpy.random.default_rng(seed)
-
-        # Draw from uniform distribution
-        if method == 'mc':
-            u = rng.random(size)
-
-        elif method == 'qmc':
-            try:
-                engine = qmc.Halton(d=1, scramble=True, rng=rng)
-            except TypeError:
-                # Older scipy versions
-                engine = qmc.Halton(d=1, scramble=True, seed=rng)
-            u = engine.random(size).ravel()
-
-        else:
-            raise NotImplementedError('"method" is invalid.')
-
-        # Draw from distribution by mapping from inverse CDF
-        samples = inv_cdf(u).ravel()
-
-        if plot:
-            radius = 0.5 * (self.lam_p - self.lam_m)
-            center = 0.5 * (self.lam_p + self.lam_m)
-            scale = 1.25
-            x_min = numpy.floor(center - radius * scale)
-            x_max = numpy.ceil(center + radius * scale)
-            x = numpy.linspace(x_min, x_max, 500)
-            rho = self.density(x)
-            plot_samples(x, rho, x_min, x_max, samples, latex=latex, save=save)
-
-        return samples
-
     # ======
     # matrix
     # ======