feat: add main feat (#11)

34j · web-flow · commit 43a01bb398c7 · 2026-01-04T14:18:35.000+09:00
diff --git a/.codespellrc b/.codespellrc
@@ -1,2 +1,2 @@
 [codespell]
-ignore-words-list = socio-economic
+ignore-words-list = socio-economic, numer
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -36,8 +36,6 @@ jobs:
       fail-fast: false
       matrix:
         python-version:
-          - "3.10"
-          - "3.11"
           - "3.12"
           - "3.13"
           - "3.14"
diff --git a/.gitignore b/.gitignore
@@ -173,3 +173,4 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+*.png
diff --git a/README.md b/README.md
@@ -44,7 +44,66 @@ Monte-Carlo Nystrom method in NumPy / PyTorch
 
 Install this via pip (or your favourite package manager):
 
-`pip install montecarlo-nystrom`
+```shell
+pip install montecarlo-nystrom
+```
+
+## Usage
+
+Solve integral equations of the second kind of the following form.
+
+$\forall d \in \mathbb{N}.$
+$\forall \Omega \in \mathbb{R}^d [\Omega \text{ is bounded Lipschitz}].$
+$`\forall p \in L^\infty(\Omega, {\mathbb{R}}_{\geq 0}) [\int_\Omega p(y) dy = 1].`$
+$\forall f, z \in L^2(\Omega,\mathbb{C}).$
+$\forall k \in L^2(\Omega,L^2(\Omega,\mathbb{C}))$
+$[z(y) + \int_\Omega k(y, y') z(y') p(y') dy' = f(y)].$
+
+Let $N \in \mathbb{N}$, $(y_i)_{i=1}^N$ be i.i.d. samples drawn from $p$.
+
+Let $(z_{N,i})_{i=1}^N$ be the solution of the linear system
+
+$$
+z_{N,i} + \frac{1}{N} \sum_{j=1}^N k(y_i, y_j) z_{N,j} = f(y_i)
+\quad i \in \{1, \ldots, N\}
+$$
+
+and
+
+$$
+z_N(y) := f(y) - \frac{1}{N} \sum_{i=1}^N k(y, y_i) z_{N,i}
+\quad y \in \Omega
+$$
+
+Then $z_N$ would approximate $z$ as $N \to \infty$.
+
+The below example solves the case where $d = 1$, $\Omega = [0, 1]$, $p$ (`random_samples`) is the uniform distribution on $[0, 1]$, $k(x, y) = \|x - y\|^{-0.4}$ (`kernel`), and $f(x) = 1$ (`rhs`), and evaluates the solution at $x = (0.5,)$.
+
+```python
+>>> import numpy as np
+>>> from montecarlo_nystrom import montecarlo_nystrom
+>>> rng = np.random.default_rng(0)
+>>> def random_samples(n):
+...     return rng.uniform(0, 1, size=(n, 1))
+>>> def kernel(x, y):
+...     return np.linalg.vector_norm(x - y, axis=-1) ** -0.4
+>>> def rhs(x):
+...     x0 = x[..., 0]
+...     return np.ones_like(x0)
+>>> z_N = montecarlo_nystrom(
+...     random_samples=random_samples,
+...     kernel=kernel,
+...     rhs=rhs,
+...     n=100,
+...     n_mean=10,
+... )
+>>> np.round(z_N(np.asarray((0.5,))), 6)  # Evaluate at x=0.5
+np.float64(0.272957)
+```
+
+## References
+
+- Feppon, F., & Ammari, H. (2022). Analysis of a Monte-Carlo Nystrom Method. SIAM J. Numer. Anal. Retrieved from https://epubs.siam.org/doi/10.1137/21M1432338
 
 ## Contributors ✨
 
diff --git a/conftest.py b/conftest.py
@@ -0,0 +1,28 @@
+from sybil import Sybil
+from sybil.evaluators.doctest import NUMBER
+from sybil.parsers.myst import DocTestDirectiveParser as MarkdownDocTestParser
+from sybil.parsers.myst import PythonCodeBlockParser as MarkdownPythonCodeBlockParser
+from sybil.parsers.myst import SkipParser as MarkdownSkipParser
+from sybil.parsers.rest import DocTestParser as ReSTDocTestParser
+from sybil.parsers.rest import PythonCodeBlockParser as ReSTPythonCodeBlockParser
+from sybil.parsers.rest import SkipParser as ReSTSkipParser
+
+markdown_examples = Sybil(
+    parsers=[
+        MarkdownDocTestParser(NUMBER),
+        MarkdownPythonCodeBlockParser(doctest_optionflags=NUMBER),
+        MarkdownSkipParser(),
+    ],
+    patterns=["*.md"],
+)
+
+rest_examples = Sybil(
+    parsers=[
+        ReSTDocTestParser(NUMBER),
+        ReSTPythonCodeBlockParser(),
+        ReSTSkipParser(),
+    ],
+    patterns=["*.py", "*.rst"],
+)
+
+pytest_collect_file = (markdown_examples + rest_examples).pytest()
diff --git a/docs/conf.py b/docs/conf.py
@@ -25,11 +25,17 @@
 # ones.
 extensions = [
     "myst_parser",
-    "sphinx.ext.napoleon",
+    "sphinx.ext.duration",
     "sphinx.ext.autodoc",
     "sphinx.ext.viewcode",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.napoleon",
+    "sphinx_math_dollar",
+]
+myst_enable_extensions = [
+    "amsmath",
+    "dollarmath",
 ]
-napoleon_google_docstring = False
 
 # The suffix of source filenames.
 source_suffix = [
diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -0,0 +1,49 @@
+{
+  description = "A flake providing a dev shell for Numba with CUDA without installing Numba via nix. Also supports PyTorch yet being minimal for Numba with CUDA.";
+
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+  };
+
+  outputs =
+    { self, nixpkgs }:
+    let
+      system = "x86_64-linux"; # Adjust if needed
+      pkgs = import nixpkgs {
+        system = system;
+        config.allowUnfree = true;
+      };
+      cudatookit-with-cudart-to-lib64 = pkgs.symlinkJoin {
+        name = "cudatoolkit";
+        paths = with pkgs.cudaPackages; [
+          cudatoolkit
+          (pkgs.lib.getStatic cuda_cudart)
+        ];
+        postBuild = ''
+          ln -s $out/lib $out/lib64
+        '';
+      };
+    in
+    {
+      devShells.${system}.default = pkgs.mkShell {
+        shellHook = ''
+          # Required for both PyTorch and Numba to find CUDA
+          export CUDA_PATH=${cudatookit-with-cudart-to-lib64}
+
+          # Required for numba-cuda Python package
+          export NUMBA_CUDA_INCLUDE_PATH=${cudatookit-with-cudart-to-lib64}/include
+
+          # cuda.bindings.nvjitlink.nvJitLinkError: ERROR_INTERNAL (6)
+          # Linker error log: ERROR 4 in nvvmAddNVVMContainerToProgram, may need newer version of nvJitLink library
+          export NUMBA_CUDA_USE_NVIDIA_BINDING=0
+
+          # Required for both PyTorch and Numba, adds necessary paths for dynamic linking
+          export LD_LIBRARY_PATH=${
+            pkgs.lib.makeLibraryPath [
+              "/run/opengl-driver" # Needed to find libGL.so, required by both PyTorch and Numba
+            ]
+          }:$LD_LIBRARY_PATH
+        '';
+      };
+    };
+}
diff --git a/pyproject.toml b/pyproject.toml
@@ -11,7 +11,7 @@ license = "MIT"
 authors = [
   { name = "ultrasphere-dev", email = "34j.95a2p@simplelogin.com" },
 ]
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 classifiers = [
   "Development Status :: 2 - Pre-Alpha",
   "Intended Audience :: Developers",
@@ -27,10 +27,12 @@ classifiers = [
 
 dependencies = [
     "array-api-compat>=1.12.0",
+    "array-api-extra>=0.9.0",
     "cyclopts>=3.24.0",
     "matplotlib>=3.10.7",
     "numpy>=2.2.6",
     "types-array-api>=1.1.4",
+    "ultrasphere>=2.0.4",
 ]
 urls."Bug Tracker" = "https://github.com/ultrasphere-dev/montecarlo-nystrom/issues"
 urls.Changelog = "https://github.com/ultrasphere-dev/montecarlo-nystrom/blob/main/CHANGELOG.md"
@@ -43,12 +45,14 @@ dev = [
   "coverage",
   "pytest>=9,<10",
   "pytest-cov>=7,<8",
+  "sybil>=9.3.0",
 ]
 docs = [
   "furo>=2023.5.20; python_version>='3.11'",
   "myst-parser>=0.16; python_version>='3.11'",
   "sphinx>=4; python_version>='3.11'",
   "sphinx-autobuild>=2025,<2026; python_version>='3.11'",
+  "sphinx-math-dollar>=1.2.1",
 ]
 
 [tool.ruff]
diff --git a/src/montecarlo_nystrom/__init__.py b/src/montecarlo_nystrom/__init__.py
@@ -1 +1,4 @@
 __version__ = "1.0.0"
+from ._main import montecarlo_nystrom
+
+__all__ = ["montecarlo_nystrom"]
diff --git a/src/montecarlo_nystrom/_main.py b/src/montecarlo_nystrom/_main.py
@@ -10,17 +10,18 @@ def montecarlo_nystrom(
     kernel: Callable[[Array, Array], Array],
     rhs: Callable[[Array], Array],
     n: int,
+    n_mean: int = 1,
+    solve: Callable[[Array, Array], Array] | None = None,
 ) -> Callable[[Array], Array]:
     r"""
     Solve integral equations of the second kind of the following form.
 
-    $\forall d \in \mathbb{N}.
-    \forall \Omega \in \mathbb{R}^d [\Omega \text{is Lipschitz}].
-    \forall p \in L^\infty(\Omega, \mathbb{R}_{\geq 0}) [\int_\Omega p(y) dy = 1].
-    \forall f, z \in L^2(\Omega,\mathbb{C}).
-    \forall k \in L^2(\Omega,L^2(\Omega,\mathbb{C})
-    [z(y) + \int k(y, y') z(y') p(y') dy' = f(y)].
-    $
+    $\forall d \in \mathbb{N}.$
+    $\forall \Omega \in \mathbb{R}^d [\Omega \text{ is bounded Lipschitz}].$
+    $\forall p \in L^\infty(\Omega, {\mathbb{R}}_{\geq 0}) [\int_\Omega p(y) dy = 1].$
+    $\forall f, z \in L^2(\Omega,\mathbb{C}).$
+    $\forall k \in L^2(\Omega,L^2(\Omega,\mathbb{C}))$
+    $[z(y) + \int_\Omega k(y, y') z(y') p(y') dy' = f(y)].$
 
     Let $N \in \mathbb{N}$, $(y_i)_{i=1}^N$ be i.i.d. samples drawn from $p$.
 
@@ -47,28 +48,87 @@ def montecarlo_nystrom(
         returns n i.i.d. samples from the distribution p of shape (..., n, d).
     kernel : Callable[[Array, Array], Array]
         Kernel function k of (..., :, :, d), (..., :, :, d) -> (..., :, :).
+
+        The diagonal part will be automatically dropped
+        and can be of any value.
     rhs : Callable[[Array], Array]
         Right-hand side function f of (..., :, d) -> (..., :).
     n : int
         The number of random samples to draw from the distribution p.
+    n_mean : int
+        The number of independent runs to average over, by default 1.
+    solve : Callable[[Array, Array], Array] | None, optional
+        A function that takes a matrix A of shape (..., n, n)
+        and a right-hand side b of shape (..., n, 1)
+        and returns the solution of shape (..., n, 1).
+        If None, uses `xp.linalg.solve`, by default None.
 
     Returns
     -------
     Callable[[Array], Array]
         Approximate solution function z_N that takes an array of shape (...(x), ..., d)
         and returns an array of shape (...(x), ...).
 
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from montecarlo_nystrom import montecarlo_nystrom
+    >>> rng = np.random.default_rng(0)
+    >>> def random_samples(n):
+    ...     return rng.uniform(0, 1, size=(n, 1))
+    >>> def kernel(x, y):
+    ...     return np.linalg.vector_norm(x - y, axis=-1) ** -0.4
+    >>> def rhs(x):
+    ...     x0 = x[..., 0]
+    ...     return np.ones_like(x0)
+    >>> z_N = montecarlo_nystrom(
+    ...     random_samples=random_samples,
+    ...     kernel=kernel,
+    ...     rhs=rhs,
+    ...     n=100,
+    ...     n_mean=10,
+    ... )
+    >>> np.round(z_N(np.asarray((0.5,))), 6)  # Evaluate at x=0.5
+    np.float64(0.272957)
+
     """
-    y = random_samples(n)
-    K = kernel(y[..., :, None, :], y[..., None, :, :])
-    b = rhs(y)
-    xp = array_namespace(K, b)
-    K[..., xp.arange(n), xp.arange(n)] = 0
-    A = xp.eye(n) + K / n
-    z_N_samples = xp.linalg.solve(A, b)
+    y = random_samples(n * n_mean)
+    if y.shape[-2] != n * n_mean:
+        raise ValueError(
+            f"random_samples returned array of shape {y.shape}, "
+            f"expected (..., {n * n_mean}, d)"
+        )
+    xp = array_namespace(y)
+    y = xp.reshape(y, (*y.shape[:-2], n_mean, n, y.shape[-1]))  # (..., n_mean, n, d)
+    K = kernel(y[..., :, :, None, :], y[..., :, None, :, :])  # (..., n_mean, n, n)
+    b = rhs(y)  # (..., n_mean, n)
+    if K.shape[-3:] != (n_mean, n, n):
+        raise ValueError(
+            f"kernel returned array of shape {K.shape}, "
+            f"expected (..., {n_mean}, {n}, {n})"
+        )
+    if b.shape[-2:] != (n_mean, n):
+        raise ValueError(
+            f"rhs returned array of shape {b.shape}, expected (..., {n_mean}, {n})"
+        )
+    K[..., :, xp.arange(n), xp.arange(n)] = 0  # drop diagonal
+    A = xp.eye(n, dtype=K.dtype, device=K.device) + K / n
+    z_N_samples = (solve or xp.linalg.solve)(A, b[..., None])
+    if z_N_samples.shape[-3:] != (n_mean, n, 1):
+        raise ValueError(
+            f"solve returned array of shape {z_N_samples.shape}, "
+            f"expected (..., {n_mean}, {n}, 1)"
+        )
+    z_N_samples = z_N_samples[..., 0]  # (..., n_mean, n)
 
     def z_N(x: Array) -> Array:
-        K_x = kernel(x[..., None, :], y)
-        return rhs(x) - K_x @ z_N_samples / n
+        if x.shape[-1] != y.shape[-1]:
+            raise ValueError(
+                f"x has shape {x.shape}, expected (...(x), ..., {y.shape[-1]})"
+            )
+        K_x = kernel(x[..., None, None, :], y)  # (...(x), ..., n_mean, n)
+        return (
+            rhs(x) - xp.mean(xp.vecdot(xp.conj(K_x), z_N_samples, axis=-1), axis=-1) / n
+        )
 
     return z_N
diff --git a/src/montecarlo_nystrom/cli.py b/src/montecarlo_nystrom/cli.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`[codespell]`
`2`		`-ignore-words-list = socio-economic`
	`2`	`+ignore-words-list = socio-economic, numer`
Original file line number	Diff line number	Diff line change
`@@ -173,3 +173,4 @@ cython_debug/`
`173`	`173`
`174`	`174`	`# PyPI configuration file`
`175`	`175`	`.pypirc`
	`176`	`+*.png`