diff --git a/.github/workflows/alltests.yml b/.github/workflows/alltests.yml index 97486f78f..b94c8cfab 100644 --- a/.github/workflows/alltests.yml +++ b/.github/workflows/alltests.yml @@ -312,6 +312,15 @@ jobs: run: | make check_colab_notebooks make check_colab_notebooks_smoke + - name: Check test-suite conventions (Linux) + if: runner.os == 'Linux' + shell: bash -l {0} + run: | + set -e + make check_test_style STRICT=--strict + python -m pip install -q "pydoclint>=0.5.0" + make check_docstring # informational (exits 0 without STRICT) + make check_baseline # real gate: docstring/pydoclint/annotation issues must not increase # ----------------------------------------------------------- # Install minimal LaTeX required by Jupyter notebooks (OS-specific) # ----------------------------------------------------------- diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 486686ac7..6c0fe3ae5 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -163,6 +163,10 @@ make tests Please see the targets in the makefile for more granular control over tests. +### Test file layout + +Unit tests live flat in `test/`, named `test__.py` where `` is a short code for the `qmcpy` subpackage under test (`tm` true_measure, `dd` discrete_distribution, `sc` stopping_criterion, `ig` integrand, ...) or a cross-cutting bucket (`ee`, `sr`). So `pytest test/ -k test_tm_` runs every true-measure test. A test that spans two areas goes under the component actually under test, with the other named in `` (e.g. `test_sc_cubbayes_kernels.py`); use `ee` only when neither side is the clear subject, and never coin a new code — `STRICT=--strict` rejects anything outside the table. New files should also be written as a `unittest.TestCase` subclass rather than bare `def test_*` functions. `make check_test_style` lists any file that breaks either convention (informational; also runs inside `make format`; `STRICT=--strict` makes it fail). The full area table is in [`test/README.md`](test/README.md#test-file-organization). + ## Documentation ### Ensure `pyreverse` Is On Your PATH diff --git a/demos/acceptance_rejection.ipynb b/demos/acceptance_rejection.ipynb index 8bed5e925..b6dfbba55 100644 --- a/demos/acceptance_rejection.ipynb +++ b/demos/acceptance_rejection.ipynb @@ -494,7 +494,7 @@ "$$D^*_N(\\psi) = \\sup_{t \\in [0,1]} \\left| \\frac{\\#\\{i : x_i \\leq t\\}}{N} - \\frac{1}{C}\\int_0^t \\psi(x)\\,dx \\right|$$\n", "\n", "Theorem 1 of Zhu & Dick (2014) predicts:\n", - "- **Sobol driver (DAR):** $D^*_N = O(N^{-1/s})$ with $s=2$ for $d=1$, i.e. $O(N^{-1/2})$ in the worst case — but empirically faster due to the $\\log N$ factors.\n", + "- **Sobol driver (DAR):** $D^*_N = O(N^{-1/s})$ with $s=2$ for $d=1$, i.e., $O(N^{-1/2})$ in the worst case — but empirically faster due to the $\\log N$ factors.\n", "- **Random driver (standard A-R):** $D^*_N = O(N^{-1/2})$ always.\n", "\n", "The plot below measures this empirically for $\\psi(x) = 2x$ on $[0,1]$, where the CDF is $F(t) = t^2$ and $C = 1$." diff --git a/demos/lebesgue_integration.ipynb b/demos/lebesgue_integration.ipynb index 4db928774..3894898b1 100644 --- a/demos/lebesgue_integration.ipynb +++ b/demos/lebesgue_integration.ipynb @@ -3,7 +3,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "# QMCPy for Lebesgue Integration\nThis notebook will give examples of how to use QMCPy for integration problems that are not defined in terms of a standard measure. i.e. Uniform or Gaussian. " + "source": "# QMCPy for Lebesgue Integration\nThis notebook will give examples of how to use QMCPy for integration problems that are not defined in terms of a standard measure. i.e., Uniform or Gaussian. " }, { "cell_type": "markdown", diff --git a/demos/makefile_dev_tools.ipynb b/demos/makefile_dev_tools.ipynb new file mode 100644 index 000000000..d78afacb8 --- /dev/null +++ b/demos/makefile_dev_tools.ipynb @@ -0,0 +1,692 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# `make format` and `make check`\n", + "\n", + "You run two commands around every change to QMCPy:\n", + "\n", + "| command | when | what it does |\n", + "|---|---|---|\n", + "| **`make format`** | before you commit | **changes your files**: imports, whitespace, asserts, docstring types |\n", + "| **`make check`** | before you open a PR | **only reads**: the same rules CI enforces |\n", + "\n", + "Each command runs a few small tools in order. This notebook walks through both, shows what each tool does on a tiny example, then lists the rest.\n", + "\n", + "The code cells run the real scripts when you are inside a QMCSoftware checkout. From a plain `pip install qmcpy` they print the same before/after as text." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/QMCSoftware/QMCSoftware/blob/develop/demos/makefile_dev_tools.ipynb)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "# @title Execute this cell to install dependencies\n", + "try:\n", + " import google.colab\n", + " IN_COLAB = True\n", + "except ImportError:\n", + " IN_COLAB = False\n", + "if IN_COLAB:\n", + " !pip install -q qmcpy" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "# Every import the notebook uses, in one place.\n", + "import json, os, pathlib, re, shutil, subprocess, sys, tempfile\n", + "\n", + "# The QMCSoftware checkout, or None when this notebook runs on its own.\n", + "REPO = next((d for d in (pathlib.Path.cwd(), *pathlib.Path.cwd().parents)\n", + " if (d / \"scripts/convert_asserts.py\").exists()), None)\n", + "\n", + "def demo(tool, before, after, target, name=\"snippet.py\"):\n", + " \"\"\"Print `Before`, then `After` running `tool` on it.\n", + "\n", + " `tool` is the argument list for a script under scripts/. Inside a checkout\n", + " the real script runs and its result replaces `after`; otherwise the canned\n", + " `after` text is used.\n", + " \"\"\"\n", + " if REPO:\n", + " f = pathlib.Path(tempfile.mkdtemp()) / name\n", + " f.write_text(before)\n", + " try:\n", + " subprocess.run([sys.executable, *tool, str(f)], cwd=REPO,\n", + " capture_output=True, text=True, check=True)\n", + " after = f.read_text()\n", + " except (FileNotFoundError, subprocess.CalledProcessError):\n", + " pass\n", + " print(\"Before:\\n-------\\n\" + before)\n", + " after_str = f\"After (make {target}):\"\n", + " print(f\"{after_str}\\n\" + \"-\"*len(after_str) + \"\\n\" + after)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. The two pipelines\n", + "\n", + "`make format` and `make check` are just ordered lists of smaller targets. Here they are, read straight from the `makefile`:" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "make format\n", + " flatten_qmcpy_imports\n", + " markdown-unwrap\n", + " rm_trailing_whitespace\n", + " harden_colab_notebook\n", + " convert_asserts_changed\n", + " add_docstring_arg_types_changed\n", + "make check\n", + " check_test_style\n", + " check_docstring_changed\n", + " check_baseline\n", + " check_asserts_changed\n", + " check_links\n" + ] + } + ], + "source": [ + "DEFAULT = {\n", + " \"format\": \"flatten_qmcpy_imports markdown-unwrap rm_trailing_whitespace \"\n", + " \"harden_colab_notebook convert_asserts_changed add_docstring_arg_types_changed\".split(),\n", + " \"check\": \"check_test_style check_docstring_changed check_baseline \"\n", + " \"check_asserts_changed check_links\".split(),\n", + "}\n", + "\n", + "def steps(target):\n", + " \"\"\"The sub-targets `make ` runs, read from the makefile.\"\"\"\n", + " makefile = REPO / \"makefile\" if REPO else None\n", + " if makefile and makefile.exists():\n", + " block = re.search(rf\"^{target}:(?:.*\\n)((?:[ \\t].*\\n|\\n)+)\", makefile.read_text(), re.M)\n", + " names = re.findall(r\"\\$\\(MAKE\\)\\s+(\\S+)\", block.group(1)) if block else []\n", + " if names:\n", + " return names\n", + " return DEFAULT[target]\n", + "\n", + "for target in (\"format\", \"check\"):\n", + " print(f\"make {target}\")\n", + " for name in steps(target):\n", + " print(\" \", name)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Every step ends with **one summary line**, so a run is easy to scan:\n", + "\n", + "- `clean (0/42 files)` — nothing to do\n", + "- `3 changed (3/42 files)` — `make format` rewrote 3 files\n", + "- `WARNING: 2 problem(s) (2/42 files)` — issues found, but this step does not fail the build\n", + "- `ERROR: 2 problem(s) (2/42 files)` — issues found and this step fails (same as CI)\n", + "\n", + "Any details are listed just above that line, one `-` bullet each." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. `make format`: changes your files\n", + "\n", + "Run it before you commit. Every step is safe to run again and touches only what it needs to.\n", + "\n", + "| step | what it does |\n", + "|---|---|\n", + "| `flatten_qmcpy_imports` | rewrite `from qmcpy.sub.mod import X` as `import qmcpy as qp` then `qp.X` |\n", + "| `markdown-unwrap` | join hard-wrapped Markdown lines back into one line per paragraph |\n", + "| `rm_trailing_whitespace` | remove trailing spaces across the repo |\n", + "| `harden_colab_notebook` | add the *Open in Colab* badge to any demo notebook that lacks one |\n", + "| `convert_asserts_changed` | turn `assert` into a real `raise`, in changed files (next cell) |\n", + "| `add_docstring_arg_types_changed` | copy signature types into `Args:` lines, in changed files (cell after) |\n", + "\n", + "`format` has no docstring reformatter. `format-docstring` was tried and dropped: on this code it deletes `Returns:` and `Yields:` types and turns `**References:**` into `**References: **`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.1 `flatten_qmcpy_imports`\n", + "\n", + "Collapse a deep import path to the one public name. Anything exported by `qmcpy` should be reached as `qp.`, so imports stay stable when internal modules move." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Before:\n", + "-------\n", + "from qmcpy.discrete_distribution.lattice.lattice import Lattice\n", + "\n", + "After (make flatten_qmcpy_imports):\n", + "-----------------------------------\n", + "from qmcpy import Lattice\n", + "\n" + ] + } + ], + "source": [ + "demo([\"scripts/flatten_qmcpy_imports.py\"],\n", + "\"\"\"from qmcpy.discrete_distribution.lattice.lattice import Lattice\n", + "\"\"\",\n", + "\"\"\"from qmcpy import Lattice\n", + "\"\"\", \"flatten_qmcpy_imports\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.2 `markdown_unwrap`\n", + "\n", + "Join each hard-wrapped Markdown paragraph back onto one line (code fences and math are left alone), so a later edit shows as a word change, not a whole-paragraph rewrap." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Before:\n", + "-------\n", + "QMCPy estimates an integral as the mean of an\n", + "integrand sampled on a low-discrepancy point\n", + "set, and stops once the error is small enough.\n", + "\n", + "After (make markdown_unwrap):\n", + "-----------------------------\n", + "QMCPy estimates an integral as the mean of an integrand sampled on a low-discrepancy point set, and stops once the error is small enough.\n", + "\n" + ] + } + ], + "source": [ + "demo([\"scripts/unwrap_markdown.py\"],\n", + "\"QMCPy estimates an integral as the mean of an\\n\"\n", + "\"integrand sampled on a low-discrepancy point\\n\"\n", + "\"set, and stops once the error is small enough.\\n\",\n", + "\"QMCPy estimates an integral as the mean of an integrand sampled on a \"\n", + "\"low-discrepancy point set, and stops once the error is small enough.\\n\",\n", + "\"markdown_unwrap\", name=\"snippet.md\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.3 `rm_trailing_whitespace`\n", + "\n", + "Strip spaces and tabs at end of line across every git-tracked text file. Nothing to run on a snippet here — it asks git which files exist. Effect: `\"muhat = 1.80 \\n\"` becomes `\"muhat = 1.80\\n\"`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.4 `harden_colab_notebook`\n", + "\n", + "Add the *Open in Colab* badge cell and the dependency-install cell to any notebook under `demos/` that is listed in the Colab manifest but missing them. It works on real notebook files, not snippets; `make format` runs it for every still-unclassified notebook." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.5 `convert_asserts`\n", + "\n", + "`python -O` removes every `assert`, so a check written that way is gone when Python runs with `-O`. This rewrites it as a real `raise`, keeping comments and layout." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Before:\n", + "-------\n", + "def clip(x, lo, hi):\n", + " assert lo <= hi, \"empty interval\"\n", + "\n", + "After (make convert_asserts_changed):\n", + "-------------------------------------\n", + "def clip(x, lo, hi):\n", + " if not (lo <= hi):\n", + " raise AssertionError(\"empty interval\")\n", + "\n" + ] + } + ], + "source": [ + "demo([\"scripts/convert_asserts.py\", \"--exception\", \"AssertionError\"],\n", + "\"\"\"def clip(x, lo, hi):\n", + " assert lo <= hi, \"empty interval\"\n", + "\"\"\",\n", + "\"\"\"def clip(x, lo, hi):\n", + " if not (lo <= hi):\n", + " raise AssertionError(\"empty interval\")\n", + "\"\"\", \"convert_asserts_changed\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.6 `add_docstring_arg_types`\n", + "\n", + "The type is already in the signature. This copies it into the matching `Args:` line (`name` becomes `name (type)`). It never makes up a type or a description." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Before:\n", + "-------\n", + "def disc_area(radius: float, n_sectors: int = 4) -> float:\n", + " \"\"\"Area of a disc.\n", + "\n", + " Args:\n", + " radius: distance to the edge.\n", + " n_sectors: wedge count.\n", + " \"\"\"\n", + " return 3.14159 * radius ** 2\n", + "\n", + "After (make add_docstring_arg_types_changed):\n", + "---------------------------------------------\n", + "def disc_area(radius: float, n_sectors: int = 4) -> float:\n", + " \"\"\"Area of a disc.\n", + "\n", + " Args:\n", + " radius (float): distance to the edge.\n", + " n_sectors (int): wedge count.\n", + " \"\"\"\n", + " return 3.14159 * radius ** 2\n", + "\n" + ] + } + ], + "source": [ + "src = \"\"\"def disc_area(radius: float, n_sectors: int = 4) -> float:\n", + " \\\"\\\"\\\"Area of a disc.\n", + "\n", + " Args:\n", + " radius: distance to the edge.\n", + " n_sectors: wedge count.\n", + " \\\"\\\"\\\"\n", + " return 3.14159 * radius ** 2\n", + "\"\"\"\n", + "\n", + "demo([\"scripts/add_docstring_arg_types.py\"], src,\n", + " src.replace(\"radius:\", \"radius (float):\").replace(\"n_sectors:\", \"n_sectors (int):\"),\n", + " \"add_docstring_arg_types_changed\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. `make check`: only reads\n", + "\n", + "Run it before you open a PR. Nothing here changes your files. It runs the same checks as the *Check test-suite conventions* job in CI.\n", + "\n", + "| step | what it checks |\n", + "|---|---|\n", + "| `check_test_style` | each `test/` file is named `test__*.py` and uses a `unittest.TestCase` class (next cell) |\n", + "| `check_docstring_changed` | Google-style format plus `pydoclint`, on your changed files |\n", + "| `check_baseline` | the count of known issues did not go up (cell after) |\n", + "| `check_asserts_changed` | a dry run of `convert_asserts` on changed source |\n", + "| `check_links` | internal doc links and anchors resolve |\n", + "\n", + "Left out on purpose: `check_links_external` (needs the network) and `check_pep8_changed` (too many old violations to be a useful gate today)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.1 `check_test_style`\n", + "\n", + "One file that follows both rules, one that breaks both:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + " - 1 of 2 files use a unittest.TestCase class\n", + " - 1 file(s) use bare pytest functions (no unittest.TestCase class):\n", + " - test_zz_bad.py\n", + " - 1 of 2 files use a test__ prefix (dd, ee, ft, ig, kn, sc, sr, tm, ut)\n", + " - 1 file(s) have no recognized test__ prefix:\n", + " - test_zz_bad.py\n", + "WARNING: 1 problem(s) (1 of 2 files)\n", + "\n" + ] + } + ], + "source": [ + "folder = pathlib.Path(tempfile.mkdtemp())\n", + "(folder / \"test_tm_ok.py\").write_text(\n", + " \"import unittest\\nclass T(unittest.TestCase):\\n def test_x(self): self.assertEqual(2, 2)\\n\")\n", + "(folder / \"test_zz_bad.py\").write_text(\"def test_x():\\n assert 2 == 2\\n\")\n", + "\n", + "if REPO:\n", + " # run from inside `folder` so the report shows plain file names\n", + " result = subprocess.run([sys.executable, str(REPO / \"scripts/check_test_style.py\"), \".\"],\n", + " cwd=folder, capture_output=True, text=True)\n", + " print(result.stdout)\n", + "else:\n", + " print(\"test_zz_bad.py is flagged twice: it uses a bare function, and 'zz' is not a known area.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.2 `check_docstring`\n", + "\n", + "Two passes over public docstrings under `qmcpy/`: `check_docstring.py` for Google-style **format** (a one-line summary first, a blank line before each `Args:` / `Returns:`, canonical `Name:` headers) and `pydoclint` for **content** (every parameter and return documented). Informational unless `STRICT=--strict`." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + " - mod.py:3: missing-summary: docstring opens with `Args:`; add a one-line summary first\n", + " - 1 file(s) scanned: 1 issue(s) across 1 file(s): 1 missing-summary\n", + "WARNING: 1 problem(s) (1 of 1 files)\n", + "\n", + "With STRICT=--strict the last line becomes 'ERROR: ...' and the exit is 1.\n" + ] + } + ], + "source": [ + "mod = pathlib.Path(tempfile.mkdtemp()) / \"mod.py\"\n", + "mod.write_text(\n", + " 'def disc_area(radius):\\n'\n", + " ' \"\"\"\\n'\n", + " ' Args:\\n'\n", + " ' radius (float): distance to the edge.\\n'\n", + " ' \"\"\"\\n'\n", + " ' return 3.14159 * radius ** 2\\n')\n", + "\n", + "if REPO:\n", + " print(subprocess.run([sys.executable, str(REPO / \"scripts/check_docstring.py\"), mod.name],\n", + " cwd=mod.parent, capture_output=True, text=True).stdout)\n", + "else:\n", + " print(\" - mod.py:3: missing-summary: docstring opens with `Args:`; add a one-line summary first\")\n", + " print(\"WARNING: 1 problem(s) (1/1 files)\")\n", + "print(\"With STRICT=--strict the last line becomes 'ERROR: ...' and the exit is 1.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.3 `check_baseline`\n", + "\n", + "Some checks have many old violations, so they cannot fail the build yet. Their counts are stored in a file, and the build fails only if a count goes up." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "stored counts: {'check_docstring': 0, 'pydoclint': 0, 'unsafe_annotations': 0}\n", + "make check_baseline fails if any count goes above this.\n", + "make check_baseline_update saves new counts after you fix or accept a change.\n" + ] + } + ], + "source": [ + "path = REPO / \"scripts/baseline_counts.json\" if REPO else None\n", + "counts = (json.loads(path.read_text()) if path and path.exists()\n", + " else {\"check_docstring\": 0, \"pydoclint\": 0, \"unsafe_annotations\": 0})\n", + "print(\"stored counts:\", counts)\n", + "print(\"make check_baseline fails if any count goes above this.\")\n", + "print(\"make check_baseline_update saves new counts after you fix or accept a change.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.4 `check_asserts`\n", + "\n", + "The read-only half of `convert_asserts` (2.5): same detection, but it writes nothing and exits non-zero when an `assert` in changed code could be converted. This is the variant `make check` runs." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3.5 `check_links`\n", + "\n", + "Build the docs site (`mkdocs build`) and check that every internal link and heading anchor resolves. No inline example — it needs the built `site/`. A failure reads:\n", + "\n", + "```\n", + " - path/page.html: broken internal link '../missing.md'\n", + "ERROR: 3 problem(s) (2/90 pages)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Other tools\n", + "\n", + "Things `format` and `check` do not call directly:\n", + "\n", + "| tool | what it is for |\n", + "|---|---|\n", + "| **`$(PYTHON)`** | the makefile finds the interpreter once (active env, then `$CONDA_PREFIX/bin/python`, then `conda run -n qmcpy`, then `python3`) and every rule uses it (next cell) |\n", + "| **the `_changed` suffix** | most tools have a whole-repo form and a `_changed` form that looks only at files changed from a base branch. On a large codebase, only the `_changed` form is practical day to day |\n", + "| **`annotate_public_api_types_changed`** and **`sync_docstring_types_changed`** | the reverse of `add_docstring_arg_types`: copy a documented type into a missing signature annotation, then copy annotations back into the docstrings. `check_public_api_types_changed` only reports (last cell) |\n", + "| **`tests_fast`** | run doctests, unit tests, and notebook tests at the same time; the rule records each exit code so a failing suite fails the target |\n", + "| **whole-repo forms** | `convert_asserts`, `add_docstring_arg_types`, `check_docstring` also exist without `_changed`, for a full pass |" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "interpreter: /Users/terrya/miniconda3/envs/qmcpy/bin/python\n" + ] + } + ], + "source": [ + "def discover_python():\n", + " \"\"\"The makefile's PYTHON ?= $(shell ...) cascade, written in Python.\"\"\"\n", + " env = os.environ.get(\"CONDA_PREFIX\")\n", + " if shutil.which(\"python\"): # 1. active env on PATH\n", + " return shutil.which(\"python\")\n", + " if env and shutil.which(\"python\", path=env + \"/bin\"): # 2. active env, not on PATH\n", + " return shutil.which(\"python\", path=env + \"/bin\")\n", + " if shutil.which(\"conda\"): # 3. the repo's qmcpy env\n", + " out = subprocess.run([\"conda\", \"run\", \"-n\", \"qmcpy\", \"python\",\n", + " \"-c\", \"import sys; print(sys.executable)\"],\n", + " capture_output=True, text=True)\n", + " if out.stdout.strip():\n", + " return out.stdout.strip()\n", + " return shutil.which(\"python3\") or sys.executable # 4. system python3\n", + "\n", + "\n", + "print(\"interpreter:\", discover_python())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4.1 `annotate_public_api_types`\n", + "\n", + "The other direction: the types are only in the docstring and the signature is bare." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Before:\n", + "-------\n", + "def disc_area(radius, n_sectors=4):\n", + " \"\"\"Area of a disc.\n", + "\n", + " Args:\n", + " radius (float): distance to the edge.\n", + " n_sectors (int): wedge count.\n", + "\n", + " Returns:\n", + " float: the area.\n", + " \"\"\"\n", + " return 3.14159 * radius ** 2\n", + "\n", + "After (make annotate_public_api_types_changed):\n", + "-----------------------------------------------\n", + "def disc_area(radius: float, n_sectors: int = 4) -> float:\n", + " \"\"\"Area of a disc.\n", + "\n", + " Args:\n", + " radius (float): distance to the edge.\n", + " n_sectors (int): wedge count.\n", + "\n", + " Returns:\n", + " float: the area.\n", + " \"\"\"\n", + " return 3.14159 * radius ** 2\n", + "\n" + ] + } + ], + "source": [ + "src = \"\"\"def disc_area(radius, n_sectors=4):\n", + " \\\"\\\"\\\"Area of a disc.\n", + "\n", + " Args:\n", + " radius (float): distance to the edge.\n", + " n_sectors (int): wedge count.\n", + "\n", + " Returns:\n", + " float: the area.\n", + " \\\"\\\"\\\"\n", + " return 3.14159 * radius ** 2\n", + "\"\"\"\n", + "\n", + "demo([\"-m\", \"scripts.annotate_public_api_types\"], src,\n", + " src.replace(\"(radius, n_sectors=4)\", \"(radius: float, n_sectors: int = 4) -> float\"),\n", + " \"annotate_public_api_types_changed\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Takeaways\n", + "\n", + "- Run `make format` before you commit, `make check` before you open a PR. One changes files, one only reads.\n", + "- Find the interpreter once. `$(PYTHON)` as a resolved variable removes a class of \"wrong Python\" bugs.\n", + "- Scope checks to the diff. A `_changed` target makes a whole-repo linter usable.\n", + "- Add a strict check to old code with a baseline count, not a hard gate.\n", + "- Use a codemod, not hand edits, for the assert and type changes, and keep a careless reformatter out of the pipeline." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "qmcpy", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.13" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/demos/talk_paper_demos/MCQMC_Tutorial_2020/MCQMC_2020_QMC_Software_Tutorial.ipynb b/demos/talk_paper_demos/MCQMC_Tutorial_2020/MCQMC_2020_QMC_Software_Tutorial.ipynb index 82136d126..aa7a8da55 100644 --- a/demos/talk_paper_demos/MCQMC_Tutorial_2020/MCQMC_2020_QMC_Software_Tutorial.ipynb +++ b/demos/talk_paper_demos/MCQMC_Tutorial_2020/MCQMC_2020_QMC_Software_Tutorial.ipynb @@ -2780,7 +2780,7 @@ " | - Pass in `generating_matrices` *without* interlacing and supply `alpha`>1 to apply interlacing, or\n", " | - Pass in `generating_matrices` *with* interlacing and set `alpha=1` to avoid additional interlacing\n", " |\n", - " | i.e. do *not* pass in interlaced `generating_matrices` and set `alpha>1`, this will apply additional interlacing.\n", + " | i.e., do *not* pass in interlaced `generating_matrices` and set `alpha>1`, this will apply additional interlacing.\n", " |\n", " | Examples:\n", " | >>> discrete_distrib = DigitalNetB2(2,seed=7)\n", diff --git a/demos/talk_paper_demos/SorokinThesis2025/sorokin_thesis_2025.ipynb b/demos/talk_paper_demos/SorokinThesis2025/sorokin_thesis_2025.ipynb index 95e13d4fe..125d952c6 100644 --- a/demos/talk_paper_demos/SorokinThesis2025/sorokin_thesis_2025.ipynb +++ b/demos/talk_paper_demos/SorokinThesis2025/sorokin_thesis_2025.ipynb @@ -95,7 +95,7 @@ " dimension = 52, \n", " randomize = \"LMS DS\", # Matousek's LMS then a digital shift\n", " # other options [\"NUS\", \"DS\", \"LMS\", None]\n", - " t = 64, # number of LMS bits i.e. number of rows in S_j\n", + " t = 64, # number of LMS bits i.e., number of rows in S_j\n", " alpha = 2, # interlacing factor for higher order digital nets\n", " replications = 16, # R\n", " order = \"radical inverse\", # also supports \"Gray code\"\n", @@ -124,7 +124,7 @@ " dimension = 52, \n", " randomize = \"LMS DP\", # Matousek's LMS then a digital permutation\n", " # other options [\"LMS DS\", \"LMS\", \"DP\", \"DS\", \"NUS\", \"QRNG\", None]\n", - " t = 64, # number of LMS digits i.e. number of rows in S_j\n", + " t = 64, # number of LMS digits i.e., number of rows in S_j\n", " replications = 16, # R\n", " seed = None) # pass integer seed for reproducibility\n", "x = halton(2**10) # a numpy.ndarray with shape 16 x 1024 x 52" diff --git a/docs/api/discrete_distributions.md b/docs/api/discrete_distributions.md index daa782cb6..49b60fbb4 100644 --- a/docs/api/discrete_distributions.md +++ b/docs/api/discrete_distributions.md @@ -70,10 +70,7 @@ python -m pip install "qmcpy[mpmc]" qmcpy-install-mpmc ``` -The second command selects the `pyg_lib` wheel page matching the installed -PyTorch and accelerator builds. For GPU support or platform-specific wheels, -see the [PyTorch installation guide](https://pytorch.org/get-started/locally/) -and the [PyTorch Geometric installation guide](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html). +The second command selects the `pyg_lib` wheel page matching the installed PyTorch and accelerator builds. For GPU support or platform-specific wheels, see the [PyTorch installation guide](https://pytorch.org/get-started/locally/) and the [PyTorch Geometric installation guide](https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html). ::: qmcpy.discrete_distribution.mpmc.mpmc.MPMC diff --git a/docs/good_practices.md b/docs/good_practices.md index d87ebc827..e4da6ed56 100644 --- a/docs/good_practices.md +++ b/docs/good_practices.md @@ -38,10 +38,27 @@ When notebook-backed content changes: QMCPy documentation is built from docstrings, so public APIs should document their behavior clearly and consistently. - Use **Google-style docstrings** for public classes, methods, and functions. -- Document parameters, return values, shapes, assumptions, and any stochastic behavior. +- Start every docstring with a one-line summary before any section header. +- Document every parameter and return value, plus shapes, assumptions, and any stochastic behavior. Constructor arguments go in the `__init__` method's own docstring, with the type in the docstring (`name (type): ...`). +- Put a blank line before every section header (`Args:`, `Returns:`, `Raises:`, `Examples:`, ...) and write the header as `Name:` — not a NumPy-style `Name` followed by an `-----` underline. - Include short doctestable examples when they clarify expected use. - Update docstrings at the same time as the implementation so the rendered API docs do not drift from the code. +`make check_docstring` runs two checks over public objects under `qmcpy/`: + +- `scripts/check_docstring.py` for **formatting** — a one-line summary before the first section (`missing-summary`), no NumPy-style sections, a blank line before every section header, canonical `Name:` headers, and public objects with no docstring. After the overall count it prints a second summary restricted to files changed relative to `DOCSTRING_BASE` (default `develop`), so you can see your branch's contribution to the backlog. +- `pydoclint` (configured in `pyproject.toml` under `[tool.pydoclint]`) for **content** — every parameter and return value is documented and matches the signature, in Google form. + +It is informational by default; `STRICT=--strict make check_docstring` makes both parts fail the build. Pass `CHECK_DOCSTRING_ARGS=--skip-missing` to skip the "no docstring" formatting check, or `DOCSTRING_PATH=qmcpy/true_measure` to narrow the scan. `make check_docstring_changed` runs the same two checks on just the `qmcpy/*.py` files that changed relative to `DOCSTRING_BASE` — the quick check to run before opening a PR (it is also part of `make format`). + +For annotated public APIs, `make add_docstring_arg_types` inserts missing Google-style argument types into existing `Args:` entries from the function signature. For example, `distance: float` becomes `distance (float): ...` in the docstring. Use `DOCSTRING_TYPE_PATH=path/to/file.py` to narrow the scan, or run `make add_docstring_arg_types_changed` to apply it only to Python files reported by `git diff --name-only develop -- '*.py'`. Use `DOCSTRING_TYPE_DIFF_BASE=origin/develop` to compare against a different base, and use `make check_docstring_arg_types_changed` to fail when changed files still need annotation-derived updates. The helper does not infer types for unannotated functions and does not invent missing scientific argument descriptions. + +For changed public APIs, `make annotate_public_api_types_changed` performs the reverse operation conservatively: it copies explicit, valid Google `Args:` and `Returns:` types into missing function annotations and adds `-> None` to constructors. It never replaces an existing annotation. Types that are prose, use syntax unsafe for Python 3.9, or reference names not already available in the module are reported and skipped. Then `make sync_docstring_types_changed` copies the resulting signature annotations back into existing `Args:`, `Returns:`, and `Yields:` descriptions. Run the annotation target before the synchronization target, review the complete diff, and run `make check_public_api_types_changed` for a non-mutating consistency check. All three targets default to files under `qmcpy/` changed relative to `develop`; override this with `PUBLIC_API_TYPE_PATH` or `PUBLIC_API_TYPE_DIFF_BASE`. + +These helpers synchronize explicit type information; they do not infer a scientific API contract from default values, implementation expressions, or one observed runtime type. They also do not invent missing docstring descriptions or sections. Resolve every reported conflict manually, especially scalar-versus-array inputs, optional values, shape conventions, and abstract interfaces. + +There is intentionally no full third-party docstring reformatter in the Makefile. `format-docstring` was evaluated and rejected: on this codebase it strips `Returns:`/`Yields:` types and rewrites `**References:**` to `**References: **`. If wrapping/whitespace normalization is ever wanted, prefer a tool that leaves section structure and type hints untouched (for example `docformatter` or `pydocstringformatter`), and still review the diff. + ## Extend the Existing Object Model New functionality should fit the existing QMCPy class hierarchy instead of introducing parallel designs without discussion. @@ -82,6 +99,10 @@ Several reviews focused on avoidable cleanup that is easy to catch before reques - Remove unused imports, trailing whitespace, and other style-only churn before requesting review. - Use explicit runtime exceptions such as `ParameterError` for invalid user inputs instead of relying on `assert` statements in production code. +For a mechanical first pass, `make check_asserts_changed` reports standalone assertions in production Python files changed relative to `ASSERT_DIFF_BASE` (default `develop`) and returns nonzero when conversions are available. `make convert_asserts_changed` uses the open-source [LibCST](https://libcst.readthedocs.io/) codemod library to convert those assertions to explicit `AssertionError` raises while preserving comments and formatting. Use `make convert_asserts ASSERT_PATH=path/to/file.py` for a specific file or directory. + +`AssertionError` is the conservative default because it preserves the original exception class and message while making validation active under `python -O`. For a reviewed set of input checks, a developer may select an exception already imported by every target file, for example `make convert_asserts ASSERT_PATH=path/to/file.py ASSERT_EXCEPTION=ParameterError`. The tool does not infer whether a condition represents invalid input, a dimension mismatch, or an internal invariant; choose `ParameterError`, `DimensionError`, `ValueError`, or another public exception only after reviewing the API contract. Assertions sharing a semicolon-delimited line with another statement, or appearing in a one-line compound suite such as `if condition: assert invariant`, are reported but skipped. Always inspect the complete diff and run the focused tests after conversion. + ## Add Demos or Blogs as Notebooks User-facing methods, new workflows, and mathematically important additions should usually come with an executable notebook. diff --git a/docs/tests.md b/docs/tests.md index a1f20762e..73cc08d92 100644 --- a/docs/tests.md +++ b/docs/tests.md @@ -25,6 +25,52 @@ This document describes the available test targets in the Makefile for QMCSoftwa | `make delcoverage` | Reset coverage tracking | Instant | Start fresh coverage analysis | +## Test File Organization + +Unit tests live flat in `test/` (no subpackage subfolders). Every file is named: + +``` +test__.py +``` + +`` is a short code for the `qmcpy` subpackage under test, or a cross-cutting bucket: + +| area | scope | +|------|-------| +| `dd` | `qmcpy/discrete_distribution` | +| `ft` | `qmcpy/fast_transform` | +| `ig` | `qmcpy/integrand` | +| `kn` | `qmcpy/kernel` | +| `sc` | `qmcpy/stopping_criterion` | +| `tm` | `qmcpy/true_measure` | +| `ut` | `qmcpy/util` | +| `ee` | end-to-end / cross-cutting pipeline (`integrate()`, worked problems such as Keister and pi) | +| `sr` | `scripts/` tooling, packaging, and docs checks | + +This keeps related tests adjacent when the directory is sorted, and lets you run one area at a time: + +```bash +python -m pytest test/ -k test_tm_ # every true_measure test +make unittests PYTEST_EXTRA_ARGS="-k test_sc_" +``` + +When a test spans two areas (say a stopping criterion exercised against a particular kernel), file it under the component actually under test and name the other in `` — e.g. `test_sc_cubbayes_kernels.py`. Reserve `ee` for cases where neither side is the clear subject. Do not invent new area codes: only the prefixes in the table are accepted, and `make check_test_style STRICT=--strict` fails on anything else. + +Notebook tests are separate: they live in `test/booktests/` as `tb_*.py` and are generated from `demos/` (see `test/booktests/README.md`). + +### Conventions checked by `make check_test_style` + +1. **Area prefix** — the filename must start with a recognized `test__` prefix from the table above. +2. **Object class** — write a test file as one or more `unittest.TestCase` subclasses rather than bare `def test_*` pytest functions. A class groups related assertions under a name (so `pytest -k TestCubMCG` selects them and a failure report names the group), shares construction through `setUp` / `setUpClass` / `self.addCleanup`, and runs identically under `pytest`, `python -m unittest`, and the coverage and booktest runners without depending on pytest fixtures. Most of the suite already follows this; a few legacy files still use bare functions and new files should not. + +`make check_test_style` lists any violation and is informational (exit 0). It also runs as part of `make format`. To make it fail instead — for a pre-commit hook or CI gate — pass `--strict`: + +```bash +STRICT=--strict make check_test_style +``` + +`STRICT=--strict make check_test_style` also runs in CI (the `alltests` workflow), so both conventions are enforced on every pull request. + ## Detailed Descriptions ## Scope @@ -153,6 +199,7 @@ Runs notebook tests with **Parsl distributed parallelization** for compute-heavy - **Dependencies**: Parsl must be installed and configured - **Use when**: Running large notebook suites with distributed compute resources + --- ### Helper / Internal Targets @@ -242,7 +289,6 @@ Displays the current coverage report (must run other targets first to accumulate Deletes `.coverage` and `coverage.json` files to reset coverage tracking. - **Use before**: Running a fresh coverage report without accumulated data - --- ## Currently Active Targets: Justification diff --git a/makefile b/makefile index 8f30cb146..812b8cda9 100644 --- a/makefile +++ b/makefile @@ -1,10 +1,13 @@ +# Prefer an active environment, then the repository's conventional qmcpy Conda +# environment, before falling back to a system interpreter. Override with +# ``make PYTHON=/path/to/python `` when needed. +PYTHON ?= $(shell command -v python 2>/dev/null || { [ -n "$$CONDA_PREFIX" ] && command -v "$$CONDA_PREFIX/bin/python" 2>/dev/null; } || { command -v conda >/dev/null 2>&1 && conda run -n qmcpy python -c 'import sys; print(sys.executable)' 2>/dev/null; } || command -v python3 2>/dev/null) # Emit pytest-xdist argument if available; can be overridden on the make command line -PYTEST_XDIST ?= $(shell python scripts/pytest_xdist.py 2>/dev/null) +PYTEST_XDIST ?= $(shell $(PYTHON) scripts/pytest_xdist.py 2>/dev/null) PYTEST ?= -PYTHON ?= python3 SMOKE_CODE_CELLS ?= 2 WITH_MPMC ?= 0 -HAS_MPMC ?= $(shell python -c "import importlib.util; mods=('torch','pyg_lib','torch_geometric'); print(int(all(importlib.util.find_spec(m) is not None for m in mods)))" 2>/dev/null || echo 0) +HAS_MPMC ?= $(shell $(PYTHON) -c "import importlib.util; mods=('torch','pyg_lib','torch_geometric'); print(int(all(importlib.util.find_spec(m) is not None for m in mods)))" 2>/dev/null || echo 0) # set environment variable for documentation export JUPYTER_PLATFORM_DIRS=1 @@ -41,20 +44,135 @@ clean_local_only_files: clean_coverage: rm -fr artifacts/coverage/ .coverage* test/booktests/.coverage* +TEST_STYLE_PATH ?= test +# Check test/test_*.py against two suite conventions: (1) written as a +# unittest.TestCase subclass ("object class"), not bare pytest functions; +# (2) named test__*.py where is the qmcpy subpackage under test +# (dd ft ig kn sc tm ut) or a cross-cutting bucket (ee sr). +# Informational by default; pass --strict to make it fail +# (e.g. STRICT=--strict make check_test_style). +check_test_style: + @$(PYTHON) scripts/check_test_style.py $(TEST_STYLE_PATH) $(STRICT) + +ASSERT_PATH ?= qmcpy +ASSERT_DIFF_BASE ?= develop +ASSERT_EXCEPTION ?= AssertionError +ASSERT_CONVERT_ARGS ?= + +check_libcst_dependency: + @$(PYTHON) -c "import libcst" 2>/dev/null || { \ + echo 'Missing LibCST. Install the test tools with: $(PYTHON) -m pip install -e ".[test]"'; \ + exit 127; \ + } + +check_assert_codemod_dependency: check_libcst_dependency + +convert_asserts: check_assert_codemod_dependency + $(PYTHON) scripts/convert_asserts.py --exception "$(ASSERT_EXCEPTION)" $(ASSERT_CONVERT_ARGS) $(ASSERT_PATH) + +convert_asserts_changed: check_assert_codemod_dependency + @$(PYTHON) scripts/convert_asserts.py --diff "$(ASSERT_DIFF_BASE)" --exception "$(ASSERT_EXCEPTION)" $(ASSERT_CONVERT_ARGS) + +check_asserts_changed: check_assert_codemod_dependency + @$(PYTHON) scripts/convert_asserts.py --diff "$(ASSERT_DIFF_BASE)" --exception "$(ASSERT_EXCEPTION)" --check $(ASSERT_CONVERT_ARGS) + +DOCSTRING_PATH ?= qmcpy +DOCSTRING_BASE ?= origin/develop +PYDOCLINT ?= pydoclint +PYDOCLINT_ARGS ?= -q +DOCSTRING_TYPE_PATH ?= qmcpy +DOCSTRING_TYPE_DIFF_BASE ?= develop +DOCSTRING_TYPE_ARGS ?= +PUBLIC_API_TYPE_PATH ?= qmcpy +PUBLIC_API_TYPE_DIFF_BASE ?= develop +PUBLIC_API_ANNOTATE_ARGS ?= +DOCSTRING_SYNC_ARGS ?= +# Two-part docstring check for public APIs under qmcpy/: +# 1. scripts/check_docstring.py -- formatting: a one-line summary before the +# first section, no NumPy-style "-----" section underlines, a blank line +# before every Args:/Returns:/... header, canonical "Name:" headers, and +# public objects with no docstring (pass --skip-missing via +# CHECK_DOCSTRING_ARGS to check style only). It also prints a second summary +# restricted to files changed relative to DOCSTRING_BASE. +# 2. pydoclint (config in pyproject.toml [tool.pydoclint]) -- content: every +# parameter and return value is documented and matches the signature, in +# Google form. +# Informational by default; pass --strict (STRICT=--strict make check_docstring) +# to make both parts fail the build. +check_docstring: + @$(PYTHON) scripts/check_docstring.py $(DOCSTRING_PATH) --diff $(DOCSTRING_BASE) $(CHECK_DOCSTRING_ARGS) $(STRICT) + @out="$$($(PYDOCLINT) $(PYDOCLINT_ARGS) $(DOCSTRING_PATH) 2>&1)"; rc=$$?; \ + [ -z "$$out" ] || printf '\n%s\n' "$$out"; \ + $(if $(STRICT),exit $$rc,true) + +# Ratchet gate: check_docstring/pydoclint/annotate_public_api_types are +# informational (existing backlog is large, see PR #613 review F9/F10), but +# this fails if a change increases any of their full-tree violation counts +# above scripts/baseline_counts.json. Run with --update after intentionally +# reducing (or, with justification, increasing) one of the counts. +check_baseline: + @$(PYTHON) scripts/check_baseline.py + +check_baseline_update: + @$(PYTHON) scripts/check_baseline.py --update + +add_docstring_arg_types: + $(PYTHON) scripts/add_docstring_arg_types.py $(DOCSTRING_TYPE_ARGS) $(DOCSTRING_TYPE_PATH) + +add_docstring_arg_types_changed: + @$(PYTHON) scripts/add_docstring_arg_types.py --diff "$(DOCSTRING_TYPE_DIFF_BASE)" $(DOCSTRING_TYPE_ARGS) + +check_docstring_arg_types_changed: + $(PYTHON) scripts/add_docstring_arg_types.py --diff "$(DOCSTRING_TYPE_DIFF_BASE)" --check $(DOCSTRING_TYPE_ARGS) + +annotate_public_api_types_changed: check_libcst_dependency + $(PYTHON) -m scripts.annotate_public_api_types --diff "$(PUBLIC_API_TYPE_DIFF_BASE)" --root "$(PUBLIC_API_TYPE_PATH)" $(PUBLIC_API_ANNOTATE_ARGS) + +sync_docstring_types_changed: + $(PYTHON) scripts/add_docstring_arg_types.py --diff "$(PUBLIC_API_TYPE_DIFF_BASE)" --root "$(PUBLIC_API_TYPE_PATH)" --include-outputs --overwrite-existing $(DOCSTRING_SYNC_ARGS) + +check_public_api_types_changed: check_libcst_dependency + @status=0; \ + $(PYTHON) -m scripts.annotate_public_api_types --diff "$(PUBLIC_API_TYPE_DIFF_BASE)" --root "$(PUBLIC_API_TYPE_PATH)" --check $(PUBLIC_API_ANNOTATE_ARGS) || status=$$?; \ + $(PYTHON) scripts/add_docstring_arg_types.py --diff "$(PUBLIC_API_TYPE_DIFF_BASE)" --root "$(PUBLIC_API_TYPE_PATH)" --include-outputs --overwrite-existing --check $(DOCSTRING_SYNC_ARGS) || { code=$$?; if [ $$code -gt $$status ]; then status=$$code; fi; }; \ + exit $$status + +# Same checks as check_docstring, but only on qmcpy/*.py files that changed +# relative to DOCSTRING_BASE (committed, staged/unstaged, and untracked). +check_docstring_changed: + @set -e; \ + changed_files="$$( \ + { \ + git diff --name-only --diff-filter=ACMR "$(DOCSTRING_BASE)...HEAD" -- 'qmcpy/*.py' 2>/dev/null || true; \ + git diff --name-only --diff-filter=ACMR HEAD -- 'qmcpy/*.py'; \ + git ls-files --others --exclude-standard -- 'qmcpy/*.py'; \ + } | sort -u \ + )"; \ + if [ -z "$$changed_files" ]; then \ + echo " - No changed qmcpy/*.py files relative to $(DOCSTRING_BASE)."; \ + else \ + file_count=$$(printf '%s\n' "$$changed_files" | wc -l | tr -d ' '); \ + echo " - Checking docstrings on $$file_count changed qmcpy file(s) relative to $(DOCSTRING_BASE)."; \ + $(PYTHON) scripts/check_docstring.py $$changed_files $(CHECK_DOCSTRING_ARGS) $(STRICT); \ + out="$$($(PYDOCLINT) $(PYDOCLINT_ARGS) $$changed_files 2>&1)"; rc=$$?; \ + [ -z "$$out" ] || printf '\n%s\n' "$$out"; \ + $(if $(STRICT),test $$rc -eq 0,true); \ + fi + ########################################################## # Doctests ########################################################## doctests_minimal: ensure_artifacts @mkdir -p $(DOCTEST_COV_DIR)/minimal COVERAGE_FILE=$(DOCTEST_COV_DIR)/minimal/.coverage \ - python -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/minimal/coverage.json --no-header --cov-append \ + $(PYTHON) -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/minimal/coverage.json --no-header --cov-append \ --doctest-modules qmcpy/ \ --ignore qmcpy/fast_transform/ft_pytorch.py \ --ignore qmcpy/stopping_criterion/pf_gp_ci.py \ --ignore qmcpy/kernel/ \ --ignore qmcpy/util/dig_shift_invar_ops.py \ --ignore qmcpy/util/shift_invar_ops.py \ - --ignore qmcpy/util/exact_gpytorch_gression_model.py \ + --ignore qmcpy/util/exact_gpytorch_regression_model.py \ --ignore qmcpy/integrand/umbridge_wrapper.py \ --ignore qmcpy/integrand/hartmann6d.py \ --ignore qmcpy/discrete_distribution/mpmc/ \ @@ -62,7 +180,7 @@ doctests_minimal: ensure_artifacts doctests_torch: ensure_artifacts @mkdir -p $(DOCTEST_COV_DIR)/torch COVERAGE_FILE=$(DOCTEST_COV_DIR)/torch/.coverage \ - python -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/torch/coverage.json --no-header --cov-append \ + $(PYTHON) -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/torch/coverage.json --no-header --cov-append \ --doctest-modules qmcpy/fast_transform/ft_pytorch.py \ --doctest-modules qmcpy/kernel/*.py \ --doctest-modules qmcpy/util/dig_shift_invar_ops.py \ @@ -71,26 +189,26 @@ doctests_torch: ensure_artifacts doctests_gpytorch: ensure_artifacts @mkdir -p $(DOCTEST_COV_DIR)/gpytorch COVERAGE_FILE=$(DOCTEST_COV_DIR)/gpytorch/.coverage \ - python -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/gpytorch/coverage.json --no-header --cov-append \ + $(PYTHON) -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/gpytorch/coverage.json --no-header --cov-append \ --doctest-modules qmcpy/stopping_criterion/pf_gp_ci.py \ doctests_botorch: ensure_artifacts @mkdir -p $(DOCTEST_COV_DIR)/botorch COVERAGE_FILE=$(DOCTEST_COV_DIR)/botorch/.coverage \ - python -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/botorch/coverage.json --no-header --cov-append \ + $(PYTHON) -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/botorch/coverage.json --no-header --cov-append \ --doctest-modules qmcpy/integrand/hartmann6d.py \ doctests_mpmc: @mkdir -p $(DOCTEST_COV_DIR)/mpmc COVERAGE_FILE=$(DOCTEST_COV_DIR)/mpmc/.coverage \ - python -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/mpmc/coverage.json --no-header --cov-append \ + $(PYTHON) -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/mpmc/coverage.json --no-header --cov-append \ --doctest-modules qmcpy/discrete_distribution/mpmc/*.py \ doctests_umbridge: ensure_artifacts # https://github.com/UM-Bridge/umbridge/issues/96 @mkdir -p $(DOCTEST_COV_DIR)/umbridge @docker --version COVERAGE_FILE=$(DOCTEST_COV_DIR)/umbridge/.coverage \ - python -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/umbridge/coverage.json --no-header --cov-append \ + $(PYTHON) -m pytest $(PYTEST_XDIST) -x --cov qmcpy/ --cov-report term --cov-report json:$(DOCTEST_COV_DIR)/umbridge/coverage.json --no-header --cov-append \ --doctest-modules qmcpy/integrand/umbridge_wrapper.py \ doctests_markdown: @@ -110,13 +228,8 @@ doctests: doctests_markdown doctests_minimal doctests_torch doctests_gpytorch do ########################################################## unittests: ensure_artifacts @mkdir -p $(UNIT_COV_DIR) - @PYTHON_BIN=$$(command -v python 2>/dev/null || { [ -n "$$CONDA_PREFIX" ] && command -v "$$CONDA_PREFIX/bin/python" 2>/dev/null; } || { command -v conda >/dev/null 2>&1 && conda run -n qmcpy python -c 'import sys; print(sys.executable)' 2>/dev/null; } || command -v python3 2>/dev/null); \ - if [ -z "$$PYTHON_BIN" ]; then \ - echo "No Python interpreter found (tried: python, $$CONDA_PREFIX/bin/python, python3)."; \ - exit 127; \ - fi; \ - COVERAGE_FILE=$(UNIT_COV_DIR)/.coverage \ - "$$PYTHON_BIN" -m pytest $(PYTEST_XDIST) -x $(PYTEST_EXTRA_ARGS) \ + @COVERAGE_FILE=$(UNIT_COV_DIR)/.coverage \ + $(PYTHON) -m pytest $(PYTEST_XDIST) -x $(PYTEST_EXTRA_ARGS) \ --cov=qmcpy \ --cov-report term \ --cov-report json:$(UNIT_COV_DIR)/coverage.json \ @@ -130,7 +243,7 @@ unittests: ensure_artifacts unittests_core: ensure_artifacts @mkdir -p $(UNIT_COV_DIR) COVERAGE_FILE=$(UNIT_COV_DIR)/.coverage \ - python -m pytest $(PYTEST_XDIST) $(PYTEST_EXTRA_ARGS) \ + $(PYTHON) -m pytest $(PYTEST_XDIST) $(PYTEST_EXTRA_ARGS) \ --cov=qmcpy \ --cov-report term \ --cov-report json:$(UNIT_COV_DIR)/coverage.json \ @@ -145,7 +258,7 @@ tests_no_docker_no_mpmc: doctests_no_docker_no_mpmc unittests coverage ########################################################## generate_booktests: @echo "\nGenerating missing booktest files..." - cd test/booktests/ && python generate_test.py --check-missing + cd test/booktests/ && $(PYTHON) generate_test.py --check-missing check_colab_notebooks: # faster $(PYTHON) -m scripts.check_colab_notebooks --strict @@ -254,11 +367,11 @@ booktests_no_docker: check_booktests generate_booktests clean_local_only_files e if [ -z "$(TESTS)" ]; then \ PYTHONWARNINGS="ignore::UserWarning,ignore::DeprecationWarning,ignore::FutureWarning,ignore::ImportWarning" \ COVERAGE_FILE=../../$(BOOKTEST_COV_DIR)/.coverage \ - python -W ignore -m coverage run --append --source=../../qmcpy/ -m unittest discover -s . -p "*.py" -v --failfast; \ + $(PYTHON) -W ignore -m coverage run --append --source=../../qmcpy/ -m unittest discover -s . -p "*.py" -v --failfast; \ else \ PYTHONWARNINGS="ignore::UserWarning,ignore::DeprecationWarning,ignore::FutureWarning,ignore::ImportWarning" \ COVERAGE_FILE=../../$(BOOKTEST_COV_DIR)/.coverage \ - python -W ignore -m coverage run --append --source=../../qmcpy/ -m unittest $(TESTS) -v --failfast; \ + $(PYTHON) -W ignore -m coverage run --append --source=../../qmcpy/ -m unittest $(TESTS) -v --failfast; \ fi && \ cd ../.. @@ -268,7 +381,7 @@ booktests_parallel_no_docker: check_booktests generate_booktests clean_local_onl cd test/booktests/ && \ rm -fr *.eps *.jpg *.pdf *.png *.part *.txt *.log && rm -fr logs && rm -fr runinfo prob_failure_gp_ci_plots && \ PYTHONWARNINGS="ignore::UserWarning,ignore::DeprecationWarning,ignore::FutureWarning,ignore::ImportWarning" \ - python parsl_test_runner.py $(TESTS) -v --failfast && \ + $(PYTHON) parsl_test_runner.py $(TESTS) -v --failfast && \ cd ../.. # Windows-compatible parallel booktests using pytest-xdist instead of Parsl @@ -277,7 +390,7 @@ booktests_parallel_pytest: check_booktests generate_booktests clean_local_only_f cd test/booktests/ && \ PYTHONWARNINGS="ignore::UserWarning,ignore::DeprecationWarning,ignore::FutureWarning,ignore::ImportWarning" \ COVERAGE_FILE=../../$(BOOKTEST_COV_DIR)/.coverage \ - python -W ignore -m pytest $(PYTEST_XDIST) $(PYTEST) -v tb_*.py \ + $(PYTHON) -W ignore -m pytest $(PYTEST_XDIST) $(PYTEST) -v tb_*.py \ --cov=qmcpy \ --cov-append \ --cov-report=term \ @@ -304,19 +417,23 @@ tests_no_docker: # Fast test target: run doctests, unittests, booktests concurrently tests_fast: @echo "Running fast tests: doctests and unittests concurrently (splitting CPU cores)." - @make clean_local_only_files clean_coverage && \ + @set -e; \ + $(MAKE) clean_local_only_files clean_coverage; \ if [ "$(WITH_MPMC)" = "1" ] || [ "$(HAS_MPMC)" = "1" ]; then \ DOCTESTS_TARGET=doctests_no_docker; \ UNITTESTS_ARGS=""; \ else \ DOCTESTS_TARGET=doctests_no_docker_no_mpmc; \ UNITTESTS_ARGS="--ignore=test/test_dd_mpmc.py"; \ - fi && \ - set -e && \ - $(MAKE) $$DOCTESTS_TARGET & \ - $(MAKE) unittests PYTEST_EXTRA_ARGS="$$UNITTESTS_ARGS" & \ - $(MAKE) booktests_parallel_no_docker & \ - wait + fi; \ + $(MAKE) $$DOCTESTS_TARGET & doctests_pid=$$!; \ + $(MAKE) unittests PYTEST_EXTRA_ARGS="$$UNITTESTS_ARGS" & unittests_pid=$$!; \ + $(MAKE) booktests_parallel_no_docker & booktests_pid=$$!; \ + status=0; \ + wait $$doctests_pid || status=$$?; \ + wait $$unittests_pid || status=$$?; \ + wait $$booktests_pid || status=$$?; \ + exit $$status $(MAKE) coverage ########################################################## @@ -331,7 +448,7 @@ coverage: ensure_artifacts # https://github.com/marketplace/actions/coverage-bad @echo "============================================================" @echo "" COVERAGE_FILE=$(UNIT_COV_DIR)/.coverage \ - python -m coverage report -m + $(PYTHON) -m coverage report -m combine-coverage-local: ensure_artifacts # Combine coverage files and build reports locally (NOT official) @echo "Combining coverage files from $(COV_DIR)/ into coverage-data/ and generating reports" @@ -350,7 +467,7 @@ combine-coverage-local: ensure_artifacts # Combine coverage files and build rep echo "No coverage data found. Run tests first (e.g., make unittests / make doctests / make booktests_*)"; \ exit 1; \ fi; \ - python scripts/combine_coverage.py --dir coverage-data --outdir coverage_html --keep + $(PYTHON) scripts/combine_coverage.py --dir coverage-data --outdir coverage_html --keep coverage_html: ensure_artifacts @mkdir -p $(UNIT_COV_DIR)/html @@ -361,7 +478,7 @@ coverage_html: ensure_artifacts @echo "============================================================" @echo "" COVERAGE_FILE=$(UNIT_COV_DIR)/.coverage \ - python -m coverage html -d $(UNIT_COV_DIR)/html + $(PYTHON) -m coverage html -d $(UNIT_COV_DIR)/html delcoverage: @rm -f .coverage coverage.json test/booktests/.coverage @@ -416,6 +533,8 @@ copydocs: # mkdocs only looks for content in the docs/ folder, so we have to co @# Rewrite repo-root-relative link for the copied MkDocs page. @perl -0pi -e 's!\(docs/good_practices\.md\)!\(good_practices.md\)!g' docs/CONTRIBUTING.md @perl -0pi -e 's!\(docs/ai-assisted-contributions\.md\)!\(ai-assisted-contributions.md\)!g' docs/CONTRIBUTING.md + @perl -0pi -e 's!\(docs/tests\.md\)!\(tests.md\)!g' docs/CONTRIBUTING.md + @perl -0pi -e 's!\(test/README\.md(#[^)]*)?\)!\(tests.md$$1\)!g' docs/CONTRIBUTING.md @cp community.md docs/community.md @cp -r demos docs @find docs/demos -mindepth 2 -name README.md -delete @@ -424,7 +543,7 @@ copydocs: # mkdocs only looks for content in the docs/ folder, so we have to co @./scripts/render_paper_for_mkdocs.sh @cp test/booktests/README.md docs/booktests.md @cp test/README.md docs/tests.md - @python scripts/make_qmc_software_page.py + @$(PYTHON) scripts/make_qmc_software_page.py @mkdir -p docs/stats @cp stats/pypi_downloads.md docs/stats/pypi_downloads.md @cp docs/assets/logos/qmcpy_logo.png docs/apple-touch-icon.png @@ -447,19 +566,19 @@ docnouml: copydocs runmkdocserve check_links: copydocs # internal links + anchors only; fast, no network, safe for CI @NO_MKDOCS_2_WARNING=1 mkdocs build -q -d site - @python scripts/check_links.py site + @$(PYTHON) scripts/check_links.py site check_links_external: copydocs # also checks http/https links; slow and network-flaky, run locally @NO_MKDOCS_2_WARNING=1 mkdocs build -q -d site - @python scripts/check_links.py site --external + @$(PYTHON) scripts/check_links.py site --external # The targets above check links inside the new site; these check the other # direction -- already-published URLs that would 404 after the next deploy. check_removed_urls: copydocs # fetches the deployed sitemap.xml; needs network - @python scripts/check_removed_urls.py + @$(PYTHON) scripts/check_removed_urls.py check_removed_urls_verify: copydocs # also HTTP-checks every redirect target - @python scripts/check_removed_urls.py --verify-redirects + @$(PYTHON) scripts/check_removed_urls.py --verify-redirects ########################################################## # PEP8 @@ -492,7 +611,7 @@ pep8: update_pep8_badge update_pep8_badge: @mkdir -p $(LOG_DIR) docs/assets @make check_pep8 > $(LOG_DIR)/pylint.out - @python3 scripts/update_pep8_badge.py $(LOG_DIR)/pylint.out docs/assets/pep8-badge.json docs/assets/pep8-badge.svg + @$(PYTHON) scripts/update_pep8_badge.py $(LOG_DIR)/pylint.out docs/assets/pep8-badge.json docs/assets/pep8-badge.svg ########################################################## @@ -502,17 +621,86 @@ update_pep8_badge: FORMAT_PATH ?= . MARKDOWN_UNWRAP_PATH ?= $(FORMAT_PATH) +RULE := ========================================================================== +RULE2 := $(subst =,-,$(RULE)) + +# `make format` rewrites files in place. Every step ends with one summary line: +# : clean (0/N files) -- nothing changed +# : 3 changed (3/N files) -- 3 files were rewritten +# Review the result with `git diff` before committing. format: - $(MAKE) flatten_qmcpy_imports - $(MAKE) markdown-unwrap MARKDOWN_UNWRAP_PATH="$(MARKDOWN_UNWRAP_PATH)" - $(MAKE) rm_trailing_whitespace FORMAT_PATH="$(FORMAT_PATH)" - $(MAKE) harden_colab_notebook + @echo "$(RULE)" + @echo "make format: rewriting files in place -- review with 'git diff' afterwards" + @echo "$(RULE)" + @echo + @echo "> flatten_qmcpy_imports" + @$(MAKE) flatten_qmcpy_imports + @echo + @echo "> markdown_unwrap" + @$(MAKE) markdown-unwrap MARKDOWN_UNWRAP_PATH="$(MARKDOWN_UNWRAP_PATH)" + @echo + @echo "> trailing_whitespace" + @$(MAKE) rm_trailing_whitespace FORMAT_PATH="$(FORMAT_PATH)" + @echo + @echo "> harden_colab_notebook" + @$(MAKE) harden_colab_notebook + @echo + @echo "> convert_asserts_changed" + @$(MAKE) convert_asserts_changed + @echo + @echo "> add_docstring_arg_types_changed" + @$(MAKE) add_docstring_arg_types_changed + @echo + @echo "$(RULE2)" + @echo "make format: done -- a 'clean' line for every step means nothing changed" + @echo "$(RULE2)" + @# No third-party docstring reformatter here on purpose: format-docstring + @# (tried on this codebase) strips Returns:/Yields: types under + @# --include-return-and-yield-types=False and rewrites `**References:**` to + @# `**References: **`. Wrapping/whitespace-only tools like docformatter are + @# safe to add later if wanted; a full reflow pass is not. + +# `make check` only reads -- it never edits the tree. Every step ends with one +# summary line: +# : clean (0/N files) -- nothing to fix +# : 2 problem(s) (2/N files) -- 2 files need attention +# Same conventions as alltests.yml's "Check test-suite conventions" step. It +# stops at the first step that fails; fix that step and rerun. +check: + @echo "$(RULE)" + @echo "make check: read-only, same rules as CI -- nothing here edits the tree" + @echo "$(RULE)" + @echo + @echo "> check_test_style" + @$(MAKE) check_test_style + @echo + @echo "> check_docstring_changed" + @$(MAKE) check_docstring_changed + @echo + @echo "> check_baseline" + @$(MAKE) check_baseline + @echo + @echo "> check_asserts_changed" + @$(MAKE) check_asserts_changed + @echo + @echo "> check_links" + @$(MAKE) check_links + @echo + @echo "$(RULE2)" + @echo "make check: every step above is clean" + @echo "$(RULE2)" + @# check_links_external deliberately NOT included: its own comment already + @# says "slow and network-flaky, run locally" -- not something `check` + @# should depend on. check_pep8_changed also deliberately excluded: 664 + @# existing violations in currently-changed files would break `check` + @# immediately (same shape as F9/F10's docstring backlog; would need the + @# check_baseline ratchet, not a hard gate, if added later). flatten_qmcpy_imports: - $(PYTHON) scripts/flatten_qmcpy_imports.py + @$(PYTHON) scripts/flatten_qmcpy_imports.py markdown-unwrap: - $(PYTHON) scripts/unwrap_markdown.py "$(MARKDOWN_UNWRAP_PATH)" + @$(PYTHON) scripts/unwrap_markdown.py "$(MARKDOWN_UNWRAP_PATH)" rm_trailing_whitespace: - $(PYTHON) scripts/remove_trailing_whitespace.py "$(FORMAT_PATH)" + @$(PYTHON) scripts/remove_trailing_whitespace.py "$(FORMAT_PATH)" diff --git a/mkdocs.yml b/mkdocs.yml index 792731c47..f448b8b29 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -89,6 +89,7 @@ nav: - QMCPy MPMC compatibility matrix: mpmc-compatibility.md - Unit tests on Jupyter Notebooks: booktests.md - Coding Agents: AGENTS.md + - Makefile Developer Tooling: demos/makefile_dev_tools.ipynb - JOSS 2026 Paper: paper/paper.md @@ -131,6 +132,7 @@ plugins: docstring_style: google docstring_options: ignore_init_summary: false + returns_named_value: false # idiomatic Google style: `Type: description`, no invented name merge_init_into_class: true - glightbox: touchNavigation: true diff --git a/pyproject.toml b/pyproject.toml index 570ba971d..a3744cebe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,7 @@ mpmc = [ test = [ "pytest >= 9.0.3", "pytest-cov >= 6.1.1", + "libcst >= 1.9.0, < 2.0", # formatting-preserving source codemods "phmutest >= 1.0.1", "pytest-accept >= 0.1.10", "testbook >= 0.4.2", @@ -109,6 +110,7 @@ test = [ test_core = [ "pytest >= 7.0", "pytest-cov >= 4.0", + "libcst >= 1.9.0, < 2.0", # required by source-codemod unit tests "pytest-xdist >= 3.0", "scikit-learn >= 1.0.0", "pandas >= 1.3.0", @@ -150,6 +152,7 @@ docs = [ # brew install weasyprint "mkdocs-print-site-plugin >= 2.7.2", "mkdocs-exclude >= 1.0.2", "pylint >= 4.0.5", + "pydoclint >= 0.5.0", ] dev = [ # brew install weasyprint "qmcpy[docs,test,torch,gpytorch,botorch,umbridge,mpmc]", @@ -184,6 +187,24 @@ class = [ "networkx >= 3.0", ] +[tool.pydoclint] +# `make check_docstring` reads this. QMCPy convention: constructor arguments are +# documented in the __init__ method's own docstring, and parameter types are +# given in both the signature and the docstring (`name (type): ...`), kept in +# sync by scripts/annotate_public_api_types.py. Optional heavy dependencies +# (torch, matplotlib, pandas, ...) that a module only imports lazily are +# referenced via `from __future__ import annotations` + `if TYPE_CHECKING:` +# so the annotation never forces an eager import. +style = "google" +allow-init-docstring = true +arg-type-hints-in-signature = true +arg-type-hints-in-docstring = true +check-return-types = false +check-yield-types = false +check-class-attributes = false +skip-checking-short-docstrings = true +skip-checking-raises = true + [tool.pylint.typecheck] # Members that live on compiled/C-extension objects (numpy, scipy, matplotlib) # which pylint's static inference cannot see, so they should not trigger diff --git a/qmcpy/accumulate_data/__init__.py b/qmcpy/accumulate_data/__init__.py deleted file mode 100644 index 6f913ef81..000000000 --- a/qmcpy/accumulate_data/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Accumulate data module.""" diff --git a/qmcpy/discrete_distribution/abstract_discrete_distribution.py b/qmcpy/discrete_distribution/abstract_discrete_distribution.py index 708141c76..1ac23439a 100644 --- a/qmcpy/discrete_distribution/abstract_discrete_distribution.py +++ b/qmcpy/discrete_distribution/abstract_discrete_distribution.py @@ -1,3 +1,4 @@ +from typing import Union from ..util import ( ParameterError, MethodImplementationError, @@ -7,8 +8,14 @@ class AbstractDiscreteDistribution(object): + """Abstract base class for QMCPy discrete distributions (samplers). - def __init__(self, dimension, replications, seed, d_limit, n_limit): + Every concrete discrete distribution (e.g. `DigitalNetB2`, `Lattice`, + `IIDStdUniform`) subclasses this and implements `_gen_samples` and + `_spawn`. + """ + + def __init__(self, dimension, replications, seed, d_limit, n_limit) -> None: self.mimics = "StdUniform" if not hasattr(self, "parameters"): self.parameters = [] @@ -26,7 +33,7 @@ def __init__(self, dimension, replications, seed, d_limit, n_limit): self.no_replications = replications is None self.replications = 1 if self.no_replications else int(replications) if self.replications < 0: - raise ParameterError("replications must be None or a postive int") + raise ParameterError("replications must be None or a positive int") if ( isinstance(dimension, list) or isinstance(dimension, tuple) @@ -52,7 +59,7 @@ def __init__(self, dimension, replications, seed, d_limit, n_limit): self.spawn_key = self._base_seed.spawn_key self.rng = np.random.Generator(np.random.SFC64(self._base_seed)) - def __call__(self, n=None, n_min=None, n_max=None, return_binary=False, warn=True): + def __call__(self, n: Union[None, int] = None, n_min: Union[None, int] = None, n_max: Union[None, int] = None, return_binary: bool = False, warn: bool = True) -> np.ndarray: r""" - If just `n` is supplied, generate samples from the sequence at indices 0,...,`n`-1. - If `n_min` and `n_max` are supplied, generate samples from the sequence at indices `n_min`,...,`n_max`-1. @@ -62,17 +69,19 @@ def __call__(self, n=None, n_min=None, n_max=None, return_binary=False, warn=Tru n (Union[None, int]): Number of points to generate. n_min (Union[None, int]): Starting index of sequence. n_max (Union[None, int]): Final index of sequence. - return_binary (bool): Only used for `DigitalNetB2`. - If `True`, *only* return the integer representation `x_integer` of base 2 digital net. + return_binary (bool): Only used for `DigitalNetB2`. If `True`, + *only* return the integer representation `x_integer` of base 2 + digital net. warn (bool): If `False`, disable warnings when generating samples. Returns: - x (np.ndarray): Samples from the sequence. + np.ndarray: Samples from the sequence. - If `replications` is `None` then this will be of size (`n_max`-`n_min`) $\times$ `dimension` - If `replications` is a positive int, then `x` will be of size `replications` $\times$ (`n_max`-`n_min`) $\times$ `dimension` - Note that if `return_binary=True` then `x` is returned where `x` are integer representations of the digital net points. + Note that if `return_binary=True` then `x` is returned where `x` + are integer representations of the digital net points. """ return self.gen_samples( n=n, n_min=n_min, n_max=n_max, return_binary=return_binary, warn=warn @@ -81,6 +90,9 @@ def __call__(self, n=None, n_min=None, n_max=None, return_binary=False, warn=Tru def gen_samples( self, n=None, n_min=None, n_max=None, return_binary=False, warn=True ): + r"""Generate samples from the sequence. Called by `__call__`; see its + docstring for the full `Args:`/`Returns:` description. + """ if n is not None and n_min is None and n_max is None: n_min = 0 n_max = int(n) @@ -122,20 +134,22 @@ def gen_samples( def _gen_samples(self, *args, **kwargs): raise MethodImplementationError(self, "_gen_samples") - def spawn(self, s=1, dimensions=None): - r""" - Spawn new instances of the current discrete distribution but with new seeds and dimensions. - Used by multi-level QMC algorithms which require different seeds and dimensions on each level. + def spawn(self, s: int = 1, dimensions: Union[None, np.ndarray] = None) -> list: + r"""Spawn new instances of the current discrete distribution but with + new seeds and dimensions. Used by multi-level QMC algorithms which + require different seeds and dimensions on each level. - Note: - Use `replications` instead of using `spawn` when possible, e.g., when spawning copies which all have the same dimension. + Notes: + Use `replications` instead of using `spawn` when possible, e.g., + when spawning copies which all have the same dimension. Args: s (int): Number of copies to spawn - dimensions (np.ndarray): Length `s` array of dimensions for each copy. Defaults to the current dimension. + dimensions (Union[None, np.ndarray]): Length `s` array of dimensions for each + copy. Defaults to the current dimension. Returns: - spawned_discrete_distribs (list): Discrete distributions with new seeds and dimensions. + list: Discrete distributions with new seeds and dimensions. """ s = int(s) if s <= 0: @@ -161,7 +175,17 @@ def spawn(self, s=1, dimensions=None): def _spawn(self, child_seed, dimension): raise MethodImplementationError(self, "_spawn") - def pdf(self, x): + def pdf(self, x: np.ndarray) -> np.ndarray: + """Probability density function of the distribution this sampler mimics. + + Args: + x (np.ndarray): Points at which to evaluate the density, shape `(*batch_shape, d)`. + + Returns: + np.ndarray: Density values with shape `batch_shape`. The base + implementation is uniform on `[0,1]^d` (density 1 everywhere); + subclasses that mimic a different distribution override this. + """ return np.ones_like(x[..., 0]) def __repr__(self, abc_class_name): @@ -171,14 +195,18 @@ def __repr__(self, abc_class_name): class AbstractLDDiscreteDistribution(AbstractDiscreteDistribution): - """Low discrepancy sequence. Alias for `AbstractDiscreteDistribution` used for compatibility checks.""" + """Low discrepancy sequence. Alias for `AbstractDiscreteDistribution` + used for compatibility checks. + """ def __repr__(self): return super().__repr__("AbstractLDDiscreteDistribution") class AbstractIIDDiscreteDistribution(AbstractDiscreteDistribution): - """IID sequence. Alias for `AbstractDiscreteDistribution` used for compatibility checks.""" + """IID sequence. Alias for `AbstractDiscreteDistribution` used for + compatibility checks. + """ def __repr__(self): return super().__repr__("AbstractIIDDiscreteDistribution") diff --git a/qmcpy/discrete_distribution/digital_net_any_bases/digital_net_any_bases.py b/qmcpy/discrete_distribution/digital_net_any_bases/digital_net_any_bases.py index 0504deb81..010b11b0d 100644 --- a/qmcpy/discrete_distribution/digital_net_any_bases/digital_net_any_bases.py +++ b/qmcpy/discrete_distribution/digital_net_any_bases/digital_net_any_bases.py @@ -1,3 +1,4 @@ +from typing import Union import warnings from ..abstract_discrete_distribution import AbstractLDDiscreteDistribution from ...util import ParameterError,ParameterWarning @@ -9,26 +10,29 @@ class DigitalNetAnyBases(AbstractLDDiscreteDistribution): - r""" - Low discrepancy digital net with arbitrary bases for each dimension. - - Note: - - Digital net samples sizes should be products of powers of bases, - i.e., a digital net with bases $(b_1,\dots,b_d)$ - will prefer sample sizes $n = b_1^{p_1} \cdots b_d^{p_d}$ for some $p_1,\dots,p_d \in \mathbb{N}_0$. - - The first point of an unrandomized digital net is the origin. - - The construction of higher order digital nets requires the same base for each dimension. - To construct higher order digital nets, either: - - - Pass in `generating_matrices` *without* interlacing and supply `alpha>1` to apply interlacing, or - - Pass in `generating_matrices` *with* interlacing and set `alpha=1` to avoid additional interlacing. - - i.e. do *not* pass in interlaced `generating_matrices` and set `alpha>1`, this will apply additional interlacing. - - A few examples below showcase how to pass in custom bases and generating matrices. Many other examples can be found in the Halton and Faure implementations - + r"""Low discrepancy digital net with arbitrary bases for each dimension. + + Notes: + - Digital net samples sizes should be products of powers of bases, + i.e., a digital net with bases $(b_1,\dots,b_d)$ will prefer sample + sizes $n = b_1^{p_1} \cdots b_d^{p_d}$ for some $p_1,\dots,p_d \in + \mathbb{N}_0$. + - The first point of an unrandomized digital net is the origin. + - The construction of higher order digital nets requires the same base for each dimension. + To construct higher order digital nets, either: + + - Pass in `generating_matrices` *without* interlacing and supply `alpha>1` to apply interlacing, or + - Pass in `generating_matrices` *with* interlacing and set `alpha=1` to avoid additional interlacing. + + i.e., do *not* pass in interlaced `generating_matrices` and set + `alpha>1`, this will apply additional interlacing. + + A few examples below showcase how to pass in custom bases and generating + matrices. Many other examples can be found in the Halton and Faure + implementations + Examples: - >>> bases = 3 + >>> bases = 3 >>> generating_matrices = np.array( ... [ ... [[1, 0, 0], @@ -59,7 +63,7 @@ class DigitalNetAnyBases(AbstractLDDiscreteDistribution): [0.30864198], [0.5308642 ], [0.75308642]]) - + >>> rng = np.random.Generator(np.random.PCG64(7)) >>> bases = np.array( ... [[2,5,7,23], @@ -107,7 +111,7 @@ class DigitalNetAnyBases(AbstractLDDiscreteDistribution): [0.16158904, 0.74257456, 0.19604142, 0.98366484], [0.33578478, 0.31746094, 0.35948446, 0.75911922], [0.64593022, 0.11007165, 0.63174328, 0.55910368]]]) - + >>> bases = 2 >>> generating_matrices = np.array([ ... [[1, 0, 0], @@ -139,41 +143,44 @@ class DigitalNetAnyBases(AbstractLDDiscreteDistribution): [0.375, 0.125]]) >>> bool((x==x_b2).all()) True - + **References:** - 1. Dick, Josef, and Friedrich Pillichshammer. - Digital nets and sequences: discrepancy theory and quasi–Monte Carlo integration. + 1. Dick, Josef, and Friedrich Pillichshammer. + Digital nets and sequences: discrepancy theory and quasi–Monte Carlo integration. Cambridge University Press, 2010. - - 2. Sorokin, Aleksei. - "QMCPy: A Python Software for Randomized Low-Discrepancy Sequences, Quasi-Monte Carlo, and Fast Kernel Methods" + + 2. Sorokin, Aleksei. + "QMCPy: A Python Software for Randomized Low-Discrepancy Sequences, Quasi-Monte Carlo, and Fast Kernel Methods" arXiv preprint arXiv:2502.14256 (2025). """ DEFAULT_GENERATING_MATRICES = None def __init__(self, - dimension = 1, - replications = None, - seed = None, - randomize = 'LMS DP', - bases_generating_matrices = None, - t = None, - alpha = 1, - n_lim = 2**32, - warn = True): - r""" + dimension: Union[int, np.ndarray] = 1, + replications: Union[None, int] = None, + seed: Union[None, int, np.random.SeedSequence] = None, + randomize: str = 'LMS DP', + bases_generating_matrices: Union[None, str, tuple] = None, + t: Union[None, int] = None, + alpha: int = 1, + n_lim: int = 2**32, + warn: bool = True) -> None: + r"""Initialize a DigitalNetAnyBases discrete distribution. + Args: - dimension (Union[int,np.ndarray]): Dimension of the generator. + dimension (Union[int, np.ndarray]): Dimension of the generator. - If an `int` is passed in, use generating vector components at indices 0,...,`dimension`-1. - If an `np.ndarray` is passed in, use generating vector components at these indices. - - replications (int): Number of independent randomizations of a pointset. - seed (Union[None,int,np.random.SeedSeq]): Seed the random number generator for reproducibility. + + replications (Union[None, int]): Number of independent randomizations of a + pointset. + seed (Union[None, int, np.random.SeedSequence]): Seed the random number + generator for reproducibility. randomize (str): Options are - + - `'LMS DP'`: Linear matrix scramble with digital permutation. - `'LMS DS'`: Linear matrix scramble with digital shift. - `'LMS'`: Linear matrix scramble only. @@ -181,24 +188,28 @@ def __init__(self, - `'DS'`: Digital shift only. - `'NUS'`: Nested uniform scrambling. - `'QRNG'`: Deterministic permutation scramble and random digital shift from QRNG [1] (with `generalize=True`). Does *not* support replications>1. - - `None`: No randomization. In this case the first point will be the origin. - - bases_generating_matrices (Union[str, tuple]): Specify the bases and the generating matrices. - + - `None`: No randomization. In this case the first point will be the origin. + + bases_generating_matrices (Union[None, str, tuple]): Specify the bases + and the generating matrices. + - `"HALTON"` will use Halton generating matrices. - `"FAURE"` will use Faure generating matrices . - - `bases,generating_matrices` requires - + - `bases,generating_matrices` requires + - `bases` is an `np.ndarray` of integers with shape $(,d)$ or $(r,d)$ where $d$ is the number of dimensions and $r$ is the number of replications. - `generating_matrices` is an `np.ndarray` of integers with shape $(d,m_\mathrm{max},t_\mathrm{max})$ or $(r,d,m_\mathrm{max},t_\mathrm{max})$ where $d$ is the number of dimensions, $r$ is the number of replications, and $2^{m_\mathrm{max}}$ is the maximum number of supported points. - - t (int): Number of digits *after* randomization. The number of digits in the generating matrices is inferred. - alpha (int): Interlacing factor for higher order nets. - When `alpha`>1, interlacing is performed regardless of the generating matrices, - i.e., for `alpha`>1 do *not* pass in generating matrices which are already interlaced. - The Note for this class contains more info. - n_lim (int): Maximum number of compatible points, determines the number of rows in the generating matrices. - warn (bool): If `False`, suppress warnings in construction + + t (Union[None, int]): Number of digits *after* randomization. The number of + digits in the generating matrices is inferred. + alpha (int): Interlacing factor for higher order nets. When + `alpha`>1, interlacing is performed regardless of the + generating matrices, i.e., for `alpha`>1 do *not* pass in + generating matrices which are already interlaced. The Note for + this class contains more info. + n_lim (int): Maximum number of compatible points, determines the + number of rows in the generating matrices. + warn (bool): If `False`, suppress warnings in construction """ self.parameters = ['randomize','t','n_limit'] self.all_primes = np.array([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311, 313, 317, 331, 337, 347, 349, 353, 359, 367, 373, 379, 383, 389, 397, 401, 409, 419, 421, 431, 433, 439, 443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503, 509, 521, 523, 541, 547, 557, 563, 569, 571, 577, 587, 593, 599, 601, 607, 613, 617, 619, 631, 641, 643, 647, 653, 659, 661, 673, 677, 683, 691, 701, 709, 719, 727, 733, 739, 743, 751, 757, 761, 769, 773, 787, 797, 809, 811, 821, 823, 827, 829, 839, 853, 857, 859, 863, 877, 881, 883, 887, 907, 911, 919, 929, 937, 941, 947, 953, 967, 971, 977, 983, 991, 997, 1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069, 1087, 1091, 1093, 1097, 1103, 1109, 1117, 1123, 1129, 1151, 1153, 1163, 1171, 1181, 1187, 1193, 1201, 1213, 1217, 1223, 1229, 1231, 1237, 1249, 1259, 1277, 1279, 1283, 1289, 1291, 1297, 1301, 1303, 1307, 1319, 1321, 1327, 1361, 1367, 1373, 1381, 1399, 1409, 1423, 1427, 1429, 1433, 1439, 1447, 1451, 1453, 1459, 1471, 1481, 1483, 1487, 1489, 1493, 1499, 1511, 1523, 1531, 1543, 1549, 1553, 1559, 1567, 1571, 1579, 1583, 1597, 1601, 1607, 1609, 1613, 1619, 1621, 1627, 1637, 1657, 1663, 1667, 1669, 1693, 1697, 1699, 1709, 1721, 1723, 1733, 1741, 1747, 1753, 1759, 1777, 1783, 1787, 1789, 1801, 1811, 1823, 1831, 1847, 1861, 1867, 1871, 1873, 1877, 1879, 1889, 1901, 1907, 1913, 1931, 1933, 1949, 1951, 1973, 1979, 1987, 1993, 1997, 1999, 2003, 2011, 2017, 2027, 2029, 2039, 2053, 2063, 2069, 2081, 2083, 2087, 2089, 2099, 2111, 2113, 2129, 2131, 2137, 2141, 2143, 2153, 2161, 2179, 2203, 2207, 2213, 2221, 2237, 2239, 2243, 2251, 2267, 2269, 2273, 2281, 2287, 2293, 2297, 2309, 2311, 2333, 2339, 2341, 2347, 2351, 2357, 2371, 2377, 2381, 2383, 2389, 2393, 2399, 2411, 2417, 2423, 2437, 2441, 2447, 2459, 2467, 2473, 2477, 2503, 2521, 2531, 2539, 2543, 2549, 2551, 2557, 2579, 2591, 2593, 2609, 2617, 2621, 2633, 2647, 2657, 2659, 2663, 2671, 2677, 2683, 2687, 2689, 2693, 2699, 2707, 2711, 2713, 2719, 2729, 2731, 2741, 2749, 2753, 2767, 2777, 2789, 2791, 2797, 2801, 2803, 2819, 2833, 2837, 2843, 2851, 2857, 2861, 2879, 2887, 2897, 2903, 2909, 2917, 2927, 2939, 2953, 2957, 2963, 2969, 2971, 2999, 3001, 3011, 3019, 3023, 3037, 3041, 3049, 3061, 3067, 3079, 3083, 3089, 3109, 3119, 3121, 3137, 3163, 3167, 3169, 3181, 3187, 3191, 3203, 3209, 3217, 3221, 3229, 3251, 3253, 3257, 3259, 3271, 3299, 3301, 3307, 3313, 3319, 3323, 3329, 3331, 3343, 3347, 3359, 3361, 3371, 3373, 3389, 3391, 3407, 3413, 3433, 3449, 3457, 3461, 3463, 3467, 3469, 3491, 3499, 3511, 3517, 3527, 3529, 3533, 3539, 3541, 3547, 3557, 3559, 3571, 3581, 3583, 3593, 3607, 3613, 3617, 3623, 3631, 3637, 3643, 3659, 3671, 3673, 3677, 3691, 3697, 3701, 3709, 3719, 3727, 3733, 3739, 3761, 3767, 3769, 3779, 3793, 3797, 3803, 3821, 3823, 3833, 3847, 3851, 3853, 3863, 3877, 3881, 3889, 3907, 3911, 3917, 3919, 3923, 3929, 3931, 3943, 3947, 3967, 3989, 4001, 4003, 4007, 4013, 4019, 4021, 4027, 4049, 4051, 4057, 4073, 4079, 4091, 4093, 4099, 4111, 4127, 4129, 4133, 4139, 4153, 4157, 4159, 4177, 4201, 4211, 4217, 4219, 4229, 4231, 4241, 4243, 4253, 4259, 4261, 4271, 4273, 4283, 4289, 4297, 4327, 4337, 4339, 4349, 4357, 4363, 4373, 4391, 4397, 4409, 4421, 4423, 4441, 4447, 4451, 4457, 4463, 4481, 4483, 4493, 4507, 4513, 4517, 4519, 4523, 4547, 4549, 4561, 4567, 4583, 4591, 4597, 4603, 4621, 4637, 4639, 4643, 4649, 4651, 4657, 4663, 4673, 4679, 4691, 4703, 4721, 4723, 4729, 4733, 4751, 4759, 4783, 4787, 4789, 4793, 4799, 4801, 4813, 4817, 4831, 4861, 4871, 4877, 4889, 4903, 4909, 4919, 4931, 4933, 4937, 4943, 4951, 4957, 4967, 4969, 4973, 4987, 4993, 4999, 5003, 5009, 5011, 5021, 5023, 5039, 5051, 5059, 5077, 5081, 5087, 5099, 5101, 5107, 5113, 5119, 5147, 5153, 5167, 5171, 5179, 5189, 5197, 5209, 5227, 5231, 5233, 5237, 5261, 5273, 5279, 5281, 5297, 5303, 5309, 5323, 5333, 5347, 5351, 5381, 5387, 5393, 5399, 5407, 5413, 5417, 5419, 5431, 5437, 5441, 5443, 5449, 5471, 5477, 5479, 5483, 5501, 5503, 5507, 5519, 5521, 5527, 5531, 5557, 5563, 5569, 5573, 5581, 5591, 5623, 5639, 5641, 5647, 5651, 5653, 5657, 5659, 5669, 5683, 5689, 5693, 5701, 5711, 5717, 5737, 5741, 5743, 5749, 5779, 5783, 5791, 5801, 5807, 5813, 5821, 5827, 5839, 5843, 5849, 5851, 5857, 5861, 5867, 5869, 5879, 5881, 5897, 5903, 5923, 5927, 5939, 5953, 5981, 5987, 6007, 6011, 6029, 6037, 6043, 6047, 6053, 6067, 6073, 6079, 6089, 6091, 6101, 6113, 6121, 6131, 6133, 6143, 6151, 6163, 6173, 6197, 6199, 6203, 6211, 6217, 6221, 6229, 6247, 6257, 6263, 6269, 6271, 6277, 6287, 6299, 6301, 6311, 6317, 6323, 6329, 6337, 6343, 6353, 6359, 6361, 6367, 6373, 6379, 6389, 6397, 6421, 6427, 6449, 6451, 6469, 6473, 6481, 6491, 6521, 6529, 6547, 6551, 6553, 6563, 6569, 6571, 6577, 6581, 6599, 6607, 6619, 6637, 6653, 6659, 6661, 6673, 6679, 6689, 6691, 6701, 6703, 6709, 6719, 6733, 6737, 6761, 6763, 6779, 6781, 6791, 6793, 6803, 6823, 6827, 6829, 6833, 6841, 6857, 6863, 6869, 6871, 6883, 6899, 6907, 6911, 6917, 6947, 6949, 6959, 6961, 6967, 6971, 6977, 6983, 6991, 6997, 7001, 7013, 7019, 7027, 7039, 7043, 7057, 7069, 7079, 7103, 7109, 7121, 7127, 7129, 7151, 7159, 7177, 7187, 7193, 7207, 7211, 7213, 7219, 7229, 7237, 7243, 7247, 7253, 7283, 7297, 7307, 7309, 7321, 7331, 7333, 7349, 7351, 7369, 7393, 7411, 7417, 7433, 7451, 7457, 7459, 7477, 7481, 7487, 7489, 7499, 7507, 7517, 7523, 7529, 7537, 7541, 7547, 7549, 7559, 7561, 7573, 7577, 7583, 7589, 7591, 7603, 7607, 7621, 7639, 7643, 7649, 7669, 7673, 7681, 7687, 7691, 7699, 7703, 7717, 7723, 7727, 7741, 7753, 7757, 7759, 7789, 7793, 7817, 7823, 7829, 7841, 7853, 7867, 7873, 7877, 7879, 7883, 7901, 7907, 7919],dtype=np.uint64) @@ -213,16 +224,22 @@ def __init__(self, raise ParameterError("must supply bases_generating_matrices") else: self.type_bases_generating_matrices = "CUSTOM" - assert len(bases_generating_matrices)==2 + if not (len(bases_generating_matrices)==2): + raise AssertionError bases,generating_matrices = bases_generating_matrices - assert isinstance(generating_matrices,np.ndarray) - assert generating_matrices.ndim==3 or generating_matrices.ndim==4 + if not (isinstance(generating_matrices,np.ndarray)): + raise AssertionError + if not (generating_matrices.ndim==3 or generating_matrices.ndim==4): + raise AssertionError d_limit = generating_matrices.shape[1] if np.isscalar(bases): - assert bases>0 - assert bases%1==0 + if not (bases>0): + raise AssertionError + if not (bases%1==0): + raise AssertionError bases = int(bases)*np.ones(d_limit,dtype=int) - assert bases.ndim==1 or bases.ndim==2 + if not (bases.ndim==1 or bases.ndim==2): + raise AssertionError self.input_t = deepcopy(t) self.input_bases_generating_matrices = deepcopy(bases_generating_matrices) super(DigitalNetAnyBases,self).__init__(dimension,replications,seed,d_limit,n_lim) @@ -233,16 +250,22 @@ def __init__(self, if self.randomize=="OWEN": self.randomize = "NUS" if self.randomize=="NONE": self.randomize = "FALSE" if self.randomize=="NO": self.randomize = "FALSE" - assert self.randomize in ["LMS DP","LMS DS","LMS","DP","DS","NUS","QRNG","FALSE"] + if not (self.randomize in ["LMS DP","LMS DS","LMS","DP","DS","NUS","QRNG","FALSE"]): + raise AssertionError if self.randomize=="QRNG": - assert self.type_bases_generating_matrices=="HALTON", "QRNG randomization is only applicable for the Halton generator." - assert self.replications==1, "QRNG requires replications=1" + if not (self.type_bases_generating_matrices=="HALTON"): + raise AssertionError("QRNG randomization is only applicable for the Halton generator.") + if not (self.replications==1): + raise AssertionError("QRNG requires replications=1") self.randu_d_32 = self.rng.uniform(size=(self.d,32)) self.alpha = alpha - assert self.alpha>=1 - assert self.alpha%1==0 + if not (self.alpha>=1): + raise AssertionError + if not (self.alpha%1==0): + raise AssertionError if self.alpha>1: - assert (self.dvec==np.arange(self.d)).all(), "digital interlacing requires dimension is an int" + if not ((self.dvec==np.arange(self.d)).all()): + raise AssertionError("digital interlacing requires dimension is an int") self.dtalpha = self.alpha*self.d if self.type_bases_generating_matrices=="HALTON": self.bases = self.all_primes[self.dvec][None,:] @@ -252,7 +275,8 @@ def __init__(self, self.t = self.m_max if self.m_max>t else t self.C = qmctoolscl.gdn_get_halton_generating_matrix(np.uint64(1),np.uint64(self.d),np.uint64(self._t_curr)) elif self.type_bases_generating_matrices=="FAURE": - assert (self.dvec==np.arange(self.d)).all(), "Faure requires dimension is an int" + if not ((self.dvec==np.arange(self.d)).all()): + raise AssertionError("Faure requires dimension is an int") p = self.all_primes[np.argmax(self.all_primes>=self.d)] self.bases = p*np.ones((1,self.dtalpha),dtype=np.uint64) self.m_max = int(np.ceil(np.log(self.n_limit)/np.log(p))) @@ -274,14 +298,16 @@ def __init__(self, else: self.bases = bases.astype(np.uint64) if self.bases.ndim==1: self.bases = self.bases[None,:] - assert self.bases.shape[1]>=self.dtalpha + if not (self.bases.shape[1]>=self.dtalpha): + raise AssertionError if self.alpha==1: self.bases = self.bases[:,self.dvec] else: self.bases = self.bases[:,:self.dtalpha] self.C = generating_matrices.astype(np.uint64) if self.C.ndim==3: self.C = self.C[None,:,:,:] - assert self.C.shape[1]>=self.dtalpha + if not (self.C.shape[1]>=self.dtalpha): + raise AssertionError if self.alpha==1: self.C = self.C[:,self.dvec,:,:] else: @@ -289,20 +315,29 @@ def __init__(self, self.m_max,self._t_curr = self.C.shape[-2:] if t is None: t = int(np.ceil(-np.log(2**(-63))/np.log(self.bases.min()))) self.t = self.m_max if self.m_max>t else t - assert (0<=self.C).all() - assert (self.C1: - assert (self.bases==self.bases[0,0]).all(), "alpha>1 performs digital interlacing which requires the same base across dimensions and replications." + if not ((self.bases==self.bases[0,0]).all()): + raise AssertionError("alpha>1 performs digital interlacing which requires the same base across dimensions and replications.") if warn and self.m_max!=self._t_curr: warnings.warn("Digital interlacing is often performed on generating matrices with the number of columns (m_max = %d) equal to the number of rows (_t_curr = %d), but this is not the case. Ensure you are NOT setting alpha>1 when generating matrices are already interlaced."%(self.m_max,self._t_curr),ParameterWarning) - assert self.bases.ndim==2 - assert self.bases.shape[-1]==self.dtalpha - assert self.bases.shape[0]==1 or self.bases.shape[0]==self.replications - assert self.C.ndim==4 - assert self.C.shape[-3:]==(self.dtalpha,self.m_max,self._t_curr) - assert self.C.shape[0]==1 or self.C.shape[0]==self.replications + if not (self.bases.ndim==2): + raise AssertionError + if not (self.bases.shape[-1]==self.dtalpha): + raise AssertionError + if not (self.bases.shape[0]==1 or self.bases.shape[0]==self.replications): + raise AssertionError + if not (self.C.ndim==4): + raise AssertionError + if not (self.C.shape[-3:]==(self.dtalpha,self.m_max,self._t_curr)): + raise AssertionError + if not (self.C.shape[0]==1 or self.C.shape[0]==self.replications): + raise AssertionError r_b = self.bases.shape[0] r_C = self.C.shape[0] if self.randomize=="FALSE": @@ -349,10 +384,15 @@ def __init__(self, new_seeds = self._base_seed.spawn(self.replications*self.dtalpha) self.rngs = np.array([np.random.Generator(np.random.SFC64(new_seeds[j])) for j in range(self.replications*self.dtalpha)]).reshape(self.replications,self.dtalpha) self.root_nodes = np.array([qmctoolscl.NUSNode_gdn() for i in range(self.replications*self.dtalpha)]).reshape(self.replications,self.dtalpha) - assert self.C.ndim==4 and (self.C.shape[0]==1 or self.C.shape[0]==self.replications) and self.C.shape[1]==(self.dtalpha if self.randomize=="NUS" else self.d) and self.C.shape[2]==self.m_max and self.C.shape[3]==self._t_curr - assert self.bases.ndim==2 and (self.bases.shape[0]==1 or self.bases.shape[0]==self.replications) and self.bases.shape[1]==(self.dtalpha if self.randomize=="NUS" else self.d) - assert 0>> discrete_distrib = Faure(4,seed=7) >>> discrete_distrib(25) @@ -44,8 +43,8 @@ class Faure(DigitalNetAnyBases): t 28 n_limit 2^(32) entropy 7 - - Replications of independent randomizations + + Replications of independent randomizations >>> x = Faure(3,seed=7,replications=2)(9) >>> x.shape @@ -71,7 +70,7 @@ class Faure(DigitalNetAnyBases): [0.30097968, 0.36957094, 0.23358374], [0.99369356, 0.78380717, 0.74090153]]]) - Unrandomized Faure + Unrandomized Faure >>> Faure(4,randomize="FALSE",seed=7)(25,warn=False) array([[0. , 0. , 0. , 0. ], @@ -99,8 +98,8 @@ class Faure(DigitalNetAnyBases): [0.56, 0.36, 0.16, 0.96], [0.76, 0.56, 0.36, 0.16], [0.96, 0.76, 0.56, 0.36]]) - - All randomizations + + All randomizations >>> Faure(3,randomize="LMS DP",seed=7)(9) array([[0.60869072, 0.76096155, 0.79807281], @@ -162,8 +161,8 @@ class Faure(DigitalNetAnyBases): [0.25089638, 0.17805972, 0.95988146], [0.68344029, 0.77065782, 0.26676153], [0.4322891 , 0.40799837, 0.34911626]]) - - Replications of randomizations + + Replications of randomizations >>> Faure(3,randomize="LMS DP",seed=7,replications=2)(9) array([[[0.46995809, 0.81347921, 0.84921511], @@ -287,7 +286,7 @@ class Faure(DigitalNetAnyBases): [0.59326363, 0.50120469, 0.9906825 ]]]) Higher order Faure - + >>> Faure(3,randomize="LMS DP",seed=7,alpha=2)(9) array([[0.07060326, 0.24965078, 0.49971375], [0.9104272 , 0.77359118, 0.02813304], @@ -338,9 +337,9 @@ class Faure(DigitalNetAnyBases): [0.32098765, 0.43209877, 0.87654321], [0.43209877, 0.87654321, 0.32098765], [0.87654321, 0.32098765, 0.43209877]]) - + Replications of higher order Faure - + >>> Faure(3,randomize="LMS DP",seed=7,alpha=2,replications=2)(9) array([[[0.65006542, 0.84004771, 0.39377772], [0.73541117, 0.25289783, 0.11639162], diff --git a/qmcpy/discrete_distribution/digital_net_any_bases/halton.py b/qmcpy/discrete_distribution/digital_net_any_bases/halton.py index 0fd3eb1e3..b4630bb4a 100644 --- a/qmcpy/discrete_distribution/digital_net_any_bases/halton.py +++ b/qmcpy/discrete_distribution/digital_net_any_bases/halton.py @@ -2,13 +2,12 @@ class Halton(DigitalNetAnyBases): - r""" - Low discrepancy Halton points. + r"""Low discrepancy Halton points. - Note: + Notes: - The first point of an unrandomized Halton sequence is the origin. - QRNG does *not* support multiple replications (independent randomizations). - + Examples: >>> discrete_distrib = Halton(2,seed=7) >>> discrete_distrib(4) @@ -24,8 +23,8 @@ class Halton(DigitalNetAnyBases): t 63 n_limit 2^(32) entropy 7 - - Replications of independent randomizations + + Replications of independent randomizations >>> x = Halton(3,seed=7,replications=2)(4) >>> x.shape @@ -41,15 +40,15 @@ class Halton(DigitalNetAnyBases): [0.89132308, 0.12030255, 0.35715804], [0.04025218, 0.44304244, 0.10724799]]]) - Unrandomized Halton + Unrandomized Halton >>> Halton(2,randomize="FALSE",seed=7)(4,warn=False) array([[0. , 0. ], [0.5 , 0.33333333], [0.25 , 0.66666667], [0.75 , 0.11111111]]) - - All randomizations + + All randomizations >>> Halton(2,randomize="LMS DP",seed=7)(4) array([[0.83790457, 0.89981478], @@ -86,8 +85,8 @@ class Halton(DigitalNetAnyBases): [0.85362988, 0.72066823], [0.10362988, 0.05400156], [0.60362988, 0.498446 ]]) - - Replications of randomizations + + Replications of randomizations >>> Halton(3,randomize="LMS DP",seed=7,replications=2)(4) array([[[0.70988236, 0.18180876, 0.54073621], @@ -151,19 +150,19 @@ class Halton(DigitalNetAnyBases): [0.71866903, 0.23852281, 0.80431142]]]) **References:** - - 1. Marius Hofert and Christiane Lemieux. - qrng: (Randomized) Quasi-Random Number Generators. - R package version 0.0-7. (2019). + + 1. Marius Hofert and Christiane Lemieux. + qrng: (Randomized) Quasi-Random Number Generators. + R package version 0.0-7. (2019). [https://CRAN.R-project.org/package=qrng](https://CRAN.R-project.org/package=qrng). - - 2. A. B. Owen. - A randomized Halton algorithm in R. - [arXiv:1706.02808](https://arxiv.org/abs/1706.02808) [stat.CO]. 2017. - - 3. A. B. Owen and Z. Pan. - Gain coefficients for scrambled Halton points. - [arXiv:2308.08035](https://arxiv.org/abs/2308.08035) [stat.CO]. 2023. + + 2. A. B. Owen. + A randomized Halton algorithm in R. + [arXiv:1706.02808](https://arxiv.org/abs/1706.02808) [stat.CO]. 2017. + + 3. A. B. Owen and Z. Pan. + Gain coefficients for scrambled Halton points. + [arXiv:2308.08035](https://arxiv.org/abs/2308.08035) [stat.CO]. 2023. """ DEFAULT_GENERATING_MATRICES = "HALTON" diff --git a/qmcpy/discrete_distribution/digital_net_any_bases/hammersley.py b/qmcpy/discrete_distribution/digital_net_any_bases/hammersley.py index 2e765b549..3fdfaf4e0 100644 --- a/qmcpy/discrete_distribution/digital_net_any_bases/hammersley.py +++ b/qmcpy/discrete_distribution/digital_net_any_bases/hammersley.py @@ -1,3 +1,4 @@ +from typing import Union from qmcpy.util import ParameterError,ParameterWarning import numpy as np from .halton import Halton @@ -8,21 +9,21 @@ class Hammersley(DigitalNetAnyBases): - r""" - Hammersley point set: a deterministic, 'closed' low discrepancy point set. + r"""Hammersley point set: a deterministic, 'closed' low discrepancy point + set. With $p_1,\dots,p_{d-1}$ the first $d-1$ prime numbers, the point set - $\{t_0,\dots,t_{n-1}\}$ with $n$ points in $d$ dimensions is given by - $t_i = (i/n,\ \varphi_{p_1}(i),\ \dots,\ \varphi_{p_{d-1}}(i))$ - for $i=0,\dots,n-1$, where $\varphi_p$ denotes the radical inverse - function in base $p$. + $\{t_0,\dots,t_{n-1}\}$ with $n$ points in $d$ dimensions is given by $t_i + = (i/n,\ \varphi_{p_1}(i),\ \dots,\ \varphi_{p_{d-1}}(i))$ for + $i=0,\dots,n-1$, where $\varphi_p$ denotes the radical inverse function in + base $p$. Being a 'closed' point set (n must be fixed in advance, unlike an - extensible sequence such as Halton), the QMC error bound gains one - fewer power of $\log n$ than the corresponding Halton bound: - $|I_d(f)-Q_{n,d}(f)| \le C_d\, (\log n)^{d-1}/n\, V(f)$. + extensible sequence such as Halton), the QMC error bound gains one fewer + power of $\log n$ than the corresponding Halton bound: $|I_d(f)-Q_{n,d}(f)| + \le C_d\, (\log n)^{d-1}/n\, V(f)$. - Note: + Notes: - This class is fully deterministic: no randomization is supported, and the `seed` argument has no effect on the generated points. - The first point is always the origin. @@ -66,30 +67,31 @@ class Hammersley(DigitalNetAnyBases): """ def __init__(self, - dimension=1, - seed=None, - t=None, - n_lim=2**32, - warn = True - ): - r""" + dimension: int = 1, + seed: Union[None, int, np.random.SeedSequence] = None, + t: Union[None, int] = None, + n_lim: int = 2**32, + warn: bool = True + ) -> None: + r"""Initialize a Hammersley discrete distribution. + Args: - dimension (int): Dimension of the samples. Must be a scalar - `int` (unlike `Halton`, an array of indices is not - supported -- see class Notes). + dimension (int): Dimension of the samples. Must be a scalar `int` + (unlike `Halton`, an array of indices is not supported -- see + class Notes). - seed (Union[None, int, np.random.SeedSequence]): Unused; kept - for API consistency with the other discrete distributions. - This point set is fully deterministic, so `seed` has no - effect on the generated points. + seed (Union[None, int, np.random.SeedSequence]): Unused; kept for + API consistency with the other discrete distributions. This + point set is fully deterministic, so `seed` has no effect on + the generated points. t (Union[None, int]): Passed through to the internal `Halton` - generator used for dimensions 2,...,`dimension` (ignored - when `dimension` is 1). See `Halton`'s docstring for - details. + generator used for dimensions 2,...,`dimension` (ignored when + `dimension` is 1). See `Halton`'s docstring for details. - n_lim (int): Maximum number of points `n` this distribution - can be asked to generate. + n_lim (int): Maximum number of points `n` this distribution can be + asked to generate. + warn (bool): If `False`, disable warnings when generating samples. """ if not np.isscalar(dimension): diff --git a/qmcpy/discrete_distribution/digital_net_b2/digital_net_b2.py b/qmcpy/discrete_distribution/digital_net_b2/digital_net_b2.py index 89887d578..240223708 100644 --- a/qmcpy/discrete_distribution/digital_net_b2/digital_net_b2.py +++ b/qmcpy/discrete_distribution/digital_net_b2/digital_net_b2.py @@ -1,3 +1,4 @@ +from typing import Union from ..abstract_discrete_distribution import AbstractLDDiscreteDistribution from ...util import ParameterError, ParameterWarning import qmctoolscl @@ -9,10 +10,9 @@ import platform class DigitalNetB2(AbstractLDDiscreteDistribution): - r""" - Low discrepancy digital net in base 2. + r"""Low discrepancy digital net in base 2. - Note: + Notes: - Digital net sample sizes should be powers of $2$ e.g. $1$, $2$, $4$, $8$, $16$, $\dots$. - The first point of an unrandomized digital nets is the origin. - `Sobol` is an alias for `DigitalNetB2`. @@ -21,7 +21,8 @@ class DigitalNetB2(AbstractLDDiscreteDistribution): - Pass in `generating_matrices` *without* interlacing and supply `alpha`>1 to apply interlacing, or - Pass in `generating_matrices` *with* interlacing and set `alpha=1` to avoid additional interlacing - i.e. do *not* pass in interlaced `generating_matrices` and set `alpha>1`, this will apply additional interlacing. + i.e., do *not* pass in interlaced `generating_matrices` and set + `alpha>1`, this will apply additional interlacing. Examples: >>> discrete_distrib = DigitalNetB2(2,seed=7) @@ -69,7 +70,8 @@ class DigitalNetB2(AbstractLDDiscreteDistribution): array([[0.25, 0.75], [0.75, 0.25]]) - Generating matrices from [https://github.com/QMCSoftware/LDData/tree/main/dnet](https://github.com/QMCSoftware/LDData/tree/main/dnet) + Generating matrices from + [https://github.com/QMCSoftware/LDData/tree/main/dnet](https://github.com/QMCSoftware/LDData/tree/main/dnet) >>> DigitalNetB2(dimension=3,randomize=False,generating_matrices="mps.nx_s5_alpha2_m32.txt")(8,warn=False) array([[0. , 0. , 0. ], @@ -214,30 +216,33 @@ class DigitalNetB2(AbstractLDDiscreteDistribution): def __init__( self, - dimension=1, - replications=None, - seed=None, - randomize="LMS DS", - generating_matrices="joe_kuo.6.21201.txt", - order="RADICAL INVERSE", - t=63, - alpha=1, - msb=None, - _verbose=False, + dimension: Union[int, np.ndarray] = 1, + replications: Union[None, int] = None, + seed: Union[None, int, np.random.SeedSequence] = None, + randomize: str = "LMS DS", + generating_matrices: Union[str, np.ndarray, int] = "joe_kuo.6.21201.txt", + order: str = "RADICAL INVERSE", + t: int = 63, + alpha: int = 1, + msb: Union[None, bool] = None, + _verbose: bool = False, # deprecated - graycode=None, - t_max=None, - t_lms=None, - ): - r""" + graycode: Union[None, bool] = None, + t_max: Union[None, int] = None, + t_lms: Union[None, int] = None, + ) -> None: + r"""Initialize a DigitalNetB2 discrete distribution. + Args: dimension (Union[int, np.ndarray]): Dimension of the generator. - If an `int` is passed in, use generating vector components at indices 0,...,`dimension`-1. - If an `np.ndarray` is passed in, use generating vector components at these indices. - replications (int): Number of independent randomizations of a pointset. - seed (Union[None, int, np.random.SeedSeq): Seed the random number generator for reproducibility. + replications (Union[None, int]): Number of independent randomizations of a + pointset. + seed (Union[None, int, np.random.SeedSequence]): Seed the random + number generator for reproducibility. randomize (str): Options are - `'LMS DS'`: Linear matrix scramble with digital shift. @@ -246,19 +251,33 @@ def __init__( - `'NUS'`: Nested uniform scrambling. Also known as Owen scrambling. - `'FALSE'`: No randomization. In this case the first point will be the origin. - generating_matrices (Union[str, np.ndarray, int]): Specify the generating matrices. + generating_matrices (Union[str, np.ndarray, int]): Specify the + generating matrices. - A `str` should be the name (or path) of a file from the LDData repo at [https://github.com/QMCSoftware/LDData/tree/main/dnet](https://github.com/QMCSoftware/LDData/tree/main/dnet). - An `np.ndarray` of integers with shape $(d,m_\mathrm{max})$ or $(r,d,m_\mathrm{max})$ where $d$ is the number of dimensions, $r$ is the number of replications, and $2^{m_\mathrm{max}}$ is the maximum number of supported points. Setting `msb=False` will flip the bits of ints in the generating matrices. - order (str): `'RADICAL INVERSE'`, or `'GRAY'` ordering. See the doctest example above. - t (int): Number of bits in integer represetation of points *after* randomization. The number of bits in the generating matrices is inferred based on the largest value. - alpha (int): Interlacing factor for higher order nets. - When `alpha`>1, interlacing is performed regardless of the generating matrices, - i.e., for `alpha`>1 do *not* pass in generating matrices which are already interlaced. - The Note for this class contains more info. - msb (bool): Flag for Most Significant Bit (MSB) vs Least Significant Bit (LSB) integer representations in generating matrices. If `msb=False` (LSB order), then integers in generating matrices will be bit-reversed. - _verbose (bool): If `True`, print linear matrix scrambling matrices. + order (str): `'RADICAL INVERSE'`, or `'GRAY'` ordering. See the + doctest example above. + t (int): Number of bits in integer representation of points *after* + randomization. The number of bits in the generating matrices is + inferred based on the largest value. + alpha (int): Interlacing factor for higher order nets. When + `alpha`>1, interlacing is performed regardless of the + generating matrices, i.e., for `alpha`>1 do *not* pass in + generating matrices which are already interlaced. The Note for + this class contains more info. + msb (Union[None, bool]): Flag for Most Significant Bit (MSB) vs Least + Significant Bit (LSB) integer representations in generating + matrices. If `msb=False` (LSB order), then integers in + generating matrices will be bit-reversed. + _verbose (bool): If `True`, print linear matrix scrambling + matrices. + graycode (Union[None, bool]): Deprecated; set `order='GRAY'` or + `order='RADICAL INVERSE'` instead. + t_max (Union[None, int]): Deprecated; has no effect, as it can be inferred + from the generating matrices. + t_lms (Union[None, int]): Deprecated; set `t` instead. """ if graycode is not None: order = "GRAY" if graycode else "RADICAL INVERSE" @@ -318,7 +337,8 @@ def __init__( gen_mats = gen_mats >> compat_shift elif isinstance(generating_matrices, str): self.gen_mats_source = generating_matrices - assert generating_matrices[-4:] == ".txt" + if not (generating_matrices[-4:] == ".txt"): + raise AssertionError local_root = dirname(abspath(__file__)) + "/generating_matrices/" repos = DataSource() if repos.exists(local_root + generating_matrices): @@ -362,7 +382,8 @@ def __init__( contents = [line.split("#", 1)[0] for line in contents if line[0] != "#"] datafile.close() msb = True - assert int(contents[0]) == 2, "DigitalNetB2 requires base=2 " # base 2 + if not (int(contents[0]) == 2): # base 2 + raise AssertionError("DigitalNetB2 requires base=2 ") d_limit = int(contents[1]) n_limit = int(contents[2]) self._t_curr = int(contents[3]) @@ -381,17 +402,20 @@ def __init__( )[None, :] elif isinstance(generating_matrices, np.ndarray): self.gen_mats_source = "custom" - assert generating_matrices.ndim == 2 or generating_matrices.ndim == 3 + if not (generating_matrices.ndim == 2 or generating_matrices.ndim == 3): + raise AssertionError gen_mats = ( generating_matrices[None, :, :] if generating_matrices.ndim == 2 else generating_matrices ) - assert isinstance( + if not (isinstance( msb, bool - ), "when generating_matrices is a np.ndarray you must set either msb=True (for most significant bit ordering) or msb=False (for least significant bit ordering which will require a bit reversal)" + )): + raise AssertionError("when generating_matrices is a np.ndarray you must set either msb=True (for most significant bit ordering) or msb=False (for least significant bit ordering which will require a bit reversal)") gen_mat_max = gen_mats.max() - assert gen_mat_max > 0, "generating matrix must have positive ints" + if not (gen_mat_max > 0): + raise AssertionError("generating matrix must have positive ints") self._t_curr = int(np.ceil(np.log2(gen_mat_max + 1))) d_limit = gen_mats.shape[1] n_limit = int(2 ** (gen_mats.shape[2])) @@ -402,12 +426,13 @@ def __init__( super(DigitalNetB2, self).__init__( dimension, replications, seed, d_limit, n_limit ) - assert ( + if not ( gen_mats.ndim == 3 and gen_mats.shape[1] >= self.d and (gen_mats.shape[0] == 1 or gen_mats.shape[0] == self.replications) and gen_mats.shape[2] > 0 - ), "invalid gen_mats.shape = %s" % str(gen_mats.shape) + ): + raise AssertionError("invalid gen_mats.shape = %s" % str(gen_mats.shape)) self.m_max = int(gen_mats.shape[-1]) if isinstance(generating_matrices, np.ndarray) and (not msb): qmctoolscl.dnb2_gmat_lsb_to_msb( @@ -424,18 +449,23 @@ def __init__( self.order = "GRAY" if self.order == "NATURAL": self.order = "RADICAL INVERSE" - assert self.order in ["RADICAL INVERSE", "GRAY"] - assert isinstance(t, int) and t > 0 - assert self._t_curr <= t <= 64, ( - "t must no more than 64 and no less than %d (the number of bits used to represent the generating matrices)" - % (self._t_curr) - ) - assert isinstance(alpha, int) and alpha > 0 + if not (self.order in ["RADICAL INVERSE", "GRAY"]): + raise AssertionError + if not (isinstance(t, int) and t > 0): + raise AssertionError + if not (self._t_curr <= t <= 64): + raise AssertionError( + "t must no more than 64 and no less than %d (the number of bits used to represent the generating matrices)" + % (self._t_curr) + ) + if not (isinstance(alpha, int) and alpha > 0): + raise AssertionError self.alpha = alpha if self.alpha > 1: - assert ( + if not (( self.dvec == np.arange(self.d) - ).all(), "digital interlacing requires dimension is an int" + ).all()): + raise AssertionError("digital interlacing requires dimension is an int") if self.m_max != self._t_curr: warnings.warn( "Digital interlacing is often performed on matrices with the number of columns (m_max = %d) equal to the number of bits in each int (%d), but this is not the case. Ensure you are NOT setting alpha>1 when generating matrices are already interlaced." @@ -452,7 +482,8 @@ def __init__( self.randomize = "FALSE" if self.randomize == "NO": self.randomize = "FALSE" - assert self.randomize in ["LMS DS", "LMS", "DS", "NUS", "FALSE"] + if not (self.randomize in ["LMS DS", "LMS", "DS", "NUS", "FALSE"]): + raise AssertionError self.dtalpha = self.alpha * self.d if self.randomize == "FALSE": if self.alpha == 1: @@ -615,19 +646,23 @@ def __init__( raise ParameterError("self.randomize parsing error") self.gen_mats = np.ascontiguousarray(self.gen_mats) gen_mat_max = self.gen_mats.max() - assert gen_mat_max > 0, "generating matrix must have positive ints" - assert self._t_curr == int(np.ceil(np.log2(gen_mat_max + 1))) - assert ( + if not (gen_mat_max > 0): + raise AssertionError("generating matrix must have positive ints") + if not (self._t_curr == int(np.ceil(np.log2(gen_mat_max + 1)))): + raise AssertionError + if not ( 0 < self._t_curr <= self.t <= 64 - ), "invalid 0 <= self._t_curr (%d) <= self.t (%d) <= 64" % ( - self._t_curr, - self.t, - ) + ): + raise AssertionError("invalid 0 <= self._t_curr (%d) <= self.t (%d) <= 64" % ( + self._t_curr, + self.t, + )) if self.randomize == "FALSE": - assert self.gen_mats.shape[0] == self.replications, ( - "randomize='FALSE' but replications = %d does not equal the number of sets of generating matrices %d" - % (self.replications, self.gen_mats.shape[0]) - ) + if not (self.gen_mats.shape[0] == self.replications): + raise AssertionError( + "randomize='FALSE' but replications = %d does not equal the number of sets of generating matrices %d" + % (self.replications, self.gen_mats.shape[0]) + ) def _try_gen_samples_float(self, r, n, d, n_start, mmax, r_x, return_binary): if return_binary or "NUS" in self.randomize: @@ -673,7 +708,7 @@ def _try_gen_samples_float(self, r, n, d, n_start, mmax, r_x, return_binary): def _gen_samples(self, n_min, n_max, return_binary, warn): if n_min == 0 and self.randomize in ["FALSE", "LMS"] and warn: warnings.warn( - "Without randomization, the first digtial net point is the origin", + "Without randomization, the first digital net point is the origin", ParameterWarning, ) r_x = np.uint64(self.gen_mats.shape[0]) diff --git a/qmcpy/discrete_distribution/dummy_sampler.py b/qmcpy/discrete_distribution/dummy_sampler.py index cf33720e5..3190d94b4 100644 --- a/qmcpy/discrete_distribution/dummy_sampler.py +++ b/qmcpy/discrete_distribution/dummy_sampler.py @@ -1,10 +1,12 @@ +from typing import Union +import numpy as np from .abstract_discrete_distribution import AbstractLDDiscreteDistribution from ..util import ParameterError class DummySampler(AbstractLDDiscreteDistribution): - r""" - Placeholder discrete distribution for constructing true-measure marginals. + r"""Placeholder discrete distribution for constructing true-measure + marginals. ``DummySampler`` is useful when a true measure is needed only for its dimension, transform, range, and weight behavior. QMCPy's current @@ -15,8 +17,7 @@ class DummySampler(AbstractLDDiscreteDistribution): Direct calls to ``DummySampler`` raise an error because the sampler is only a construction placeholder and cannot generate meaningful QMC points. - Examples - -------- + Examples: >>> from qmcpy import DummySampler >>> sampler = DummySampler(2) >>> sampler.d @@ -29,7 +30,18 @@ class DummySampler(AbstractLDDiscreteDistribution): qmcpy.util.exceptions_warnings.ParameterError: DummySampler is only a construction placeholder for ProductMeasure child true measures and cannot generate samples. """ - def __init__(self, dimension=1, replications=None, seed=None, warn=True): + def __init__(self, dimension: int = 1, replications: Union[None, int] = None, seed: Union[None, int, np.random.SeedSequence] = None, warn: bool = True) -> None: + """Initialize a DummySampler discrete distribution. + + Args: + dimension (int): Dimension of the placeholder sampler. + replications (Union[None, int]): Number of independent randomizations, kept + for API consistency with the other discrete distributions. + seed (Union[None, int, np.random.SeedSequence]): Unused; kept for + API consistency with the other discrete distributions. + warn (bool): Unused; kept for API consistency with the other + discrete distributions. + """ # Keep the same constructor as other discrete distributions. del warn diff --git a/qmcpy/discrete_distribution/iid_std_uniform.py b/qmcpy/discrete_distribution/iid_std_uniform.py index 195108ca8..50a04ae1f 100644 --- a/qmcpy/discrete_distribution/iid_std_uniform.py +++ b/qmcpy/discrete_distribution/iid_std_uniform.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_discrete_distribution import AbstractIIDDiscreteDistribution from ..util import ParameterError, ParameterWarning import numpy as np @@ -5,10 +6,10 @@ class IIDStdUniform(AbstractIIDDiscreteDistribution): - r""" - IID standard uniform points, a wrapper around [`numpy.random.rand`](https://numpy.org/doc/stable/reference/random/generated/numpy.random.rand.html). + r"""IID standard uniform points, a wrapper around + [`numpy.random.rand`](https://numpy.org/doc/stable/reference/random/generated/numpy.random.rand.html). - Note: + Notes: - Unlike low discrepancy sequence, calling an `IIDStdUniform` instance gives new samples every time, e.g., running the first doctest below with `dd = Lattice(dimension=2)` would give the same 4 points in both calls, but since we are using an `IIDStdUniform` instance it gives different points every call. @@ -49,18 +50,22 @@ class IIDStdUniform(AbstractIIDDiscreteDistribution): [0.6171181 , 0.1239209 , 0.16809479]]]) """ - def __init__(self, dimension=1, replications=None, seed=None): - r""" + def __init__(self, dimension: int = 1, replications: Union[None, int] = None, seed: Union[None, int, np.random.SeedSequence] = None) -> None: + r"""Initialize an IIDStdUniform discrete distribution. + Args: dimension (int): Dimension of the samples. - replications (Union[None, int]): Number of randomizations. This is implemented only for API consistency. Equivalent to reshaping samples. - seed (Union[None, int, np.random.SeedSeq): Seed the random number generator for reproducibility. + replications (Union[None, int]): Number of randomizations. This is + implemented only for API consistency. Equivalent to reshaping + samples. + seed (Union[None, int, np.random.SeedSequence]): Seed the random + number generator for reproducibility. """ super(IIDStdUniform, self).__init__( int(dimension), replications, seed, d_limit=np.inf, n_limit=np.inf ) if not (self.dvec == np.arange(self.d)).all(): - warnings.warn("IIDStdUniform does not accomodate dvec", ParameterWarning) + warnings.warn("IIDStdUniform does not accommodate dvec", ParameterWarning) def _gen_samples(self, n_min, n_max, return_binary, warn): if n_min > 0 and warn: diff --git a/qmcpy/discrete_distribution/korobov.py b/qmcpy/discrete_distribution/korobov.py index b3ae08e99..10f1b0b50 100644 --- a/qmcpy/discrete_distribution/korobov.py +++ b/qmcpy/discrete_distribution/korobov.py @@ -1,3 +1,4 @@ +from typing import Union import numpy as np from qmcpy.util import ParameterError, ParameterWarning from pathlib import Path @@ -11,8 +12,9 @@ def load_korobov_table( npz_path=Path(__file__).resolve().parent / "generating_params" / "korobov_p2_table.npz" ): """Load the Korobov table from the compressed .npz file. Cached via - lru_cache: the file is only actually read once per process, with no - explicit module-level global variable.""" + lru_cache: the file is only actually read once per process, with no + explicit module-level global variable. + """ with np.load(npz_path) as data: raw = data["raw"] lut = { @@ -24,7 +26,18 @@ def load_korobov_table( } return raw, lut -def get_a(lut, n, d): +def get_a(lut: dict, n: int, d: int) -> int: + """Look up the tabulated Korobov generator `a` for a given `n` and `d`. + + Args: + lut (dict): Lookup table returned by the module's table loader, with + keys `n_values`, `d_values`, and `a`. + n (int): Number of points; must be one of `lut["n_values"]`. + d (int): Dimension; must be one of `lut["d_values"]`. + + Returns: + int: The tabulated generator value `a` for this `(n, d)` pair. + """ i = np.searchsorted(lut["n_values"], n) if i >= len(lut["n_values"]) or lut["n_values"][i] != n: raise ParameterError( @@ -42,21 +55,22 @@ def get_a(lut, n, d): class KorobovLattice(AbstractLDDiscreteDistribution): - r""" - Korobov lattice rule with a tabulated, quality-optimized generating parameter. + r"""Korobov lattice rule with a tabulated, quality-optimized generating + parameter. - A rank-1 lattice rule with $n$ points and generating vector $z\in\mathbb{Z}^d$ is - $P_n(z) = \{(\{k z_1/n\},\dots,\{k z_d/n\}) : k=0,\dots,n-1\}$. The Korobov - construction restricts $z$ to a single integer parameter $a$: - $z(a) = (1,a,a^2,\dots,a^{d-1}) \bmod n$, with $\gcd(a,n)=1$. + A rank-1 lattice rule with $n$ points and generating vector + $z\in\mathbb{Z}^d$ is $P_n(z) = \{(\{k z_1/n\},\dots,\{k z_d/n\}) : + k=0,\dots,n-1\}$. The Korobov construction restricts $z$ to a single + integer parameter $a$: $z(a) = (1,a,a^2,\dots,a^{d-1}) \bmod n$, with + $\gcd(a,n)=1$. Rather than searching for $a$ at construction time, this class looks up $a$ in a precomputed table, for every $(n,d)$ pair in the table, minimizing the weighted $P_2$ figure of merit (the squared worst-case integration error in - the weighted Korobov space of smoothness 2) with product weights - $\gamma_j = 1/j^2$. + the weighted Korobov space of smoothness 2) with product weights $\gamma_j + = 1/j^2$. - Note: + Notes: - Because the optimal $a$ depends on the *total* number of points $n$, a Korobov lattice cannot be incrementally extended the way `Lattice` can: `n_min` must be 0, and `n` must be one of the values in the @@ -135,18 +149,19 @@ class KorobovLattice(AbstractLDDiscreteDistribution): """ def __init__( self, - dimension=1, - replications=None, - seed=None, - randomize="SHIFT", - ): - r""" + dimension: int = 1, + replications: Union[None, int] = None, + seed: Union[None, int, np.random.SeedSequence] = None, + randomize: str = "SHIFT", + ) -> None: + r"""Initialize a KorobovLattice discrete distribution. + Args: dimension (int): Dimension of the samples. Must be between 1 and 250 (the range covered by the precomputed table). - replications (int): Number of independent Cranley-Patterson - shifts of the same underlying deterministic lattice. + replications (Union[None, int]): Number of independent Cranley-Patterson shifts + of the same underlying deterministic lattice. seed (Union[None, int, np.random.SeedSequence]): Seed the random number generator for reproducibility. @@ -156,7 +171,7 @@ def __init__( - `'SHIFT'` or `'TRUE'`: Random Cranley-Patterson shift (the default). - `'FALSE'`, `'NONE'`, or `'NO'`: No randomization. In this case the first point will be the origin. - """ + """ super().__init__(dimension, replications, seed, d_limit = 250, n_limit = 131072) self.randomize = str(randomize).upper() @@ -166,7 +181,8 @@ def __init__( self.randomize = "FALSE" if self.randomize == "NO": self.randomize = "FALSE" - assert self.randomize in ["SHIFT", "FALSE"] + if not (self.randomize in ["SHIFT", "FALSE"]): + raise AssertionError if self.randomize not in ("SHIFT", "FALSE"): raise ParameterError( f"randomize must be one of 'SHIFT', 'TRUE', 'FALSE', 'NONE', or 'NO' (case-insensitive), got {randomize!r}." diff --git a/qmcpy/discrete_distribution/kronecker.py b/qmcpy/discrete_distribution/kronecker.py index 89311121a..4ebcf0617 100644 --- a/qmcpy/discrete_distribution/kronecker.py +++ b/qmcpy/discrete_distribution/kronecker.py @@ -1,3 +1,4 @@ +from typing import Union, Tuple, Callable from .abstract_discrete_distribution import AbstractLDDiscreteDistribution from ..util import ParameterError import numpy as np @@ -48,25 +49,26 @@ def _richtmyer_generating_vector(dimension): PRIMES = np.array([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311, 313, 317, 331, 337, 347, 349, 353, 359, 367, 373, 379, 383, 389, 397, 401, 409, 419, 421, 431, 433, 439, 443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503, 509, 521, 523, 541, 547, 557, 563, 569, 571, 577, 587, 593, 599, 601, 607, 613, 617, 619, 631, 641, 643, 647, 653, 659, 661, 673, 677, 683, 691, 701, 709, 719, 727, 733, 739, 743, 751, 757, 761, 769, 773, 787, 797, 809, 811, 821, 823, 827, 829, 839, 853, 857, 859, 863, 877, 881, 883, 887, 907, 911, 919, 929, 937, 941, 947, 953, 967, 971, 977, 983, 991, 997, 1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069, 1087, 1091, 1093, 1097, 1103, 1109, 1117, 1123, 1129, 1151, 1153, 1163, 1171, 1181, 1187, 1193, 1201, 1213, 1217, 1223, 1229, 1231, 1237, 1249, 1259, 1277, 1279, 1283, 1289, 1291, 1297, 1301, 1303, 1307, 1319, 1321, 1327, 1361, 1367, 1373, 1381, 1399, 1409, 1423, 1427, 1429, 1433, 1439, 1447, 1451, 1453, 1459, 1471, 1481, 1483, 1487, 1489, 1493, 1499, 1511, 1523, 1531, 1543, 1549, 1553, 1559, 1567, 1571, 1579, 1583, 1597, 1601, 1607, 1609, 1613, 1619, 1621, 1627, 1637, 1657, 1663, 1667, 1669, 1693, 1697, 1699, 1709, 1721, 1723, 1733, 1741, 1747, 1753, 1759, 1777, 1783, 1787, 1789, 1801, 1811, 1823, 1831, 1847, 1861, 1867, 1871, 1873, 1877, 1879, 1889, 1901, 1907, 1913, 1931, 1933, 1949, 1951, 1973, 1979, 1987, 1993, 1997, 1999, 2003, 2011, 2017, 2027, 2029, 2039, 2053, 2063, 2069, 2081, 2083, 2087, 2089, 2099, 2111, 2113, 2129, 2131, 2137, 2141, 2143, 2153, 2161, 2179, 2203, 2207, 2213, 2221, 2237, 2239, 2243, 2251, 2267, 2269, 2273, 2281, 2287, 2293, 2297, 2309, 2311, 2333, 2339, 2341, 2347, 2351, 2357, 2371, 2377, 2381, 2383, 2389, 2393, 2399, 2411, 2417, 2423, 2437, 2441, 2447, 2459, 2467, 2473, 2477, 2503, 2521, 2531, 2539, 2543, 2549, 2551, 2557, 2579, 2591, 2593, 2609, 2617, 2621, 2633, 2647, 2657, 2659, 2663, 2671, 2677, 2683, 2687, 2689, 2693, 2699, 2707, 2711, 2713, 2719, 2729, 2731, 2741, 2749, 2753, 2767, 2777, 2789, 2791, 2797, 2801, 2803, 2819, 2833, 2837, 2843, 2851, 2857, 2861, 2879, 2887, 2897, 2903, 2909, 2917, 2927, 2939, 2953, 2957, 2963, 2969, 2971, 2999, 3001, 3011, 3019, 3023, 3037, 3041, 3049, 3061, 3067, 3079, 3083, 3089, 3109, 3119, 3121, 3137, 3163, 3167, 3169, 3181, 3187, 3191, 3203, 3209, 3217, 3221, 3229, 3251, 3253, 3257, 3259, 3271, 3299, 3301, 3307, 3313, 3319, 3323, 3329, 3331, 3343, 3347, 3359, 3361, 3371, 3373, 3389, 3391, 3407, 3413, 3433, 3449, 3457, 3461, 3463, 3467, 3469, 3491, 3499, 3511, 3517, 3527, 3529, 3533, 3539, 3541, 3547, 3557, 3559, 3571, 3581, 3583, 3593, 3607, 3613, 3617, 3623, 3631, 3637, 3643, 3659, 3671, 3673, 3677, 3691, 3697, 3701, 3709, 3719, 3727, 3733, 3739, 3761, 3767, 3769, 3779, 3793, 3797, 3803, 3821, 3823, 3833, 3847, 3851, 3853, 3863, 3877, 3881, 3889, 3907, 3911, 3917, 3919, 3923, 3929, 3931, 3943, 3947, 3967, 3989, 4001, 4003, 4007, 4013, 4019, 4021, 4027, 4049, 4051, 4057, 4073, 4079, 4091, 4093, 4099, 4111, 4127, 4129, 4133, 4139, 4153, 4157, 4159, 4177, 4201, 4211, 4217, 4219, 4229, 4231, 4241, 4243, 4253, 4259, 4261, 4271, 4273, 4283, 4289, 4297, 4327, 4337, 4339, 4349, 4357, 4363, 4373, 4391, 4397, 4409, 4421, 4423, 4441, 4447, 4451, 4457, 4463, 4481, 4483, 4493, 4507, 4513, 4517, 4519, 4523, 4547, 4549, 4561, 4567, 4583, 4591, 4597, 4603, 4621, 4637, 4639, 4643, 4649, 4651, 4657, 4663, 4673, 4679, 4691, 4703, 4721, 4723, 4729, 4733, 4751, 4759, 4783, 4787, 4789, 4793, 4799, 4801, 4813, 4817, 4831, 4861, 4871, 4877, 4889, 4903, 4909, 4919, 4931, 4933, 4937, 4943, 4951, 4957, 4967, 4969, 4973, 4987, 4993, 4999, 5003, 5009, 5011, 5021, 5023, 5039, 5051, 5059, 5077, 5081, 5087, 5099, 5101, 5107, 5113, 5119, 5147, 5153, 5167, 5171, 5179, 5189, 5197, 5209, 5227, 5231, 5233, 5237, 5261, 5273, 5279, 5281, 5297, 5303, 5309, 5323, 5333, 5347, 5351, 5381, 5387, 5393, 5399, 5407, 5413, 5417, 5419, 5431, 5437, 5441, 5443, 5449, 5471, 5477, 5479, 5483, 5501, 5503, 5507, 5519, 5521, 5527, 5531, 5557, 5563, 5569, 5573, 5581, 5591, 5623, 5639, 5641, 5647, 5651, 5653, 5657, 5659, 5669, 5683, 5689, 5693, 5701, 5711, 5717, 5737, 5741, 5743, 5749, 5779, 5783, 5791, 5801, 5807, 5813, 5821, 5827, 5839, 5843, 5849, 5851, 5857, 5861, 5867, 5869, 5879, 5881, 5897, 5903, 5923, 5927, 5939, 5953, 5981, 5987, 6007, 6011, 6029, 6037, 6043, 6047, 6053, 6067, 6073, 6079, 6089, 6091, 6101, 6113, 6121, 6131, 6133, 6143, 6151, 6163, 6173, 6197, 6199, 6203, 6211, 6217, 6221, 6229, 6247, 6257, 6263, 6269, 6271, 6277, 6287, 6299, 6301, 6311, 6317, 6323, 6329, 6337, 6343, 6353, 6359, 6361, 6367, 6373, 6379, 6389, 6397, 6421, 6427, 6449, 6451, 6469, 6473, 6481, 6491, 6521, 6529, 6547, 6551, 6553, 6563, 6569, 6571, 6577, 6581, 6599, 6607, 6619, 6637, 6653, 6659, 6661, 6673, 6679, 6689, 6691, 6701, 6703, 6709, 6719, 6733, 6737, 6761, 6763, 6779, 6781, 6791, 6793, 6803, 6823, 6827, 6829, 6833, 6841, 6857, 6863, 6869, 6871, 6883, 6899, 6907, 6911, 6917, 6947, 6949, 6959, 6961, 6967, 6971, 6977, 6983, 6991, 6997, 7001, 7013, 7019, 7027, 7039, 7043, 7057, 7069, 7079, 7103, 7109, 7121, 7127, 7129, 7151, 7159, 7177, 7187, 7193, 7207, 7211, 7213, 7219, 7229, 7237, 7243, 7247, 7253, 7283, 7297, 7307, 7309, 7321, 7331, 7333, 7349, 7351, 7369, 7393, 7411, 7417, 7433, 7451, 7457, 7459, 7477, 7481, 7487, 7489, 7499, 7507, 7517, 7523, 7529, 7537, 7541, 7547, 7549, 7559, 7561, 7573, 7577, 7583, 7589, 7591, 7603, 7607, 7621, 7639, 7643, 7649, 7669, 7673, 7681, 7687, 7691, 7699, 7703, 7717, 7723, 7727, 7741, 7753, 7757, 7759, 7789, 7793, 7817, 7823, 7829, 7841, 7853, 7867, 7873, 7877, 7879, 7883, 7901, 7907, 7919]) - assert dimension>> x = Kronecker(3,seed=7,replications=2)(4) @@ -111,23 +113,24 @@ class Kronecker(AbstractLDDiscreteDistribution): [[0.49700422, 0.41789272, 0.80339779], [0.91944141, 0.77848924, 0.15206993]]]) - - Switch from CBC to Richtmyer generating vector when the dimension is too large. + + Switch from CBC to Richtmyer generating vector when the dimension is + too large. >>> Kronecker(15,seed=7,warn=False)(4).shape (4, 15) >>> Kronecker(15,replications=2,seed=7,warn=False)(4).shape (2, 4, 15) - CBC unrandomized - + CBC unrandomized + >>> Kronecker(3,generating_vector="CBC",randomize=False)(4) array([[0. , 0. , 0. ], [0.42243719, 0.36059652, 0.34867214], [0.84487437, 0.72119304, 0.69734427], [0.26731156, 0.08178956, 0.04601641]]) - - Richtmyer construction + + Richtmyer construction >>> Kronecker(3,generating_vector="RICHTMYER",randomize=False)(4) array([[0. , 0. , 0. ], @@ -145,7 +148,7 @@ class Kronecker(AbstractLDDiscreteDistribution): [0.48055697, 0.16080129, 0.57818947], [0.89477054, 0.8928521 , 0.81425745]]]) - Suzuki construction + Suzuki construction >>> Kronecker(3,generating_vector="SUZUKI",randomize=False)(4) array([[0. , 0. , 0. ], @@ -181,7 +184,7 @@ class Kronecker(AbstractLDDiscreteDistribution): [0.77841423, 0.32842712, 0.96358566], [0.96762135, 0.74264069, 0.64537849]]]) - Custom generating vectors + Custom generating vectors >>> Kronecker(3,generating_vector=2**(np.arange(1,4)/(3 + 1)),randomize=False)(4) array([[0. , 0. , 0. ], @@ -199,8 +202,8 @@ class Kronecker(AbstractLDDiscreteDistribution): [0.84133696, 0.11091324, 0.78784635], [0.03054408, 0.5251268 , 0.46963918], [0.21975119, 0.93934037, 0.15143201]]]) - - Subset dimensions + + Subset dimensions >>> Kronecker([0,2],generating_vector=2**(np.arange(1,4)/(3 + 1)),randomize=False)(4) array([[0. , 0. ], @@ -211,42 +214,47 @@ class Kronecker(AbstractLDDiscreteDistribution): **References** 1. Richtmyer, R. D. (1951). "The evaluation of definite integrals and a quasi-Monte Carlo method." - + 2. Niederreiter, H. (1992). *Random Number Generation and Quasi-Monte Carlo Methods*. """ def __init__(self, - dimension=1, - replications=None, - seed=None, - randomize="SHIFT", - generating_vector="CBC", - shift=None, - warn=True, - ): - r""" + dimension: Union[int, np.ndarray] = 1, + replications: Union[None, int] = None, + seed: Union[None, int, np.random.SeedSequence] = None, + randomize: str = "SHIFT", + generating_vector: Union[str, np.ndarray] = "CBC", + shift: Union[None, np.ndarray] = None, + warn: bool = True, + ) -> None: + r"""Initialize a Kronecker discrete distribution. + Args: dimension (Union[int, np.ndarray]): Dimension of the generator. - If an `int` is passed in, use generating vector components at indices 0,...,`dimension`-1. - If an `np.ndarray` is passed in, use generating vector components at these indices. - - replications (int): Number of independent randomizations. - seed (Union[None, int, np.random.SeedSeq): Seed the random number generator for reproducibility. + + replications (Union[None, int]): Number of independent randomizations. + seed (Union[None, int, np.random.SeedSequence]): Seed the random + number generator for reproducibility. randomize (str): Options are - `'SHIFT'`: use `shift` if supplied, otherwise use a random shift $\boldsymbol{\delta} \sim \mathrm{Uniform}([0,1)^d)$. - `'FALSE'`: zero shift. - - generating_vector (Union[str,np.ndarray]): Generating vector $\boldsymbol{\alpha}$. - + + generating_vector (Union[str, np.ndarray]): Generating vector + $\boldsymbol{\alpha}$. + - `"CBC"`: uses the first $d$ components of a known good Component-by-Component (CBC) generating vector. - `"RICHTMYER"`: uses $\boldsymbol{\alpha}_j = \sqrt{p_j} \bmod 1$, where $p_j$ are primes. This is the classical Richtmyer construction. - `"SUZUKI"`: uses a deterministic construction $\boldsymbol{\alpha}_j = 2^{j/(d+1)}$. - np.array: user-specified generating vector. - shift (np.ndarray): Shift vector $\boldsymbol{\delta}$. If `randomize=True`, this is ignored and a random shift is generated. Otherwise, a fixed shift is used. - warn (bool): If False, suppress warnings during construction + shift (Union[None, np.ndarray]): Shift vector $\boldsymbol{\delta}$. If + `randomize=True`, this is ignored and a random shift is + generated. Otherwise, a fixed shift is used. + warn (bool): If False, suppress warnings during construction """ self.parameters = ["randomize", "gen_vec_source"] self.input_generating_vector = generating_vector @@ -291,10 +299,14 @@ def __init__(self, if gen_vec.ndim >2: raise ParameterError("generating_vector must be a 1D or 2D np.ndarray") gen_vec = np.atleast_2d(gen_vec).astype(float) - assert gen_vec.ndim==2, "gen_vec must be a 2D array" - assert gen_vec.shape[1]>=self.d - assert (gen_vec.shape[0] == 1 or gen_vec.shape[0] == self.replications) - assert gen_vec.shape[1]>self.dvec.max() + if not (gen_vec.ndim==2): + raise AssertionError("gen_vec must be a 2D array") + if not (gen_vec.shape[1]>=self.d): + raise AssertionError + if not (gen_vec.shape[0] == 1 or gen_vec.shape[0] == self.replications): + raise AssertionError + if not (gen_vec.shape[1]>self.dvec.max()): + raise AssertionError self.gen_vec = gen_vec[:,self.dvec].copy() self.randomize = str(randomize).upper() if self.randomize == "TRUE": @@ -303,8 +315,11 @@ def __init__(self, self.randomize = "FALSE" if self.randomize == "NO": self.randomize = "FALSE" - assert self.randomize in ["SHIFT", "FALSE"] - if shift is not None: assert self.randomize=="SHIFT", "require randomize='SHIFT' when shift is not None" + if not (self.randomize in ["SHIFT", "FALSE"]): + raise AssertionError + if shift is not None: + if not (self.randomize=="SHIFT"): + raise AssertionError("require randomize='SHIFT' when shift is not None") if self.randomize=="SHIFT": if shift is not None: self.shift = np.atleast_2d(shift).astype(float) @@ -312,9 +327,12 @@ def __init__(self, self.shift = self.rng.uniform(size=(self.replications, self.d)) else: # self.randomize=="FALSE": self.shift = np.zeros((self.replications, self.d)) - assert self.shift.ndim==2 - assert self.shift.shape[1]==self.d - assert (self.shift.shape[0] == 1 or self.shift.shape[0] == self.replications) + if not (self.shift.ndim==2): + raise AssertionError + if not (self.shift.shape[1]==self.d): + raise AssertionError + if not (self.shift.shape[0] == 1 or self.shift.shape[0] == self.replications): + raise AssertionError def _gen_samples(self, n_min, n_max, return_binary, warn): if return_binary: @@ -323,23 +341,24 @@ def _gen_samples(self, n_min, n_max, return_binary, warn): points = ((i[:,None] * self.gen_vec[:,None,:]) + self.shift[:, None, :]) % 1 return points - def periodic_discrepancy(self, n, k_tilde=None, gamma=None): - # """ - # Calculates the discrepancy for a periodic kernel. - - # Args: - # n (int): the number of sample points - # k_tilde (Tuple[function, float]): the function takes in 2 arguments: the sample points and the coordinate weights. - # The float is the integral over the unit hypercube. - # gamma (np.ndarray): shape (1xd) + def periodic_discrepancy(self, n: int, k_tilde: Union[None, Tuple[Callable, float]] = None, gamma: Union[None, np.ndarray] = None) -> np.ndarray: + """Calculate the discrepancy for a periodic kernel. - # Returns: - # discrep (np.ndarray): discrepancy - - # Notes: - # - If k_tilde is not specified, the second Bernoulli polynomial is used. - # - If gamma is not specified, the coordinate weights will be just all ones. - # """ + Args: + n (int): The number of sample points. + k_tilde (Union[None, Tuple[Callable, float]]): A `(function, integral)` pair + where the function takes the sample points and coordinate + weights and returns kernel values, and `integral` is that + function's integral over the unit hypercube. + gamma (Union[None, np.ndarray]): Coordinate weights, shape `(d,)`. + + Returns: + np.ndarray: The discrepancy. + + Note: + - If `k_tilde` is not specified, the second Bernoulli polynomial is used. + - If `gamma` is not specified, the coordinate weights are all ones. + """ if gamma is None: gamma = np.ones(self.d) @@ -349,8 +368,19 @@ def periodic_discrepancy(self, n, k_tilde=None, gamma=None): return np.sqrt(self._square_periodic_discrepancies(n, k_tilde, gamma)) - def wssd_discrepancy(self, n, weights, k_tilde = None, gamma = None): - # calculates the weighted sum of square discrepancy + def wssd_discrepancy(self, n: int, weights: np.ndarray, k_tilde: Union[None, Tuple[Callable, float]] = None, gamma: Union[None, np.ndarray] = None) -> np.ndarray: + """Calculate the weighted sum of squared discrepancies. + + Args: + n (int): The number of sample points. + weights (np.ndarray): Weights applied to each squared discrepancy + before summing. + k_tilde (Union[None, Tuple[Callable, float]]): Same as in `periodic_discrepancy`. + gamma (Union[None, np.ndarray]): Coordinate weights, shape `(d,)`. + + Returns: + np.ndarray: The weighted sum of squared discrepancies. + """ if gamma is None: gamma = np.ones(self.d) @@ -375,7 +405,8 @@ def _square_periodic_discrepancies(self, n, k_tilde, gamma): def _spawn(self, child_seed, dimension): - assert self.input_shift is None, "spawn requires shift=None" + if not (self.input_shift is None): + raise AssertionError("spawn requires shift=None") return Kronecker( dimension=dimension, replications=None if self.no_replications else self.replications, diff --git a/qmcpy/discrete_distribution/latin_hypercube.py b/qmcpy/discrete_distribution/latin_hypercube.py index 95aed4c8a..2c1581fe9 100644 --- a/qmcpy/discrete_distribution/latin_hypercube.py +++ b/qmcpy/discrete_distribution/latin_hypercube.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_discrete_distribution import AbstractDiscreteDistribution import numpy as np from qmcpy.util import ParameterError, ParameterWarning @@ -5,23 +6,22 @@ class LatinHypercube(AbstractDiscreteDistribution): - r""" - Latin Hypercube Sampler for quasi-Monte Carlo and experimental design. + r"""Latin Hypercube Sampler for quasi-Monte Carlo and experimental design. Latin Hypercube Sampling (LHS) generates points with excellent univariate stratification: splitting $[0,1)$ into `n` equal strata along *any* single coordinate axis places exactly one point in each stratum. Introduced by McKay, Beckman, and Conover as a variance-reduction alternative to simple - random sampling for computer experiments, LHS is asymptotically at least - as accurate as Monte Carlo for the additive part of an integrand, with the + random sampling for computer experiments, LHS is asymptotically at least as + accurate as Monte Carlo for the additive part of an integrand, with the rate of improvement characterized by Stein and later by Loh via a multivariate central limit theorem. - Note: + Notes: - Unlike the low discrepancy sequences in this package (e.g. `Lattice`, `Halton`, `DigitalNetB2`), `LatinHypercube` points are *not* extensible in `n`: the entire point set must be regenerated whenever `n` changes, - since the strata boundaries themselves depend on `n`. + since the strata boundaries themselves depend on `n`. Consequently `LatinHypercube` requires `n_min=0`, it cannot be generated starting from a nonzero offset. - `replications` produces independent randomizations (independent random permutations, and independent within-stratum jitter when `randomize` @@ -69,50 +69,53 @@ class LatinHypercube(AbstractDiscreteDistribution): **References:** - 1. M. D. McKay, R. J. Beckman, and W. J. Conover. - A Comparison of Three Methods for Selecting Values of Input Variables in the Analysis of Output from a Computer Code. - Technometrics, 21(2):239-245, 1979. + 1. M. D. McKay, R. J. Beckman, and W. J. Conover. + A Comparison of Three Methods for Selecting Values of Input Variables in the Analysis of Output from a Computer Code. + Technometrics, 21(2):239-245, 1979. [https://doi.org/10.1080/00401706.1979.10489755](https://doi.org/10.1080/00401706.1979.10489755). - 2. M. Stein. - Large Sample Properties of Simulations Using Latin Hypercube Sampling. - Technometrics, 29(2):143-151, 1987. + 2. M. Stein. + Large Sample Properties of Simulations Using Latin Hypercube Sampling. + Technometrics, 29(2):143-151, 1987. [https://doi.org/10.1080/00401706.1987.10488205](https://doi.org/10.1080/00401706.1987.10488205). - 3. A. B. Owen. - Controlling Correlations in Latin Hypercube Samples. - Journal of the American Statistical Association, 89(428):1517-1522, 1994. + 3. A. B. Owen. + Controlling Correlations in Latin Hypercube Samples. + Journal of the American Statistical Association, 89(428):1517-1522, 1994. [https://doi.org/10.1080/01621459.1994.10476891](https://doi.org/10.1080/01621459.1994.10476891). - 4. W.-L. Loh. - On Latin Hypercube Sampling. - The Annals of Statistics, 24(5):2058-2080, 1996. + 4. W.-L. Loh. + On Latin Hypercube Sampling. + The Annals of Statistics, 24(5):2058-2080, 1996. [https://doi.org/10.1214/aos/1069362310](https://doi.org/10.1214/aos/1069362310). - 5. B. Tang. - Orthogonal Array-Based Latin Hypercubes. - Journal of the American Statistical Association, 88(424):1392-1397, 1993. + 5. B. Tang. + Orthogonal Array-Based Latin Hypercubes. + Journal of the American Statistical Association, 88(424):1392-1397, 1993. [https://doi.org/10.1080/01621459.1993.10476423](https://doi.org/10.1080/01621459.1993.10476423). """ def __init__( - self, dimension, replications, seed, randomize="TRUE" - ): - r""" + self, dimension: int, replications: Union[None, int], seed: Union[None, int, np.random.SeedSequence], randomize: str = "TRUE" + ) -> None: + r"""Initialize a LatinHypercube discrete distribution. + Args: dimension (int): Dimension of the samples. replications (Union[None, int]): Number of independent LHS designs - to generate. Each replication is its own independently permuted, - independently jittered stratification into `n` strata. + to generate. Each replication is its own independently + permuted, independently jittered stratification into `n` + strata. - seed (Union[None, int, np.random.SeedSequence]): Seed for the random - number generator to ensure reproducibility. + seed (Union[None, int, np.random.SeedSequence]): Seed for the + random number generator to ensure reproducibility. randomize (str): Whether to jitter each point uniformly within its stratum (`True`, the default) or place it at the stratum's - center (`False`), must be one of 'TRUE', 'FALSE', 'NONE', or 'NO' (case-insensitive). - """ + center (`False`), must be one of 'TRUE', 'FALSE', 'NONE', or + 'NO' (case-insensitive). + """ super().__init__(dimension=dimension, replications=replications, seed=seed, d_limit=np.inf, n_limit=np.inf) self.randomize = str(randomize).upper() if self.randomize in ("NONE", "NO", "FALSE"): @@ -127,7 +130,7 @@ def __init__( def _gen_samples( self, n=None, n_min=None, n_max=None, return_binary=False, warn=True ): - r"""...""" # (inchangee) + r"""...""" # (unchanged) if return_binary: raise ParameterError("LatinHypercube does not support return_binary=True") if n_min != 0: @@ -164,4 +167,4 @@ def _spawn(self, child_seed, dimension): ) def __repr__(self): - return super().__repr__("LatinHypercube") \ No newline at end of file + return super().__repr__("LatinHypercube") diff --git a/qmcpy/discrete_distribution/lattice/lattice.py b/qmcpy/discrete_distribution/lattice/lattice.py index 248692c32..d769c7d3c 100644 --- a/qmcpy/discrete_distribution/lattice/lattice.py +++ b/qmcpy/discrete_distribution/lattice/lattice.py @@ -1,3 +1,4 @@ +from typing import Union from ..abstract_discrete_distribution import AbstractLDDiscreteDistribution from ...util import ParameterError, ParameterWarning import qmctoolscl @@ -9,10 +10,9 @@ class Lattice(AbstractLDDiscreteDistribution): - r""" - Low discrepancy lattice sequence. + r"""Low discrepancy lattice sequence. - Note: + Notes: - Lattice sample sizes should be powers of $2$ e.g. $1$, $2$, $4$, $8$, $16$, $\dots$. - The first point of an unrandomized lattice is the origin. @@ -52,7 +52,8 @@ class Lattice(AbstractLDDiscreteDistribution): [0.40212985, 0.94669968, 0.35605352]]]) - Different orderings (avoid warnings that the first point is the origin). + Different orderings (avoid warnings that the first point is the + origin). >>> Lattice(dimension=2,randomize=False,order='RADICAL INVERSE')(4,warn=False) array([[0. , 0. ], @@ -70,7 +71,8 @@ class Lattice(AbstractLDDiscreteDistribution): [0.5 , 0.5 ], [0.75, 0.25]]) - Generating vector from [https://github.com/QMCSoftware/LDData/tree/main/lattice](https://github.com/QMCSoftware/LDData/tree/main/lattice) + Generating vector from + [https://github.com/QMCSoftware/LDData/tree/main/lattice](https://github.com/QMCSoftware/LDData/tree/main/lattice) >>> Lattice(dimension=3,randomize=False,generating_vector="mps.exod2_base2_m20_CKN.txt")(8,warn=False) array([[0. , 0. , 0. ], @@ -93,7 +95,8 @@ class Lattice(AbstractLDDiscreteDistribution): [0.25, 0.75, 0.75], [0.75, 0.25, 0.25]]) - Two random generating vectors both supporting $2^{25}$ points along with independent random shifts + Two random generating vectors both supporting $2^{25}$ points along + with independent random shifts >>> discrete_distrib = Lattice(3,seed=7,generating_vector=25,replications=2) >>> discrete_distrib.gen_vec @@ -138,40 +141,46 @@ class Lattice(AbstractLDDiscreteDistribution): def __init__( self, - dimension=1, - replications=None, - seed=None, - randomize="SHIFT", - generating_vector="kuo.lattice-33002-1024-1048576.9125.txt", - order="RADICAL INVERSE", - m_max=None, - ): - r""" + dimension: Union[int, np.ndarray] = 1, + replications: Union[None, int] = None, + seed: Union[None, int, np.random.SeedSequence] = None, + randomize: str = "SHIFT", + generating_vector: Union[str, np.ndarray, int] = "kuo.lattice-33002-1024-1048576.9125.txt", + order: str = "RADICAL INVERSE", + m_max: Union[None, int] = None, + ) -> None: + r"""Initialize a Lattice discrete distribution. + Args: dimension (Union[int, np.ndarray]): Dimension of the generator. - If an `int` is passed in, use generating vector components at indices 0,...,`dimension`-1. - If an `np.ndarray` is passed in, use generating vector components at these indices. - replications (int): Number of independent randomizations. - seed (Union[None, int, np.random.SeedSeq): Seed the random number generator for reproducibility. + replications (Union[None, int]): Number of independent randomizations. + seed (Union[None, int, np.random.SeedSequence]): Seed the random + number generator for reproducibility. randomize (str): Options are - `'SHIFT'`: Random shift. - `'FALSE'`: No randomization. In this case the first point will be the origin. - generating_vector (Union[str, np.ndarray, int]): Specify the generating vector. + generating_vector (Union[str, np.ndarray, int]): Specify the + generating vector. - A `str` should be the name (or path) of a file from the LDData repo at [https://github.com/QMCSoftware/LDData/tree/main/lattice](https://github.com/QMCSoftware/LDData/tree/main/lattice). - A `np.ndarray` of integers with shape $(d,)$ or $(r,d)$ where $d$ is the number of dimensions and $r$ is the number of replications. Must supply `m_max` where $2^{m_\mathrm{max}}$ is the max number of supported samples. - An `int`, call it $M$, gives the random generating vector $(1,v_1,\dots,v_{d-1})^T$ - where $d$ is the dimension and $v_i$ are randomly selected from $\{3,5,\dots,2^M-1\}$ uniformly and independently. - We require require $1 < M < 27$. - - order (str): `'LINEAR'`, `'RADICAL INVERSE'`, or `'GRAY'` ordering. See the doctest example above. - m_max (int): $2^{m_\mathrm{max}}$ is the maximum number of supported samples. + where $d$ is the dimension and $v_i$ are randomly selected from + $\{3,5,\dots,2^M-1\}$ uniformly and independently. We require + require $1 < M < 27$. + + order (str): `'LINEAR'`, `'RADICAL INVERSE'`, or `'GRAY'` ordering. + See the doctest example above. + m_max (Union[None, int]): $2^{m_\mathrm{max}}$ is the maximum number of + supported samples. """ self.parameters = ["randomize", "gen_vec_source", "order", "n_limit"] self.input_generating_vector = deepcopy(generating_vector) @@ -189,7 +198,8 @@ def __init__( n_limit = 1048576 elif isinstance(generating_vector, str): self.gen_vec_source = generating_vector - assert generating_vector[-4:] == ".txt" + if not (generating_vector[-4:] == ".txt"): + raise AssertionError local_root = dirname(abspath(__file__)) + "/generating_vectors/" repos = DataSource() if repos.exists(local_root + generating_vector): @@ -247,11 +257,13 @@ def __init__( n_limit = int(2**m_max) d_limit = int(gen_vec.shape[-1]) elif isinstance(generating_vector, int): - assert 1 < generating_vector < 27, "int generating vector out of range" + if not (1 < generating_vector < 27): + raise AssertionError("int generating vector out of range") n_limit = 2**generating_vector - assert isinstance( + if not (isinstance( dimension, int - ), "random generating vector requires int dimension" + )): + raise AssertionError("random generating vector requires int dimension") d_limit = dimension else: raise ParameterError( @@ -274,20 +286,23 @@ def __init__( + 1, ] ).copy() - assert isinstance(gen_vec, np.ndarray) + if not (isinstance(gen_vec, np.ndarray)): + raise AssertionError gen_vec = np.atleast_2d(gen_vec) - assert ( + if not ( gen_vec.ndim == 2 and gen_vec.shape[1] >= self.d and (gen_vec.shape[0] == 1 or gen_vec.shape[0] == self.replications) - ), "invalid gen_vec.shape = %s" % str(gen_vec.shape) + ): + raise AssertionError("invalid gen_vec.shape = %s" % str(gen_vec.shape)) self.gen_vec = gen_vec[:, self.dvec].copy() self.order = str(order).upper().strip().replace("_", " ") if self.order == "GRAY CODE": self.order = "GRAY" if self.order == "NATURAL": self.order = "RADICAL INVERSE" - assert self.order in ["LINEAR", "RADICAL INVERSE", "GRAY"] + if not (self.order in ["LINEAR", "RADICAL INVERSE", "GRAY"]): + raise AssertionError self.randomize = str(randomize).upper() if self.randomize == "TRUE": self.randomize = "SHIFT" @@ -295,14 +310,16 @@ def __init__( self.randomize = "FALSE" if self.randomize == "NO": self.randomize = "FALSE" - assert self.randomize in ["SHIFT", "FALSE"] + if not (self.randomize in ["SHIFT", "FALSE"]): + raise AssertionError if self.randomize == "SHIFT": self.shift = self.rng.uniform(size=(self.replications, self.d)) if self.randomize == "FALSE": - assert self.gen_vec.shape[0] == self.replications, ( - "randomize='FALSE' but replications = %d does not equal the number of sets of generating vectors %d" - % (self.replications, self.gen_vec.shape[0]) - ) + if not (self.gen_vec.shape[0] == self.replications): + raise AssertionError( + "randomize='FALSE' but replications = %d does not equal the number of sets of generating vectors %d" + % (self.replications, self.gen_vec.shape[0]) + ) def _gen_samples(self, n_min, n_max, return_binary, warn): if return_binary: @@ -318,14 +335,16 @@ def _gen_samples(self, n_min, n_max, return_binary, warn): n_start = np.uint64(n_min) x = np.empty((r_x, n, d), dtype=np.float64) if self.order == "LINEAR": - assert ( + if not ( r_x == 1 - ), "lattice linear currently requires there be only 1 generating matrix" + ): + raise AssertionError("lattice linear currently requires there be only 1 generating matrix") x = self._gail_linear(n_min, n_max)[None, :, :] elif self.order == "RADICAL INVERSE": - assert (n_min == 0 or (n_min & (n_min - 1)) == 0) and ( + if not ((n_min == 0 or (n_min & (n_min - 1)) == 0) and ( n_max == 0 or (n_max & (n_max - 1)) == 0 - ), "lattice in natural order requires n_min and n_max be 0 or powers of 2" + )): + raise AssertionError("lattice in natural order requires n_min and n_max be 0 or powers of 2") _ = qmctoolscl.lat_gen_natural( r_x, n, d, n_start, self.gen_vec, x, backend="c" ) @@ -334,7 +353,8 @@ def _gen_samples(self, n_min, n_max, return_binary, warn): r_x, n, d, n_start, self.gen_vec, x, backend="c" ) else: - assert False, "invalid lattice order" + if not (False): + raise AssertionError("invalid lattice order") if self.randomize == "FALSE": xr = x elif self.randomize == "SHIFT": @@ -352,7 +372,22 @@ def _gen_block_linear(self, m_next, first=True): x = np.outer(y, self.gen_vec) % 1 return x - def calculate_y(self, m_low, m_high, y): + def calculate_y(self, m_low: int, m_high: int, y: np.ndarray) -> np.ndarray: + """Refine 1D interval midpoints from level `m_low` up to `m_high`. + + At each level, interleaves the current midpoints `y` with the new + midpoints introduced at that level, doubling the length of `y` each + step. Used internally by `_gail_linear` to build up linear-order + lattice coordinates level-by-level. + + Args: + m_low (int): Starting level (`y` must already hold the midpoints for this level). + m_high (int): Final level (exclusive) to refine up to. + y (np.ndarray): Interval midpoints at level `m_low`, shape `(2**(m_low-1), 1)`. + + Returns: + np.ndarray: Interval midpoints at level `m_high`, shape `(2**(m_high-1), 1)`. + """ for m in range(m_low, m_high): n = 2**m y_next = np.arange(1 / n, 1, 2 / n).reshape((int(n / 2), 1)) diff --git a/qmcpy/discrete_distribution/mpmc/__init__.py b/qmcpy/discrete_distribution/mpmc/__init__.py index 9ceea1244..8502009f7 100644 --- a/qmcpy/discrete_distribution/mpmc/__init__.py +++ b/qmcpy/discrete_distribution/mpmc/__init__.py @@ -1,23 +1,23 @@ -""" -Message Passing Monte Carlo (MPMC) discrete distribution. +"""Message Passing Monte Carlo (MPMC) discrete distribution. -This module implements MPMC using PyTorch and PyTorch Geometric for -generating low-discrepancy point sets through neural message passing. +This module implements MPMC using PyTorch and PyTorch Geometric for generating +low-discrepancy point sets through neural message passing. -Installation Requirements --------------------------- -MPMC requires PyTorch and PyTorch Geometric. Install with: +Installation Requirements -------------------------- MPMC requires PyTorch and +PyTorch Geometric. Install with: - python -m pip install "qmcpy[mpmc]" - qmcpy-install-mpmc +python -m pip install "qmcpy[mpmc]" qmcpy-install-mpmc -For GPU support (NVIDIA CUDA), see https://pytorch.org/get-started/locally/ -For torch-geometric wheels, see https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html +For GPU support (NVIDIA CUDA), see https://pytorch.org/get-started/locally/ For +torch-geometric wheels, see +https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html -If these dependencies are not installed, attempting to use MPMC will raise an ImportError -with installation instructions. You can check availability by running: +If these dependencies are not installed, attempting to use MPMC will raise an +ImportError with installation instructions. You can check availability by +running: - python -c "import torch; import pyg_lib; import torch_geometric; print('MPMC dependencies ready')" +python -c "import torch; import pyg_lib; import torch_geometric; print('MPMC +dependencies ready')" """ try: @@ -29,7 +29,9 @@ _missing_dep = str(e) class MPMC(object): - """Placeholder MPMC class shown when PyTorch dependencies are missing.""" + """Placeholder MPMC class shown when PyTorch dependencies are + missing. + """ def __init__(self, *args, **kwargs): raise ImportError( f"MPMC requires PyTorch, pyg_lib, and PyTorch Geometric, but they are not installed.\n" diff --git a/qmcpy/discrete_distribution/mpmc/models.py b/qmcpy/discrete_distribution/mpmc/models.py index 48deb56cf..c88ff9111 100644 --- a/qmcpy/discrete_distribution/mpmc/models.py +++ b/qmcpy/discrete_distribution/mpmc/models.py @@ -10,7 +10,14 @@ class MPNN_layer(MessagePassing): - def __init__(self, ninp, nhid): + """One message-passing neural network layer used by `MPMC_net`. + + Implements `torch_geometric.nn.MessagePassing`'s `message`/`update` + interface: each node aggregates messages from its neighbors (from the + `edge_index` graph built in `MPMC_net`) and updates its own features. + """ + + def __init__(self, ninp, nhid) -> None: super(MPNN_layer, self).__init__() self.ninp = ninp self.nhid = nhid @@ -29,24 +36,66 @@ def __init__(self, ninp, nhid): ) self.norm = InstanceNorm(nhid) - def forward(self, x, edge_index, batch): + def forward(self, x: torch.Tensor, edge_index: torch.Tensor, batch: torch.Tensor) -> torch.Tensor: + """Propagate messages over the graph and instance-normalize the result. + + Args: + x (torch.Tensor): Node features, shape `(num_nodes, ninp)`. + edge_index (torch.Tensor): Graph connectivity, shape `(2, num_edges)`. + batch (torch.Tensor): Batch assignment for each node, shape `(num_nodes,)`. + + Returns: + torch.Tensor: Updated, normalized node features, shape `(num_nodes, nhid)`. + """ x = self.propagate(edge_index, x=x) x = self.norm(x, batch) return x - def message(self, x_i, x_j): + def message(self, x_i: torch.Tensor, x_j: torch.Tensor) -> torch.Tensor: + """Compute the message sent from neighbor `x_j` to node `x_i`. + + Called internally by `MessagePassing.propagate`. + + Args: + x_i (torch.Tensor): Features of the target node, shape `(num_edges, ninp)`. + x_j (torch.Tensor): Features of the source (neighbor) node, shape `(num_edges, ninp)`. + + Returns: + torch.Tensor: Message for each edge, shape `(num_edges, nhid)`. + """ message = self.message_net_1(torch.cat((x_i, x_j), dim=-1)) message = self.message_net_2(message) return message - def update(self, message, x): + def update(self, message: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """Combine a node's aggregated message with its own features. + + Called internally by `MessagePassing.propagate`. + + Args: + message (torch.Tensor): Aggregated incoming message, shape `(num_nodes, nhid)`. + x (torch.Tensor): The node's own features, shape `(num_nodes, ninp)`. + + Returns: + torch.Tensor: Updated node features, shape `(num_nodes, nhid)`. + """ update = self.update_net_1(torch.cat((x, message), dim=-1)) update = self.update_net_2(update) return update class MPMC_net(nn.Module): - def __init__(self, dim, nhid, nlayers, nsamples, nbatch, radius, loss_fn, weights): + """Graph neural network that transforms random points into a + low-discrepancy point set by minimizing a discrepancy-based loss. + + Encodes `nbatch` independent batches of `nsamples` random points in + `dim` dimensions, passes them through `nlayers` `MPNN_layer`s connected + by a radius graph, decodes back to `dim` dimensions, and squashes to + `[0,1]^dim` via a sigmoid. Trained (elsewhere, e.g. `MPMC._train`) to + minimize `loss_fn` evaluated on the resulting points. + """ + + def __init__(self, dim, nhid, nlayers, nsamples, nbatch, radius, loss_fn, weights) -> None: super(MPMC_net, self).__init__() self.enc = nn.Linear(dim,nhid) self.convs = nn.ModuleList() @@ -78,7 +127,15 @@ def __init__(self, dim, nhid, nlayers, nsamples, nbatch, radius, loss_fn, weight else: raise ValueError(f"Loss function DNE: {loss_fn}") - def forward(self): + def forward(self) -> tuple[torch.Tensor, torch.Tensor]: + """Transform the stored random points and compute the discrepancy loss. + + Returns: + tuple[torch.Tensor, torch.Tensor]: `(loss, X)` where `loss` is + the scalar mean discrepancy loss (weighted by `self.weights` + if given) and `X` is the transformed point set, shape + `(nbatch, nsamples, dim)`. + """ X = self.x edge_index = self.edge_index diff --git a/qmcpy/discrete_distribution/mpmc/mpmc.py b/qmcpy/discrete_distribution/mpmc/mpmc.py index 59000b29b..2dc32a2f1 100644 --- a/qmcpy/discrete_distribution/mpmc/mpmc.py +++ b/qmcpy/discrete_distribution/mpmc/mpmc.py @@ -1,3 +1,4 @@ +from typing import Union from types import SimpleNamespace from io import BytesIO import os @@ -26,16 +27,16 @@ } class MPMC(AbstractLDDiscreteDistribution): - """ - Low-discrepancy generator trained by MPMC. Produces nbatch independent pointsets of size n in [0,1]^d. - + """Low-discrepancy generator trained by MPMC. Produces nbatch independent + pointsets of size n in [0,1]^d. + Requires PyTorch and PyTorch Geometric. Install with: - python -m pip install "qmcpy[mpmc]" - qmcpy-install-mpmc - - For GPU support or platform-specific details, see https://pytorch.org/get-started/locally/ - + python -m pip install "qmcpy[mpmc]" qmcpy-install-mpmc + + For GPU support or platform-specific details, see + https://pytorch.org/get-started/locally/ + Examples: >>> mpmc = MPMC( ... dimension=2, @@ -66,26 +67,69 @@ class MPMC(AbstractLDDiscreteDistribution): def __init__( self, - randomize='shift', - seed=None, - dimension=2, - replications=1, - d_max=None, - lr=1e-3, - nlayers=3, - weight_decay=1e-6, - nhid=32, - epochs=50_000, - start_reduce=40_000, - radius=0.35, - nbatch=1, - loss_fn='L2star', - weights=None, - use_pretrained=True, - pretrained_local_dir=None, - pretrained_base_url='https://github.com/QMCSoftware/LDData/tree/main/pregenerated_pointsets/mpmc', - prompt_on_missing=True, - ): + randomize: str = 'shift', + seed: Union[None, int, np.random.SeedSequence] = None, + dimension: int = 2, + replications: int = 1, + d_max: Union[None, int] = None, + lr: float = 1e-3, + nlayers: int = 3, + weight_decay: float = 1e-6, + nhid: int = 32, + epochs: int = 50_000, + start_reduce: int = 40_000, + radius: float = 0.35, + nbatch: int = 1, + loss_fn: str = 'L2star', + weights: Union[None, list, np.ndarray, torch.Tensor] = None, + use_pretrained: bool = True, + pretrained_local_dir: Union[None, str] = None, + pretrained_base_url: str = 'https://github.com/QMCSoftware/LDData/tree/main/pregenerated_pointsets/mpmc', + prompt_on_missing: bool = True, + ) -> None: + """Initialize an MPMC discrete distribution. + + Args: + randomize (str): `'shift'`/`'true'` for a random shift, or + `'false'`/`'none'`/`'no'` for no randomization. + seed (Union[None, int, np.random.SeedSequence]): Seed the random + number generator for reproducibility. + dimension (int): Dimension of the generated pointsets. + replications (int): Number of independent pointsets to + generate. Ignored if `nbatch` is set. + d_max (Union[None, int]): Unused; kept for backward compatibility. + `self.d_max` always mirrors `dimension`. + lr (float): Learning rate for the MPMC network optimizer. + nlayers (int): Number of message-passing layers in the MPMC + network. + weight_decay (float): Weight decay (L2 regularization) for the + optimizer. + nhid (int): Hidden dimension of the MPMC network layers. + epochs (int): Number of training epochs. + start_reduce (int): Epoch at which learning-rate reduction + begins. + radius (float): Radius parameter for the discrepancy loss. + nbatch (int): Number of independent pointsets to train and + generate (overrides `replications` when not `None`). + loss_fn (str): Name of the discrepancy loss to train against; one + of the keys in `qmcpy.discrete_distribution.mpmc.utils`'s + discrepancy registry (e.g. `'L2star'`), optionally suffixed + `'_weighted'`. + weights (Union[None, list, np.ndarray, torch.Tensor]): Per- + coordinate weights, required when `loss_fn` names a weighted + discrepancy (or supplying them switches `loss_fn` to its + weighted variant automatically). + use_pretrained (bool): If `True`, load a pretrained pointset + generator instead of training a new one, when one is + available for the requested `dimension`/`nbatch`. + pretrained_local_dir (Union[None, str]): Local directory to search for (and + cache) pretrained generators. Defaults to a package cache + directory when `None`. + pretrained_base_url (str): Base URL to download pretrained + generators from when not already cached locally. + prompt_on_missing (bool): If `True`, prompt interactively before + training a new generator when no pretrained one is found. + """ self.mimics = 'StdUniform' self.low_discrepancy = True @@ -298,9 +342,15 @@ def _spawn(self, child_seed, dimension): # Training # -------------------------- def _train(self, args: SimpleNamespace): - """ + """Train an MPMC network and return its generated pointsets. + + Args: + args (SimpleNamespace): Training configuration, carrying `dim`, + `nhid`, `nlayers`, `nsamples`, `nbatch`, `radius`, `loss_fn`, + `weights`, `lr`, `weight_decay`, `epochs`, and `start_reduce`. + Returns: - x (np.ndarray): shape `(nbatch, nsamples, dim)` + np.ndarray: shape `(nbatch, nsamples, dim)` """ model = MPMC_net( dim=args.dim, nhid=args.nhid, nlayers=args.nlayers, diff --git a/qmcpy/discrete_distribution/mpmc/utils.py b/qmcpy/discrete_distribution/mpmc/utils.py index ded0beb4e..9b684de97 100644 --- a/qmcpy/discrete_distribution/mpmc/utils.py +++ b/qmcpy/discrete_distribution/mpmc/utils.py @@ -1,9 +1,7 @@ import torch def _check_inputs(x, gamma=None): - """ - x: (B, N, d) in [0,1] - gamma: (d,) nonnegative weights (optional) + """x: (B, N, d) in [0,1] gamma: (d,) nonnegative weights (optional) """ if x.dim() != 3: raise ValueError(f"x must be (batch,N,d); got {tuple(x.shape)}") @@ -25,6 +23,14 @@ def _sqrt_safe(v): # L2 STAR (Warnock) # ---------------------------- def L2star(x: torch.Tensor) -> torch.Tensor: + """Warnock $L_2$ star discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x) t1 = (1.0 / 3.0) ** d p = torch.prod(1.0 - x**2, dim=2) @@ -35,6 +41,15 @@ def L2star(x: torch.Tensor) -> torch.Tensor: return _sqrt_safe(t1 - t2 + t3) def L2star_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + """Coordinate-weighted $L_2$ star discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + gamma (torch.Tensor): Non-negative coordinate weights of shape ``(d,)``, one per dimension. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x, gamma) g = gamma t1 = torch.prod(1.0 + g / 3.0) @@ -49,6 +64,14 @@ def L2star_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: # L2 EXTREME # ----------------------------------------- def L2ext(x: torch.Tensor) -> torch.Tensor: + """$L_2$ extreme discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x) t1 = (1.0 / 12.0) ** d p = torch.prod(0.5 * (x - x**2), dim=2) @@ -59,6 +82,15 @@ def L2ext(x: torch.Tensor) -> torch.Tensor: return _sqrt_safe(t1 - t2 + t3) def L2ext_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + """Coordinate-weighted $L_2$ extreme discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + gamma (torch.Tensor): Non-negative coordinate weights of shape ``(d,)``, one per dimension. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x, gamma) g = gamma t1 = torch.prod(1.0 + g / 12.0) @@ -73,6 +105,14 @@ def L2ext_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: # L2 PERIODIC # ----------------------------------------- def L2per(x: torch.Tensor) -> torch.Tensor: + """$L_2$ periodic discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x) t1 = (1.0 / 3.0) ** d xi, xj = _pairwise(x) @@ -82,6 +122,15 @@ def L2per(x: torch.Tensor) -> torch.Tensor: return _sqrt_safe(-t1 + t3) def L2per_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + """Coordinate-weighted $L_2$ periodic discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + gamma (torch.Tensor): Non-negative coordinate weights of shape ``(d,)``, one per dimension. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x, gamma) g = gamma t1 = torch.prod(1.0 + g / 3.0) @@ -95,6 +144,14 @@ def L2per_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: # L2 CENTERED # ----------------------------------------- def L2ctr(x: torch.Tensor) -> torch.Tensor: + """$L_2$ centered discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x) t1 = (1.0 / 12.0) ** d u = torch.abs(x - 0.5) @@ -106,6 +163,15 @@ def L2ctr(x: torch.Tensor) -> torch.Tensor: return _sqrt_safe(t1 - t2 + t3) def L2ctr_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + """Coordinate-weighted $L_2$ centered discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + gamma (torch.Tensor): Non-negative coordinate weights of shape ``(d,)``, one per dimension. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x, gamma) g = gamma t1 = torch.prod(1.0 + g / 12.0) @@ -121,6 +187,14 @@ def L2ctr_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: # L2 SYMMETRIC # ----------------------------------------- def L2sym(x: torch.Tensor) -> torch.Tensor: + """$L_2$ symmetric discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x) t1 = (1.0 / 12.0) ** d p = torch.prod(0.5 * (x - x**2), dim=2) @@ -131,6 +205,15 @@ def L2sym(x: torch.Tensor) -> torch.Tensor: return _sqrt_safe(t1 - t2 + t3) def L2sym_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + """Coordinate-weighted $L_2$ symmetric discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + gamma (torch.Tensor): Non-negative coordinate weights of shape ``(d,)``, one per dimension. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x, gamma) g = gamma t1 = torch.prod(1.0 + g / 12.0) @@ -145,6 +228,14 @@ def L2sym_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: # L2 MIXTURE # ----------------------------------------- def L2mix(x: torch.Tensor) -> torch.Tensor: + """$L_2$ mixture discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x) t1 = (7.0 / 12.0) ** d u = x - 0.5 @@ -158,6 +249,15 @@ def L2mix(x: torch.Tensor) -> torch.Tensor: return _sqrt_safe(t1 - t2 + t3) def L2mix_weighted(x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor: + """Coordinate-weighted $L_2$ mixture discrepancy of each point set in a batch. + + Args: + x (torch.Tensor): Points of shape ``(batch, N, d)`` with entries in $[0,1]$. + gamma (torch.Tensor): Non-negative coordinate weights of shape ``(d,)``, one per dimension. + + Returns: + torch.Tensor: Discrepancy of shape ``(batch,)``, one value per point set. + """ _, N, d = _check_inputs(x, gamma) g = gamma t1 = torch.prod(1.0 + (7.0 / 12.0) * g) diff --git a/qmcpy/fast_transform/ft.py b/qmcpy/fast_transform/ft.py index 3aada030e..b29eaa7d2 100644 --- a/qmcpy/fast_transform/ft.py +++ b/qmcpy/fast_transform/ft.py @@ -3,11 +3,11 @@ import itertools -def fftbr(x): - r""" - 1 dimensional Bit-Reversed-Order (BRO) Fast Fourier Transform (FFT) along the last dimension. - Requires the last dimension of x is already in BRO, so we can skip the first step of the decimation-in-time FFT. - Requires the size of the last dimension is a power of 2. +def fftbr(x: np.ndarray) -> np.ndarray: + r"""1 dimensional Bit-Reversed-Order (BRO) Fast Fourier Transform (FFT) + along the last dimension. Requires the last dimension of x is already in + BRO, so we can skip the first step of the decimation-in-time FFT. Requires + the size of the last dimension is a power of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -23,10 +23,11 @@ def fftbr(x): x (np.ndarray): Array of samples at which to run BRO-FFT. Returns: - y (np.ndarray): BRO-FFT values. + np.ndarray: BRO-FFT values. """ n = x.shape[-1] - assert n & (n - 1) == 0 # require n is a power of 2 + if not (n & (n - 1) == 0): # require n is a power of 2 + raise AssertionError m = int(np.log2(n)) shape = list(x.shape) ndim = x.ndim @@ -40,11 +41,11 @@ def fftbr(x): return scipy.fft.fft(xr, norm="ortho") -def ifftbr(x): - r""" - 1 dimensional Bit-Reversed-Order (BRO) Inverse Fast Fourier Transform (IFFT) along the last dimension. - Outputs an array in bit-reversed order, so we can skip the last step of the decimation-in-time IFFT. - Requires the size of the last dimension is a power of 2. +def ifftbr(x: np.ndarray) -> np.ndarray: + r"""1 dimensional Bit-Reversed-Order (BRO) Inverse Fast Fourier Transform + (IFFT) along the last dimension. Outputs an array in bit-reversed order, so + we can skip the last step of the decimation-in-time IFFT. Requires the size + of the last dimension is a power of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -60,10 +61,11 @@ def ifftbr(x): x (np.ndarray): Array of samples at which to run BRO-IFFT. Returns: - y (np.ndarray): BRO-IFFT values. + np.ndarray: BRO-IFFT values. """ n = x.shape[-1] - assert n & (n - 1) == 0 # require n is a power of 2 + if not (n & (n - 1) == 0): # require n is a power of 2 + raise AssertionError m = int(np.log2(n)) shape = list(x.shape) ndim = x.ndim @@ -76,10 +78,9 @@ def ifftbr(x): return xr -def fwht(x): - r""" - 1 dimensional Fast Walsh Hadamard Transform (FWHT) along the last dimension. - Requires the size of the last dimension is a power of 2. +def fwht(x: np.ndarray) -> np.ndarray: + r"""1 dimensional Fast Walsh Hadamard Transform (FWHT) along the last + dimension. Requires the size of the last dimension is a power of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -93,13 +94,14 @@ def fwht(x): x (np.ndarray): Array of samples at which to run FWHT. Returns: - y (np.ndarray): FWHT values. + np.ndarray: FWHT values. """ y = x.copy() + 0.0 n = x.shape[-1] if n <= 1: return y - assert n & (n - 1) == 0 # require n is a power of 2 + if not (n & (n - 1) == 0): # require n is a power of 2 + raise AssertionError m = int(np.log2(n)) it = np.arange(n, dtype=np.int64).reshape( [2] * m @@ -114,9 +116,9 @@ def fwht(x): return y -def omega_fwht(m): - r""" - A useful when efficiently updating FWHT values after doubling the sample size. +def omega_fwht(m: int) -> np.ndarray: + r"""A useful when efficiently updating FWHT values after doubling the + sample size. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -136,14 +138,14 @@ def omega_fwht(m): m (int): Size $2^m$ output. Returns: - y (np.ndarray): $\left(1\right)_{k=0}^{2^m}$. + np.ndarray: $\left(1\right)_{k=0}^{2^m}$. """ return np.ones(2**m) -def omega_fftbr(m): - r""" - A useful when efficiently updating FFT values after doubling the sample size. +def omega_fftbr(m: int) -> np.ndarray: + r"""A useful when efficiently updating FFT values after doubling the + sample size. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -163,6 +165,6 @@ def omega_fftbr(m): m (int): Size $2^m$ output. Returns: - y (np.ndarray): $\left(e^{- \pi \mathrm{i} k / 2^m}\right)_{k=0}^{2^m}$. + np.ndarray: $\left(e^{- \pi \mathrm{i} k / 2^m}\right)_{k=0}^{2^m}$. """ return np.exp(-np.pi * 1j * np.arange(2**m) / 2**m) diff --git a/qmcpy/fast_transform/ft_pytorch.py b/qmcpy/fast_transform/ft_pytorch.py index b06b2a7da..0819423ad 100644 --- a/qmcpy/fast_transform/ft_pytorch.py +++ b/qmcpy/fast_transform/ft_pytorch.py @@ -1,13 +1,15 @@ +from typing import Union import torch import numpy as np import itertools -def fftbr_torch(x): - r""" - Torch implementation of the 1 dimensional Bit-Reversed-Order (BRO) Fast Fourier Transform (FFT) along the last dimension. - Requires the last dimension of x is already in BRO, so we can skip the first step of the decimation-in-time FFT. - Requires the size of the last dimension is a power of 2. +def fftbr_torch(x: torch.Tensor) -> torch.Tensor: + r"""Torch implementation of the 1 dimensional Bit-Reversed-Order (BRO) + Fast Fourier Transform (FFT) along the last dimension. Requires the last + dimension of x is already in BRO, so we can skip the first step of the + decimation-in-time FFT. Requires the size of the last dimension is a power + of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -36,10 +38,11 @@ def fftbr_torch(x): x (torch.Tensor): Array of samples at which to run BRO-FFT. Returns: - y (torch.Tensor): BRO-FFT values. + torch.Tensor: BRO-FFT values. """ n = x.size(-1) - assert n & (n - 1) == 0 # require n is a power of 2 + if not (n & (n - 1) == 0): # require n is a power of 2 + raise AssertionError m = int(np.log2(n)) shape = list(x.shape) ndim = x.ndim @@ -53,11 +56,12 @@ def fftbr_torch(x): return torch.fft.fft(xr, norm="ortho") -def ifftbr_torch(x): - r""" - Torch implementation of the 1 dimensional Bit-Reversed-Order (BRO) Inverse Fast Fourier Transform (IFFT) along the last dimension. - Outputs an array in bit-reversed order, so we can skip the last step of the decimation-in-time IFFT. - Requires the size of the last dimension is a power of 2. +def ifftbr_torch(x: torch.Tensor) -> torch.Tensor: + r"""Torch implementation of the 1 dimensional Bit-Reversed-Order (BRO) + Inverse Fast Fourier Transform (IFFT) along the last dimension. Outputs an + array in bit-reversed order, so we can skip the last step of the + decimation-in-time IFFT. Requires the size of the last dimension is a power + of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -86,10 +90,11 @@ def ifftbr_torch(x): x (torch.Tensor): Array of samples at which to run BRO-IFFT. Returns: - y (torch.Tensor): BRO-IFFT values. + torch.Tensor: BRO-IFFT values. """ n = x.size(-1) - assert n & (n - 1) == 0 # require n is a power of 2 + if not (n & (n - 1) == 0): # require n is a power of 2 + raise AssertionError m = int(np.log2(n)) shape = list(x.shape) ndim = x.ndim @@ -107,7 +112,8 @@ def _fwht_torch(x): n = x.size(-1) if n <= 1: return y - assert n & (n - 1) == 0 # require n is a power of 2 + if not (n & (n - 1) == 0): # require n is a power of 2 + raise AssertionError m = int(np.log2(n)) it = torch.arange(n, dtype=torch.int64, device=x.device).reshape( [2] * m @@ -134,10 +140,10 @@ def backward(ctx, dx): return _fwht_torch(dx) -def fwht_torch(x): - r""" - Torch implementation of the 1 dimensional Fast Walsh Hadamard Transform (FWHT) along the last dimension. - Requires the size of the last dimension is a power of 2. +def fwht_torch(x: torch.Tensor) -> torch.Tensor: + r"""Torch implementation of the 1 dimensional Fast Walsh Hadamard + Transform (FWHT) along the last dimension. Requires the size of the last + dimension is a power of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -163,14 +169,14 @@ def fwht_torch(x): x (torch.Tensor): Array of samples at which to run FWHT. Returns: - y (torch.Tensor): FWHT values. + torch.Tensor: FWHT values. """ return _FWHTB2Ortho.apply(x) -def omega_fwht_torch(m, device=None): - r""" - Torch implementation useful when efficiently updating FWHT values after doubling the sample size. +def omega_fwht_torch(m: int, device: Union[None, torch.device] = None) -> np.ndarray: + r"""Torch implementation useful when efficiently updating FWHT values + after doubling the sample size. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -188,18 +194,20 @@ def omega_fwht_torch(m, device=None): Args: m (int): Size $2^m$ output. + device (Union[None, torch.device]): Device to place the output tensor on. + Defaults to CPU. Returns: - y (np.ndarray): $\left(1\right)_{k=0}^{2^m}$. + np.ndarray: $\left(1\right)_{k=0}^{2^m}$. """ if device is None: device = "cpu" return torch.ones(2**m, device=device) -def omega_fftbr_torch(m, device=None): - r""" - Torch implementation useful when efficiently updating FFT values after doubling the sample size. +def omega_fftbr_torch(m: int, device: Union[None, torch.device] = None) -> np.ndarray: + r"""Torch implementation useful when efficiently updating FFT values after + doubling the sample size. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -217,9 +225,11 @@ def omega_fftbr_torch(m, device=None): Args: m (int): Size $2^m$ output. + device (Union[None, torch.device]): Device to place the output tensor on. + Defaults to CPU. Returns: - y (np.ndarray): $\left(e^{- \pi \mathrm{i} k / 2^m}\right)_{k=0}^{2^m}$. + np.ndarray: $\left(e^{- \pi \mathrm{i} k / 2^m}\right)_{k=0}^{2^m}$. """ if device is None: device = "cpu" diff --git a/qmcpy/fast_transform/ft_qmctoolscl.py b/qmcpy/fast_transform/ft_qmctoolscl.py index b36163c49..32e6a176e 100644 --- a/qmcpy/fast_transform/ft_qmctoolscl.py +++ b/qmcpy/fast_transform/ft_qmctoolscl.py @@ -15,15 +15,17 @@ def _parse_ft_input(x): n = shape[-1] x = x.reshape(-1, n) d = x.shape[0] - assert (n & (n - 1)) == 0 # require n is 0 or a power of 2 + if not ((n & (n - 1)) == 0): # require n is 0 or a power of 2 + raise AssertionError return x, shape, d, n, n // 2 -def fftbr_qmctoolscl(x): - r""" - QMCToolsCL implementation of the 1 dimensional Bit-Reversed-Order (BRO) Fast Fourier Transform (FFT) along the last dimension. - Requires the last dimension of x is already in BRO, so we can skip the first step of the decimation-in-time FFT. - Requires the size of the last dimension is a power of 2. +def fftbr_qmctoolscl(x: np.ndarray) -> np.ndarray: + r"""QMCToolsCL implementation of the 1 dimensional Bit-Reversed-Order + (BRO) Fast Fourier Transform (FFT) along the last dimension. Requires the + last dimension of x is already in BRO, so we can skip the first step of the + decimation-in-time FFT. Requires the size of the last dimension is a power + of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -39,7 +41,7 @@ def fftbr_qmctoolscl(x): x (np.ndarray): Array of samples at which to run BRO-FFT. Returns: - y (np.ndarray): BRO-FFT values. + np.ndarray: BRO-FFT values. """ x, shape, d, n, n_half = _parse_ft_input(x) if n <= 1: @@ -53,11 +55,12 @@ def fftbr_qmctoolscl(x): return xc.reshape(shape) -def ifftbr_qmctoolscl(x): - r""" - QMCToolsCL implementation of the 1 dimensional Bit-Reversed-Order (BRO) Inverse Fast Fourier Transform (IFFT) along the last dimension. - Outputs an array in bit-reversed order, so we can skip the last step of the decimation-in-time IFFT. - Requires the size of the last dimension is a power of 2. +def ifftbr_qmctoolscl(x: np.ndarray) -> np.ndarray: + r"""QMCToolsCL implementation of the 1 dimensional Bit-Reversed-Order + (BRO) Inverse Fast Fourier Transform (IFFT) along the last dimension. + Outputs an array in bit-reversed order, so we can skip the last step of the + decimation-in-time IFFT. Requires the size of the last dimension is a power + of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -73,7 +76,7 @@ def ifftbr_qmctoolscl(x): x (np.ndarray): Array of samples at which to run BRO-IFFT. Returns: - y (np.ndarray): BRO-IFFT values. + np.ndarray: BRO-IFFT values. """ x, shape, d, n, n_half = _parse_ft_input(x) if n <= 1: @@ -87,10 +90,10 @@ def ifftbr_qmctoolscl(x): return xc.reshape(shape) -def fwht_qmctoolscl(x): - r""" - QMCToolsCL implementation of the 1 dimensional Fast Walsh Hadamard Transform (FWHT) along the last dimension. - Requires the size of the last dimension is a power of 2. +def fwht_qmctoolscl(x: np.ndarray) -> np.ndarray: + r"""QMCToolsCL implementation of the 1 dimensional Fast Walsh Hadamard + Transform (FWHT) along the last dimension. Requires the size of the last + dimension is a power of 2. Examples: >>> rng = np.random.Generator(np.random.SFC64(11)) @@ -104,7 +107,7 @@ def fwht_qmctoolscl(x): x (np.ndarray): Array of samples at which to run FWHT. Returns: - y (np.ndarray): FWHT values. + np.ndarray: FWHT values. """ x, shape, d, n, n_half = _parse_ft_input(x) if n <= 1: diff --git a/qmcpy/integrand/abstract_integrand.py b/qmcpy/integrand/abstract_integrand.py index 4d838e81f..44a9495d4 100644 --- a/qmcpy/integrand/abstract_integrand.py +++ b/qmcpy/integrand/abstract_integrand.py @@ -1,3 +1,4 @@ +from typing import Union from ..util import MethodImplementationError, _univ_repr, ParameterError from ..true_measure.abstract_true_measure import AbstractTrueMeasure from ..discrete_distribution.abstract_discrete_distribution import ( @@ -11,9 +12,16 @@ class AbstractIntegrand(object): + """Base class for integrands. + + An integrand pairs a function $g$ with the true measure its argument is + distributed by, and exposes $f$, the composition that a stopping criterion + samples. Subclasses implement ``g``. + """ + + def __init__(self, dimension_indv: tuple, dimension_comb: tuple, parallel: int, threadpool: bool = False) -> None: + r"""Initialize an AbstractIntegrand integrand. - def __init__(self, dimension_indv, dimension_comb, parallel, threadpool=False): - r""" Args: dimension_indv (tuple): Individual solution shape. dimension_comb (tuple): Combined solution shape. @@ -22,7 +30,8 @@ def __init__(self, dimension_indv, dimension_comb, parallel, threadpool=False): - When `parallel = 0` or `parallel = 1` then function evaluation is done in serial fashion. - `parallel > 1` specifies the number of processes used by `multiprocessing.Pool` or `multiprocessing.pool.ThreadPool`. - Setting `parallel=True` is equivalent to `parallel = os.cpu_count()`. + Setting `parallel=True` is equivalent to `parallel = + os.cpu_count()`. threadpool (bool): When `parallel > 1`: - Setting `threadpool = True` will use `multiprocessing.pool.ThreadPool`. @@ -65,7 +74,8 @@ def __init__(self, dimension_indv, dimension_comb, parallel, threadpool=False): self.parameters = [] if not hasattr(self, "multilevel"): self.multilevel = False - assert isinstance(self.multilevel, bool) + if not (isinstance(self.multilevel, bool)): + raise AssertionError if not hasattr(self, "max_level"): self.max_level = np.inf if not hasattr(self, "discrete_distrib"): @@ -79,7 +89,7 @@ def __init__(self, dimension_indv, dimension_comb, parallel, threadpool=False): ) self.EPS = np.finfo(float).eps - def __call__(self, n=None, n_min=None, n_max=None, warn=True): + def __call__(self, n: Union[None, int] = None, n_min: Union[None, int] = None, n_max: Union[None, int] = None, warn: bool = True) -> np.ndarray: r""" - If just `n` is supplied, generate samples from the sequence at indices 0,...,`n`-1. - If `n_min` and `n_max` are supplied, generate samples from the sequence at indices `n_min`,...,`n_max`-1. @@ -92,62 +102,82 @@ def __call__(self, n=None, n_min=None, n_max=None, warn=True): warn (bool): If `False`, disable warnings when generating samples. Returns: - t (np.ndarray): Samples from the sequence. + np.ndarray: Samples from the sequence. - If `replications` is `None` then this will be of size (`n_max`-`n_min`) $\times$ `dimension` - If `replications` is a positive int, then `t` will be of size `replications` $\times$ (`n_max`-`n_min`) $\times$ `dimension` - weights (np.ndarray): Only returned when `return_weights=True`. The Jacobian weights for the transformation """ return self.gen_samples(n=n, n_min=n_min, n_max=n_max, warn=warn) def gen_samples( - self, n=None, n_min=None, n_max=None, return_weights=False, warn=True - ): + self, n: Union[None, int] = None, n_min: Union[None, int] = None, n_max: Union[None, int] = None, return_weights: bool = False, warn: bool = True + ) -> np.ndarray: + """Generate discrete distribution samples and evaluate the integrand at them. + + Args: + n (Union[None, int]): Number of points, taken from index ``0`` to ``n``. + n_min (Union[None, int]): Starting index of the sequence. + n_max (Union[None, int]): Final index of the sequence. + return_weights (bool): Accepted for API consistency; unused here. + warn (bool): If ``False``, disable warnings while generating samples. + + Returns: + np.ndarray: Integrand values at the generated points. + """ x = self.discrete_distrib(n=n, n_min=n_min, n_max=n_max, warn=warn) y = self.f(x) return y - def g(self, t, *args, **kwargs): - r""" - *Abstract method* implementing the integrand as a function of the true measure. + def g(self, t: np.ndarray, *args: tuple, **kwargs: dict) -> np.ndarray: + r"""*Abstract method* implementing the integrand as a function of the + true measure. Args: t (np.ndarray): Inputs with shape `(*batch_shape, d)`. - args (tuple): positional arguments to `g`. - kwargs (dict): keyword arguments to `g`. - - Some algorithms will additionally try to pass in a `compute_flags` keyword argument. - This `np.ndarray` are flags indicating which outputs require evaluation. - For example, if the vector function has 3 outputs and `compute_flags = [False, True, False]`, - then the function is only required to evaluate the second output and may leave the remaining outputs as `np.nan` values, - i.e., the outputs corresponding to `compute_flags` which are `False` will not be used in the computation. + *args (tuple): positional arguments to `g`. + **kwargs (dict): keyword arguments to `g`. + + Some algorithms will additionally try to pass in a + `compute_flags` keyword argument. This `np.ndarray` are flags + indicating which outputs require evaluation. For example, if + the vector function has 3 outputs and `compute_flags = [False, + True, False]`, then the function is only required to evaluate + the second output and may leave the remaining outputs as + `np.nan` values, i.e., the outputs corresponding to + `compute_flags` which are `False` will not be used in the + computation. Returns: - y (np.ndarray): function evaluations with shape `(*batch_shape, *dimension_indv)` where `dimension_indv` is the shape of the function outputs. + np.ndarray: function evaluations with shape `(*batch_shape, *dimension_indv)` + where `dimension_indv` is the shape of the function outputs. """ raise MethodImplementationError(self, "g") - def f(self, x, *args, **kwargs): - r""" - Function to evaluate the transformed integrand as a function of the discrete distribution. - Automatically applies the transformation determined by the true measure. + def f(self, x: np.ndarray, *args: tuple, **kwargs: dict) -> np.ndarray: + r"""Function to evaluate the transformed integrand as a function of + the discrete distribution. Automatically applies the transformation + determined by the true measure. Args: x (np.ndarray): Inputs with shape `(*batch_shape, d)`. - args (tuple): positional arguments to `g`. - kwargs (dict): keyword arguments to `g`. - - Some algorithms will additionally try to pass in a `compute_flags` keyword argument. - This `np.ndarray` are flags indicating which outputs require evaluation. - For example, if the vector function has 3 outputs and `compute_flags = [False, True, False]`, - then the function is only required to evaluate the second output and may leave the remaining outputs as `np.nan` values, - i.e., the outputs corresponding to `compute_flags` which are `False` will not be used in the computation. - - The keyword argument `periodization_transform`, a string, specifies a periodization transform. - Options are: + *args (tuple): positional arguments to `g`. + **kwargs (dict): keyword arguments to `g`. + + Some algorithms will additionally try to pass in a + `compute_flags` keyword argument. This `np.ndarray` are flags + indicating which outputs require evaluation. For example, if + the vector function has 3 outputs and `compute_flags = [False, + True, False]`, then the function is only required to evaluate + the second output and may leave the remaining outputs as + `np.nan` values, i.e., the outputs corresponding to + `compute_flags` which are `False` will not be used in the + computation. + + The keyword argument `periodization_transform`, a string, + specifies a periodization transform. Options are: - `False`: No periodizing transform, $\psi(x) = x$. - - `'BAKER'`: Baker tansform $\psi(x) = 1-2\lvert x-1/2 \rvert$. + - `'BAKER'`: Baker transform $\psi(x) = 1-2\lvert x-1/2 \rvert$. - `'C0'`: $C^0$ transform $\psi(x) = 3x^2-2x^3$. - `'C1'`: $C^1$ transform $\psi(x) = x^3(10-15x+6x^2)$. - `'C1SIN'`: Sidi $C^1$ transform $\psi(x) = x-\sin(2 \pi x)/(2 \pi)$. @@ -155,7 +185,8 @@ def f(self, x, *args, **kwargs): - `'C3SIN'`: Sidi $C^3$ transform $\psi(x) = (12\pi x-8\sin(2 \pi x) + \sin(4 \pi x))/(12 \pi)$. Returns: - y (np.ndarray): function evaluations with shape `(*batch_shape, *dimension_indv)` where `dimension_indv` is the shape of the function outputs. + np.ndarray: function evaluations with shape `(*batch_shape, *dimension_indv)` + where `dimension_indv` is the shape of the function outputs. """ if "periodization_transform" in kwargs: periodization_transform = kwargs["periodization_transform"] @@ -219,8 +250,10 @@ def f(self, x, *args, **kwargs): if periodization_transform in ["C1", "C1SIN", "C2SIN", "C3SIN"]: xp[xp <= 0] = self.EPS xp[xp >= 1] = 1 - self.EPS - assert wp.shape == batch_shape - assert xp.shape == x.shape + if not (wp.shape == batch_shape): + raise AssertionError + if not (xp.shape == x.shape): + raise AssertionError # function evaluation with chain rule i = (None,) * d_indv_ndim + (...,) if self.true_measure == self.true_measure.transform: @@ -228,25 +261,33 @@ def f(self, x, *args, **kwargs): xtf = self.true_measure._jacobian_transform_r( xp, return_weights=False ) # get transformed samples, equivalent to self.true_measure._transform_r(x) - assert xtf.shape == xp.shape + if not (xtf.shape == xp.shape): + raise AssertionError y = self._g(xtf, *args, **kwargs) else: # using importance sampling --> need to compute pdf, jacobian(s), and weight explicitly pdf = self.discrete_distrib.pdf(xp) # pdf of samples - assert pdf.shape == batch_shape + if not (pdf.shape == batch_shape): + raise AssertionError xtf, jacobians = self.true_measure.transform._jacobian_transform_r( xp, return_weights=True ) # compute recursive transform+jacobian - assert xtf.shape == xp.shape - assert jacobians.shape == batch_shape + if not (xtf.shape == xp.shape): + raise AssertionError + if not (jacobians.shape == batch_shape): + raise AssertionError weight = self.true_measure._weight(xtf) # weight based on the true measure - assert weight.shape == batch_shape + if not (weight.shape == batch_shape): + raise AssertionError gvals = self._g(xtf, *args, **kwargs) - assert gvals.shape == (self.d_indv + batch_shape) + if not (gvals.shape == (self.d_indv + batch_shape)): + raise AssertionError y = gvals * weight[i] / pdf[i] * jacobians[i] - assert y.shape == (self.d_indv + batch_shape) + if not (y.shape == (self.d_indv + batch_shape)): + raise AssertionError # account for periodization weight y = y * wp[i] - assert y.shape == (self.d_indv + batch_shape) + if not (y.shape == (self.d_indv + batch_shape)): + raise AssertionError return y def _g(self, t, *args, **kwargs): @@ -263,10 +304,11 @@ def _g(self, t, *args, **kwargs): else: y = self._g2(t, comb_args=(args, kwargs)) expected_y_shape = self.d_indv + t.shape[:-1] - assert y.shape == expected_y_shape, "expected y.shape to be %s but got %s" % ( - str(expected_y_shape), - str(y.shape), - ) + if not (y.shape == expected_y_shape): + raise AssertionError("expected y.shape to be %s but got %s" % ( + str(expected_y_shape), + str(y.shape), + )) return y def _g2(self, t, comb_args=((), {})): @@ -282,22 +324,22 @@ def _g2(self, t, comb_args=((), {})): raise e return y - def bound_fun(self, bound_low, bound_high): - """ - Compute the bounds on the combined function based on bounds for the - individual functions. + def bound_fun(self, bound_low: np.ndarray, bound_high: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Compute the bounds on the combined function based on bounds for + the individual functions. - Defaults to the identity where we essentially - do not combine integrands, but instead integrate each function - individually. + Defaults to the identity where we essentially do not combine + integrands, but instead integrate each function individually. Args: - bound_low (np.ndarray): Lower bounds on individual estimates with shape `integrand.d_indv`. - bound_high (np.ndarray): Upper bounds on individual estimates with shape `integrand.d_indv`. + bound_low (np.ndarray): Lower bounds on individual estimates with + shape `integrand.d_indv`. + bound_high (np.ndarray): Upper bounds on individual estimates with + shape `integrand.d_indv`. Returns: - comb_bound_low (np.ndarray): Lower bounds on combined estimates with shape `integrand.d_comb`. - comb_bound_high (np.ndarray): Upper bounds on combined estimates with shape `integrand.d_comb`. + tuple[np.ndarray, np.ndarray]: Lower and upper bounds on the + combined estimates, respectively, each with shape `integrand.d_comb`. """ if self.d_indv != self.d_comb: raise ParameterError( @@ -310,19 +352,25 @@ def bound_fun(self, bound_low, bound_high): ) return bound_low, bound_high - def dependency(self, comb_flags): - """ - Takes a vector of indicators of weather of not the error bound is satisfied for combined integrands and returns flags for individual integrands. + def dependency(self, comb_flags: np.ndarray) -> np.ndarray: + """Takes a vector of indicators of weather of not the error bound is + satisfied for combined integrands and returns flags for individual + integrands. - For example, if we are taking the ratio of 2 individual integrands, then getting `comb_flags=True` means the ratio - has not been approximated to within the tolerance, so the dependency function should return `indv_flags=[True,True]` - indicating that both the numerator and denominator integrands need to be better approximated. + For example, if we are taking the ratio of 2 individual integrands, + then getting `comb_flags=True` means the ratio has not been + approximated to within the tolerance, so the dependency function should + return `indv_flags=[True,True]` indicating that both the numerator and + denominator integrands need to be better approximated. Args: - comb_flags (np.ndarray): Flags of shape `integrand.d_comb` indicating whether the combined outputs are insufficiently approximated. + comb_flags (np.ndarray): Flags of shape `integrand.d_comb` + indicating whether the combined outputs are insufficiently + approximated. Returns: - indv_flags (np.ndarray): Flags of shape `integrand.d_indv` indicating whether the individual integrands require additional sampling. + np.ndarray: Flags of shape `integrand.d_indv` indicating whether the individual + integrands require additional sampling. """ return ( comb_flags @@ -330,19 +378,20 @@ def dependency(self, comb_flags): else np.tile((comb_flags == False).any(), self.d_indv) ) - def spawn(self, levels): - r""" - Spawn new instances of the current integrand at different levels with new seeds. - Used by multi-level QMC algorithms which require integrands at multiple levels. + def spawn(self, levels: np.ndarray) -> list: + r"""Spawn new instances of the current integrand at different levels + with new seeds. Used by multi-level QMC algorithms which require + integrands at multiple levels. - Note: - Use `replications` instead of using `spawn` when possible, e.g., when spawning copies which all have the same level. + Notes: + Use `replications` instead of using `spawn` when possible, e.g., + when spawning copies which all have the same level. Args: levels (np.ndarray): Levels at which to spawn new integrands. Returns: - spawned_integrand (list): Integrands with new true measures and discrete distributions. + list: Integrands with new true measures and discrete distributions. """ levels = np.array([levels]) if np.isscalar(levels) else np.array(levels) if (levels > self.max_level).any(): @@ -355,18 +404,18 @@ def spawn(self, levels): spawned_integrand[l] = self._spawn(level, tm_spawns[l]) return spawned_integrand - def dimension_at_level(self, level): - """ - *Abstract method* which returns the dimension of the generator required for a given level. + def dimension_at_level(self, level: int) -> int: + """*Abstract method* which returns the dimension of the generator + required for a given level. - Note: + Notes: Only used for multilevel problems. Args: level (int): Level at which to return the dimension. Returns: - d (int): Dimension at the given input level. + int: Dimension at the given input level. """ return self.d diff --git a/qmcpy/integrand/bayesian_lr_coeffs.py b/qmcpy/integrand/bayesian_lr_coeffs.py index e9003fbb3..c732cf357 100644 --- a/qmcpy/integrand/bayesian_lr_coeffs.py +++ b/qmcpy/integrand/bayesian_lr_coeffs.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 #pylint: disable=unused-import from ..true_measure import Gaussian @@ -6,8 +11,8 @@ class BayesianLRCoeffs(AbstractIntegrand): - r""" - Logistic Regression Coefficients computed as the posterior mean in a Bayesian framework. + r"""Logistic Regression Coefficients computed as the posterior mean in a + Bayesian framework. Examples: >>> integrand = BayesianLRCoeffs(DigitalNetB2(3,seed=7),feature_array=np.arange(8).reshape((4,2)),response_vector=[0,0,1,1]) @@ -33,21 +38,28 @@ class BayesianLRCoeffs(AbstractIntegrand): """ def __init__( - self, sampler, feature_array, response_vector, prior_mean=0, prior_covariance=10 - ): - r""" + self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], feature_array: np.ndarray, response_vector: np.ndarray, prior_mean: Union[float, np.ndarray] = 0, prior_covariance: Union[float, np.ndarray] = 10 + ) -> None: + r"""Initialize a BayesianLRCoeffs integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. - feature_array (np.ndarray): Array of features with shape $(N,d-1)$ where $N$ is the number of observations and $d$ is the dimension. - response_vector (np.ndarray): Binary responses vector of length $N$. - prior_mean (np.ndarray): Length $d$ vector of prior means, one for each coefficient. + feature_array (np.ndarray): Array of features with shape $(N,d-1)$ + where $N$ is the number of observations and $d$ is the + dimension. + response_vector (np.ndarray): Binary responses vector of length + $N$. + prior_mean (Union[float, np.ndarray]): Length $d$ vector of prior means, one for + each coefficient. - The first $d-1$ inputs correspond to the $d-1$ features. - The last input corresponds to the intercept coefficient. - prior_covariance (np.ndarray): Prior covariance array with shape $(d,d)$ d x d where indexing is consistent with the prior mean. + prior_covariance (Union[float, np.ndarray]): Prior covariance array with shape + $(d,d)$ d x d where indexing is consistent with the prior mean. """ self.prior_mean = prior_mean self.prior_covariance = prior_covariance @@ -77,7 +89,16 @@ def __init__( parallel=False, ) - def g(self, x): + def g(self, x: np.ndarray) -> np.ndarray: + """Evaluate the unnormalized posterior numerator and denominator. + + Args: + x (np.ndarray): Coefficient vectors, coefficients along the last axis. + + Returns: + np.ndarray: Stacked numerator (coefficient-weighted likelihood) and + denominator (likelihood), whose ratio is the posterior mean. + """ z = np.einsum("...j,ij->...i", x, self.feature_array) z1 = z * self.response_vector with np.errstate(over="ignore"): @@ -96,7 +117,17 @@ def _spawn(self, level, sampler): prior_covariance=self.prior_covariance, ) - def bound_fun(self, bound_low, bound_high): + def bound_fun(self, bound_low: np.ndarray, bound_high: np.ndarray) -> tuple: + """Combine numerator and denominator bounds into bounds on their ratio. + + Args: + bound_low (np.ndarray): Lower bounds on the numerator and denominator. + bound_high (np.ndarray): Upper bounds on the numerator and denominator. + + Returns: + tuple: Lower and upper bounds on the ratio, infinite where the + denominator interval straddles zero. + """ num_bounds_low, den_bounds_low = bound_low[0], bound_low[1] num_bounds_high, den_bounds_high = bound_high[0], bound_high[1] comb_bounds_low = np.minimum.reduce( @@ -119,5 +150,13 @@ def bound_fun(self, bound_low, bound_high): comb_bounds_low[violated], comb_bounds_high[violated] = -np.inf, np.inf return comb_bounds_low, comb_bounds_high - def dependency(self, comb_flags): + def dependency(self, comb_flags: np.ndarray) -> np.ndarray: + """Map combined-output flags onto the individual outputs they require. + + Args: + comb_flags (np.ndarray): Flags for the combined outputs. + + Returns: + np.ndarray: Flags for the numerator and denominator outputs. + """ return np.vstack((comb_flags, comb_flags)) diff --git a/qmcpy/integrand/box_integral.py b/qmcpy/integrand/box_integral.py index 5b930408f..aaf93720b 100644 --- a/qmcpy/integrand/box_integral.py +++ b/qmcpy/integrand/box_integral.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 from ..true_measure import Uniform @@ -5,10 +10,10 @@ class BoxIntegral(AbstractIntegrand): - r""" - Box integral from [1], see also + r"""Box integral from [1], see also - $$B_s(\boldsymbol{t}) = \left(\sum_{j=1}^d t_j^2 \right)^{s/2}, \qquad \boldsymbol{T} \sim \mathcal{U}[0,1]^d.$$ + $$B_s(\boldsymbol{t}) = \left(\sum_{j=1}^d t_j^2 \right)^{s/2}, \qquad + \boldsymbol{T} \sim \mathcal{U}[0,1]^d.$$ Examples: Scalar `s` @@ -64,18 +69,22 @@ class BoxIntegral(AbstractIntegrand): [https://www.davidhbailey.com/dhbpapers/boxintegrals.pdf](https://www.davidhbailey.com/dhbpapers/boxintegrals.pdf) """ - def __init__(self, sampler, s=1): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], s: Union[float, np.ndarray] = 1) -> None: + r"""Initialize a BoxIntegral integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. - s (Union[float, np.ndarray]): `s` parameter or parameters. The output shape of `g` is the shape of `s`. + s (Union[float, np.ndarray]): `s` parameter or parameters. The + output shape of `g` is the shape of `s`. """ self.parameters = ["s"] self.s = np.array(s) - assert self.s.size > 0 + if not (self.s.size > 0): + raise AssertionError self.sampler = sampler self.true_measure = Uniform(self.sampler) self.s_over_2 = self.s / 2 @@ -83,7 +92,16 @@ def __init__(self, sampler, s=1): dimension_indv=self.s.shape, dimension_comb=self.s.shape, parallel=False ) - def g(self, t, **kwargs): + def g(self, t: np.ndarray, **kwargs: dict) -> np.ndarray: + r"""Evaluate the box integral function. + + Args: + t (np.ndarray): Points in the unit cube, dimensions along the last axis. + **kwargs (dict): Unused; accepted for API consistency. + + Returns: + np.ndarray: $\lVert t \rVert_2^s$ for each exponent $s$. + """ sum_squares = (t**2).sum(-1) y = sum_squares ** self.s_over_2[(...,) + (None,) * sum_squares.ndim] return y diff --git a/qmcpy/integrand/custom_fun.py b/qmcpy/integrand/custom_fun.py index b7d730e2b..a8b0fbbfb 100644 --- a/qmcpy/integrand/custom_fun.py +++ b/qmcpy/integrand/custom_fun.py @@ -1,3 +1,5 @@ +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union, Callable from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 #pylint: disable=unused-import from ..true_measure import Gaussian, Uniform #pylint: disable=unused-import @@ -5,13 +7,13 @@ class CustomFun(AbstractIntegrand): - r""" - User supplied integrand $g$. In the following example we implement + r"""User supplied integrand $g$. In the following example we implement Examples: First we will implement - $$g(\boldsymbol{t}) = t_1^2t_2, \qquad \boldsymbol{T}=(T_1,T_2) \sim \mathcal{N}((1,2)^T,\mathsf{I}).$$ + $$g(\boldsymbol{t}) = t_1^2t_2, \qquad \boldsymbol{T}=(T_1,T_2) \sim + \mathcal{N}((1,2)^T,\mathsf{I}).$$ >>> integrand = CustomFun( ... true_measure = Gaussian(DigitalNetB2(2,seed=7),mean=[1,2]), @@ -36,7 +38,10 @@ class CustomFun(AbstractIntegrand): Next we will implement the multi-output function - $$g(\boldsymbol{t}) = \begin{pmatrix} \sin(t_1)\cos(t_2) \\ \cos(t_1)\sin(t_2) \\ \sin(t_1)+\cos(t_2) \\ \cos(t_1)+\sin(t_2) \end{pmatrix} \qquad \boldsymbol{T}=(T_1,T_2) \sim \mathcal{U}[0,2\pi]^2.$$ + $$g(\boldsymbol{t}) = \begin{pmatrix} \sin(t_1)\cos(t_2) \\ + \cos(t_1)\sin(t_2) \\ \sin(t_1)+\cos(t_2) \\ \cos(t_1)+\sin(t_2) + \end{pmatrix} \qquad \boldsymbol{T}=(T_1,T_2) \sim + \mathcal{U}[0,2\pi]^2.$$ >>> def g(t): ... t1,t2 = t[...,0],t[...,1] @@ -59,8 +64,11 @@ class CustomFun(AbstractIntegrand): ... y.mean(-1) array([8.18e-04, 1.92e-06, -2.26e-10, 5.05e-07]) - Stopping criterion which supporting vectorized outputs may pass in Boolean `compute_flags` with `dimension_indv` shape indicating which output need to evaluated, - i.e. where `compute_flags` is `False` we do not need to evaluate the integrand. We have not used this in inexpensive example above. + Stopping criterion which supporting vectorized outputs may pass in + Boolean `compute_flags` with `dimension_indv` shape indicating which + output need to evaluated, + i.e., where `compute_flags` is `False` we do not need to evaluate + the integrand. We have not used this in inexpensive example above. With independent replications @@ -80,24 +88,27 @@ class CustomFun(AbstractIntegrand): >>> with np.printoptions(formatter={"float": lambda x: "%.2e"%x}): ... muhats.mean(-1) array([3.83e-03, -6.78e-03, -1.56e-03, -5.65e-04]) - """ - def __init__(self, true_measure, g, dimension_indv=(), parallel=False): - """ + def __init__(self, true_measure: AbstractTrueMeasure, g: Callable, dimension_indv: tuple = (), parallel: Union[bool, int] = False) -> None: + """Initialize a CustomFun integrand. + Args: true_measure (AbstractTrueMeasure): The true measure. - g (callable): A function handle. - dimension_indv (tuple): Shape of individual solution outputs from `g`. - parallel (int): Parallelization flag. + g (Callable): A function handle. + dimension_indv (tuple): Shape of individual solution outputs from + `g`. + parallel (Union[bool, int]): Parallelization flag. - When `parallel = 0` or `parallel = 1` then function evaluation is done in serial fashion. - `parallel > 1` specifies the number of processes used by `multiprocessing.Pool` or `multiprocessing.pool.ThreadPool`. - Setting `parallel=True` is equivalent to `parallel = os.cpu_count()`. + Setting `parallel=True` is equivalent to `parallel = + os.cpu_count()`. - Note: - For `parallel > 1` do *not* set `g` to be anonymous function (i.e. a `lambda` function) + Notes: + For `parallel > 1` do *not* set `g` to be anonymous function (i.e. + a `lambda` function) """ self.parameters = [] self.true_measure = true_measure @@ -109,7 +120,17 @@ def __init__(self, true_measure, g, dimension_indv=(), parallel=False): parallel=parallel, ) - def g(self, t, *args, **kwargs): + def g(self, t: np.ndarray, *args: tuple, **kwargs: dict) -> np.ndarray: + """Evaluate the user-supplied function. + + Args: + t (np.ndarray): Points distributed by the true measure. + *args (tuple): Positional arguments forwarded to the user function. + **kwargs (dict): Keyword arguments forwarded to the user function. + + Returns: + np.ndarray: Function values. + """ return self.__g(t, *args, **kwargs) def _spawn(self, level, sampler): diff --git a/qmcpy/integrand/financial_option.py b/qmcpy/integrand/financial_option.py index 2459540e1..3676b0e9d 100644 --- a/qmcpy/integrand/financial_option.py +++ b/qmcpy/integrand/financial_option.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 from ..true_measure import GeometricBrownianMotion @@ -7,8 +12,7 @@ class FinancialOption(AbstractIntegrand): - r""" - Financial options. + r"""Financial options. - Start price $S_0$ - Strike price $K$ @@ -17,11 +21,15 @@ class FinancialOption(AbstractIntegrand): - Drift $\gamma$ - Equidistant monitoring times $\boldsymbol{\tau} = (\tau_1,\dots,\tau_d)^T$ with $\tau_d$ the final (exercise) time and $\tau_j = \tau_d j/d$. - Define the [geometric brownian motion](https://en.wikipedia.org/wiki/Geometric_Brownian_motion) as + Define the [geometric brownian + motion](https://en.wikipedia.org/wiki/Geometric_Brownian_motion) as - $$\boldsymbol{S}(\boldsymbol{t}) = S_0 e^{(\gamma-\sigma^2/2)\boldsymbol{\tau}+\sigma\boldsymbol{t}}, \qquad \boldsymbol{T} \sim \mathcal{N}(\boldsymbol{0},\mathsf{\Sigma})$$ + $$\boldsymbol{S}(\boldsymbol{t}) = S_0 + e^{(\gamma-\sigma^2/2)\boldsymbol{\tau}+\sigma\boldsymbol{t}}, \qquad + \boldsymbol{T} \sim \mathcal{N}(\boldsymbol{0},\mathsf{\Sigma})$$ - where $\boldsymbol{T}$ is a standard Brownian motion so $\mathsf{\Sigma} = \left(\min\{\tau_j,\tau_{j'}\}\right)_{j,j'=1}^d$. + where $\boldsymbol{T}$ is a standard Brownian motion so $\mathsf{\Sigma} = + \left(\min\{\tau_j,\tau_{j'}\}\right)_{j,j'=1}^d$. The discounted payoff is @@ -29,56 +37,77 @@ class FinancialOption(AbstractIntegrand): where the payoff function $P$ will be defined depending on the option. - Below we will use $S_{-1}$ to denote the final element of $\boldsymbol{S}$, the value of the path at exercise time. + Below we will use $S_{-1}$ to denote the final element of $\boldsymbol{S}$, + the value of the path at exercise time. # European Options *European Call and Put Options* have respective payoffs - $$P(\boldsymbol{S}) = \max\{S_{-1}-K,0\}, \qquad P(\boldsymbol{S}) = \max\{K-S_{-1},0\}.$$ + $$P(\boldsymbol{S}) = \max\{S_{-1}-K,0\}, \qquad P(\boldsymbol{S}) = + \max\{K-S_{-1},0\}.$$ # Asian Options - An asian option considers the average value of an asset path across time. We use the trapezoidal rule to approximate either the *arithmetic mean* by + An asian option considers the average value of an asset path across time. + We use the trapezoidal rule to approximate either the *arithmetic mean* by - $$A(\boldsymbol{S}) = \frac{1}{d}\left[\frac{1}{2} S_0 + \sum_{j=1}^{d-1} S_j + \frac{1}{2} S_{-1}\right]$$ + $$A(\boldsymbol{S}) = \frac{1}{d}\left[\frac{1}{2} S_0 + \sum_{j=1}^{d-1} + S_j + \frac{1}{2} S_{-1}\right]$$ or the *geometric mean* by - $$A(\boldsymbol{S}) = \left[\sqrt{S_0} \prod_{j=1}^{d-1} S_j \sqrt{S_{-1}}\right]^{1/d}.$$ + $$A(\boldsymbol{S}) = \left[\sqrt{S_0} \prod_{j=1}^{d-1} S_j + \sqrt{S_{-1}}\right]^{1/d}.$$ *Asian Call and Put Option* have respective payoffs - $$P(\boldsymbol{S}) = \max\{A(\boldsymbol{S})-K,0\}, \qquad P(\boldsymbol{S}) = \max\{K-A(\boldsymbol{S}),0\}.$$ + $$P(\boldsymbol{S}) = \max\{A(\boldsymbol{S})-K,0\}, \qquad + P(\boldsymbol{S}) = \max\{K-A(\boldsymbol{S}),0\}.$$ # Barrier Options - Barrier $B$. - *In* options are activate when the path crosses the barrier $B$, while *out* options are activated only if the path never crosses the barrier $B$. - An *up* option satisfies $S_0B$, both indicating the direction of the barrier from the start price. + *In* options are activate when the path crosses the barrier $B$, while + *out* options are activated only if the path never crosses the barrier $B$. + An *up* option satisfies $S_0B$, + both indicating the direction of the barrier from the start price. *Barrier Up-In Call and Put Options* have respective payoffs - $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{any } \boldsymbol{S} \geq B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{any } \boldsymbol{S} \geq B \\ 0, & \mathrm{otherwise} \end{cases}.$$ + $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{any } + \boldsymbol{S} \geq B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad + P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{any } + \boldsymbol{S} \geq B \\ 0, & \mathrm{otherwise} \end{cases}.$$ *Barrier Up-Out Call and Put Options* have respective payoffs - $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{all } \boldsymbol{S} < B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{all } \boldsymbol{S} < B \\ 0, & \mathrm{otherwise} \end{cases}.$$ + $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{all } + \boldsymbol{S} < B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad + P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{all } + \boldsymbol{S} < B \\ 0, & \mathrm{otherwise} \end{cases}.$$ *Barrier Down-In Call and Put Options* have respective payoffs - $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{any } \boldsymbol{S} \leq B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{any } \boldsymbol{S} \leq B \\ 0, & \mathrm{otherwise} \end{cases}.$$ + $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{any } + \boldsymbol{S} \leq B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad + P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{any } + \boldsymbol{S} \leq B \\ 0, & \mathrm{otherwise} \end{cases}.$$ *Barrier Down-Out Call and Put Options* have respective payoffs - $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{all } \boldsymbol{S} > B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{all } \boldsymbol{S} > B \\ 0, & \mathrm{otherwise} \end{cases}.$$ + $$P(\boldsymbol{S}) = \begin{cases} \max\{S_{-1})-K,0\}, & \text{all } + \boldsymbol{S} > B \\ 0, & \mathrm{otherwise} \end{cases}, \qquad + P(\boldsymbol{S}) = \begin{cases} \max\{K-S_{-1}),0\}, & \text{all } + \boldsymbol{S} > B \\ 0, & \mathrm{otherwise} \end{cases}.$$ # Lookback Options *Lookback Call and Put Options* have respective payoffs - $$P(\boldsymbol{S}) = S_{-1}-\min(S_0, \ldots S_{-1}), \qquad P(\boldsymbol{S}) = \max(S_0, \ldots S_{-1})-S_{-1}.$$ + $$P(\boldsymbol{S}) = S_{-1}-\min(S_0, \ldots S_{-1}), \qquad + P(\boldsymbol{S}) = \max(S_0, \ldots S_{-1})-S_{-1}.$$ # Digital Option @@ -86,21 +115,30 @@ class FinancialOption(AbstractIntegrand): *Digital Call and Put Options* have respective payoffs - $$P(\boldsymbol{S}) = \begin{cases} \rho, & S_{-1} \geq K \\ 0, & \mathrm{otherwise} \end{cases}, \qquad P(\boldsymbol{S}) = \begin{cases} \rho, & S_{-1} \leq K \\ 0, & \mathrm{otherwise} \end{cases}.$$ + $$P(\boldsymbol{S}) = \begin{cases} \rho, & S_{-1} \geq K \\ 0, & + \mathrm{otherwise} \end{cases}, \qquad P(\boldsymbol{S}) = \begin{cases} + \rho, & S_{-1} \leq K \\ 0, & \mathrm{otherwise} \end{cases}.$$ # Multilevel Options - Initial level $\ell_0 \geq 0$. - Level $\ell \geq \ell_0$. - Let $\boldsymbol{S}_\mathrm{fine}=\boldsymbol{S}$ be the *fine* full path. For $\ell>\ell_0$ write the *coarse* path as $\boldsymbol{S}_\mathrm{coarse} = (S_j)_{j \text{ even}}$ which only considers every other element of $\boldsymbol{S}$. - In this multilevel setting the payoff is + Let $\boldsymbol{S}_\mathrm{fine}=\boldsymbol{S}$ be the *fine* full path. + For $\ell>\ell_0$ write the *coarse* path as + $\boldsymbol{S}_\mathrm{coarse} = (S_j)_{j \text{ even}}$ which only + considers every other element of $\boldsymbol{S}$. In this multilevel + setting the payoff is - $$P_\ell(\boldsymbol{S}) = \begin{cases} P(\boldsymbol{S}_\mathrm{fine}), & \ell = \ell_0, \\ P(\boldsymbol{S}_\mathrm{fine})-P(\boldsymbol{S}_\mathrm{coarse}), & \ell > \ell_0 \end{cases}.$$ + $$P_\ell(\boldsymbol{S}) = \begin{cases} P(\boldsymbol{S}_\mathrm{fine}), & + \ell = \ell_0, \\ + P(\boldsymbol{S}_\mathrm{fine})-P(\boldsymbol{S}_\mathrm{coarse}), & \ell > + \ell_0 \end{cases}.$$ Cancellations from the telescoping sum allow us to write - $$\lim_{\ell \to \infty} P_\ell = P_{\ell_0} + \sum_{\ell=\ell_0+1}^\infty P_\ell.$$ + $$\lim_{\ell \to \infty} P_\ell = P_{\ell_0} + \sum_{\ell=\ell_0+1}^\infty + P_\ell.$$ Examples: >>> integrand = FinancialOption(DigitalNetB2(dimension=3,seed=7),option="EUROPEAN") @@ -205,43 +243,48 @@ class FinancialOption(AbstractIntegrand): def __init__( self, - sampler, - option="ASIAN", - call_put="CALL", - volatility=0.5, - start_price=30, - strike_price=35, - interest_rate=0, - t_final=1, - decomp_type="PCA", - level=None, - d_coarsest=2, - asian_mean="ARITHMETIC", - asian_mean_quadrature_rule="TRAPEZOIDAL", - barrier_in_out="IN", - barrier_price=38, - digital_payout=10, - ): - r""" + sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], + option: str = "ASIAN", + call_put: str = "CALL", + volatility: float = 0.5, + start_price: float = 30, + strike_price: float = 35, + interest_rate: float = 0, + t_final: float = 1, + decomp_type: str = "PCA", + level: Union[None, int] = None, + d_coarsest: Union[None, int] = 2, + asian_mean: str = "ARITHMETIC", + asian_mean_quadrature_rule: str = "TRAPEZOIDAL", + barrier_in_out: str = "IN", + barrier_price: float = 38, + digital_payout: float = 10, + ) -> None: + r"""Initialize a FinancialOption integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. - option (str): Option type in `['ASIAN', 'EUROPEAN', 'BARRIER', 'LOOKBACK', 'DIGITAL']` + option (str): Option type in `['ASIAN', 'EUROPEAN', 'BARRIER', + 'LOOKBACK', 'DIGITAL']` call_put (str): Either `'CALL'` or `'PUT'`. volatility (float): $\sigma$. start_price (float): $S_0$. strike_price (float): $K$. interest_rate (float): $r$. t_final (float): $\tau_d$. - decomp_type (str): Method for decomposition for covariance matrix. Options include + decomp_type (str): Method for decomposition for covariance matrix. + Options include - `'PCA'` for principal component analysis, - `'Cholesky'` for cholesky decomposition, or - `'BrownianBridge'` or `'Bridge'` for brownian bridge construction. level (Union[None, int]): Level for multilevel problems - d_coarsest (Union[None, int]): Dimension of the problem on the coarsest level. + d_coarsest (Union[None, int]): Dimension of the problem on the + coarsest level. asian_mean (str): Either `'ARITHMETIC'` or `'GEOMETRIC'`. asian_mean_quadrature_rule (str): Either 'TRAPEZOIDAL' or 'RIGHT'. barrier_in_out (str): Either `'IN'` or `'OUT'`. @@ -278,34 +321,40 @@ def __init__( if self.level is not None: self.multilevel = True self.parameters += ["level", "d_coarsest"] - assert np.isscalar(self.level) and self.level % 1 == 0 - assert ( + if not (np.isscalar(self.level) and self.level % 1 == 0): + raise AssertionError + if not ( np.isscalar(self.d_coarsest) and self.d_coarsest % 1 == 0 and d_coarsest > 0 and np.log2(d_coarsest) % 1 == 0 - ), "d_coarsest must be an integer power of 2" + ): + raise AssertionError("d_coarsest must be an integer power of 2") self.level = int(self.level) self.d_coarsest = int(self.d_coarsest) - assert ( + if not ( self.sampler.d == self.d_coarsest * 2**self.level - ), "the dimension of the sampler must equal d_coarsest*2^level = %d" % ( - d_coarsest * 2**self.level - ) + ): + raise AssertionError("the dimension of the sampler must equal d_coarsest*2^level = %d" % ( + d_coarsest * 2**self.level + )) self.cost = self.d_coarsest * 2**self.level dim_shape = (2,) else: self.multilevel = False dim_shape = () self.call_put = str(call_put).upper() - assert self.call_put in ["CALL", "PUT"], "invalid call_put = %s" % self.call_put + if not (self.call_put in ["CALL", "PUT"]): + raise AssertionError("invalid call_put = %s" % self.call_put) self.option = str(option).upper() self.asian_mean = str(asian_mean).upper() self.asian_mean_quadrature_rule = str(asian_mean_quadrature_rule).upper() self.barrier_in_out = str(barrier_in_out).upper() - assert np.isscalar(barrier_price) + if not (np.isscalar(barrier_price)): + raise AssertionError self.barrier_price = float(barrier_price) - assert np.isscalar(digital_payout) and digital_payout > 0 + if not (np.isscalar(digital_payout) and digital_payout > 0): + raise AssertionError self.digital_payout = float(digital_payout) if self.option == "EUROPEAN": self.payoff = ( @@ -315,13 +364,15 @@ def __init__( ) elif self.option == "ASIAN": self.parameters += ["asian_mean"] - assert self.asian_mean in ["ARITHMETIC", "GEOMETRIC"], ( - "invalid asian_mean = %s" % self.asian_mean - ) - assert self.asian_mean_quadrature_rule in ["TRAPEZOIDAL", "RIGHT"], ( - "invalid asian_mean_quadrature_rule = %s" - % self.asian_mean_quadrature_rule - ) + if not (self.asian_mean in ["ARITHMETIC", "GEOMETRIC"]): + raise AssertionError( + "invalid asian_mean = %s" % self.asian_mean + ) + if not (self.asian_mean_quadrature_rule in ["TRAPEZOIDAL", "RIGHT"]): + raise AssertionError( + "invalid asian_mean_quadrature_rule = %s" + % self.asian_mean_quadrature_rule + ) if self.asian_mean == "ARITHMETIC": if self.asian_mean_quadrature_rule == "TRAPEZOIDAL": self.payoff = ( @@ -401,7 +452,17 @@ def __init__( dimension_indv=dim_shape, dimension_comb=dim_shape, parallel=False ) - def g(self, t, **kwargs): + def g(self, t: np.ndarray, **kwargs: dict) -> np.ndarray: + """Evaluate the discounted option payoff along each price path. + + Args: + t (np.ndarray): Geometric Brownian motion paths from the true measure. + **kwargs (dict): Unused; accepted for API consistency. + + Returns: + np.ndarray: Discounted payoffs; for a multilevel problem, the coarse + and fine payoffs stacked together. + """ gbm = t # GeometricBrownianMotion already provides GBM paths directly discounted_payoffs = self.payoff(gbm) * self.discount_factor if self.multilevel: @@ -416,13 +477,37 @@ def g(self, t, **kwargs): ) return discounted_payoffs - def payoff_european_call(self, gbm): + def payoff_european_call(self, gbm: np.ndarray) -> np.ndarray: + """European call payoff at maturity. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum(gbm[..., -1] - self.strike_price, 0) - def payoff_european_put(self, gbm): + def payoff_european_put(self, gbm: np.ndarray) -> np.ndarray: + """European put payoff at maturity. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum(self.strike_price - gbm[..., -1], 0) - def payoff_asian_arithmetic_trap_call(self, gbm): + def payoff_asian_arithmetic_trap_call(self, gbm: np.ndarray) -> np.ndarray: + """Asian arithmetic-mean call payoff, trapezoidal averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum( (self.start_price / 2 + gbm[..., :-1].sum(-1) + gbm[..., -1] / 2) / gbm.shape[-1] @@ -430,7 +515,15 @@ def payoff_asian_arithmetic_trap_call(self, gbm): 0, ) - def payoff_asian_arithmetic_trap_put(self, gbm): + def payoff_asian_arithmetic_trap_put(self, gbm: np.ndarray) -> np.ndarray: + """Asian arithmetic-mean put payoff, trapezoidal averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum( self.strike_price - (self.start_price / 2 + gbm[..., :-1].sum(-1) + gbm[..., -1] / 2) @@ -438,7 +531,15 @@ def payoff_asian_arithmetic_trap_put(self, gbm): 0, ) - def payoff_asian_geometric_trap_call(self, gbm): + def payoff_asian_geometric_trap_call(self, gbm: np.ndarray) -> np.ndarray: + """Asian geometric-mean call payoff, trapezoidal averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum( np.exp( ( @@ -452,7 +553,15 @@ def payoff_asian_geometric_trap_call(self, gbm): 0, ) - def payoff_asian_geometric_trap_put(self, gbm): + def payoff_asian_geometric_trap_put(self, gbm: np.ndarray) -> np.ndarray: + """Asian geometric-mean put payoff, trapezoidal averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum( self.strike_price - np.exp( @@ -466,101 +575,229 @@ def payoff_asian_geometric_trap_put(self, gbm): 0, ) - def payoff_asian_arithmetic_right_call(self, gbm): + def payoff_asian_arithmetic_right_call(self, gbm: np.ndarray) -> np.ndarray: + """Asian arithmetic-mean call payoff, right-endpoint averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum(gbm.sum(-1) / gbm.shape[-1] - self.strike_price, 0) - def payoff_asian_arithmetic_right_put(self, gbm): + def payoff_asian_arithmetic_right_put(self, gbm: np.ndarray) -> np.ndarray: + """Asian arithmetic-mean put payoff, right-endpoint averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum((self.strike_price - gbm.sum(-1)) / gbm.shape[-1], 0) - def payoff_asian_geometric_right_call(self, gbm): + def payoff_asian_geometric_right_call(self, gbm: np.ndarray) -> np.ndarray: + """Asian geometric-mean call payoff, right-endpoint averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum( np.exp(np.log(gbm).sum(-1) / gbm.shape[-1]) - self.strike_price, 0 ) - def payoff_asian_geometric_right_put(self, gbm): + def payoff_asian_geometric_right_put(self, gbm: np.ndarray) -> np.ndarray: + """Asian geometric-mean put payoff, right-endpoint averaging. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.maximum( self.strike_price - np.exp(np.log(gbm).sum(-1) / gbm.shape[-1]), 0 ) - def payoff_barrier_in_up_call(self, gbm): + def payoff_barrier_in_up_call(self, gbm: np.ndarray) -> np.ndarray: + """Up-and-in barrier call payoff; pays only if the barrier is reached from below. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm >= self.barrier_price).any(-1) v[~flag] = 0 v[flag] = np.maximum(v[flag] - self.strike_price, 0) return v - def payoff_barrier_out_up_call(self, gbm): + def payoff_barrier_out_up_call(self, gbm: np.ndarray) -> np.ndarray: + """Up-and-out barrier call payoff; pays only if the barrier is never reached. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm < self.barrier_price).all(-1) v[~flag] = 0 v[flag] = np.maximum(v[flag] - self.strike_price, 0) return v - def payoff_barrier_in_down_call(self, gbm): + def payoff_barrier_in_down_call(self, gbm: np.ndarray) -> np.ndarray: + """Down-and-in barrier call payoff; pays only if the barrier is reached from above. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm <= self.barrier_price).any(-1) v[~flag] = 0 v[flag] = np.maximum(v[flag] - self.strike_price, 0) return v - def payoff_barrier_out_down_call(self, gbm): + def payoff_barrier_out_down_call(self, gbm: np.ndarray) -> np.ndarray: + """Down-and-out barrier call payoff; pays only if the barrier is never reached. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm > self.barrier_price).all(-1) v[~flag] = 0 v[flag] = np.maximum(v[flag] - self.strike_price, 0) return v - def payoff_barrier_in_up_put(self, gbm): + def payoff_barrier_in_up_put(self, gbm: np.ndarray) -> np.ndarray: + """Up-and-in barrier put payoff; pays only if the barrier is reached from below. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm >= self.barrier_price).any(-1) v[~flag] = 0 v[flag] = np.maximum(self.strike_price - v[flag], 0) return v - def payoff_barrier_out_up_put(self, gbm): + def payoff_barrier_out_up_put(self, gbm: np.ndarray) -> np.ndarray: + """Up-and-out barrier put payoff; pays only if the barrier is never reached. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm < self.barrier_price).all(-1) v[~flag] = 0 v[flag] = np.maximum(self.strike_price - v[flag], 0) return v - def payoff_barrier_in_down_put(self, gbm): + def payoff_barrier_in_down_put(self, gbm: np.ndarray) -> np.ndarray: + """Down-and-in barrier put payoff; pays only if the barrier is reached from above. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm <= self.barrier_price).any(-1) v[~flag] = 0 v[flag] = np.maximum(self.strike_price - v[flag], 0) return v - def payoff_barrier_out_down_put(self, gbm): + def payoff_barrier_out_down_put(self, gbm: np.ndarray) -> np.ndarray: + """Down-and-out barrier put payoff; pays only if the barrier is never reached. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ v = gbm[..., -1].copy() flag = (gbm > self.barrier_price).all(-1) v[~flag] = 0 v[flag] = np.maximum(self.strike_price - v[flag], 0) return v - def payoff_lookback_call(self, gbm): # include start price in min + def payoff_lookback_call(self, gbm: np.ndarray) -> np.ndarray: # include start price in min + """Lookback call payoff: final price less the running minimum, including the start price. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ min_path = np.minimum(gbm.min(-1), self.start_price) return gbm[..., -1] - min_path - def payoff_lookback_put(self, gbm): # include start price in max + def payoff_lookback_put(self, gbm: np.ndarray) -> np.ndarray: # include start price in max + """Lookback put payoff: the running maximum, including the start price, less the final price. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ max_path = np.maximum(gbm.max(-1), self.start_price) return max_path - gbm[..., -1] - def payoff_digital_call(self, gbm): + def payoff_digital_call(self, gbm: np.ndarray) -> np.ndarray: + """Digital call payoff: a fixed payout when the final price is at or above the strike. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. + + Returns: + np.ndarray: Payoff of each path. + """ return np.where(gbm[..., -1] >= self.strike_price, self.digital_payout, 0) - def payoff_digital_put(self, gbm): - return np.where(gbm[..., -1] <= self.strike_price, self.digital_payout, 0) + def payoff_digital_put(self, gbm: np.ndarray) -> np.ndarray: + """Digital put payoff: a fixed payout when the final price is at or below the strike. + + Args: + gbm (np.ndarray): Geometric Brownian motion paths, monitoring times last. - def get_exact_value(self): + Returns: + np.ndarray: Payoff of each path. """ - Compute the exact analytic fair price of the option in finite dimensions. Supports + return np.where(gbm[..., -1] <= self.strike_price, self.digital_payout, 0) + + def get_exact_value(self) -> float: + """Compute the exact analytic fair price of the option in finite + dimensions. Supports - `option='EUROPEAN'` - `option='ASIAN'` with `asian_mean='GEOMETRIC'` and `asian_mean_quadrature_rule='RIGHT'` Returns: - mean (float): Exact value of the integral. + float: Exact value of the integral. """ if self.option == "EUROPEAN": denom = self.volatility * np.sqrt(self.t_final) @@ -590,10 +827,11 @@ def get_exact_value(self): term2 / denom ) elif self.option == "ASIAN": - assert ( + if not ( self.asian_mean == "GEOMETRIC" and self.asian_mean_quadrature_rule == "RIGHT" - ), "exact value for Asian options only implemented for self.asian_mean=='GEOMETRIC' and self.asian_mean_quadrature_rule=='RIGHT'" + ): + raise AssertionError("exact value for Asian options only implemented for self.asian_mean=='GEOMETRIC' and self.asian_mean_quadrature_rule=='RIGHT'") Tbar = (1 + 1 / self.d) * self.t_final / 2 sigmabar = self.volatility * np.sqrt((2 + 1 / self.d) / 3) rbar = self.interest_rate + (sigmabar**2 - self.volatility**2) / 2 @@ -611,19 +849,20 @@ def get_exact_value(self): ) return fp - def get_exact_value_inf_dim(self): - r""" - Get the exact analytic fair price of the option in infinite dimensions. Supports + def get_exact_value_inf_dim(self) -> float: + r"""Get the exact analytic fair price of the option in infinite + dimensions. Supports - `option='ASIAN'` with `asian_mean='GEOMETRIC'` Returns: - mean (float): Exact value of the integral. + float: Exact value of the integral. """ if self.option == "ASIAN": - assert ( + if not ( self.asian_mean == "GEOMETRIC" - ), "get_exact_value_inf_dim for the Asian option only available for self.asian_mean=='GEOMETRIC'" + ): + raise AssertionError("get_exact_value_inf_dim for the Asian option only available for self.asian_mean=='GEOMETRIC'") sigma_g = self.volatility / np.sqrt(3) b = 1 / 2 * (self.interest_rate - 1 / 2 * sigma_g**2) d1 = ( @@ -643,7 +882,15 @@ def get_exact_value_inf_dim(self): ) return val - def dimension_at_level(self, level): + def dimension_at_level(self, level: int) -> int: + """Return the number of monitoring times used at a multilevel level. + + Args: + level (int): Multilevel level index. + + Returns: + int: Monitoring times at that level, doubling with each level. + """ return self.d_coarsest * 2**level def _spawn(self, level, sampler): @@ -677,7 +924,12 @@ def _eurogbmprice(S0, r, T, sigma, K): class AsianOption(FinancialOption): - def __init__(self, *args, **kwargs): + """Asian option. + + Deprecated, please use :class:`FinancialOption` with ``option="ASIAN"``. + """ + + def __init__(self, *args, **kwargs) -> None: """Deprecated, please use FinancialOption""" if "option" in kwargs: raise ParameterError("please do not pass 'option' to AsianOption") @@ -685,7 +937,12 @@ def __init__(self, *args, **kwargs): class EuropeanOption(FinancialOption): - def __init__(self, *args, **kwargs): + """European option. + + Deprecated, please use :class:`FinancialOption` with ``option="EUROPEAN"``. + """ + + def __init__(self, *args, **kwargs) -> None: """Deprecated, please use FinancialOption""" if "option" in kwargs: raise ParameterError("please do not pass 'option' to EuropeanOption") @@ -693,7 +950,12 @@ def __init__(self, *args, **kwargs): class BarrierOption(FinancialOption): - def __init__(self, *args, **kwargs): + """Barrier option. + + Deprecated, please use :class:`FinancialOption` with ``option="BARRIER"``. + """ + + def __init__(self, *args, **kwargs) -> None: """Deprecated, please use FinancialOption""" if "option" in kwargs: raise ParameterError("please do not pass 'option' to BarrierOption") @@ -701,7 +963,12 @@ def __init__(self, *args, **kwargs): class LookbackOption(FinancialOption): - def __init__(self, *args, **kwargs): + """Lookback option. + + Deprecated, please use :class:`FinancialOption` with ``option="LOOKBACK"``. + """ + + def __init__(self, *args, **kwargs) -> None: """Deprecated, please use FinancialOption""" if "option" in kwargs: raise ParameterError("please do not pass 'option' to LookbackOption") @@ -709,7 +976,12 @@ def __init__(self, *args, **kwargs): class DigitalOption(FinancialOption): - def __init__(self, *args, **kwargs): + """Digital option. + + Deprecated, please use :class:`FinancialOption` with ``option="DIGITAL"``. + """ + + def __init__(self, *args, **kwargs) -> None: """Deprecated, please use FinancialOption""" if "option" in kwargs: raise ParameterError("please do not pass 'option' to DigitalOption") diff --git a/qmcpy/integrand/fourbranch2d.py b/qmcpy/integrand/fourbranch2d.py index 04123c4f6..18f77094c 100644 --- a/qmcpy/integrand/fourbranch2d.py +++ b/qmcpy/integrand/fourbranch2d.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union import numpy as np from .abstract_integrand import AbstractIntegrand from ..true_measure import Uniform @@ -5,10 +10,13 @@ class FourBranch2d(AbstractIntegrand): - r""" - Four Branch function in $d=2$. + r"""Four Branch function in $d=2$. - $$g(\boldsymbol{t}) = \min \begin{cases} 3+0.1(t_0-t_1)^2-\frac{t_0-t_1}{\sqrt{2}} \\ 3+0.1(t_0-t_1)^2+\frac{t_0-t_1}{\sqrt{2}} \\ t_0-t_1 + 7/\sqrt{2} \\ t_1-t_0 + 7/\sqrt{2}\end{cases}, \qquad \boldsymbol{T}=(T_0,T_1) \sim \mathcal{U}[-8,8]^2.$$ + $$g(\boldsymbol{t}) = \min \begin{cases} + 3+0.1(t_0-t_1)^2-\frac{t_0-t_1}{\sqrt{2}} \\ + 3+0.1(t_0-t_1)^2+\frac{t_0-t_1}{\sqrt{2}} \\ t_0-t_1 + 7/\sqrt{2} \\ + t_1-t_0 + 7/\sqrt{2}\end{cases}, \qquad \boldsymbol{T}=(T_0,T_1) \sim + \mathcal{U}[-8,8]^2.$$ Examples: >>> integrand = FourBranch2d(DigitalNetB2(2,seed=7)) @@ -41,22 +49,33 @@ class FourBranch2d(AbstractIntegrand): -2.5042 """ - def __init__(self, sampler): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure]) -> None: + r"""Initialize a FourBranch2d integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. """ self.sampler = sampler - assert self.sampler.d == 2 + if not (self.sampler.d == 2): + raise AssertionError self.true_measure = Uniform(self.sampler, lower_bound=-8, upper_bound=8) super(FourBranch2d, self).__init__( dimension_indv=(), dimension_comb=(), parallel=False ) - def g(self, t): + def g(self, t: np.ndarray) -> np.ndarray: + """Evaluate the four-branch function. + + Args: + t (np.ndarray): Two-dimensional points. + + Returns: + np.ndarray: Minimum of the four branches at each point. + """ t0, t1 = t[..., 0], t[..., 1] return np.minimum.reduce( [ diff --git a/qmcpy/integrand/genz.py b/qmcpy/integrand/genz.py index 73973b62f..9898442e2 100644 --- a/qmcpy/integrand/genz.py +++ b/qmcpy/integrand/genz.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 from ..true_measure import Uniform @@ -6,14 +11,16 @@ class Genz(AbstractIntegrand): - r""" - Genz function following the [`DAKOTA` implementation](https://snl-dakota.github.io/docs/6.17.0/users/usingdakota/examples/additionalexamples.html?highlight=genz#genz-functions). + r"""Genz function following the [`DAKOTA` + implementation](https://snl-dakota.github.io/docs/6.17.0/users/usingdakota/examples/additionalexamples.html?highlight=genz#genz-functions). - $$g_\mathrm{oscillatory}(\boldsymbol{t}) = \cos\left(-\sum_{j=1}^d c_j t_j\right)$$ + $$g_\mathrm{oscillatory}(\boldsymbol{t}) = \cos\left(-\sum_{j=1}^d c_j + t_j\right)$$ or - $$g_\mathrm{corner-peak}(\boldsymbol{t}) = \left(1+\sum_{j=1}^d c_j t_j\right)^{-(d+1)}$$ + $$g_\mathrm{corner-peak}(\boldsymbol{t}) = \left(1+\sum_{j=1}^d c_j + t_j\right)^{-(d+1)}$$ where @@ -21,7 +28,9 @@ class Genz(AbstractIntegrand): and the coefficients $\boldsymbol{c}$ are have three kinds - $$c_k^{(1)} = \frac{k-1/2}{d}, \qquad c_k^{(2)} = \frac{1}{k^2}, \qquad c_k^{(3)} = \exp\left(\frac{k \log(10^{-8})}{d}\right), \qquad k=1,\dots,d.$$ + $$c_k^{(1)} = \frac{k-1/2}{d}, \qquad c_k^{(2)} = \frac{1}{k^2}, \qquad + c_k^{(3)} = \exp\left(\frac{k \log(10^{-8})}{d}\right), \qquad + k=1,\dots,d.$$ Examples: >>> for kind_func in ['OSCILLATORY','CORNER PEAK']: @@ -50,10 +59,12 @@ class Genz(AbstractIntegrand): 0.7200 """ - def __init__(self, sampler, kind_func="OSCILLATORY", kind_coeff=1): - """ + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], kind_func: str = "OSCILLATORY", kind_coeff: int = 1) -> None: + """Initialize a Genz integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -92,10 +103,26 @@ def __init__(self, sampler, kind_func="OSCILLATORY", kind_coeff=1): self.parameters = ["kind_func", "kind_coeff"] super(Genz, self).__init__(dimension_indv=(), dimension_comb=(), parallel=False) - def g_oscillatory(self, t): + def g_oscillatory(self, t: np.ndarray) -> np.ndarray: + r"""Evaluate the oscillatory Genz function. + + Args: + t (np.ndarray): Points in the unit cube. + + Returns: + np.ndarray: $\cos(-c \cdot t)$ at each point. + """ return np.cos(-(self.c * t).sum(-1)) - def g_corner_peak(self, t): + def g_corner_peak(self, t: np.ndarray) -> np.ndarray: + r"""Evaluate the corner-peak Genz function. + + Args: + t (np.ndarray): Points in the unit cube. + + Returns: + np.ndarray: $(1 + c \cdot t)^{-(d+1)}$ at each point. + """ return (1 + (self.c * t).sum(-1)) ** (-(self.d + 1)) def _spawn(self, level, sampler): diff --git a/qmcpy/integrand/hartmann6d.py b/qmcpy/integrand/hartmann6d.py index c1fe57d13..ba304836a 100644 --- a/qmcpy/integrand/hartmann6d.py +++ b/qmcpy/integrand/hartmann6d.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union import numpy as np from .abstract_integrand import AbstractIntegrand from ..true_measure import Uniform @@ -5,8 +10,9 @@ class Hartmann6d(AbstractIntegrand): - r""" - Wrapper around [`BoTorch`'s implementation of the Augmented Hartmann function](https://botorch.readthedocs.io/en/stable/test_functions.html#botorch.test_functions.multi_fidelity.AugmentedHartmann) in dimension $d=6$. + r"""Wrapper around [`BoTorch`'s implementation of the Augmented Hartmann + function](https://botorch.readthedocs.io/en/stable/test_functions.html#botorch.test_functions.multi_fidelity.AugmentedHartmann) + in dimension $d=6$. Examples: >>> integrand = Hartmann6d(DigitalNetB2(6,seed=7)) @@ -29,7 +35,7 @@ class Hartmann6d(AbstractIntegrand): (3, 3) 0.08333333333333333 (4, 4) 0.08333333333333333 (5, 5) 0.08333333333333333 - + With independent replications >>> integrand = Hartmann6d(DigitalNetB2(6,seed=7,replications=2**4)) @@ -43,16 +49,19 @@ class Hartmann6d(AbstractIntegrand): -0.2599 """ - def __init__(self, sampler): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure]) -> None: + r"""Initialize a Hartmann6d integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. """ self.sampler = sampler - assert self.sampler.d == 6 + if not (self.sampler.d == 6): + raise AssertionError self.true_measure = Uniform(self.sampler, lower_bound=0, upper_bound=1) super(Hartmann6d, self).__init__( dimension_indv=(), dimension_comb=(), parallel=False @@ -61,7 +70,15 @@ def __init__(self, sampler): self.ah = AugmentedHartmann(negate=False) - def g(self, t): + def g(self, t: np.ndarray) -> np.ndarray: + """Evaluate the six-dimensional augmented Hartmann function. + + Args: + t (np.ndarray): Six-dimensional points. + + Returns: + np.ndarray: Function values, via BoTorch's ``AugmentedHartmann``. + """ import torch t = np.concatenate([t, np.ones(tuple(t.shape[:-1]) + (1,))], axis=-1) diff --git a/qmcpy/integrand/ishigami.py b/qmcpy/integrand/ishigami.py index 8ff3aa429..b8e8c6504 100644 --- a/qmcpy/integrand/ishigami.py +++ b/qmcpy/integrand/ishigami.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union import numpy as np from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 @@ -6,10 +11,11 @@ class Ishigami(AbstractIntegrand): - r""" - Ishigami function in $d=3$ dimensions from [1] and [https://www.sfu.ca/~ssurjano/ishigami.html](https://www.sfu.ca/~ssurjano/ishigami.html). + r"""Ishigami function in $d=3$ dimensions from [1] and + [https://www.sfu.ca/~ssurjano/ishigami.html](https://www.sfu.ca/~ssurjano/ishigami.html). - $$g(\boldsymbol{t}) = (1+bt_2^4)\sin(t_0)+a\sin^2(t_1), \qquad \boldsymbol{T} = (T_0,T_1,T_2) \sim \mathcal{U}(-\pi,\pi)^3.$$ + $$g(\boldsymbol{t}) = (1+bt_2^4)\sin(t_0)+a\sin^2(t_1), \qquad + \boldsymbol{T} = (T_0,T_1,T_2) \sim \mathcal{U}(-\pi,\pi)^3.$$ Examples: >>> integrand = Ishigami(DigitalNetB2(3,seed=7)) @@ -52,10 +58,12 @@ class Ishigami(AbstractIntegrand): Proceedings, First International Symposium on (pp. 398-403). IEEE. """ - def __init__(self, sampler, a=7, b=0.1): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], a: float = 7, b: float = 0.1) -> None: + r"""Initialize an Ishigami integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -72,7 +80,15 @@ def __init__(self, sampler, a=7, b=0.1): dimension_indv=(), dimension_comb=(), parallel=False ) - def g(self, t): + def g(self, t: np.ndarray) -> np.ndarray: + r"""Evaluate the Ishigami function. + + Args: + t (np.ndarray): Three-dimensional points. + + Returns: + np.ndarray: $(1 + b t_3^4)\sin(t_1) + a \sin^2(t_2)$. + """ y = (1 + self.b * t[..., 2] ** 4) * np.sin(t[..., 0]) + self.a * np.sin( t[..., 1] ) ** 2 @@ -84,7 +100,8 @@ def _spawn(self, level, sampler): @staticmethod def _exact_sensitivity_indices(indices, a, b): a, b = np.atleast_1d(a), np.atleast_1d(b) - assert a.shape == b.shape and a.ndim == 1 and b.ndim == 1 + if not (a.shape == b.shape and a.ndim == 1 and b.ndim == 1): + raise AssertionError mu = a / 2 m2 = 1 / 2 + 3 / 8 * a**2 + np.pi**4 / 5 * b + np.pi**8 / 18 * b**2 tau_closed = { @@ -123,7 +140,8 @@ def _exact_fu_functions(x, indices, a, b): x = np.atleast_2d(x) n = len(x) a, b = np.atleast_1d(a), np.atleast_1d(b) - assert x.ndim == 2 and x.shape == (n, 3) and a.shape == (1,) and b.shape == (1,) + if not (x.ndim == 2 and x.shape == (n, 3) and a.shape == (1,) and b.shape == (1,)): + raise AssertionError x0, x1, x2 = x[:, 0], x[:, 1], x[:, 2] fus = { repr([]): a / 2, diff --git a/qmcpy/integrand/keister.py b/qmcpy/integrand/keister.py index 949b3f644..54399f8ad 100644 --- a/qmcpy/integrand/keister.py +++ b/qmcpy/integrand/keister.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 from ..true_measure import Gaussian @@ -6,10 +11,10 @@ class Keister(AbstractIntegrand): - r""" - Keister function from [1]. + r"""Keister function from [1]. - $$f(\boldsymbol{t}) = \pi^{d/2} \cos(\lVert \boldsymbol{t} \rVert_2) \qquad \boldsymbol{T} \sim \mathcal{N}(\boldsymbol{0},\mathsf{I}/2).$$ + $$f(\boldsymbol{t}) = \pi^{d/2} \cos(\lVert \boldsymbol{t} \rVert_2) \qquad + \boldsymbol{T} \sim \mathcal{N}(\boldsymbol{0},\mathsf{I}/2).$$ Examples: >>> integrand = Keister(DigitalNetB2(2,seed=7)) @@ -44,10 +49,12 @@ class Keister(AbstractIntegrand): Computers in Physics, 10, pp. 119-122, 1996. """ - def __init__(self, sampler): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure]) -> None: + r"""Initialize a Keister integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -58,7 +65,15 @@ def __init__(self, sampler): dimension_indv=(), dimension_comb=(), parallel=False ) - def g(self, t): + def g(self, t: np.ndarray) -> np.ndarray: + r"""Evaluate the Keister function. + + Args: + t (np.ndarray): Points, dimensions along the last axis. + + Returns: + np.ndarray: $\pi^{d/2}\cos(\lVert t \rVert_2)$. + """ d = t.shape[-1] norm = np.linalg.norm(t, axis=-1) k = np.pi ** (d / 2) * np.cos(norm) @@ -68,15 +83,15 @@ def _spawn(self, level, sampler): return Keister(sampler=sampler) @classmethod - def get_exact_value(self, d): - """ - Compute the exact analytic value of the Keister integral with dimension $d$. + def get_exact_value(cls, d: int) -> float: + """Compute the exact analytic value of the Keister integral with + dimension $d$. Args: d (int): Dimension. Returns: - mean (float): Exact value of the integral. + float: Exact value of the integral. """ cosinteg = np.zeros(shape=(d)) cosinteg[0] = np.sqrt(np.pi) / (2 * np.exp(1 / 4)) @@ -91,5 +106,16 @@ def get_exact_value(self, d): I = (2 * (np.pi ** (d / 2)) / gamma(d / 2)) * cosinteg[d - 1] return I - def exact_integ(self, *args, **kwargs): + def exact_integ(self, *args: tuple, **kwargs: dict) -> float: + """Return the exact value of the Keister integral. + + Deprecated alias for :meth:`get_exact_value`. + + Args: + *args (tuple): Forwarded to :meth:`get_exact_value`. + **kwargs (dict): Forwarded to :meth:`get_exact_value`. + + Returns: + float: The exact integral value. + """ return self.get_exact_value(*args, **kwargs) diff --git a/qmcpy/integrand/linear0.py b/qmcpy/integrand/linear0.py index b43188edf..640e39b0f 100644 --- a/qmcpy/integrand/linear0.py +++ b/qmcpy/integrand/linear0.py @@ -1,13 +1,19 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union +import numpy as np from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 from ..true_measure import Uniform class Linear0(AbstractIntegrand): - r""" - Linear Function with analytic mean $0$. + r"""Linear Function with analytic mean $0$. - $$g(\boldsymbol{t}) = \sum_{j=1}^d t_j \qquad \boldsymbol{T} \sim \mathcal{U}[0,1]^d.$$ + $$g(\boldsymbol{t}) = \sum_{j=1}^d t_j \qquad \boldsymbol{T} \sim + \mathcal{U}[0,1]^d.$$ Examples: >>> integrand = Linear0(DigitalNetB2(100,seed=7)) @@ -28,10 +34,12 @@ class Linear0(AbstractIntegrand): -9.8203e-05 """ - def __init__(self, sampler): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure]) -> None: + r"""Initialize a Linear0 integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -42,7 +50,15 @@ def __init__(self, sampler): dimension_indv=(), dimension_comb=(), parallel=False ) - def g(self, t): + def g(self, t: np.ndarray) -> np.ndarray: + """Evaluate the centered linear function. + + Args: + t (np.ndarray): Points, dimensions along the last axis. + + Returns: + np.ndarray: Sum of the coordinates of each point. + """ y = t.sum(-1) return y diff --git a/qmcpy/integrand/multimodal2d.py b/qmcpy/integrand/multimodal2d.py index 1291f1a5c..eb8e6c7f4 100644 --- a/qmcpy/integrand/multimodal2d.py +++ b/qmcpy/integrand/multimodal2d.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union import numpy as np from .abstract_integrand import AbstractIntegrand from ..true_measure import Uniform @@ -5,10 +10,10 @@ class Multimodal2d(AbstractIntegrand): - r""" - Multimodal function in $d=2$ dimensions. + r"""Multimodal function in $d=2$ dimensions. - $$g(\boldsymbol{t}) = (t_0^2+4)(t_1-1)/20-\sin(5t_0/2)-2 \qquad \boldsymbol{T} = (T_0,T_1) \sim \mathcal{U}([-4,7] \times [-3,8]).$$ + $$g(\boldsymbol{t}) = (t_0^2+4)(t_1-1)/20-\sin(5t_0/2)-2 \qquad + \boldsymbol{T} = (T_0,T_1) \sim \mathcal{U}([-4,7] \times [-3,8]).$$ Examples: >>> integrand = Multimodal2d(DigitalNetB2(2,seed=7)) @@ -41,16 +46,19 @@ class Multimodal2d(AbstractIntegrand): -0.7366 """ - def __init__(self, sampler): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure]) -> None: + r"""Initialize a Multimodal2d integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. """ self.sampler = sampler - assert self.sampler.d == 2 + if not (self.sampler.d == 2): + raise AssertionError self.true_measure = Uniform( self.sampler, lower_bound=[-4, -3], upper_bound=[7, 8] ) @@ -58,7 +66,15 @@ def __init__(self, sampler): dimension_indv=(), dimension_comb=(), parallel=False ) - def g(self, t): + def g(self, t: np.ndarray) -> np.ndarray: + """Evaluate the two-dimensional multimodal function. + + Args: + t (np.ndarray): Two-dimensional points. + + Returns: + np.ndarray: Function values. + """ t0, t1 = t[..., 0], t[..., 1] return (t0**2 + 4) * (t1 - 1) / 20 - np.sin(5 * t0 / 2) - 2 diff --git a/qmcpy/integrand/sensitivity_indices.py b/qmcpy/integrand/sensitivity_indices.py index 9c90d300b..7dc5ab703 100644 --- a/qmcpy/integrand/sensitivity_indices.py +++ b/qmcpy/integrand/sensitivity_indices.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_integrand import AbstractIntegrand from .keister import Keister from .box_integral import BoxIntegral @@ -8,8 +9,7 @@ class SensitivityIndices(AbstractIntegrand): - r""" - Sensitivity indices i.e. normalized Sobol' Indices. + r"""Sensitivity indices i.e., normalized Sobol' Indices. Examples: Singleton indices @@ -111,11 +111,15 @@ class SensitivityIndices(AbstractIntegrand): [https://artowen.su.domains/mc/A-anova.pdf](https://artowen.su.domains/mc/A-anova.pdf). """ - def __init__(self, integrand, indices="singletons"): - r""" + def __init__(self, integrand: AbstractIntegrand, indices: Union[str, np.ndarray] = "singletons") -> None: + r"""Initialize a SensitivityIndices integrand. + Args: - integrand (AbstractIntegrand): Integrand to find sensitivity indices of. - indices (np.ndarray): Bool array with shape $(\dots,d)$ where each length $d$ vector item indicates which dimensions are active in the subset. + integrand (AbstractIntegrand): Integrand to find sensitivity + indices of. + indices (Union[str, np.ndarray]): Bool array with shape $(\dots,d)$ where each + length $d$ vector item indicates which dimensions are active in + the subset. - The default `indices='singletons'` sets `indices=np.eye(d,dtype=bool)`. - Setting `incides='all'` sets `indices = np.array([[bool(int(b)) for b in np.binary_repr(i,width=d)] for i in range(1,2**d-1)],dtype=bool)` @@ -123,7 +127,8 @@ def __init__(self, integrand, indices="singletons"): self.parameters = ["indices"] self.integrand = integrand self.dtilde = self.integrand.d - assert self.dtilde > 1, "SensitivityIndices does not make sense for d=1" + if not (self.dtilde > 1): + raise AssertionError("SensitivityIndices does not make sense for d=1") self.indices = indices if isinstance(self.indices, str) and self.indices == "singletons": self.indices = np.eye(self.dtilde, dtype=bool) @@ -137,14 +142,16 @@ def __init__(self, integrand, indices="singletons"): idxs_r[i, comb] = True self.indices = np.vstack([self.indices, idxs_r]) self.indices = np.atleast_1d(self.indices) - assert ( + if not ( self.indices.dtype == bool and self.indices.ndim >= 1 and self.indices.shape[-1] == self.dtilde - ) - assert ( + ): + raise AssertionError + if not ( not (self.indices == self.indices[..., 0, None]).all(-1).any() - ), "indices cannot include the emptyset or the set of all dimensions" + ): + raise AssertionError("indices cannot include the emptyset or the set of all dimensions") self.not_indices = ~self.indices # sensitivity_index self.true_measure = self.integrand.true_measure @@ -160,13 +167,25 @@ def __init__(self, integrand, indices="singletons"): ) self.d = 2 * self.dtilde - def f(self, x, *args, **kwargs): + def f(self, x: np.ndarray, *args: tuple, **kwargs: dict) -> np.ndarray: + r"""Evaluate the numerator and moment terms needed for the sensitivity indices. + + Args: + x (np.ndarray): Points from the discrete distribution. + *args (tuple): Forwarded to the wrapped integrand. + **kwargs (dict): Forwarded to the wrapped integrand; ``compute_flags`` + selects which outputs to evaluate. + + Returns: + np.ndarray: The $\tau$, mean, and second-moment terms. + """ if "compute_flags" in kwargs: compute_flags = kwargs["compute_flags"] del kwargs["compute_flags"] else: compute_flags = np.ones(self.d_indv, dtype=bool) - assert compute_flags.shape == self.d_indv + if not (compute_flags.shape == self.d_indv): + raise AssertionError z = x[..., self.dtilde :] x = x[..., : self.dtilde] v = np.zeros_like(x) @@ -190,16 +209,28 @@ def f(self, x, *args, **kwargs): y[(slice(None), 2) + i + self.i_slice] = ( f_x[(None,) + self.i_slice] ** 2 ) # sigma^2+mu^2 - # here we copy mu and sigma^2+mu^2 since if these these were not copied there is a chance the bounds could change - # for mu and/or sigma and then an index which was previously approximated sufficiently woulud become insufficientlly approximated - # and it would then be difficult ot go back and resample the numerator for that approximation + # Here we copy mu and sigma^2+mu^2 since, if these were not copied, + # there is a chance the bounds could change for mu and/or sigma. + # Then an index that was previously approximated sufficiently could + # become insufficiently approximated, and it would be difficult to + # go back and resample the numerator for that approximation. return y def _spawn(self, level, sampler): new_integrand = self.integrand.spawn(level, sampler) return SensitivityIndices(integrand=new_integrand, indices=self.indices) - def bound_fun(self, bound_low, bound_high): + def bound_fun(self, bound_low: np.ndarray, bound_high: np.ndarray) -> tuple: + r"""Combine bounds on the moment terms into bounds on the sensitivity indices. + + Args: + bound_low (np.ndarray): Lower bounds on $\tau$, the mean, and the second moment. + bound_high (np.ndarray): Upper bounds on the same terms. + + Returns: + tuple: Lower and upper bounds on the indices, clipped to $[0,1]$ and + widened to $[0,1]$ where the variance bound is non-positive. + """ tau_low, mu_low, f2_low = bound_low[:, 0], bound_low[:, 1], bound_low[:, 2] tau_high, mu_high, f2_high = ( bound_high[:, 0], @@ -218,5 +249,13 @@ def bound_fun(self, bound_low, bound_high): comb_bounds_low[violated], comb_bounds_high[violated] = 0, 1 return comb_bounds_low, comb_bounds_high - def dependency(self, comb_flags): + def dependency(self, comb_flags: np.ndarray) -> np.ndarray: + """Map combined-output flags onto the individual outputs they require. + + Args: + comb_flags (np.ndarray): Flags for the combined outputs. + + Returns: + np.ndarray: Flags for the three moment terms behind each index. + """ return np.repeat(comb_flags[:, None], 3, axis=1) diff --git a/qmcpy/integrand/sin1d.py b/qmcpy/integrand/sin1d.py index 9ffb69170..8854dfe6a 100644 --- a/qmcpy/integrand/sin1d.py +++ b/qmcpy/integrand/sin1d.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union import numpy as np from .abstract_integrand import AbstractIntegrand from ..true_measure import Uniform @@ -5,8 +10,7 @@ class Sin1d(AbstractIntegrand): - r""" - Sine function in $d=1$ dimension. + r"""Sine function in $d=1$ dimension. $$g(t) = \sin(t), \qquad t \sim \mathcal{U}[0,2\pi k]$$ @@ -40,18 +44,22 @@ class Sin1d(AbstractIntegrand): 7.0800e-04 """ - def __init__(self, sampler, k=1): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], k: float = 1) -> None: + r"""Initialize a Sin1d integrand. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. - k (float): The true measure will be uniform between $0$ and $2 \pi k$. + k (float): The true measure will be uniform between $0$ and $2 \pi + k$. """ self.sampler = sampler self.k = k - assert self.sampler.d == 1 + if not (self.sampler.d == 1): + raise AssertionError self.true_measure = Uniform( self.sampler, lower_bound=0, upper_bound=2 * self.k * np.pi ) @@ -59,7 +67,15 @@ def __init__(self, sampler, k=1): dimension_indv=(), dimension_comb=(), parallel=False ) - def g(self, t): + def g(self, t: np.ndarray) -> np.ndarray: + r"""Evaluate the one-dimensional sine function. + + Args: + t (np.ndarray): One-dimensional points. + + Returns: + np.ndarray: $\sin(t)$. + """ return np.sin(t[..., 0]) def _spawn(self, level, sampler): diff --git a/qmcpy/integrand/umbridge_wrapper.py b/qmcpy/integrand/umbridge_wrapper.py index 7f2caf7f8..425047ad2 100644 --- a/qmcpy/integrand/umbridge_wrapper.py +++ b/qmcpy/integrand/umbridge_wrapper.py @@ -1,3 +1,7 @@ +from __future__ import annotations + +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import TYPE_CHECKING, Union from .abstract_integrand import AbstractIntegrand from ..discrete_distribution import DigitalNetB2 from ..true_measure import Uniform @@ -5,10 +9,15 @@ import numpy as np import os +if TYPE_CHECKING: + import umbridge + class UMBridgeWrapper(AbstractIntegrand): - """ - Wrapper around a [`UM-Bridge`](https://um-bridge-benchmarks.readthedocs.io/en/docs/index.html) model. See also the [`UM-Bridge` documentation for the QMCPy client](https://um-bridge-benchmarks.readthedocs.io/en/docs/umbridge/clients.html). + """Wrapper around a + [`UM-Bridge`](https://um-bridge-benchmarks.readthedocs.io/en/docs/index.html) + model. See also the [`UM-Bridge` documentation for the QMCPy + client](https://um-bridge-benchmarks.readthedocs.io/en/docs/umbridge/clients.html). Requires [Docker](https://www.docker.com/) is installed. Examples: @@ -64,18 +73,21 @@ class UMBridgeWrapper(AbstractIntegrand): [['-1.59e-08', '1.49e-04', '1.49e-04'], ['8.20e-06', '-1.38e-04'], ['-8.14e-06']] """ - def __init__(self, true_measure, model, config=None, parallel=False): - """ + def __init__(self, true_measure: AbstractTrueMeasure, model: umbridge.HTTPModel, config: Union[None, dict] = None, parallel: Union[bool, int] = False) -> None: + """Initialize a UMBridgeWrapper integrand. + Args: true_measure (AbstractTrueMeasure): The true measure. model (umbridge.HTTPModel): A `UM-Bridge` model. - config (dict): Configuration keyword argument to `umbridge.HTTPModel(url,name).__call__`. - parallel (int): Parallelization flag. + config (Union[None, dict]): Configuration keyword argument to + `umbridge.HTTPModel(url,name).__call__`. + parallel (Union[bool, int]): Parallelization flag. - When `parallel = 0` or `parallel = 1` then function evaluation is done in serial fashion. - `parallel > 1` specifies the number of processes used by `multiprocessing.Pool` or `multiprocessing.pool.ThreadPool`. - Setting `parallel=True` is equivalent to `parallel = os.cpu_count()`. + Setting `parallel=True` is equivalent to `parallel = + os.cpu_count()`. """ if config is None: config = {} @@ -114,7 +126,16 @@ def __init__(self, true_measure, model, config=None, parallel=False): threadpool=True, ) - def g(self, t, **kwargs): + def g(self, t: np.ndarray, **kwargs: dict) -> np.ndarray: + """Evaluate the wrapped UM-Bridge model at each point. + + Args: + t (np.ndarray): Points, model inputs along the last axis. + **kwargs (dict): Unused; accepted for API consistency. + + Returns: + np.ndarray: Model outputs, flattened across the UM-Bridge output blocks. + """ y = np.zeros((self.total_out_elements,) + tuple(t.shape[:-1]), dtype=float) idxiterator = np.ndindex(t.shape[:-1]) for i in idxiterator: @@ -139,15 +160,18 @@ def _spawn(self, _level, _sampler): parallel=self.parallel, ) - def to_umbridge_out_sizes(self, x): - """ - Convert a data attribute to `UM-Bridge` output sized list of lists. + def to_umbridge_out_sizes(self, x: np.ndarray) -> list: + """Convert a data attribute to `UM-Bridge` output sized list of + lists. Args: - x (np.ndarray): Array of length `sum(model.get_output_sizes(self.config))` where `model` is a `umbridge.HTTPModel`. + x (np.ndarray): Array of length + `sum(model.get_output_sizes(self.config))` where `model` is a + `umbridge.HTTPModel`. Returns: - x_list_list (list): List of lists with sub-list lengths specified by `model.get_output_sizes(self.config)`. + list: List of lists with sub-list lengths specified by + `model.get_output_sizes(self.config)`. """ return [ x[..., self.d_out_umbridge[j] : self.d_out_umbridge[j + 1]].tolist() diff --git a/qmcpy/kernel/abstract_kernel.py b/qmcpy/kernel/abstract_kernel.py index 94b89c3fe..27269f684 100644 --- a/qmcpy/kernel/abstract_kernel.py +++ b/qmcpy/kernel/abstract_kernel.py @@ -1,3 +1,9 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union, Tuple, Callable +if TYPE_CHECKING: + import torch + from ..util import MethodImplementationError import numpy as np from ..util.transforms import ( @@ -9,6 +15,13 @@ class AbstractKernel(object): + """Abstract base class for QMCPy kernels. + + Concrete kernels subclass this and implement `parsed___call__` and + `parsed_single_integral_01d` (and optionally `double_integral_01d`); + `AbstractKernel` handles NumPy/PyTorch backend dispatch, batched + parameters, and `torch.compile` wiring. + """ def __new__(cls, *args, **kwargs): if ( @@ -28,10 +41,11 @@ def __new__(cls, *args, **kwargs): instance = super().__new__(cls) return instance - def __init__(self, d, torchify, device, compile_call, compile_call_kwargs): + def __init__(self, d, torchify, device, compile_call, compile_call_kwargs) -> None: super().__init__() # dimension - assert d % 1 == 0 and d > 0, "dimension d must be a positive int" + if not (d % 1 == 0 and d > 0): + raise AssertionError("dimension d must be a positive int") self.d = d # torchify self.torchify = torchify @@ -51,7 +65,8 @@ def __init__(self, d, torchify, device, compile_call, compile_call_kwargs): self.nptkwargs = {} self.batch_param_names = [] if compile_call: - assert self.torchify, "compile_call requires torchify is True" + if not (self.torchify): + raise AssertionError("compile_call requires torchify is True") import torch self.compiled_parsed___call__ = torch.compile( @@ -62,6 +77,10 @@ def __init__(self, d, torchify, device, compile_call, compile_call_kwargs): @property def nbdim(self): + """int: Number of batch dimensions this kernel's output carries, + beyond the sample dimensions. Determined by calling the kernel on + empty inputs and inspecting the output's number of dimensions. + """ empty = self.npt.empty((0, self.d), **self.nptkwargs) v = self.__call__(empty, empty) nbdim = v.ndim - 1 @@ -69,44 +88,68 @@ def nbdim(self): @property def batch_params(self): + """dict: This kernel's batched parameters (e.g. `scale`, + `lengthscales`), keyed by name, at their natural (unbroadcast) shapes. + """ return {pname: getattr(self, pname) for pname in self.batch_param_names} - def get_batch_params(self, ndim): + def get_batch_params(self, ndim: int) -> dict: + """Return `batch_params` with each value reshaped to broadcast against `ndim` extra dimensions. + + Args: + ndim (int): Number of leading sample dimensions to broadcast against. + + Returns: + dict: `batch_params`, each value passed through `insert_batch_dims(value, ndim, -1)`. + """ return { pname: insert_batch_dims(batch_param, ndim, -1) for pname, batch_param in self.batch_params.items() } - def __call__(self, x0, x1, beta0=None, beta1=None, c=None, **kwargs): - r""" - Evaluate the kernel with (optional) partial derivatives + def __call__(self, x0: Union[np.ndarray, torch.Tensor], x1: Union[np.ndarray, torch.Tensor], beta0: Union[None, np.ndarray, torch.Tensor] = None, beta1: Union[None, np.ndarray, torch.Tensor] = None, c: Union[None, np.ndarray, torch.Tensor] = None, **kwargs: dict) -> Union[np.ndarray, torch.Tensor]: + r"""Evaluate the kernel with (optional) partial derivatives - $$\sum_{\ell=1}^p c_{\ell} \partial_{\boldsymbol{x}_0}^{\boldsymbol{\beta}_{\ell 0}} \partial_{\boldsymbol{x}_1}^{\boldsymbol{\beta}_{\ell 1}} K(\boldsymbol{x}_0,\boldsymbol{x}_1).$$ + $$\sum_{\ell=1}^p c_{\ell} + \partial_{\boldsymbol{x}_0}^{\boldsymbol{\beta}_{\ell 0}} + \partial_{\boldsymbol{x}_1}^{\boldsymbol{\beta}_{\ell 1}} + K(\boldsymbol{x}_0,\boldsymbol{x}_1).$$ Args: - x0 (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` first input to kernel with - x1 (Union[np.ndarray, torch.Tensor]): Shape `x1.shape=(...,d)` second input to kernel with - beta0 (Union[np.ndarray, torch.Tensor]): Shape `beta0.shape=(p,d)` derivative orders with respect to first inputs, $\boldsymbol{\beta}_0$. - beta1 (Union[np.ndarray, torch.Tensor]): Shape `beta1.shape=(p,d)` derivative orders with respect to first inputs, $\boldsymbol{\beta}_1$. - c (Union[np.ndarray, torch.Tensor]): Shape `c.shape=(p,)` coefficients of derivatives. - kwargs (dict): keyword arguments to parsed call + x0 (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` + first input to kernel with + x1 (Union[np.ndarray, torch.Tensor]): Shape `x1.shape=(...,d)` + second input to kernel with + beta0 (Union[None, np.ndarray, torch.Tensor]): Shape `beta0.shape=(p,d)` + derivative orders with respect to first inputs, + $\boldsymbol{\beta}_0$. + beta1 (Union[None, np.ndarray, torch.Tensor]): Shape `beta1.shape=(p,d)` + derivative orders with respect to first inputs, + $\boldsymbol{\beta}_1$. + c (Union[None, np.ndarray, torch.Tensor]): Shape `c.shape=(p,)` + coefficients of derivatives. + **kwargs (dict): keyword arguments to parsed call Returns: - k (Union[np.ndarray, torch.Tensor]): Shape `y.shape=(x0+x1).shape[:-1]` kernel evaluations. + Union[np.ndarray, torch.Tensor]: Shape `y.shape=(x0+x1).shape[:-1]` kernel evaluations. """ - assert isinstance(x0, self.nptarraytype) - assert isinstance(x0, self.nptarraytype) - assert ( + if not (isinstance(x0, self.nptarraytype)): + raise AssertionError + if not (isinstance(x1, self.nptarraytype)): + raise AssertionError + if not ( x0.shape[-1] == self.d - ), "the size of the last dimension of x0 must equal d=%d, got x0.shape=%s" % ( - self.d, - str(tuple(x0.shape)), - ) - assert ( + ): + raise AssertionError("the size of the last dimension of x0 must equal d=%d, got x0.shape=%s" % ( + self.d, + str(tuple(x0.shape)), + )) + if not ( x1.shape[-1] == self.d - ), "the size of the last dimension of x1 must equal d=%d, got x1.shape=%s" % ( - self.d, - str(tuple(x1.shape)), - ) + ): + raise AssertionError("the size of the last dimension of x1 must equal d=%d, got x1.shape=%s" % ( + self.d, + str(tuple(x1.shape)), + )) if beta0 is None: beta0 = self.npt.zeros((1, self.d), dtype=int, **self.nptkwargs) if beta1 is None: @@ -117,37 +160,43 @@ def __call__(self, x0, x1, beta0=None, beta1=None, c=None, **kwargs): beta1 = self.nptarray(beta1) beta0 = self.npt.atleast_2d(beta0) beta1 = self.npt.atleast_2d(beta1) - assert ( + if not ( beta0.ndim == 2 and beta1.ndim == 2 - ), "beta0 and beta1 must both be 2 dimensional" + ): + raise AssertionError("beta0 and beta1 must both be 2 dimensional") p = beta0.shape[0] - assert beta0.shape == ( - p, - self.d, - ), "expected beta0.shape=(%d,%d) but got beta0.shape=%s" % ( - p, - self.d, - str(tuple(beta0.shape)), - ) - assert beta1.shape == ( + if not (beta0.shape == ( p, self.d, - ), "expected beta1.shape=(%d,%d) but got beta1.shape=%s" % ( + )): + raise AssertionError("expected beta0.shape=(%d,%d) but got beta0.shape=%s" % ( + p, + self.d, + str(tuple(beta0.shape)), + )) + if not (beta1.shape == ( p, self.d, - str(tuple(beta1.shape)), - ) - assert (beta0 % 1 == 0).all() and (beta0 >= 0).all(), "require int beta0 >= 0" - assert (beta1 % 1 == 0).all() and (beta1 >= 0).all(), "require int beta1 >= 0" + )): + raise AssertionError("expected beta1.shape=(%d,%d) but got beta1.shape=%s" % ( + p, + self.d, + str(tuple(beta1.shape)), + )) + if not ((beta0 % 1 == 0).all() and (beta0 >= 0).all()): + raise AssertionError("require int beta0 >= 0") + if not ((beta1 % 1 == 0).all() and (beta1 >= 0).all()): + raise AssertionError("require int beta1 >= 0") if c is None: c = self.npt.ones(p, **self.nptkwargs) if not isinstance(c, self.nptarraytype): c = self.nptarray(c) c = self.npt.atleast_1d(c) - assert c.shape == (p,), "expected c.shape=(%d,) but got c.shape=%s" % ( - p, - str(tuple(c.shape)), - ) + if not (c.shape == (p,)): + raise AssertionError("expected c.shape=(%d,) but got c.shape=%s" % ( + p, + str(tuple(c.shape)), + )) if not self.AUTOGRADKERNEL: batch_params = self.get_batch_params(max(x0.ndim - 1, x1.ndim - 1)) k = self.compiled_parsed___call__( @@ -160,7 +209,8 @@ def __call__(self, x0, x1, beta0=None, beta1=None, c=None, **kwargs): x0, x1, batch_params, **kwargs ) else: # requires autograd, so self.npt=torch - assert self.torchify, "autograd requires torchify=True" + if not (self.torchify): + raise AssertionError("autograd requires torchify=True") import torch incoming_grad_enabled = torch.is_grad_enabled() @@ -250,60 +300,103 @@ def __call__(self, x0, x1, beta0=None, beta1=None, c=None, **kwargs): return k def parsed___call__(self, *args, **kwargs): + """*Abstract method* computing the kernel on already-validated, + batch-parsed inputs. Called by `__call__` after input validation and + batch-parameter preparation; subclasses implement the actual kernel + formula here. + """ raise MethodImplementationError(self, "parsed___call__") - def single_integral_01d(self, x): - r""" - Evaluate the integral of the kernel over the unit cube + def single_integral_01d(self, x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + r"""Evaluate the integral of the kernel over the unit cube - $$\tilde{K}(\boldsymbol{x}) = \int_{[0,1]^d} K(\boldsymbol{x},\boldsymbol{z}) \; \mathrm{d} \boldsymbol{z}.$$ + $$\tilde{K}(\boldsymbol{x}) = \int_{[0,1]^d} + K(\boldsymbol{x},\boldsymbol{z}) \; \mathrm{d} \boldsymbol{z}.$$ Args: - x (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` first input to kernel with + x (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` first + input to kernel with Returns: - tildek (Union[np.ndarray, torch.Tensor]): Shape `y.shape=x.shape[:-1]` integral kernel evaluations. + Union[np.ndarray, torch.Tensor]: Shape `y.shape=x.shape[:-1]` integral kernel evaluations. """ if self.npt == np: - assert isinstance(x, np.ndarray) + if not (isinstance(x, np.ndarray)): + raise AssertionError else: # self.npt==torch - assert isinstance(x, self.npt.Tensor) - assert ( + if not (isinstance(x, self.npt.Tensor)): + raise AssertionError + if not ( x.shape[-1] == self.d - ), "the size of the last dimension of x must equal d=%d, got x.shape=%s" % ( - self.d, - str(tuple(x.shape)), - ) + ): + raise AssertionError("the size of the last dimension of x must equal d=%d, got x.shape=%s" % ( + self.d, + str(tuple(x.shape)), + )) batch_params = self.get_batch_params(x.ndim - 1) return self.parsed_single_integral_01d(x, batch_params) def parsed_single_integral_01d(self, x, batch_params): + """*Abstract method* computing `single_integral_01d` on already- + validated inputs with batch parameters prepared. Called by + `single_integral_01d`; subclasses implement the actual formula here. + """ raise MethodImplementationError(self, "parsed_single_integral_01d") - def double_integral_01d(self): - r""" - Evaluate the integral of the kernel over the unit cube + def double_integral_01d(self) -> "Union[np.ndarray, torch.Tensor]": + r"""Evaluate the integral of the kernel over the unit cube - $$\tilde{K} = \int_{[0,1]^d} \int_{[0,1]^d} K(\boldsymbol{x},\boldsymbol{z}) \; \mathrm{d} \boldsymbol{x} \; \mathrm{d} \boldsymbol{z}.$$ + $$\tilde{K} = \int_{[0,1]^d} \int_{[0,1]^d} + K(\boldsymbol{x},\boldsymbol{z}) \; \mathrm{d} \boldsymbol{x} \; + \mathrm{d} \boldsymbol{z}.$$ Returns: - tildek (Union[np.ndarray, torch.Tensor]): Double integral kernel evaluations. + Union[np.ndarray, torch.Tensor]: Double integral kernel evaluations. """ raise MethodImplementationError(self, "double_integral_01d") - def rel_pairwise_dist_func(self, x0, x1, lengthscales): + def rel_pairwise_dist_func(self, x0: Union[np.ndarray, torch.Tensor], x1: Union[np.ndarray, torch.Tensor], lengthscales: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + r"""Lengthscale-normalized pairwise distance $\lVert x_0-x_1\rVert / (\sqrt{2}\boldsymbol{\gamma})$. + + A common building block for stationary/RBF-style kernels. + + Args: + x0 (Union[np.ndarray, torch.Tensor]): First input, shape `(...,d)`. + x1 (Union[np.ndarray, torch.Tensor]): Second input, shape `(...,d)`. + lengthscales (Union[np.ndarray, torch.Tensor]): Lengthscales $\boldsymbol{\gamma}$. + + Returns: + Union[np.ndarray, torch.Tensor]: Normalized pairwise distances, shape `(...,)`. + """ return self.npt.linalg.norm((x0 - x1) / (np.sqrt(2) * lengthscales), 2, -1) def parse_assign_param( self, - pname, - param, - shape_param, - requires_grad_param, - tfs_param, - endsize_ops, - constraints, - ): + pname: str, + param: Union[float, np.ndarray, torch.Tensor], + shape_param: list, + requires_grad_param: bool, + tfs_param: Tuple[Callable, Callable], + endsize_ops: list, + constraints: list, + ) -> Union[np.ndarray, torch.Tensor]: + """Validate, transform, and store a kernel hyperparameter. + + Thin wrapper around `qmcpy.util.transforms.parse_assign_param` that + fills in this kernel's backend (`torchify`, `npt`, `nptkwargs`). + + Args: + pname (str): Name of the parameter (for error messages). + param (Union[float, np.ndarray, torch.Tensor]): The raw user-supplied parameter value. + shape_param (list): Shape to broadcast `param` to when it is scalar. + requires_grad_param (bool): If `True` and `torchify`, set `requires_grad=True`. + tfs_param (Tuple[Callable, Callable]): `(to_raw, from_raw)` transform pair. + endsize_ops (list): Allowed sizes for the parameter's trailing dimension. + constraints (list): Named constraints to enforce (e.g. `["POSITIVE"]`). + + Returns: + Union[np.ndarray, torch.Tensor]: The raw (unconstrained) parameter value to store. + """ return parse_assign_param( pname=pname, param=param, @@ -319,38 +412,57 @@ def parse_assign_param( class AbstractKernelScaleLengthscales(AbstractKernel): + """Abstract base class for kernels parameterized by a scale and lengthscales. + + Adds `scale` and `lengthscales` hyperparameters (each stored internally + in an unconstrained "raw" form and exposed via a positivity-preserving + transform) on top of `AbstractKernel`. + """ def __init__( self, - d, - scale=1.0, - lengthscales=1.0, - shape_scale=None, - shape_lengthscales=None, - tfs_scale=(tf_exp_eps_inv, tf_exp_eps), - tfs_lengthscales=(tf_exp_eps_inv, tf_exp_eps), - torchify=False, - requires_grad_scale=True, - requires_grad_lengthscales=True, - device="cpu", - compile_call=False, - compile_call_kwargs=None, - ): - r""" + d: int, + scale: Union[float, np.ndarray, torch.Tensor] = 1.0, + lengthscales: Union[float, np.ndarray, torch.Tensor] = 1.0, + shape_scale: Union[None, list] = None, + shape_lengthscales: Union[None, list] = None, + tfs_scale: Tuple[Callable, Callable] = (tf_exp_eps_inv, tf_exp_eps), + tfs_lengthscales: Tuple[Callable, Callable] = (tf_exp_eps_inv, tf_exp_eps), + torchify: bool = False, + requires_grad_scale: bool = True, + requires_grad_lengthscales: bool = True, + device: Union[str, torch.device] = "cpu", + compile_call: bool = False, + compile_call_kwargs: Union[None, dict] = None, + ) -> None: + r"""Initialize an AbstractKernelScaleLengthscales kernel. + Args: d (int): Dimension. - scale (Union[np.ndarray, torch.Tensor]): Scaling factor $S$. - lengthscales (Union[np.ndarray, torch.Tensor]): Lengthscales $\boldsymbol{\gamma}$. - shape_scale (list): Shape of `scale` when `np.isscalar(scale)`. - shape_lengthscales (list): Shape of `lengthscales` when `np.isscalar(lengthscales)` - tfs_scale (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_lengthscales (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - torchify (bool): If `True`, use the `torch` backend. Set to `True` if computing gradients with respect to inputs and/or hyperparameters. - requires_grad_scale (bool): If `True` and `torchify`, set `requires_grad=True` for `scale`. - requires_grad_lengthscales (bool): If `True` and `torchify`, set `requires_grad=True` for `lengthscales`. - device (torch.device): If `torchify`, put things onto this device. - compile_call (bool): If `True`, `torch.compile` the `parsed___call__` method. - compile_call_kwargs (dict): When `compile_call` is `True`, pass these keyword arguments to `torch.compile`. + scale (Union[float, np.ndarray, torch.Tensor]): Scaling factor $S$. + lengthscales (Union[float, np.ndarray, torch.Tensor]): Lengthscales + $\boldsymbol{\gamma}$. + shape_scale (Union[None, list]): Shape of `scale` when `np.isscalar(scale)`. + shape_lengthscales (Union[None, list]): Shape of `lengthscales` when + `np.isscalar(lengthscales)` + tfs_scale (Tuple[Callable, Callable]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + tfs_lengthscales (Tuple[Callable, Callable]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + torchify (bool): If `True`, use the `torch` backend. Set to `True` + if computing gradients with respect to inputs and/or + hyperparameters. + requires_grad_scale (bool): If `True` and `torchify`, set + `requires_grad=True` for `scale`. + requires_grad_lengthscales (bool): If `True` and `torchify`, set + `requires_grad=True` for `lengthscales`. + device (Union[str, torch.device]): If `torchify`, put things onto this device. + compile_call (bool): If `True`, `torch.compile` the + `parsed___call__` method. + compile_call_kwargs (Union[None, dict]): When `compile_call` is `True`, pass + these keyword arguments to `torch.compile`. """ if shape_scale is None: shape_scale = [1] @@ -388,8 +500,15 @@ def __init__( @property def scale(self): + """Union[np.ndarray, torch.Tensor]: The scaling factor $S$, computed + from the raw (unconstrained) stored value via `tfs_scale`'s inverse transform. + """ return self.tfs_scale[1](self.raw_scale) @property def lengthscales(self): + """Union[np.ndarray, torch.Tensor]: The lengthscales + $\\boldsymbol{\\gamma}$, computed from the raw (unconstrained) stored + value via `tfs_lengthscales`'s inverse transform. + """ return self.tfs_lengthscales[1](self.raw_lengthscales) diff --git a/qmcpy/kernel/common_kernels.py b/qmcpy/kernel/common_kernels.py index 193a694e8..85bfdca20 100644 --- a/qmcpy/kernel/common_kernels.py +++ b/qmcpy/kernel/common_kernels.py @@ -1,3 +1,9 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union, Tuple, Callable +if TYPE_CHECKING: + import torch + from .abstract_kernel import AbstractKernelScaleLengthscales from ..discrete_distribution import DigitalNetB2 from ..util.transforms import tf_exp_eps, tf_exp_eps_inv, tf_identity @@ -8,10 +14,24 @@ class AbstractKernelGaussianSE(AbstractKernelScaleLengthscales): + """Abstract base class for Gaussian / squared-exponential-family kernels. + + Provides the analytic `[0,1]^d` single and double integrals shared by + this whole kernel family; subclasses need only implement `parsed___call__`. + """ AUTOGRADKERNEL = True - def parsed_single_integral_01d(self, x, batch_params): + def parsed_single_integral_01d(self, x: Union[np.ndarray, torch.Tensor], batch_params: dict) -> Union[np.ndarray, torch.Tensor]: + """Analytic single integral of the Gaussian/SE-family kernel over `[0,1]^d`. + + Args: + x (Union[np.ndarray, torch.Tensor]): Points, shape `(...,d)`. + batch_params (dict): Batch-broadcast `scale`/`lengthscales`, from `get_batch_params`. + + Returns: + Union[np.ndarray, torch.Tensor]: Shape `(...,)` integral kernel evaluations. + """ s = batch_params["scale"][..., 0] l = batch_params["lengthscales"] norm_class = ( @@ -25,7 +45,12 @@ def parsed_single_integral_01d(self, x, batch_params): ) return kint - def double_integral_01d(self): + def double_integral_01d(self) -> Union[np.ndarray, torch.Tensor]: + """Analytic double integral of the Gaussian/SE-family kernel over `[0,1]^d x [0,1]^d`. + + Returns: + Union[np.ndarray, torch.Tensor]: Double integral kernel evaluations. + """ erf = self.npt.erf if self.torchify else scipy.special.erf s = self.scale[..., 0] l = self.lengthscales @@ -39,10 +64,11 @@ def double_integral_01d(self): class KernelGaussian(AbstractKernelGaussianSE): - r""" - Gaussian / Squared Exponential kernel implemented using the product of exponentials. + r"""Gaussian / Squared Exponential kernel implemented using the product of + exponentials. - $$K(\boldsymbol{x},\boldsymbol{z}) = S \prod_{j=1}^d \exp\left(-\left(\frac{x_j-z_j}{\sqrt{2} \gamma_j}\right)^2\right)$$ + $$K(\boldsymbol{x},\boldsymbol{z}) = S \prod_{j=1}^d + \exp\left(-\left(\frac{x_j-z_j}{\sqrt{2} \gamma_j}\right)^2\right)$$ Examples: >>> rng = np.random.Generator(np.random.PCG64(7)) @@ -260,6 +286,9 @@ class KernelGaussian(AbstractKernelGaussianSE): """ def parsed___call__(self, x0, x1, batch_params): + """Gaussian / squared exponential kernel evaluation via a direct + elementwise formula; see the class docstring for the formula. + """ scale = batch_params["scale"][..., 0] lengthscales = batch_params["lengthscales"] k = scale * self.npt.exp( @@ -269,11 +298,14 @@ def parsed___call__(self, x0, x1, batch_params): class KernelSquaredExponential(AbstractKernelGaussianSE): - r""" - Gaussian / Squared Exponential kernel implemented using the pairwise distance function. - Please use `KernelGaussian` when using derivative information. + r"""Gaussian / Squared Exponential kernel implemented using the pairwise + distance function. Please use `KernelGaussian` when using derivative + information. - $$K(\boldsymbol{x},\boldsymbol{z}) = S \exp\left(-d_{\boldsymbol{\gamma}}^2(\boldsymbol{x},\boldsymbol{z})\right), \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ + $$K(\boldsymbol{x},\boldsymbol{z}) = S + \exp\left(-d_{\boldsymbol{\gamma}}^2(\boldsymbol{x},\boldsymbol{z})\right), + \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = + \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ Examples: >>> rng = np.random.Generator(np.random.PCG64(7)) @@ -355,6 +387,9 @@ class KernelSquaredExponential(AbstractKernelGaussianSE): """ def parsed___call__(self, x0, x1, batch_params): + """Gaussian / squared exponential kernel evaluation via the pairwise + distance function; see the class docstring for the formula. + """ scale = batch_params["scale"][..., 0] lengthscales = batch_params["lengthscales"] rdists = self.rel_pairwise_dist_func(x0, x1, lengthscales) @@ -363,10 +398,12 @@ def parsed___call__(self, x0, x1, batch_params): class KernelRationalQuadratic(AbstractKernelScaleLengthscales): - r""" - Rational Quadratic kernel + r"""Rational Quadratic kernel - $$K(\boldsymbol{x},\boldsymbol{z}) = S \left(1+\frac{d_{\boldsymbol{\gamma}}^2(\boldsymbol{x},\boldsymbol{z})}{\alpha}\right)^{-\alpha}, \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ + $$K(\boldsymbol{x},\boldsymbol{z}) = S + \left(1+\frac{d_{\boldsymbol{\gamma}}^2(\boldsymbol{x},\boldsymbol{z})}{\alpha}\right)^{-\alpha}, + \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = + \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ Examples: >>> rng = np.random.Generator(np.random.PCG64(7)) @@ -451,43 +488,60 @@ class KernelRationalQuadratic(AbstractKernelScaleLengthscales): def __init__( self, - d, - scale=1.0, - lengthscales=1.0, - alpha=1.0, - shape_scale=None, - shape_lengthscales=None, - shape_alpha=None, - tfs_scale=(tf_exp_eps_inv, tf_exp_eps), - tfs_lengthscales=(tf_exp_eps_inv, tf_exp_eps), - tfs_alpha=(tf_exp_eps_inv, tf_exp_eps), - torchify=False, - requires_grad_scale=True, - requires_grad_lengthscales=True, - requires_grad_alpha=True, - device="cpu", - compile_call=False, - compile_call_kwargs=None, - ): - r""" + d: int, + scale: Union[float, np.ndarray, torch.Tensor] = 1.0, + lengthscales: Union[float, np.ndarray, torch.Tensor] = 1.0, + alpha: Union[float, np.ndarray, torch.Tensor] = 1.0, + shape_scale: Union[None, list] = None, + shape_lengthscales: Union[None, list] = None, + shape_alpha: Union[None, list] = None, + tfs_scale: Tuple[Callable, Callable] = (tf_exp_eps_inv, tf_exp_eps), + tfs_lengthscales: Tuple[Callable, Callable] = (tf_exp_eps_inv, tf_exp_eps), + tfs_alpha: Tuple[Callable, Callable] = (tf_exp_eps_inv, tf_exp_eps), + torchify: bool = False, + requires_grad_scale: bool = True, + requires_grad_lengthscales: bool = True, + requires_grad_alpha: bool = True, + device: Union[str, torch.device] = "cpu", + compile_call: bool = False, + compile_call_kwargs: Union[None, dict] = None, + ) -> None: + r"""Initialize a KernelRationalQuadratic kernel. + Args: d (int): Dimension. - scale (Union[np.ndarray, torch.Tensor]): Scaling factor $S$. - lengthscales (Union[np.ndarray, torch.Tensor]): Lengthscales $\boldsymbol{\gamma}$. - alpha (Union[np.ndarray, torch.Tensor]): Scale mixture parameter $\alpha$. - shape_scale (list): Shape of `scale` when `np.isscalar(scale)`. - shape_lengthscales (list): Shape of `lengthscales` when `np.isscalar(lengthscales)` - shape_alpha (list): Shape of `alpha` when `np.isscalar(alpha)` - tfs_scale (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_lengthscales (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_alpha (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - torchify (bool): If `True`, use the `torch` backend. Set to `True` if computing gradients with respect to inputs and/or hyperparameters. - requires_grad_scale (bool): If `True` and `torchify`, set `requires_grad=True` for `scale`. - requires_grad_lengthscales (bool): If `True` and `torchify`, set `requires_grad=True` for `lengthscales`. - requires_grad_alpha (bool): If `True` and `torchify`, set `requires_grad=True` for `alpha`. - device (torch.device): If `torchify`, put things onto this device. - compile_call (bool): If `True`, `torch.compile` the `parsed___call__` method. - compile_call_kwargs (dict): When `compile_call` is `True`, pass these keyword arguments to `torch.compile`. + scale (Union[float, np.ndarray, torch.Tensor]): Scaling factor $S$. + lengthscales (Union[float, np.ndarray, torch.Tensor]): Lengthscales + $\boldsymbol{\gamma}$. + alpha (Union[float, np.ndarray, torch.Tensor]): Scale mixture parameter + $\alpha$. + shape_scale (Union[None, list]): Shape of `scale` when `np.isscalar(scale)`. + shape_lengthscales (Union[None, list]): Shape of `lengthscales` when + `np.isscalar(lengthscales)` + shape_alpha (Union[None, list]): Shape of `alpha` when `np.isscalar(alpha)` + tfs_scale (Tuple[Callable, Callable]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + tfs_lengthscales (Tuple[Callable, Callable]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + tfs_alpha (Tuple[Callable, Callable]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + torchify (bool): If `True`, use the `torch` backend. Set to `True` + if computing gradients with respect to inputs and/or + hyperparameters. + requires_grad_scale (bool): If `True` and `torchify`, set + `requires_grad=True` for `scale`. + requires_grad_lengthscales (bool): If `True` and `torchify`, set + `requires_grad=True` for `lengthscales`. + requires_grad_alpha (bool): If `True` and `torchify`, set + `requires_grad=True` for `alpha`. + device (Union[str, torch.device]): If `torchify`, put things onto this device. + compile_call (bool): If `True`, `torch.compile` the + `parsed___call__` method. + compile_call_kwargs (Union[None, dict]): When `compile_call` is `True`, pass + these keyword arguments to `torch.compile`. """ if shape_scale is None: shape_scale = [1] @@ -524,9 +578,13 @@ def __init__( @property def alpha(self): + """Union[np.ndarray, torch.Tensor]: The shape/mixture parameter + $\alpha$, computed from the raw stored value via `tfs_alpha`'s inverse transform. + """ return self.tfs_alpha[1](self.raw_alpha) def parsed___call__(self, x0, x1, batch_params): + """Rational quadratic kernel evaluation; see the class docstring for the formula.""" scale = batch_params["scale"][..., 0] lengthscales = batch_params["lengthscales"] alpha = batch_params["alpha"][..., 0] @@ -536,10 +594,12 @@ def parsed___call__(self, x0, x1, batch_params): class KernelMatern12(AbstractKernelScaleLengthscales): - r""" - Matern kernel with $\alpha=1/2$. + r"""Matern kernel with $\alpha=1/2$. - $$K(\boldsymbol{x},\boldsymbol{z}) = S \exp\left(-d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right), \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ + $$K(\boldsymbol{x},\boldsymbol{z}) = S + \exp\left(-d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right), + \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = + \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ Examples: >>> rng = np.random.Generator(np.random.PCG64(7)) @@ -623,6 +683,7 @@ class KernelMatern12(AbstractKernelScaleLengthscales): AUTOGRADKERNEL = True def parsed___call__(self, x0, x1, batch_params): + """Matern 1/2 (exponential) kernel evaluation; see the class docstring for the formula.""" scale = batch_params["scale"][..., 0] lengthscales = batch_params["lengthscales"] rdists = self.rel_pairwise_dist_func(x0, x1, lengthscales) @@ -631,10 +692,12 @@ def parsed___call__(self, x0, x1, batch_params): class KernelMatern32(AbstractKernelScaleLengthscales): - r""" - Matern kernel with $\alpha=3/2$. + r"""Matern kernel with $\alpha=3/2$. - $$K(\boldsymbol{x},\boldsymbol{z}) = S \left(1+\sqrt{3} d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right)\exp\left(-\sqrt{3}d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right), \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ + $$K(\boldsymbol{x},\boldsymbol{z}) = S \left(1+\sqrt{3} + d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right)\exp\left(-\sqrt{3}d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right), + \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = + \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ Examples: >>> rng = np.random.Generator(np.random.PCG64(7)) @@ -718,6 +781,7 @@ class KernelMatern32(AbstractKernelScaleLengthscales): AUTOGRADKERNEL = True def parsed___call__(self, x0, x1, batch_params): + """Matern 3/2 kernel evaluation; see the class docstring for the formula.""" scale = batch_params["scale"][..., 0] lengthscales = batch_params["lengthscales"] rdists = self.rel_pairwise_dist_func(x0, x1, lengthscales) @@ -726,10 +790,13 @@ def parsed___call__(self, x0, x1, batch_params): class KernelMatern52(AbstractKernelScaleLengthscales): - r""" - Matern kernel with $\alpha=5/2$. + r"""Matern kernel with $\alpha=5/2$. - $$K(\boldsymbol{x},\boldsymbol{z}) = S \left(1+\sqrt{5} d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) + \frac{5}{3} d_{\boldsymbol{\gamma}}^2(\boldsymbol{x},\boldsymbol{z})\right)\exp\left(-\sqrt{5}d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right), \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ + $$K(\boldsymbol{x},\boldsymbol{z}) = S \left(1+\sqrt{5} + d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) + \frac{5}{3} + d_{\boldsymbol{\gamma}}^2(\boldsymbol{x},\boldsymbol{z})\right)\exp\left(-\sqrt{5}d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z})\right), + \qquad d_{\boldsymbol{\gamma}}(\boldsymbol{x},\boldsymbol{z}) = + \left\lVert\frac{\boldsymbol{x}-\boldsymbol{z}}{\sqrt{2}\boldsymbol{\gamma}}\right\rVert_2.$$ Examples: >>> rng = np.random.Generator(np.random.PCG64(7)) @@ -813,6 +880,7 @@ class KernelMatern52(AbstractKernelScaleLengthscales): AUTOGRADKERNEL = True def parsed___call__(self, x0, x1, batch_params): + """Matern 5/2 kernel evaluation; see the class docstring for the formula.""" scale = batch_params["scale"][..., 0] lengthscales = batch_params["lengthscales"] rdists = self.rel_pairwise_dist_func(x0, x1, lengthscales) diff --git a/qmcpy/kernel/multitask_kernel.py b/qmcpy/kernel/multitask_kernel.py index 9f5c3219a..12854cbd5 100644 --- a/qmcpy/kernel/multitask_kernel.py +++ b/qmcpy/kernel/multitask_kernel.py @@ -1,3 +1,9 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union, Tuple, Callable +if TYPE_CHECKING: + import torch + from .abstract_kernel import AbstractKernel from .common_kernels import KernelGaussian from ..util.transforms import tf_identity, tf_exp_eps, tf_exp_eps_inv, insert_batch_dims @@ -5,14 +11,16 @@ class KernelMultiTask(AbstractKernel): - r""" - Multi-task kernel + r"""Multi-task kernel - $$K((i,\boldsymbol{x}),(j,\boldsymbol{z})) = K_{\mathrm{task}}(i,j) K_{\mathrm{base}}(\boldsymbol{x},\boldsymbol{z})$$ + $$K((i,\boldsymbol{x}),(j,\boldsymbol{z})) = K_{\mathrm{task}}(i,j) + K_{\mathrm{base}}(\boldsymbol{x},\boldsymbol{z})$$ - parameterized for $T$ tasks by a factor $\mathsf{F} \in \mathbb{R}^{T \times r}$ and a diagonal $\boldsymbol{v} \in \mathbb{R}^T$ so that + parameterized for $T$ tasks by a factor $\mathsf{F} \in \mathbb{R}^{T + \times r}$ and a diagonal $\boldsymbol{v} \in \mathbb{R}^T$ so that - $$\left[K_{\mathrm{task}}(i,j)\right]_{i,j=1}^T = \mathsf{F} \mathsf{F}^T + \mathrm{diag}(\boldsymbol{v}).$$ + $$\left[K_{\mathrm{task}}(i,j)\right]_{i,j=1}^T = \mathsf{F} \mathsf{F}^T + + \mathrm{diag}(\boldsymbol{v}).$$ Examples: >>> kmt = KernelMultiTask(KernelGaussian(d=2),num_tasks=3,diag=[1,2,3]) @@ -326,34 +334,46 @@ class KernelMultiTask(AbstractKernel): def __init__( self, - base_kernel, - num_tasks, - factor=1.0, - diag=1.0, - shape_factor=None, - shape_diag=None, - tfs_factor=(tf_identity, tf_identity), - tfs_diag=(tf_exp_eps_inv, tf_exp_eps), - requires_grad_factor=True, - requires_grad_diag=True, - rank_factor=1, - method="LOW RANK", - ): - r""" + base_kernel: AbstractKernel, + num_tasks: int, + factor: Union[float, np.ndarray, torch.Tensor] = 1.0, + diag: Union[float, np.ndarray, torch.Tensor] = 1.0, + shape_factor: Union[None, list] = None, + shape_diag: Union[None, list] = None, + tfs_factor: Tuple[Callable, Callable] = (tf_identity, tf_identity), + tfs_diag: Tuple[Callable, Callable] = (tf_exp_eps_inv, tf_exp_eps), + requires_grad_factor: bool = True, + requires_grad_diag: bool = True, + rank_factor: int = 1, + method: str = "LOW RANK", + ) -> None: + r"""Initialize a KernelMultiTask kernel. + Args: base_kernel (AbstractKernel): $K_{\mathrm{base}}$. num_tasks (int): Number of tasks $T>1$. - factor (Union[np.ndarray, torch.Tensor]): Factor $\mathsf{F}$. - diag (Union[np.ndarray, torch.Tensor]): Diagonal parameter $\boldsymbol{v}$. - shape_factor (list): Shape of `factor` when `np.isscalar(factor)`. - shape_diag (list): Shape of `diag` when `np.isscalar(diag)`. - tfs_factor (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_diag (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - requires_grad_factor (bool): If `True` and `torchify`, set `requires_grad=True` for `factor`. - requires_grad_diag (bool): If `True` and `torchify`, set `requires_grad=True` for `diag`. + factor (Union[float, np.ndarray, torch.Tensor]): Factor $\mathsf{F}$. + diag (Union[float, np.ndarray, torch.Tensor]): Diagonal parameter + $\boldsymbol{v}$. + shape_factor (Union[None, list]): Shape of `factor` when `np.isscalar(factor)`. + shape_diag (Union[None, list]): Shape of `diag` when `np.isscalar(diag)`. + tfs_factor (Tuple[Callable, Callable]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + tfs_diag (Tuple[Callable, Callable]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + requires_grad_factor (bool): If `True` and `torchify`, set + `requires_grad=True` for `factor`. + requires_grad_diag (bool): If `True` and `torchify`, set + `requires_grad=True` for `diag`. + rank_factor (int): Rank of the low-rank `factor` matrix when + `method="LOW RANK"` and `shape_factor` is not given; must + satisfy `0 <= rank_factor <= num_tasks`. method (str): `"LOW RANK"` or "CHOLESKY" """ - assert isinstance(base_kernel, AbstractKernel) + if not (isinstance(base_kernel, AbstractKernel)): + raise AssertionError super().__init__( d=base_kernel.d, torchify=base_kernel.torchify, @@ -363,13 +383,15 @@ def __init__( ) self.base_kernel = base_kernel self.AUTOGRADKERNEL = base_kernel.AUTOGRADKERNEL - assert np.isscalar(num_tasks) and num_tasks % 1 == 0 + if not (np.isscalar(num_tasks) and num_tasks % 1 == 0): + raise AssertionError self.num_tasks = num_tasks - assert ( + if not ( np.isscalar(rank_factor) and rank_factor % 1 == 0 and 0 <= rank_factor <= self.num_tasks - ) + ): + raise AssertionError self.method = str(method).upper().replace("_", " ").strip() if self.method == "LOW RANK": if shape_factor is None: @@ -400,7 +422,8 @@ def __init__( ) self.tfs_factor = tfs_factor if self.method == "LOW RANK": - assert self.raw_factor.shape[-2] == self.num_tasks + if not (self.raw_factor.shape[-2] == self.num_tasks): + raise AssertionError self.raw_diag = self.parse_assign_param( pname="diag", param=diag, @@ -417,20 +440,33 @@ def __init__( @property def nbdim_base(self): + """int: `nbdim` of the wrapped `base_kernel` (cached after first access).""" if self._nbdim_base is None: self._nbdim_base = self.base_kernel.nbdim return self._nbdim_base @property def factor(self): + """Union[np.ndarray, torch.Tensor]: Low-rank/Cholesky factor used to + build the task covariance matrix `taskmat`, computed from the raw + stored value via `tfs_factor`'s inverse transform. + """ return self.tfs_factor[1](self.raw_factor) @property def diag(self): + """Union[np.ndarray, torch.Tensor]: Diagonal term added to the task + covariance matrix `taskmat`, computed from the raw stored value via + `tfs_diag`'s inverse transform. + """ return self.tfs_diag[1](self.raw_diag) @property def taskmat(self): + """Union[np.ndarray, torch.Tensor]: The `(num_tasks, num_tasks)` task + covariance matrix, built from `factor` and `diag` using either the + `"LOW RANK"` or `"CHOLESKY"` parameterization (see `method`). + """ factor = self.factor diag = self.diag if self.method == "LOW RANK": @@ -463,67 +499,93 @@ def _parsed__call__(self, task0, task1, k_x): kmat = k_x * kmat_tasks return kmat[..., 0] - def __call__(self, task0, task1, x0, x1, beta0=None, beta1=None, c=None): - r""" - Evaluate the kernel with (optional) partial derivatives + def __call__(self, task0: Union[int, np.ndarray, torch.Tensor], task1: Union[int, np.ndarray, torch.Tensor], x0: Union[np.ndarray, torch.Tensor], x1: Union[np.ndarray, torch.Tensor], beta0: Union[None, np.ndarray, torch.Tensor] = None, beta1: Union[None, np.ndarray, torch.Tensor] = None, c: Union[None, np.ndarray, torch.Tensor] = None) -> Union[np.ndarray, torch.Tensor]: + r"""Evaluate the kernel with (optional) partial derivatives - $$\sum_{\ell=1}^p c_\ell \partial_{\boldsymbol{x}_0}^{\boldsymbol{\beta}_{\ell,0}} \partial_{\boldsymbol{x}_1}^{\boldsymbol{\beta}_{\ell,1}} K((i_0,\boldsymbol{x}_0),(i_1,\boldsymbol{x}_1)).$$ + $$\sum_{\ell=1}^p c_\ell + \partial_{\boldsymbol{x}_0}^{\boldsymbol{\beta}_{\ell,0}} + \partial_{\boldsymbol{x}_1}^{\boldsymbol{\beta}_{\ell,1}} + K((i_0,\boldsymbol{x}_0),(i_1,\boldsymbol{x}_1)).$$ Args: - task0 (Union[int, np.ndarray, torch.Tensor]): First task indices $i_0$. - task1 (Union[int, np.ndarray, torch.Tensor]): Second task indices $i_1$. - x0 (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` first input to kernel. - x1 (Union[np.ndarray, torch.Tensor]): Shape `x1.shape=(...,d)` second input to kernel. - beta0 (Union[np.ndarray, torch.Tensor]): Shape `beta0.shape=(p,d)` derivative orders with respect to first inputs, $\boldsymbol{\beta}_0$. - beta1 (Union[np.ndarray, torch.Tensor]): Shape `beta1.shape=(p,d)` derivative orders with respect to first inputs, $\boldsymbol{\beta}_1$. - c (Union[np.ndarray, torch.Tensor]): Shape `c.shape=(p,)` coefficients of derivatives. + task0 (Union[int, np.ndarray, torch.Tensor]): First task indices + $i_0$. + task1 (Union[int, np.ndarray, torch.Tensor]): Second task indices + $i_1$. + x0 (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` + first input to kernel. + x1 (Union[np.ndarray, torch.Tensor]): Shape `x1.shape=(...,d)` + second input to kernel. + beta0 (Union[None, np.ndarray, torch.Tensor]): Shape `beta0.shape=(p,d)` + derivative orders with respect to first inputs, + $\boldsymbol{\beta}_0$. + beta1 (Union[None, np.ndarray, torch.Tensor]): Shape `beta1.shape=(p,d)` + derivative orders with respect to first inputs, + $\boldsymbol{\beta}_1$. + c (Union[None, np.ndarray, torch.Tensor]): Shape `c.shape=(p,)` + coefficients of derivatives. Returns: - k (Union[np.ndarray, torch.Tensor]): Kernel evaluations with batched shape, see the doctests for examples. + Union[np.ndarray, torch.Tensor]: Kernel evaluations with batched shape, see the doctests for + examples. """ kmat_x = self.base_kernel.__call__(x0, x1, beta0, beta1, c) return self._parsed__call__(task0, task1, kmat_x) - def single_integral_01d(self, task0, task1, x): - r""" - Evaluate the integral of the kernel over the unit cube + def single_integral_01d(self, task0: Union[int, np.ndarray, torch.Tensor], task1: Union[int, np.ndarray, torch.Tensor], x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + r"""Evaluate the integral of the kernel over the unit cube - $$\tilde{K}((i_0,\boldsymbol{x}),i_1) = \int_{[0,1]^d} K((i_0,\boldsymbol{x}),(i_1,\boldsymbol{z}) \; \mathrm{d} \boldsymbol{z}.$$ + $$\tilde{K}((i_0,\boldsymbol{x}),i_1) = \int_{[0,1]^d} + K((i_0,\boldsymbol{x}),(i_1,\boldsymbol{z}) \; \mathrm{d} + \boldsymbol{z}.$$ Args: - task0 (Union[int, np.ndarray, torch.Tensor]): First task indices $i_0$. - task1 (Union[int, np.ndarray, torch.Tensor]): Second task indices $i_1$. - x (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` first input to kernel with + task0 (Union[int, np.ndarray, torch.Tensor]): First task indices + $i_0$. + task1 (Union[int, np.ndarray, torch.Tensor]): Second task indices + $i_1$. + x (Union[np.ndarray, torch.Tensor]): Shape `x0.shape=(...,d)` first + input to kernel with Returns: - tildek (Union[np.ndarray, torch.Tensor]): Shape `y.shape=x.shape[:-1]` integral kernel evaluations. + Union[np.ndarray, torch.Tensor]: Shape `y.shape=x.shape[:-1]` integral kernel evaluations. """ kint_x = self.base_kernel.single_integral_01d(x) return self._parsed__call__(task0, task1, kint_x) - def double_integral_01d(self, task0, task1): - r""" - Evaluate the integral of the kernel over the unit cube + def double_integral_01d(self, task0: Union[int, np.ndarray, torch.Tensor], task1: Union[int, np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + r"""Evaluate the integral of the kernel over the unit cube - $$\tilde{K}(i_0,i_1) = \int_{[0,1]^d} \int_{[0,1]^d} K((i_0,\boldsymbol{x}),(i_1,\boldsymbol{z})) \; \mathrm{d} \boldsymbol{x} \; \mathrm{d} \boldsymbol{z}.$$ + $$\tilde{K}(i_0,i_1) = \int_{[0,1]^d} \int_{[0,1]^d} + K((i_0,\boldsymbol{x}),(i_1,\boldsymbol{z})) \; \mathrm{d} + \boldsymbol{x} \; \mathrm{d} \boldsymbol{z}.$$ Args: - task0 (Union[int, np.ndarray, torch.Tensor]): First task indices $i_0$. - task1 (Union[int, np.ndarray, torch.Tensor]): Second task indices $i_1$. + task0 (Union[int, np.ndarray, torch.Tensor]): First task indices + $i_0$. + task1 (Union[int, np.ndarray, torch.Tensor]): Second task indices + $i_1$. Returns: - tildek (Union[np.ndarray, torch.Tensor]): Double integral kernel evaluations. + Union[np.ndarray, torch.Tensor]: Double integral kernel evaluations. """ kint_x = self.base_kernel.double_integral_01d() return self._parsed__call__(task0, task1, kint_x) class KernelMultiTaskDerivs(KernelMultiTask): + """`KernelMultiTask` specialized for taking derivatives across tasks. + + Fixes the task covariance matrix to the identity (`factor=1.0`, + `diag=0.0`, both non-trainable), so tasks are treated as independent and + the multi-task kernel reduces to `base_kernel` applied per task. + """ + def __init__( self, base_kernel, num_tasks, - ): + ) -> None: super().__init__( base_kernel=base_kernel, num_tasks=num_tasks, diff --git a/qmcpy/kernel/si_dsi_kernels.py b/qmcpy/kernel/si_dsi_kernels.py index 83ea718e9..3d8f7f3fb 100644 --- a/qmcpy/kernel/si_dsi_kernels.py +++ b/qmcpy/kernel/si_dsi_kernels.py @@ -1,3 +1,9 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union, Tuple, Callable +if TYPE_CHECKING: + import torch + from .abstract_kernel import AbstractKernelScaleLengthscales from ..util.transforms import tf_exp_eps, tf_exp_eps_inv, tf_identity from ..util.shift_invar_ops import BERNOULLIPOLYSDICT, bernoulli_poly @@ -13,6 +19,18 @@ class AbstractSIDSIKernel(AbstractKernelScaleLengthscales): + """Abstract base class for shift-invariant and digitally-shift-invariant + (Walsh) kernels, parameterized by smoothness `alpha`, `lengthscales`, and + `scale`. + + Subclasses implement `get_per_dim_components`, which builds the family- + specific per-dimension building blocks (Bernoulli polynomials for + shift-invariant kernels, weighted Walsh functions for digitally-shift- + invariant kernels); this base class handles combining them (optionally + with derivative orders `beta0`/`beta1`) into the full kernel evaluation, + plus the `[0,1]^d` single/double integrals, which are constant (`scale`) + for this whole kernel family. + """ AUTOGRADKERNEL = False @@ -40,7 +58,7 @@ def __init__( shape_weights, tfs_weights, requires_grad_weights, - ): + ) -> None: # alias lengthscales with weights if weights is not None: if lengthscales is not None: @@ -113,17 +131,45 @@ def __init__( @property def alpha(self): + """Union[np.ndarray, torch.Tensor]: The smoothness parameter + $\\boldsymbol{\\alpha}$, computed from the raw stored value via + `tfs_alpha`'s inverse transform. + """ return self.tfs_alpha[1](self.raw_alpha) def parsed_single_integral_01d(self, x, batch_params): + """Single integral of this kernel family over `[0,1]^d`, which is + the constant `scale` (a reproducing-kernel property of shift- + invariant/digitally-shift-invariant kernels). + """ return batch_params["scale"][..., 0] + 0 * x[..., 0] def double_integral_01d(self): + """Double integral of this kernel family over `[0,1]^d x [0,1]^d`, + which is the constant `scale` (same reproducing-kernel property as + `parsed_single_integral_01d`). + """ return self.scale[..., 0] def combine_per_dim_components_raw_m1( - self, kparts, beta0, beta1, c, batch_params, stable - ): + self, kparts: Union[np.ndarray, torch.Tensor], beta0: Union[np.ndarray, torch.Tensor], beta1: Union[np.ndarray, torch.Tensor], c: Union[np.ndarray, torch.Tensor], batch_params: dict, stable: bool + ) -> tuple[Union[np.ndarray, torch.Tensor], Union[np.ndarray, torch.Tensor]]: + """Combine per-dimension kernel components into `(scale_term, remainder)`. + + Args: + kparts (Union[np.ndarray, torch.Tensor]): Per-dimension components from `get_per_dim_components`. + beta0 (Union[np.ndarray, torch.Tensor]): Derivative orders for the first input. + beta1 (Union[np.ndarray, torch.Tensor]): Derivative orders for the second input. + c (Union[np.ndarray, torch.Tensor]): Coefficients of the derivative terms. + batch_params (dict): Batch-broadcast `scale`/`lengthscales`, from `get_batch_params`. + stable (bool): If `True`, use a numerically stabler (but more + expensive) product formula. + + Returns: + tuple[Union[np.ndarray, torch.Tensor], Union[np.ndarray, torch.Tensor]]: + `(sc, v)` such that the full kernel value is `sc + v`; + `combine_per_dim_components` adds these back together. + """ scale = batch_params["scale"][..., 0] lengthscales = batch_params["lengthscales"] ind = 1.0 * ((beta0 + beta1) == 0) @@ -143,9 +189,27 @@ def combine_per_dim_components_raw_m1( return sc, v def get_per_dim_components(self, x0, x1, beta0, beta1): + """*Abstract method* building this kernel family's per-dimension + components (e.g. Bernoulli polynomials or weighted Walsh functions, + depending on the subclass), with derivative orders `beta0`/`beta1` + applied. Called by `parsed___call__`. + """ raise MethodImplementationError(self, "get_per_dim_components") - def combine_per_dim_components(self, kparts, beta0, beta1, c, batch_params, stable): + def combine_per_dim_components(self, kparts: Union[np.ndarray, torch.Tensor], beta0: Union[np.ndarray, torch.Tensor], beta1: Union[np.ndarray, torch.Tensor], c: Union[np.ndarray, torch.Tensor], batch_params: dict, stable: bool) -> Union[np.ndarray, torch.Tensor]: + """Combine per-dimension kernel components into the final kernel value. + + Args: + kparts (Union[np.ndarray, torch.Tensor]): Per-dimension components from `get_per_dim_components`. + beta0 (Union[np.ndarray, torch.Tensor]): Derivative orders for the first input. + beta1 (Union[np.ndarray, torch.Tensor]): Derivative orders for the second input. + c (Union[np.ndarray, torch.Tensor]): Coefficients of the derivative terms. + batch_params (dict): Batch-broadcast `scale`/`lengthscales`, from `get_batch_params`. + stable (bool): If `True`, use a numerically stabler product formula. + + Returns: + Union[np.ndarray, torch.Tensor]: The kernel evaluation. + """ sc, v = self.combine_per_dim_components_raw_m1( kparts, beta0, beta1, c, batch_params, stable ) @@ -153,6 +217,7 @@ def combine_per_dim_components(self, kparts, beta0, beta1, c, batch_params, stab return k def parsed___call__(self, x0, x1, beta0, beta1, c, batch_params, stable=False): + """Evaluate the kernel by building then combining per-dimension components.""" kparts = self.get_per_dim_components(x0, x1, beta0, beta1) k = self.combine_per_dim_components( kparts, beta0, beta1, c, batch_params, stable @@ -161,17 +226,16 @@ def parsed___call__(self, x0, x1, beta0, beta1, c, batch_params, stable=False): class KernelShiftInvar(AbstractSIDSIKernel): - r""" - Shift invariant kernel with - smoothness $\boldsymbol{\alpha}$, product weights (lengthscales) $\boldsymbol{\gamma}$, and scale $S$: - - $$\begin{aligned} - K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d \left(1+ \gamma_j \tilde{K}_{\alpha_j}((x_j - z_j) \mod 1))\right), \\ - \tilde{K}_\alpha(x) &= (-1)^{\alpha+1}\frac{(2 \pi)^{2 \alpha}}{(2\alpha)!} B_{2\alpha}(x) - \end{aligned}$$ + r"""Shift invariant kernel with smoothness $\boldsymbol{\alpha}$, product + weights (lengthscales) $\boldsymbol{\gamma}$, and scale $S$: + + $$\begin{aligned} K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d + \left(1+ \gamma_j \tilde{K}_{\alpha_j}((x_j - z_j) \mod 1))\right), \\ + \tilde{K}_\alpha(x) &= (-1)^{\alpha+1}\frac{(2 \pi)^{2 \alpha}}{(2\alpha)!} + B_{2\alpha}(x) \end{aligned}$$ where $B_n$ is the $n^\text{th}$ Bernoulli polynomial. - + Examples: >>> from qmcpy import Lattice, fftbr, ifftbr >>> n = 8 @@ -183,7 +247,7 @@ class KernelShiftInvar(AbstractSIDSIKernel): >>> x.dtype dtype('float64') >>> kernel = KernelShiftInvar( - ... d = d, + ... d = d, ... alpha = list(range(1,d+1)), ... scale = 10, ... lengthscales = [1/j**2 for j in range(1,d+1)]) @@ -213,10 +277,10 @@ class KernelShiftInvar(AbstractSIDSIKernel): True >>> np.allclose(ifftbr(fftbr(y)/lam),np.linalg.solve(kmat,y)) True - >>> import torch + >>> import torch >>> xtorch = torch.from_numpy(x) >>> kernel_torch = KernelShiftInvar( - ... d = d, + ... d = d, ... alpha = list(range(1,d+1)), ... scale = 10, ... lengthscales = [1/j**2 for j in range(1,d+1)], @@ -229,16 +293,16 @@ class KernelShiftInvar(AbstractSIDSIKernel): >>> kernel_torch.single_integral_01d(xtorch) tensor([10., 10., 10., 10., 10., 10., 10., 10.], dtype=torch.float64, grad_fn=) - - Batch Params - + + Batch Params + >>> rng = np.random.Generator(np.random.PCG64(7)) >>> kernel = KernelShiftInvar( - ... d = 2, + ... d = 2, ... shape_scale = [4,3,1], ... shape_lengthscales = [3,2]) >>> x = rng.uniform(low=0,high=1,size=(6,5,2)) - >>> kernel(x,x).shape + >>> kernel(x,x).shape (4, 3, 6, 5) >>> kernel(x[:,:,None,:],x[:,None,:,:]).shape (4, 3, 6, 5, 5) @@ -249,7 +313,7 @@ class KernelShiftInvar(AbstractSIDSIKernel): >>> np.abs(kfast-kstable).max() np.float64(4.440892098500626e-16) - Derivatives + Derivatives >>> rng = np.random.Generator(np.random.PCG64(7)) >>> scale = rng.uniform(low=0,high=1,size=(1,)) @@ -309,54 +373,72 @@ class KernelShiftInvar(AbstractSIDSIKernel): >>> np.allclose(ynp,y.numpy()) True - **References:** - - 1. Kaarnioja, Vesa, Frances Y. Kuo, and Ian H. Sloan. - "Lattice-based kernel approximation and serendipitous weights for parametric PDEs in very high dimensions." + **References:** + + 1. Kaarnioja, Vesa, Frances Y. Kuo, and Ian H. Sloan. + "Lattice-based kernel approximation and serendipitous weights for parametric PDEs in very high dimensions." International Conference on Monte Carlo and Quasi-Monte Carlo Methods in Scientific Computing. Cham: Springer International Publishing, 2022. """ def __init__( self, - d, - scale=1.0, - lengthscales=None, - alpha=2, - shape_scale=None, - shape_lengthscales=None, - tfs_scale=None, - tfs_lengthscales=None, - torchify=False, - requires_grad_scale=None, - requires_grad_lengthscales=None, - device="cpu", - compile_call=False, - compile_call_kwargs=None, - weights=None, - shape_weights=None, - tfs_weights=None, - requires_grad_weights=None, - ): - r""" + d: int, + scale: Union[float, np.ndarray, torch.Tensor] = 1.0, + lengthscales: Union[None, np.ndarray, torch.Tensor] = None, + alpha: Union[float, np.ndarray, torch.Tensor] = 2, + shape_scale: Union[None, list] = None, + shape_lengthscales: Union[None, list] = None, + tfs_scale: Union[None, Tuple[Callable, Callable]] = None, + tfs_lengthscales: Union[None, Tuple[Callable, Callable]] = None, + torchify: bool = False, + requires_grad_scale: Union[None, bool] = None, + requires_grad_lengthscales: Union[None, bool] = None, + device: Union[str, torch.device] = "cpu", + compile_call: bool = False, + compile_call_kwargs: Union[None, dict] = None, + weights: Union[None, np.ndarray, torch.Tensor] = None, + shape_weights: Union[None, list] = None, + tfs_weights: Union[None, Tuple[Callable, Callable]] = None, + requires_grad_weights: Union[None, bool] = None, + ) -> None: + r"""Initialize a KernelShiftInvar kernel. + Args: d (int): Dimension. - scale (Union[np.ndarray, torch.Tensor]): Scaling factor $S$. - lengthscales (Union[np.ndarray, torch.Tensor]): Product weights $(\gamma_1,\dots,\gamma_d)$. - alpha (Union[np.ndarray, torch.Tensor]): Smoothness parameters $(\alpha_1,\dots,\alpha_d)$ where $\alpha_j \geq 1$ for $j=1,\dots,d$. - shape_scale (list): Shape of `scale` when `np.isscalar(scale)`. - shape_lengthscales (list): Shape of `lengthscales` when `np.isscalar(lengthscales)` - tfs_scale (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_lengthscales (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - torchify (bool): If `True`, use the `torch` backend. Set to `True` if computing gradients with respect to inputs and/or hyperparameters. - requires_grad_scale (bool): If `True` and `torchify`, set `requires_grad=True` for `scale`. - requires_grad_lengthscales (bool): If `True` and `torchify`, set `requires_grad=True` for `lengthscales`. - device (torch.device): If `torchify`, put things onto this device. - compile_call (bool): If `True`, `torch.compile` the `parsed___call__` method. - compile_call_kwargs (dict): When `compile_call` is `True`, pass these keyword arguments to `torch.compile`. - weights (Union[np.ndarray, torch.Tensor]): Alias for `lengthscales`. - shape_weights (list): Alias for `shape_lengthscales`. - tfs_weights (Tuple[callable,callable]): Alias for `tfs_lengthscales`. - requires_grad_weights (bool): Alias for `requires_grad_lengthscales`. + scale (Union[float, np.ndarray, torch.Tensor]): Scaling factor $S$. + lengthscales (Union[None, np.ndarray, torch.Tensor]): Product weights + $(\gamma_1,\dots,\gamma_d)$. + alpha (Union[float, np.ndarray, torch.Tensor]): Smoothness parameters + $(\alpha_1,\dots,\alpha_d)$ where $\alpha_j \geq 1$ for + $j=1,\dots,d$. + shape_scale (Union[None, list]): Shape of `scale` when `np.isscalar(scale)`. + shape_lengthscales (Union[None, list]): Shape of `lengthscales` when + `np.isscalar(lengthscales)` + tfs_scale (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + tfs_lengthscales (Union[None, Tuple[Callable, Callable]]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + torchify (bool): If `True`, use the `torch` backend. Set to `True` + if computing gradients with respect to inputs and/or + hyperparameters. + requires_grad_scale (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `scale`. + requires_grad_lengthscales (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `lengthscales`. + device (Union[str, torch.device]): If `torchify`, put things onto this device. + compile_call (bool): If `True`, `torch.compile` the + `parsed___call__` method. + compile_call_kwargs (Union[None, dict]): When `compile_call` is `True`, pass + these keyword arguments to `torch.compile`. + weights (Union[None, np.ndarray, torch.Tensor]): Alias for + `lengthscales`. + shape_weights (Union[None, list]): Alias for `shape_lengthscales`. + tfs_weights (Union[None, Tuple[Callable, Callable]]): Alias for + `tfs_lengthscales`. + requires_grad_weights (Union[None, bool]): Alias for + `requires_grad_lengthscales`. """ if shape_scale is None: shape_scale = [1] @@ -386,8 +468,10 @@ def __init__( tfs_weights=tfs_weights, requires_grad_weights=requires_grad_weights, ) - assert self.alpha.shape == (self.d,) - assert all(int(alphaj) in BERNOULLIPOLYSDICT for alphaj in self.alpha) + if not (self.alpha.shape == (self.d,)): + raise AssertionError + if not (all(int(alphaj) in BERNOULLIPOLYSDICT for alphaj in self.alpha)): + raise AssertionError if self.torchify: import torch @@ -396,12 +480,16 @@ def __init__( self.lgamma = scipy.special.loggamma def get_per_dim_components(self, x0, x1, beta0, beta1): + """Per-dimension Bernoulli-polynomial components; see the class + docstring for the kernel formula. + """ p = len(beta0) betasum = beta0 + beta1 order = 2 * self.alpha - betasum - assert ( + if not (( 2 <= order - ).all(), "order must all be at least 2, but got order = %s" % str(order) + ).all()): + raise AssertionError("order must all be at least 2, but got order = %s" % str(order)) coeffs = (-1) ** (self.alpha + beta1 + 1) * self.npt.exp( 2 * self.alpha * np.log(2 * np.pi) - self.lgamma(order + 1) ) @@ -423,15 +511,16 @@ def get_per_dim_components(self, x0, x1, beta0, beta1): class KernelShiftInvarCombined(AbstractSIDSIKernel): - r""" - Shift invariant kernel with - combination weights $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in \mathbb{R}_{>0}^4$, product weights (lengthscales) $\boldsymbol{\gamma}$, and scale $S$: + r"""Shift invariant kernel with combination weights + $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in \mathbb{R}_{>0}^4$, + product weights (lengthscales) $\boldsymbol{\gamma}$, and scale $S$: - $$\begin{aligned} - K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d \left(1+ \gamma_j \left(\sum_{p=1}^4 \alpha_{jp} \tilde{K}_p(x_j \mod 1 z_j)\right)\right) - \end{aligned}$$ + $$\begin{aligned} K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d + \left(1+ \gamma_j \left(\sum_{p=1}^4 \alpha_{jp} \tilde{K}_p(x_j \mod 1 + z_j)\right)\right) \end{aligned}$$ - where, $\tilde{K}_p$ are defined in `KernelShiftInvar` for $p \in \{1,2,3,4\}$ + where, $\tilde{K}_p$ are defined in `KernelShiftInvar` for $p \in + \{1,2,3,4\}$ Examples: >>> from qmcpy import Lattice, fftbr, ifftbr @@ -517,51 +606,72 @@ class KernelShiftInvarCombined(AbstractSIDSIKernel): def __init__( self, - d, - scale=1.0, - lengthscales=None, - alpha=1, - shape_scale=None, - shape_lengthscales=None, - shape_alpha=None, - tfs_scale=None, - tfs_lengthscales=None, - tfs_alpha=None, - torchify=False, - requires_grad_scale=None, - requires_grad_lengthscales=None, - requires_grad_alpha=None, - device="cpu", - compile_call=False, - compile_call_kwargs=None, - weights=None, - shape_weights=None, - tfs_weights=None, - requires_grad_weights=None, - ): - r""" + d: int, + scale: Union[float, np.ndarray, torch.Tensor] = 1.0, + lengthscales: Union[None, np.ndarray, torch.Tensor] = None, + alpha: Union[float, np.ndarray, torch.Tensor] = 1, + shape_scale: Union[None, list] = None, + shape_lengthscales: Union[None, list] = None, + shape_alpha: Union[None, list] = None, + tfs_scale: Union[None, Tuple[Callable, Callable]] = None, + tfs_lengthscales: Union[None, Tuple[Callable, Callable]] = None, + tfs_alpha: Union[None, Tuple[Callable, Callable]] = None, + torchify: bool = False, + requires_grad_scale: Union[None, bool] = None, + requires_grad_lengthscales: Union[None, bool] = None, + requires_grad_alpha: Union[None, bool] = None, + device: Union[str, torch.device] = "cpu", + compile_call: bool = False, + compile_call_kwargs: Union[None, dict] = None, + weights: Union[None, np.ndarray, torch.Tensor] = None, + shape_weights: Union[None, list] = None, + tfs_weights: Union[None, Tuple[Callable, Callable]] = None, + requires_grad_weights: Union[None, bool] = None, + ) -> None: + r"""Initialize a KernelShiftInvarCombined kernel. + Args: d (int): Dimension. - scale (Union[np.ndarray, torch.Tensor]): Scaling factor $S$. - lengthscales (Union[np.ndarray, torch.Tensor]): Product weights $(\gamma_1,\dots,\gamma_d)$. - alpha (Union[np.ndarray, torch.Tensor]): Weights $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in \mathbb{R}_{>0}^4$. - shape_scale (list): Shape of `scale` when `np.isscalar(scale)`. - shape_lengthscales (list): Shape of `lengthscales` when `np.isscalar(lengthscales)` - shape_alpha (list): Shape of `alpha` when `np.isscalar(alpha)` - tfs_scale (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_lengthscales (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_alpha (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - torchify (bool): If `True`, use the `torch` backend. Set to `True` if computing gradients with respect to inputs and/or hyperparameters. - requires_grad_scale (bool): If `True` and `torchify`, set `requires_grad=True` for `scale`. - requires_grad_lengthscales (bool): If `True` and `torchify`, set `requires_grad=True` for `lengthscales`. - requires_grad_alpha (bool): If `True` and `torchify`, set `requires_grad=True` for `alpha`. - device (torch.device): If `torchify`, put things onto this device. - compile_call (bool): If `True`, `torch.compile` the `parsed___call__` method. - compile_call_kwargs (dict): When `compile_call` is `True`, pass these keyword arguments to `torch.compile`. - weights (Union[np.ndarray, torch.Tensor]): Alias for `lengthscales`. - shape_weights (list): Alias for `shape_lengthscales`. - tfs_weights (Tuple[callable,callable]): Alias for `tfs_lengthscales`. - requires_grad_weights (bool): Alias for `requires_grad_lengthscales`. + scale (Union[float, np.ndarray, torch.Tensor]): Scaling factor $S$. + lengthscales (Union[None, np.ndarray, torch.Tensor]): Product weights + $(\gamma_1,\dots,\gamma_d)$. + alpha (Union[float, np.ndarray, torch.Tensor]): Weights + $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in + \mathbb{R}_{>0}^4$. + shape_scale (Union[None, list]): Shape of `scale` when `np.isscalar(scale)`. + shape_lengthscales (Union[None, list]): Shape of `lengthscales` when + `np.isscalar(lengthscales)` + shape_alpha (Union[None, list]): Shape of `alpha` when `np.isscalar(alpha)` + tfs_scale (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + tfs_lengthscales (Union[None, Tuple[Callable, Callable]]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + tfs_alpha (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + torchify (bool): If `True`, use the `torch` backend. Set to `True` + if computing gradients with respect to inputs and/or + hyperparameters. + requires_grad_scale (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `scale`. + requires_grad_lengthscales (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `lengthscales`. + requires_grad_alpha (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `alpha`. + device (Union[str, torch.device]): If `torchify`, put things onto this device. + compile_call (bool): If `True`, `torch.compile` the + `parsed___call__` method. + compile_call_kwargs (Union[None, dict]): When `compile_call` is `True`, pass + these keyword arguments to `torch.compile`. + weights (Union[None, np.ndarray, torch.Tensor]): Alias for + `lengthscales`. + shape_weights (Union[None, list]): Alias for `shape_lengthscales`. + tfs_weights (Union[None, Tuple[Callable, Callable]]): Alias for + `tfs_lengthscales`. + requires_grad_weights (Union[None, bool]): Alias for + `requires_grad_lengthscales`. """ if shape_scale is None: shape_scale = [1] @@ -591,7 +701,8 @@ def __init__( tfs_weights=tfs_weights, requires_grad_weights=requires_grad_weights, ) - assert self.alpha.shape[-2:] == (4, d) + if not (self.alpha.shape[-2:] == (4, d)): + raise AssertionError if self.torchify: import torch @@ -604,10 +715,15 @@ def __init__( ) def get_per_dim_components(self, x0, x1, beta0, beta1): + """Per-dimension Bernoulli-polynomial components for orders 1-4, + later combined by `combine_per_dim_components_raw_m1` weighted by + `alpha`. Does not support derivatives (`beta0`/`beta1` must be zero). + """ p = len(beta0) - assert (beta0 == 0).all() and ( + if not ((beta0 == 0).all() and ( beta1 == 0 - ).all(), "KernelDSICombined does not support derivatives" + ).all()): + raise AssertionError("KernelDSICombined does not support derivatives") delta = (x0 - x1) % 1 kparts = [None] * 4 kparts[0] = bernoulli_poly(1, delta) @@ -621,6 +737,9 @@ def get_per_dim_components(self, x0, x1, beta0, beta1): def combine_per_dim_components_raw_m1( self, kparts, beta0, beta1, c, batch_params, stable ): + """Weight the order-1-4 components by `alpha` and sum, then delegate + to the base class's combination logic. + """ kparts = (self.alpha[..., None, :, None, :] * kparts).sum(-3) return super().combine_per_dim_components_raw_m1( kparts, beta0, beta1, c, batch_params, stable @@ -628,27 +747,36 @@ def combine_per_dim_components_raw_m1( class KernelDigShiftInvar(AbstractSIDSIKernel): - r""" - Digitally shift invariant kernel in base $b=2$ with - smoothness $\boldsymbol{\alpha}$, product weights $\boldsymbol{\gamma}$, and scale $S$: - - $$\begin{aligned} - K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d \left(1+ \gamma_j \tilde{K}_{\alpha_j}(x_j \oplus z_j)\right), \qquad\mathrm{where} \\ - \tilde{K}_1(x) &= 6 \left(\frac{1}{6} - 2^{\lfloor \log_2(x) \rfloor -1}\right), \\ - \tilde{K}_2(x) &= \sum_{k \in \mathbb{N}} \frac{\mathrm{wal}_k(x)}{2^{\mu_2(k)}} = -\beta(x) x + \frac{5}{2}\left[1-t_1(x)\right]-1, \\ - \tilde{K}_3(x) &= \sum_{k \in \mathbb{N}} \frac{\mathrm{wal}_k(x)}{2^{\mu_3(k)}} = \beta(x)x^2-5\left[1-t_1(x)\right]x+\frac{43}{18}\left[1-t_2(x)\right]-1, \\ - \tilde{K}_4(x) &= \sum_{k \in \mathbb{N}} \frac{\mathrm{wal}_k(x)}{2^{\mu_4(k)}} = - \frac{2}{3}\beta(x)x^3+5\left[1-t_1(x)\right]x^2 - \frac{43}{9}\left[1-t_2(x)\right]x +\frac{701}{294}\left[1-t_3(x)\right]+\beta(x)\left[\frac{1}{48}\sum_{a=0}^\infty \frac{\mathrm{wal}_{2^a}(x)}{2^{3a}} - \frac{1}{42}\right] - 1. + r"""Digitally shift invariant kernel in base $b=2$ with smoothness + $\boldsymbol{\alpha}$, product weights $\boldsymbol{\gamma}$, and scale + $S$: + + $$\begin{aligned} K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d + \left(1+ \gamma_j \tilde{K}_{\alpha_j}(x_j \oplus z_j)\right), + \qquad\mathrm{where} \\ \tilde{K}_1(x) &= 6 \left(\frac{1}{6} - 2^{\lfloor + \log_2(x) \rfloor -1}\right), \\ \tilde{K}_2(x) &= \sum_{k \in \mathbb{N}} + \frac{\mathrm{wal}_k(x)}{2^{\mu_2(k)}} = -\beta(x) x + + \frac{5}{2}\left[1-t_1(x)\right]-1, \\ \tilde{K}_3(x) &= \sum_{k \in + \mathbb{N}} \frac{\mathrm{wal}_k(x)}{2^{\mu_3(k)}} = + \beta(x)x^2-5\left[1-t_1(x)\right]x+\frac{43}{18}\left[1-t_2(x)\right]-1, + \\ \tilde{K}_4(x) &= \sum_{k \in \mathbb{N}} + \frac{\mathrm{wal}_k(x)}{2^{\mu_4(k)}} = - + \frac{2}{3}\beta(x)x^3+5\left[1-t_1(x)\right]x^2 - + \frac{43}{9}\left[1-t_2(x)\right]x + +\frac{701}{294}\left[1-t_3(x)\right]+\beta(x)\left[\frac{1}{48}\sum_{a=0}^\infty + \frac{\mathrm{wal}_{2^a}(x)}{2^{3a}} - \frac{1}{42}\right] - 1. \end{aligned}$$ - where - - - $x \oplus z$ is XOR between bits, - - $\mathrm{wal}_k$ is the $k^\text{th}$ Walsh function, - - $\beta(x) = - \lfloor \log_2(x) \rfloor$ and $t_\nu(x) = 2^{-\nu \beta(x)}$ where $\beta(0)=t_\nu(0) = 0$, and - - and $\mu_\alpha$ is the Dick weight function which sums the first $\alpha$ largest indices of $1$ bits in the binary expansion of $k$ - e.g. $k=13=1101_2$ has 1-bit indexes $(4,3,1)$ so - - $$\mu_1(k) = 4, \mu_2(k) = 4+3, \mu_3(k) = 4+3+1 = \mu_4(k) = \mu_5(k) = \dots.$$ + where + + - $x \oplus z$ is XOR between bits, + - $\mathrm{wal}_k$ is the $k^\text{th}$ Walsh function, + - $\beta(x) = - \lfloor \log_2(x) \rfloor$ and $t_\nu(x) = 2^{-\nu \beta(x)}$ where $\beta(0)=t_\nu(0) = 0$, and + - and $\mu_\alpha$ is the Dick weight function which sums the first $\alpha$ largest indices of $1$ bits in the binary expansion of $k$ + e.g. $k=13=1101_2$ has 1-bit indexes $(4,3,1)$ so + + $$\mu_1(k) = 4, \mu_2(k) = 4+3, \mu_3(k) = 4+3+1 = \mu_4(k) = \mu_5(k) = + \dots.$$ Examples: >>> from qmcpy import DigitalNetB2, fwht @@ -661,7 +789,7 @@ class KernelDigShiftInvar(AbstractSIDSIKernel): >>> x.dtype dtype('uint64') >>> kernel = KernelDigShiftInvar( - ... d = d, + ... d = d, ... t = dnb2.t, ... alpha = list(range(1,d+1)), ... scale = 10, @@ -692,10 +820,10 @@ class KernelDigShiftInvar(AbstractSIDSIKernel): True >>> np.allclose(fwht(fwht(y)/lam),np.linalg.solve(kmat,y)) True - >>> import torch + >>> import torch >>> xtorch = bin_from_numpy_to_torch(x) >>> kernel_torch = KernelDigShiftInvar( - ... d = d, + ... d = d, ... t = dnb2.t, ... alpha = list(range(1,d+1)), ... scale = 10, @@ -719,16 +847,16 @@ class KernelDigShiftInvar(AbstractSIDSIKernel): >>> kernel_torch.single_integral_01d(xtorch) tensor([10., 10., 10., 10., 10., 10., 10., 10.], grad_fn=) - Batch Params - + Batch Params + >>> rng = np.random.Generator(np.random.PCG64(7)) >>> kernel = KernelDigShiftInvar( - ... d = 2, + ... d = 2, ... t = 10, ... shape_scale = [4,3,1], ... shape_lengthscales = [3,2]) >>> x = rng.uniform(low=0,high=1,size=(6,5,2)) - >>> kernel(x,x).shape + >>> kernel(x,x).shape (4, 3, 6, 5) >>> kernel(x[:,:,None,:],x[:,None,:,:]).shape (4, 3, 6, 5, 5) @@ -740,71 +868,90 @@ class KernelDigShiftInvar(AbstractSIDSIKernel): np.float64(4.440892098500626e-16) **References:** - - 1. Dick, Josef. - "Walsh spaces containing smooth functions and quasi-Monte Carlo rules of arbitrary high order." + + 1. Dick, Josef. + "Walsh spaces containing smooth functions and quasi-Monte Carlo rules of arbitrary high order." SIAM Journal on Numerical Analysis 46.3 (2008): 1519-1553. - 2. Dick, Josef. - "The decay of the Walsh coefficients of smooth functions." - Bulletin of the Australian Mathematical Society 80.3 (2009): 430-453. + 2. Dick, Josef. + "The decay of the Walsh coefficients of smooth functions." + Bulletin of the Australian Mathematical Society 80.3 (2009): 430-453. - 3. Jagadeeswaran, Rathinavel, and Fred J. Hickernell. - "Fast automatic Bayesian cubature using Sobol' sampling." + 3. Jagadeeswaran, Rathinavel, and Fred J. Hickernell. + "Fast automatic Bayesian cubature using Sobol' sampling." Advances in Modeling and Simulation: Festschrift for Pierre L'Ecuyer. Cham: Springer International Publishing, 2022. 301-318. - 4. Rathinavel, Jagadeeswaran. - Fast automatic Bayesian cubature using matching kernels and designs. + 4. Rathinavel, Jagadeeswaran. + Fast automatic Bayesian cubature using matching kernels and designs. Illinois Institute of Technology, 2019. - - 5. Sorokin, Aleksei. - "A Unified Implementation of Quasi-Monte Carlo Generators, Randomization Routines, and Fast Kernel Methods." + + 5. Sorokin, Aleksei. + "A Unified Implementation of Quasi-Monte Carlo Generators, Randomization Routines, and Fast Kernel Methods." arXiv preprint arXiv:2502.14256 (2025). """ def __init__( self, - d, - t=None, - scale=1.0, - lengthscales=None, - alpha=2, - shape_scale=None, - shape_lengthscales=None, - tfs_scale=None, - tfs_lengthscales=None, - torchify=False, - requires_grad_scale=None, - requires_grad_lengthscales=None, - device="cpu", - compile_call=False, - compile_call_kwargs=None, - weights=None, - shape_weights=None, - tfs_weights=None, - requires_grad_weights=None, - ): - r""" + d: int, + t: Union[None, int] = None, + scale: Union[float, np.ndarray, torch.Tensor] = 1.0, + lengthscales: Union[None, np.ndarray, torch.Tensor] = None, + alpha: Union[float, np.ndarray, torch.Tensor] = 2, + shape_scale: Union[None, list] = None, + shape_lengthscales: Union[None, list] = None, + tfs_scale: Union[None, Tuple[Callable, Callable]] = None, + tfs_lengthscales: Union[None, Tuple[Callable, Callable]] = None, + torchify: bool = False, + requires_grad_scale: Union[None, bool] = None, + requires_grad_lengthscales: Union[None, bool] = None, + device: Union[str, torch.device] = "cpu", + compile_call: bool = False, + compile_call_kwargs: Union[None, dict] = None, + weights: Union[None, np.ndarray, torch.Tensor] = None, + shape_weights: Union[None, list] = None, + tfs_weights: Union[None, Tuple[Callable, Callable]] = None, + requires_grad_weights: Union[None, bool] = None, + ) -> None: + r"""Initialize a KernelDigShiftInvar kernel. + Args: d (int): Dimension. - t (int): number of bits in binary represtnations. Typically `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. - scale (Union[np.ndarray, torch.Tensor]): Scaling factor $S$. - lengthscales (Union[np.ndarray, torch.Tensor]): Product weights $(\gamma_1,\dots,\gamma_d)$. - alpha (Union[np.ndarray, torch.Tensor]): Smoothness parameters $(\alpha_1,\dots,\alpha_d)$ where $\alpha_j \geq 1$ for $j=1,\dots,d$. - shape_scale (list): Shape of `scale` when `np.isscalar(scale)`. - shape_lengthscales (list): Shape of `lengthscales` when `np.isscalar(lengthscales)` - tfs_scale (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_lengthscales (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - torchify (bool): If `True`, use the `torch` backend. Set to `True` if computing gradients with respect to inputs and/or hyperparameters. - requires_grad_scale (bool): If `True` and `torchify`, set `requires_grad=True` for `scale`. - requires_grad_lengthscales (bool): If `True` and `torchify`, set `requires_grad=True` for `lengthscales`. - device (torch.device): If `torchify`, put things onto this device. - compile_call (bool): If `True`, `torch.compile` the `parsed___call__` method. - compile_call_kwargs (dict): When `compile_call` is `True`, pass these keyword arguments to `torch.compile`. - weights (Union[np.ndarray, torch.Tensor]): Alias for `lengthscales`. - shape_weights (list): Alias for `shape_lengthscales`. - tfs_weights (Tuple[callable,callable]): Alias for `tfs_lengthscales`. - requires_grad_weights (bool): Alias for `requires_grad_lengthscales`. + t (Union[None, int]): number of bits in binary representations. Typically + `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. + scale (Union[float, np.ndarray, torch.Tensor]): Scaling factor $S$. + lengthscales (Union[None, np.ndarray, torch.Tensor]): Product weights + $(\gamma_1,\dots,\gamma_d)$. + alpha (Union[float, np.ndarray, torch.Tensor]): Smoothness parameters + $(\alpha_1,\dots,\alpha_d)$ where $\alpha_j \geq 1$ for + $j=1,\dots,d$. + shape_scale (Union[None, list]): Shape of `scale` when `np.isscalar(scale)`. + shape_lengthscales (Union[None, list]): Shape of `lengthscales` when + `np.isscalar(lengthscales)` + tfs_scale (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + tfs_lengthscales (Union[None, Tuple[Callable, Callable]]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + torchify (bool): If `True`, use the `torch` backend. Set to `True` + if computing gradients with respect to inputs and/or + hyperparameters. + requires_grad_scale (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `scale`. + requires_grad_lengthscales (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `lengthscales`. + device (Union[str, torch.device]): If `torchify`, put things onto this device. + compile_call (bool): If `True`, `torch.compile` the + `parsed___call__` method. + compile_call_kwargs (Union[None, dict]): When `compile_call` is `True`, pass + these keyword arguments to `torch.compile`. + weights (Union[None, np.ndarray, torch.Tensor]): Alias for + `lengthscales`. + shape_weights (Union[None, list]): Alias for `shape_lengthscales`. + tfs_weights (Union[None, Tuple[Callable, Callable]]): Alias for + `tfs_lengthscales`. + requires_grad_weights (Union[None, bool]): Alias for + `requires_grad_lengthscales`. """ if shape_scale is None: shape_scale = [1] @@ -834,41 +981,62 @@ def __init__( tfs_weights=tfs_weights, requires_grad_weights=requires_grad_weights, ) - assert self.alpha.shape == (self.d,) + if not (self.alpha.shape == (self.d,)): + raise AssertionError self.set_t(t) - assert all(1 <= int(alphaj) <= 4 for alphaj in self.alpha) + if not (all(1 <= int(alphaj) <= 4 for alphaj in self.alpha)): + raise AssertionError @property def t(self): + """int: Number of bits used in the binary representation of inputs + (see `set_t`). Must be set via `set_t` before use. + """ if self._t is None: raise ParameterError("please use set_t to set the t value") return self._t - def set_t(self, t): + def set_t(self, t: Union[None, int]): + """Set the number of bits `t` used to binarize inputs via `to_bin`. + + Args: + t (Union[None, int]): Number of bits, `0 <= t <= 63` when + `torchify` (`torch.int64` limit) or `0 <= t <= 64` otherwise + (`np.uint64` limit). `None` clears the value, requiring a + later call to `set_t` before the kernel can be evaluated. + """ if t is None: self._t = t else: - assert t % 1 == 0 + if not (t % 1 == 0): + raise AssertionError if self.torchify: - assert 0 <= t <= 63 # torch only supports torch.int64 + if not (0 <= t <= 63): # torch only supports torch.int64 + raise AssertionError else: - assert 0 <= t <= 64 # numpy supports np.uint64 + if not (0 <= t <= 64): # numpy supports np.uint64 + raise AssertionError self._t = t def get_per_dim_components(self, x0, x1, beta0, beta1): + """Per-dimension weighted-Walsh-function components; see the class + docstring for the kernel formula. Inputs are first binarized to `t` bits. + """ t = self.t x0 = to_bin(x0, t) x1 = to_bin(x1, t) p = len(beta0) betasum = beta0 + beta1 order = self.alpha - betasum - assert (1 <= order).all() and (order <= 4).all(), ( - "order must all be between 2 and 4, but got order = %s. Try increasing alpha" - % str(order) - ) - assert not ( + if not ((1 <= order).all() and (order <= 4).all()): + raise AssertionError( + "order must all be between 2 and 4, but got order = %s. Try increasing alpha" + % str(order) + ) + if not (not ( (order == 1) * (self.alpha > 1) - ).any(), "taking the derivative of the order 2 digitally shift invariant kernel is not supported" + ).any()): + raise AssertionError("taking the derivative of the order 2 digitally shift invariant kernel is not supported") ind = 1.0 * (betasum > 0) delta = x0 ^ x1 kparts = [None] * p @@ -899,24 +1067,28 @@ def get_per_dim_components(self, x0, x1, beta0, beta1): class KernelDigShiftInvarAdaptiveAlpha(AbstractSIDSIKernel): - r""" - Digitally shift invariant kernel in base $b=2$ with - smoothness $\boldsymbol{\alpha} \geq \boldsymbol{0}$, product weights $\boldsymbol{\gamma}$, and scale $S$: - - $$\begin{aligned} - K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d \left(1+ \gamma_j \tilde{K}_{\alpha_j}(x_j \oplus z_j)\right), \qquad\mathrm{where} \\ - \tilde{K}_\alpha(x) &= \sum_{k \in \mathbb{N}} \frac{\mathrm{wal}_k(x)}{2^{{\alpha+1} (\mu_1(k)-1)}} = \frac{2^{\alpha+1}}{2^{\alpha+1}-2} - \left(\frac{2^{\alpha+1}}{2^{\alpha+1}-2}+1\right) 2^{\alpha(\lfloor \log_2(x) \rfloor+1)}, \\ - \end{aligned}$$ + r"""Digitally shift invariant kernel in base $b=2$ with smoothness + $\boldsymbol{\alpha} \geq \boldsymbol{0}$, product weights + $\boldsymbol{\gamma}$, and scale $S$: + + $$\begin{aligned} K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d + \left(1+ \gamma_j \tilde{K}_{\alpha_j}(x_j \oplus z_j)\right), + \qquad\mathrm{where} \\ \tilde{K}_\alpha(x) &= \sum_{k \in \mathbb{N}} + \frac{\mathrm{wal}_k(x)}{2^{{\alpha+1} (\mu_1(k)-1)}} = + \frac{2^{\alpha+1}}{2^{\alpha+1}-2} - + \left(\frac{2^{\alpha+1}}{2^{\alpha+1}-2}+1\right) 2^{\alpha(\lfloor + \log_2(x) \rfloor+1)}, \\ \end{aligned}$$ + + where - where - - - $x \oplus z$ is XOR between bits, - - $\mathrm{wal}_k$ is the $k^\text{th}$ Walsh function, - - $\beta(x) = - \lfloor \log_2(x) \rfloor$ and $t_\nu(x) = 2^{-\nu \beta(x)}$ where $\beta(0)=t_\nu(0) = 0$, and - - and $\mu_\alpha$ is the Dick weight function which sums the first $\alpha$ largest indices of $1$ bits in the binary expansion of $k$ - e.g. $k=13=1101_2$ has 1-bit indexes $(4,3,1)$ so - - $$\mu_1(k) = 4, \mu_2(k) = 4+3, \mu_3(k) = 4+3+1 = \mu_4(k) = \mu_5(k) = \dots.$$ + - $x \oplus z$ is XOR between bits, + - $\mathrm{wal}_k$ is the $k^\text{th}$ Walsh function, + - $\beta(x) = - \lfloor \log_2(x) \rfloor$ and $t_\nu(x) = 2^{-\nu \beta(x)}$ where $\beta(0)=t_\nu(0) = 0$, and + - and $\mu_\alpha$ is the Dick weight function which sums the first $\alpha$ largest indices of $1$ bits in the binary expansion of $k$ + e.g. $k=13=1101_2$ has 1-bit indexes $(4,3,1)$ so + + $$\mu_1(k) = 4, \mu_2(k) = 4+3, \mu_3(k) = 4+3+1 = \mu_4(k) = \mu_5(k) = + \dots.$$ Examples: >>> from qmcpy import DigitalNetB2, fwht @@ -929,7 +1101,7 @@ class KernelDigShiftInvarAdaptiveAlpha(AbstractSIDSIKernel): >>> x.dtype dtype('uint64') >>> kernel = KernelDigShiftInvarAdaptiveAlpha( - ... d = d, + ... d = d, ... t = dnb2.t, ... alpha = list(range(1,d+1)), ... scale = 10, @@ -960,10 +1132,10 @@ class KernelDigShiftInvarAdaptiveAlpha(AbstractSIDSIKernel): True >>> np.allclose(fwht(fwht(y)/lam),np.linalg.solve(kmat,y)) True - >>> import torch + >>> import torch >>> xtorch = bin_from_numpy_to_torch(x) >>> kernel_torch = KernelDigShiftInvarAdaptiveAlpha( - ... d = d, + ... d = d, ... t = dnb2.t, ... alpha = list(range(1,d+1)), ... scale = 10, @@ -987,16 +1159,16 @@ class KernelDigShiftInvarAdaptiveAlpha(AbstractSIDSIKernel): >>> kernel_torch.single_integral_01d(xtorch) tensor([10., 10., 10., 10., 10., 10., 10., 10.], grad_fn=) - Batch Params - + Batch Params + >>> rng = np.random.Generator(np.random.PCG64(7)) >>> kernel = KernelDigShiftInvarAdaptiveAlpha( - ... d = 2, + ... d = 2, ... t = 10, ... shape_scale = [4,3,1], ... shape_lengthscales = [3,2]) >>> x = rng.uniform(low=0,high=1,size=(6,5,2)) - >>> kernel(x,x).shape + >>> kernel(x,x).shape (4, 3, 6, 5) >>> kernel(x[:,:,None,:],x[:,None,:,:]).shape (4, 3, 6, 5, 5) @@ -1008,61 +1180,83 @@ class KernelDigShiftInvarAdaptiveAlpha(AbstractSIDSIKernel): np.float64(4.440892098500626e-16) **References:** - - 3. Dick, Josef, and Friedrich Pillichshammer. - "Multivariate integration in weighted Hilbert spaces based on Walsh functions and weighted Sobolev spaces." + + 3. Dick, Josef, and Friedrich Pillichshammer. + "Multivariate integration in weighted Hilbert spaces based on Walsh functions and weighted Sobolev spaces." Journal of Complexity 21.2 (2005): 149-195. """ def __init__( self, - d, - t=None, - scale=1.0, - lengthscales=None, - alpha=1, - shape_scale=None, - shape_lengthscales=None, - shape_alpha=None, - tfs_scale=None, - tfs_lengthscales=None, - tfs_alpha=None, - torchify=False, - requires_grad_scale=None, - requires_grad_lengthscales=None, - requires_grad_alpha=None, - device="cpu", - compile_call=False, - compile_call_kwargs=None, - weights=None, - shape_weights=None, - tfs_weights=None, - requires_grad_weights=None, - ): - r""" + d: int, + t: Union[None, int] = None, + scale: Union[float, np.ndarray, torch.Tensor] = 1.0, + lengthscales: Union[None, np.ndarray, torch.Tensor] = None, + alpha: Union[float, np.ndarray, torch.Tensor] = 1, + shape_scale: Union[None, list] = None, + shape_lengthscales: Union[None, list] = None, + shape_alpha: Union[None, list] = None, + tfs_scale: Union[None, Tuple[Callable, Callable]] = None, + tfs_lengthscales: Union[None, Tuple[Callable, Callable]] = None, + tfs_alpha: Union[None, Tuple[Callable, Callable]] = None, + torchify: bool = False, + requires_grad_scale: Union[None, bool] = None, + requires_grad_lengthscales: Union[None, bool] = None, + requires_grad_alpha: Union[None, bool] = None, + device: Union[str, torch.device] = "cpu", + compile_call: bool = False, + compile_call_kwargs: Union[None, dict] = None, + weights: Union[None, np.ndarray, torch.Tensor] = None, + shape_weights: Union[None, list] = None, + tfs_weights: Union[None, Tuple[Callable, Callable]] = None, + requires_grad_weights: Union[None, bool] = None, + ) -> None: + r"""Initialize a KernelDigShiftInvarAdaptiveAlpha kernel. + Args: d (int): Dimension. - t (int): number of bits in binary represtnations. Typically `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. - scale (Union[np.ndarray, torch.Tensor]): Scaling factor $S$. - lengthscales (Union[np.ndarray, torch.Tensor]): Product weights $(\gamma_1,\dots,\gamma_d)$. - alpha (Union[np.ndarray, torch.Tensor]): Smoothness parameters $(\alpha_1,\dots,\alpha_d)$ where $\alpha_j \geq 1$ for $j=1,\dots,d$. - shape_alpha (list): Shape of `alpha` when `np.isscalar(alpha)` - shape_scale (list): Shape of `scale` when `np.isscalar(scale)`. - shape_lengthscales (list): Shape of `lengthscales` when `np.isscalar(lengthscales)` - tfs_scale (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_lengthscales (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_alpha (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - torchify (bool): If `True`, use the `torch` backend. Set to `True` if computing gradients with respect to inputs and/or hyperparameters. - requires_grad_scale (bool): If `True` and `torchify`, set `requires_grad=True` for `scale`. - requires_grad_lengthscales (bool): If `True` and `torchify`, set `requires_grad=True` for `lengthscales`. - requires_grad_alpha (bool): If `True` and `torchify`, set `requires_grad=True` for `alpha`. - device (torch.device): If `torchify`, put things onto this device. - compile_call (bool): If `True`, `torch.compile` the `parsed___call__` method. - compile_call_kwargs (dict): When `compile_call` is `True`, pass these keyword arguments to `torch.compile`. - weights (Union[np.ndarray, torch.Tensor]): Alias for `lengthscales`. - shape_weights (list): Alias for `shape_lengthscales`. - tfs_weights (Tuple[callable,callable]): Alias for `tfs_lengthscales`. - requires_grad_weights (bool): Alias for `requires_grad_lengthscales`. + t (Union[None, int]): number of bits in binary representations. Typically + `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. + scale (Union[float, np.ndarray, torch.Tensor]): Scaling factor $S$. + lengthscales (Union[None, np.ndarray, torch.Tensor]): Product weights + $(\gamma_1,\dots,\gamma_d)$. + alpha (Union[float, np.ndarray, torch.Tensor]): Smoothness parameters + $(\alpha_1,\dots,\alpha_d)$ where $\alpha_j \geq 1$ for + $j=1,\dots,d$. + shape_scale (Union[None, list]): Shape of `scale` when `np.isscalar(scale)`. + shape_lengthscales (Union[None, list]): Shape of `lengthscales` when + `np.isscalar(lengthscales)` + shape_alpha (Union[None, list]): Shape of `alpha` when `np.isscalar(alpha)` + tfs_scale (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + tfs_lengthscales (Union[None, Tuple[Callable, Callable]]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + tfs_alpha (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + torchify (bool): If `True`, use the `torch` backend. Set to `True` + if computing gradients with respect to inputs and/or + hyperparameters. + requires_grad_scale (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `scale`. + requires_grad_lengthscales (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `lengthscales`. + requires_grad_alpha (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `alpha`. + device (Union[str, torch.device]): If `torchify`, put things onto this device. + compile_call (bool): If `True`, `torch.compile` the + `parsed___call__` method. + compile_call_kwargs (Union[None, dict]): When `compile_call` is `True`, pass + these keyword arguments to `torch.compile`. + weights (Union[None, np.ndarray, torch.Tensor]): Alias for + `lengthscales`. + shape_weights (Union[None, list]): Alias for `shape_lengthscales`. + tfs_weights (Union[None, Tuple[Callable, Callable]]): Alias for + `tfs_lengthscales`. + requires_grad_weights (Union[None, bool]): Alias for + `requires_grad_lengthscales`. """ if shape_scale is None: shape_scale = [1] @@ -1097,28 +1291,48 @@ def __init__( @property def t(self): + """int: Number of bits used in the binary representation of inputs + (see `set_t`). Must be set via `set_t` before use. + """ if self._t is None: raise ParameterError("please use set_t to set the t value") return self._t - def set_t(self, t): + def set_t(self, t: Union[None, int]): + """Set the number of bits `t` used to binarize inputs via `to_bin`. + + Args: + t (Union[None, int]): Number of bits, `0 <= t <= 63` when + `torchify` (`torch.int64` limit) or `0 <= t <= 64` otherwise + (`np.uint64` limit). `None` clears the value, requiring a + later call to `set_t` before the kernel can be evaluated. + """ if t is None: self._t = t else: - assert t % 1 == 0 + if not (t % 1 == 0): + raise AssertionError if self.torchify: - assert 0 <= t <= 63 # torch only supports torch.int64 + if not (0 <= t <= 63): # torch only supports torch.int64 + raise AssertionError else: - assert 0 <= t <= 64 # numpy supports np.uint64 + if not (0 <= t <= 64): # numpy supports np.uint64 + raise AssertionError self._t = t def get_per_dim_components(self, x0, x1, beta0, beta1): + """Per-dimension components with a per-XOR-bit-length adaptive + smoothness; see the class docstring for the kernel formula. Inputs + are first binarized to `t` bits. Does not support derivatives + (`beta0`/`beta1` must be zero). + """ t = self.t x0 = to_bin(x0, t) x1 = to_bin(x1, t) - assert (beta0 == 0).all() and ( + if not ((beta0 == 0).all() and ( beta1 == 0 - ).all(), "KernelDigShiftInvarAdaptiveAlpha does not support taking derivatives" + ).all()): + raise AssertionError("KernelDigShiftInvarAdaptiveAlpha does not support taking derivatives") p = len(beta0) delta = x0 ^ x1 flog2delta = self.npt.zeros(delta.shape, **self.nptkwargs) # should be -inf @@ -1130,6 +1344,10 @@ def get_per_dim_components(self, x0, x1, beta0, beta1): def combine_per_dim_components_raw_m1( self, flog2deltas, beta0, beta1, c, batch_params, stable ): + """Combine per-XOR-bit-length components using a smoothness `alpha` + that adapts to each bit length, then delegate to the base class's + combination logic. + """ alpha = batch_params["alpha"] p2alphap1 = 2 ** (alpha + 1) nu = p2alphap1 / (p2alphap1 - 2) @@ -1143,15 +1361,17 @@ def combine_per_dim_components_raw_m1( class KernelDigShiftInvarCombined(AbstractSIDSIKernel): - r""" - Digitally shift invariant kernel in base $b=2$ with - combination weights $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in \mathbb{R}_{>0}^4$, smoothness $\boldsymbol{\alpha}$, product weights $\boldsymbol{\gamma}$, and scale $S$: + r"""Digitally shift invariant kernel in base $b=2$ with combination + weights $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in + \mathbb{R}_{>0}^4$, smoothness $\boldsymbol{\alpha}$, product weights + $\boldsymbol{\gamma}$, and scale $S$: - $$\begin{aligned} - K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d \left(1+ \gamma_j \left(\sum_{p=1}^4 \alpha_{jp} \tilde{K}_p(x_j \oplus z_j)\right)\right) - \end{aligned}$$ + $$\begin{aligned} K(\boldsymbol{x},\boldsymbol{z}) &= S \prod_{j=1}^d + \left(1+ \gamma_j \left(\sum_{p=1}^4 \alpha_{jp} \tilde{K}_p(x_j \oplus + z_j)\right)\right) \end{aligned}$$ - where, $\oplus$ is defined in the docs for `KernelDigShiftInvar` and so are $\tilde{K}_p$ for $p \in \{1,2,3,4\}$ + where, $\oplus$ is defined in the docs for `KernelDigShiftInvar` and so are + $\tilde{K}_p$ for $p \in \{1,2,3,4\}$ Examples: >>> from qmcpy import DigitalNetB2, fwht @@ -1243,52 +1463,75 @@ class KernelDigShiftInvarCombined(AbstractSIDSIKernel): def __init__( self, - d, - t=None, - scale=1.0, - lengthscales=None, - alpha=1.0, - shape_scale=None, - shape_lengthscales=None, - shape_alpha=None, - tfs_scale=None, - tfs_lengthscales=None, - tfs_alpha=None, - torchify=False, - requires_grad_scale=None, - requires_grad_lengthscales=None, - requires_grad_alpha=None, - device="cpu", - compile_call=False, - compile_call_kwargs=None, - weights=None, - shape_weights=None, - tfs_weights=None, - requires_grad_weights=None, - ): - r""" + d: int, + t: Union[None, int] = None, + scale: Union[float, np.ndarray, torch.Tensor] = 1.0, + lengthscales: Union[None, np.ndarray, torch.Tensor] = None, + alpha: Union[float, np.ndarray, torch.Tensor] = 1.0, + shape_scale: Union[None, list] = None, + shape_lengthscales: Union[None, list] = None, + shape_alpha: Union[None, list] = None, + tfs_scale: Union[None, Tuple[Callable, Callable]] = None, + tfs_lengthscales: Union[None, Tuple[Callable, Callable]] = None, + tfs_alpha: Union[None, Tuple[Callable, Callable]] = None, + torchify: bool = False, + requires_grad_scale: Union[None, bool] = None, + requires_grad_lengthscales: Union[None, bool] = None, + requires_grad_alpha: Union[None, bool] = None, + device: Union[str, torch.device] = "cpu", + compile_call: bool = False, + compile_call_kwargs: Union[None, dict] = None, + weights: Union[None, np.ndarray, torch.Tensor] = None, + shape_weights: Union[None, list] = None, + tfs_weights: Union[None, Tuple[Callable, Callable]] = None, + requires_grad_weights: Union[None, bool] = None, + ) -> None: + r"""Initialize a KernelDigShiftInvarCombined kernel. + Args: d (int): Dimension. - t (int): number of bits in binary represtnations. Typically `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. - scale (Union[np.ndarray, torch.Tensor]): Scaling factor $S$. - lengthscales (Union[np.ndarray, torch.Tensor]): Product weights $(\gamma_1,\dots,\gamma_d)$. - alpha (Union[np.ndarray, torch.Tensor]): Weights $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in \mathbb{R}_{>0}^4$. - shape_scale (list): Shape of `scale` when `np.isscalar(scale)`. - shape_lengthscales (list): Shape of `lengthscales` when `np.isscalar(lengthscales)` - shape_alpha (list): Shape of `alpha` when `np.isscalar(alpha)` - tfs_scale (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - tfs_lengthscales (Tuple[callable,callable]): The first argument transforms to the raw value to be optimized; the second applies the inverse transform. - torchify (bool): If `True`, use the `torch` backend. Set to `True` if computing gradients with respect to inputs and/or hyperparameters. - requires_grad_scale (bool): If `True` and `torchify`, set `requires_grad=True` for `scale`. - requires_grad_lengthscales (bool): If `True` and `torchify`, set `requires_grad=True` for `lengthscales`. - requires_grad_alpha (bool): If `True` and `torchify`, set `requires_grad=True` for `alpha`. - device (torch.device): If `torchify`, put things onto this device. - compile_call (bool): If `True`, `torch.compile` the `parsed___call__` method. - compile_call_kwargs (dict): When `compile_call` is `True`, pass these keyword arguments to `torch.compile`. - weights (Union[np.ndarray, torch.Tensor]): Alias for `lengthscales`. - shape_weights (list): Alias for `shape_lengthscales`. - tfs_weights (Tuple[callable,callable]): Alias for `tfs_lengthscales`. - requires_grad_weights (bool): Alias for `requires_grad_lengthscales`. + t (Union[None, int]): number of bits in binary representations. Typically + `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. + scale (Union[float, np.ndarray, torch.Tensor]): Scaling factor $S$. + lengthscales (Union[None, np.ndarray, torch.Tensor]): Product weights + $(\gamma_1,\dots,\gamma_d)$. + alpha (Union[float, np.ndarray, torch.Tensor]): Weights + $\boldsymbol{\alpha}_1,\dots,\boldsymbol{\alpha}_d \in + \mathbb{R}_{>0}^4$. + shape_scale (Union[None, list]): Shape of `scale` when `np.isscalar(scale)`. + shape_lengthscales (Union[None, list]): Shape of `lengthscales` when + `np.isscalar(lengthscales)` + shape_alpha (Union[None, list]): Shape of `alpha` when `np.isscalar(alpha)` + tfs_scale (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + tfs_lengthscales (Union[None, Tuple[Callable, Callable]]): The first argument + transforms to the raw value to be optimized; the second applies + the inverse transform. + tfs_alpha (Union[None, Tuple[Callable, Callable]]): The first argument transforms + to the raw value to be optimized; the second applies the + inverse transform. + torchify (bool): If `True`, use the `torch` backend. Set to `True` + if computing gradients with respect to inputs and/or + hyperparameters. + requires_grad_scale (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `scale`. + requires_grad_lengthscales (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `lengthscales`. + requires_grad_alpha (Union[None, bool]): If `True` and `torchify`, set + `requires_grad=True` for `alpha`. + device (Union[str, torch.device]): If `torchify`, put things onto this device. + compile_call (bool): If `True`, `torch.compile` the + `parsed___call__` method. + compile_call_kwargs (Union[None, dict]): When `compile_call` is `True`, pass + these keyword arguments to `torch.compile`. + weights (Union[None, np.ndarray, torch.Tensor]): Alias for + `lengthscales`. + shape_weights (Union[None, list]): Alias for `shape_lengthscales`. + tfs_weights (Union[None, Tuple[Callable, Callable]]): Alias for + `tfs_lengthscales`. + requires_grad_weights (Union[None, bool]): Alias for + `requires_grad_lengthscales`. """ if shape_scale is None: shape_scale = [1] @@ -1319,33 +1562,55 @@ def __init__( requires_grad_weights=requires_grad_weights, ) self.set_t(t) - assert self.alpha.shape[-2:] == (4, d) + if not (self.alpha.shape[-2:] == (4, d)): + raise AssertionError @property def t(self): + """int: Number of bits used in the binary representation of inputs + (see `set_t`). Must be set via `set_t` before use. + """ if self._t is None: raise ParameterError("please use set_t to set the t value") return self._t - def set_t(self, t): + def set_t(self, t: Union[None, int]): + """Set the number of bits `t` used to binarize inputs via `to_bin`. + + Args: + t (Union[None, int]): Number of bits, `0 <= t <= 63` when + `torchify` (`torch.int64` limit) or `0 <= t <= 64` otherwise + (`np.uint64` limit). `None` clears the value, requiring a + later call to `set_t` before the kernel can be evaluated. + """ if t is None: self._t = t else: - assert t % 1 == 0 + if not (t % 1 == 0): + raise AssertionError if self.torchify: - assert 0 <= t <= 63 # torch only supports torch.int64 + if not (0 <= t <= 63): # torch only supports torch.int64 + raise AssertionError else: - assert 0 <= t <= 64 # numpy supports np.uint64 + if not (0 <= t <= 64): # numpy supports np.uint64 + raise AssertionError self._t = t def get_per_dim_components(self, x0, x1, beta0, beta1): + """Per-dimension weighted-Walsh-function components for orders 1-4, + later combined by `combine_per_dim_components_raw_m1` weighted by + `alpha`; see the class docstring for the kernel formula. Inputs are + first binarized to `t` bits. Does not support derivatives + (`beta0`/`beta1` must be zero). + """ t = self.t x0 = to_bin(x0, t) x1 = to_bin(x1, t) p = len(beta0) - assert (beta0 == 0).all() and ( + if not ((beta0 == 0).all() and ( beta1 == 0 - ).all(), "KernelDSICombined does not support derivatives" + ).all()): + raise AssertionError("KernelDSICombined does not support derivatives") delta = x0 ^ x1 kparts = [None] * 4 flog2deltaj = -self.npt.inf * self.npt.ones(delta.shape, **self.nptkwargs) @@ -1362,6 +1627,9 @@ def get_per_dim_components(self, x0, x1, beta0, beta1): def combine_per_dim_components_raw_m1( self, kparts, beta0, beta1, c, batch_params, stable ): + """Weight the order-1-4 components by `alpha` and sum, then delegate + to the base class's combination logic. + """ kparts = (self.alpha[..., None, :, None, :] * kparts).sum(-3) return super().combine_per_dim_components_raw_m1( kparts, beta0, beta1, c, batch_params, stable diff --git a/qmcpy/stopping_criterion/abstract_cub_bayes_ld_g.py b/qmcpy/stopping_criterion/abstract_cub_bayes_ld_g.py index 1ad2b1bc1..4fa56359e 100644 --- a/qmcpy/stopping_criterion/abstract_cub_bayes_ld_g.py +++ b/qmcpy/stopping_criterion/abstract_cub_bayes_ld_g.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_stopping_criterion import AbstractStoppingCriterion from ..util.data import Data @@ -12,6 +13,16 @@ class AbstractCubBayesLDG(AbstractStoppingCriterion): + """Abstract base class for guaranteed Bayesian low-discrepancy QMC stopping criteria. + + Implements the fast-transform Bayesian cubature error bound shared by + concrete lattice/digital-net Bayesian stopping criteria: doubling sample + counts each iteration, maintaining the running transform coefficients + (`_ytildefull`), and fitting the kernel hyperparameter + (`objective_function`, `_stopping_criterion`) to derive a + credible-interval error bound. + """ + _RESUME_REQUIRED_FIELDS = ( "solution", "comb_bound_low", "comb_bound_high", "comb_bound_diff", "comb_flags", "n", "n_max", "xfull", "yfull" ) @@ -32,7 +43,7 @@ def __init__( alpha, error_fun, errbd_type, - ): + ) -> None: self.parameters = ["abs_tol", "rel_tol", "n_init", "n_limit", "order"] # Input Checks if np.log2(n_init) % 1 != 0: @@ -48,7 +59,8 @@ def __init__( # Set Attributes self.n_init = int(n_init) self.n_limit = int(n_limit) - assert isinstance(error_fun, str) or callable(error_fun) + if not (isinstance(error_fun, str) or callable(error_fun)): + raise AssertionError # _error_fun_key stores a simple, serializable string and ensures correct state saving # in __getstate__(), bypassing serialization of complex lambda functions, which often fails. self.error_fun, self._error_fun_key = self._resolve_error_fun(error_fun) @@ -60,12 +72,14 @@ def __init__( super(AbstractCubBayesLDG, self).__init__( allowed_distribs=allowed_distribs, allow_vectorized_integrals=True ) - assert ( + if not ( self.integrand.discrete_distrib.no_replications == True - ), "Require the discrete distribution has replications=None" - assert ( + ): + raise AssertionError("Require the discrete distribution has replications=None") + if not ( self.integrand.discrete_distrib.randomize != "FALSE" - ), "Require discrete distribution is randomized" + ): + raise AssertionError("Require discrete distribution is randomized") self.alphas_indv, _ = self._compute_indv_alphas( np.full(self.integrand.d_comb, self.alpha) ) @@ -79,7 +93,8 @@ def __init__( self.use_gradient = False # If true uses gradient descent in parameter search self.one_theta = True # If true use common shape parameter for all dimensions, else allow shape parameter vary across dimensions self.errbd_type = errbd_type.upper() - assert self.errbd_type in ["MLE", "GCV", "FULL"] + if not (self.errbd_type in ["MLE", "GCV", "FULL"]): + raise AssertionError self.kernel = kernel self.debugEnable = True self.ft = ft @@ -185,10 +200,24 @@ def __setstate__(self, state): if isinstance(self.error_fun, str): self.error_fun, _ = self._resolve_error_fun(self.error_fun) - # objective function to estimate parameter theta - # MLE : Maximum likelihood estimation - # GCV : Generalized cross validation - def objective_function(self, theta, xun, ftilde): + def objective_function(self, theta: float, xun: np.ndarray, ftilde: np.ndarray) -> float: + """Compute the Bayesian cubature loss used to fit the kernel parameter theta. + + Evaluates either the negative log marginal likelihood (MLE) or the + generalized cross validation (GCV) loss, per `self.errbd_type`, along + with the kernel eigenvalues and RKHS norm needed by the error bound. + + Args: + theta (float): Kernel hyperparameter to evaluate the loss at. + xun (np.ndarray): Unique ordered node locations. + ftilde (np.ndarray): Fast-transformed function values at `xun`. + + Returns: + float: Loss value (MLE or GCV, per `self.errbd_type`). + np.ndarray: Kernel eigenvalues `vec_lambda`. + np.ndarray: Ring eigenvalues `vec_lambda_ring`, used by the error bound. + float: RKHS norm estimate of the fitted function. + """ n = len(ftilde) fudge = 100 * np.finfo(float).eps # if type(theta) != np.ndarray: @@ -250,10 +279,22 @@ def objective_function(self, theta, xun, ftilde): ) return loss, vec_lambda, vec_lambda_ring, RKHS_norm - # Computes modified kernel Km1 = K - 1 - # Useful to avoid cancellation error in the computation of (1 - n/\lambda_1) @staticmethod - def kernel_t(aconst, Bern): + def kernel_t(aconst: Union[float, np.ndarray], Bern: np.ndarray) -> np.ndarray: + r"""Compute the modified kernel ``Km1 = K - 1`` from Bernoulli polynomial values. + + Working with ``Km1`` rather than ``K`` directly avoids cancellation + error when later computing $1 - n/\lambda_1$. + + Args: + aconst (Union[float, np.ndarray]): Kernel parameter theta, scalar + or per-dimension array. + Bern (np.ndarray): Bernoulli polynomial values, shape ``(n, d)``. + + Returns: + np.ndarray: ``Km1``, the kernel minus one. + np.ndarray: ``K``, the full kernel (``1 + Km1``). + """ d = np.size(Bern, 1) if type(aconst) != np.ndarray: theta = np.ones((d, 1)) * aconst @@ -274,11 +315,16 @@ def kernel_t(aconst, Bern): K = Kj return [Km1, K] - # prints debug message if the given variable is Inf, Nan or complex, etc - # Example: alertMsg(x, 'Inf', 'Imag') - # prints if variable 'x' is either Infinite or Imaginary @staticmethod - def alert_msg(*args): + def alert_msg(*args: tuple): + """Print a debug message if a variable contains NaN, Inf, or complex values. + + Args: + *args (tuple): The variable to check, followed by one or more of + ``"Nan"``, ``"Inf"``, ``"Imag"`` naming which conditions to + report. Example: `alert_msg(x, "Inf", "Imag")` prints if `x` + contains infinite or imaginary values. + """ varargin = args nargin = len(varargin) if nargin > 1: @@ -303,7 +349,24 @@ def alert_msg(*args): else: print("unknown type check requested !") - def integrate(self, resume=None): + def integrate(self, resume: Union[None, Data] = None) -> tuple: + """Determine the samples needed to satisfy the target tolerance. + + Doubles the sample count each iteration, updates the running fast + transform (`_ytildefull`), and for each not-yet-converged output + calls `_stopping_criterion` to fit the Bayesian kernel hyperparameter + and derive a credible-interval bound on the integral. Stops once + every combined output is within tolerance or `self.n_limit` would be + exceeded. + + Args: + resume (Union[None, Data]): Existing integration state to resume from, if + supported. Defaults to None. + + Returns: + tuple: Approximation to the integral with shape ``integrand.d_comb`` + and the corresponding data object. + """ t_start = time() resume_provenance = self._capture_resume_provenance(resume) first_resume_iter = False @@ -445,8 +508,21 @@ def _validate_resume(self, data): if not self._is_power_of_two(n_total): raise ParameterError("resume data n_total must be a power of 2.") - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rmse_tol is None, "rmse_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance, broadcast to + `self.abs_tols` with shape `integrand.d_comb`. + rel_tol (Union[None, float]): Relative error tolerance, broadcast to + `self.rel_tols` with shape `integrand.d_comb`. + rmse_tol (Union[None, float]): Unsupported; must be `None`. + + Raises: + AssertionError: If `rmse_tol` is supplied. + """ + if not (rmse_tol is None): + raise AssertionError("rmse_tol not supported by this stopping criterion.") if abs_tol is not None: self.abs_tol = abs_tol self.abs_tols = np.full(self.integrand.d_comb, self.abs_tol) diff --git a/qmcpy/stopping_criterion/abstract_cub_mlmc.py b/qmcpy/stopping_criterion/abstract_cub_mlmc.py index a80db3092..99f1c2891 100644 --- a/qmcpy/stopping_criterion/abstract_cub_mlmc.py +++ b/qmcpy/stopping_criterion/abstract_cub_mlmc.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_stopping_criterion import AbstractStoppingCriterion from ..util.data import Data from ..util import ParameterError @@ -6,10 +7,19 @@ class AbstractCubMLMC(AbstractStoppingCriterion): + """Abstract base class for multilevel Monte Carlo stopping criteria. + + Shared machinery for `CubMLMC` and `CubMLMCCont`: level statistics + (`_refresh_level_statistics`), level growth (`_add_level`), and resume + checkpoint construction/validation/replay used across MLMC stopping + criteria. + """ @staticmethod def _append_level_diff_samples(data, level, dp): - """Append raw level-difference samples when checkpoint caching is enabled.""" + """Append raw level-difference samples when checkpoint caching is + enabled. + """ if not hasattr(data, "level_diffs"): return while len(data.level_diffs) <= level: @@ -63,8 +73,21 @@ def _get_next_samples(self, data): ) return ns.astype(int) - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rel_tol is None, "rel_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance, converted to an RMSE + tolerance via `self.alpha`. Ignored if `rmse_tol` is supplied. + rel_tol (Union[None, float]): Unsupported; must be `None`. + rmse_tol (Union[None, float]): Root mean squared error tolerance. Takes + precedence over `abs_tol` if both are supplied. + + Raises: + AssertionError: If `rel_tol` is supplied. + """ + if not (rel_tol is None): + raise AssertionError("rel_tol not supported by this stopping criterion.") if rmse_tol != None: self.rmse_tol = float(rmse_tol) elif abs_tol != None: @@ -155,12 +178,15 @@ def _construct_data(self): return data @staticmethod - def _validate_level_diffs(data): + def _validate_level_diffs(data: Data): """Validate the ``level_diffs`` replay cache on a resume checkpoint. Args: data (Data): Resume checkpoint to validate. + Returns: + None + Raises: ParameterError: If ``level_diffs`` is present but structurally inconsistent with ``n_level``. @@ -178,11 +204,13 @@ def _validate_level_diffs(data): % (level, level) ) - def _update_replay_data(self, data): - """Replay cached level-difference samples, falling back to fresh draws. + def _update_replay_data(self, data: Data): + """Replay cached level-difference samples, falling back to fresh + draws. Used during exact-resume replay to reconstruct the integration state by - consuming previously stored per-level samples before generating new ones. + consuming previously stored per-level samples before generating new + ones. Args: data (Data): Integration state carrying ``cached_level_diffs`` and diff --git a/qmcpy/stopping_criterion/abstract_cub_mlqmc.py b/qmcpy/stopping_criterion/abstract_cub_mlqmc.py index 9e95ed057..327906b58 100644 --- a/qmcpy/stopping_criterion/abstract_cub_mlqmc.py +++ b/qmcpy/stopping_criterion/abstract_cub_mlqmc.py @@ -1,13 +1,23 @@ +from typing import Union from .abstract_stopping_criterion import AbstractStoppingCriterion +from ..util.data import Data import numpy as np from scipy.stats import norm class AbstractCubMLQMC(AbstractStoppingCriterion): + """Abstract base class for multilevel Quasi-Monte Carlo stopping criteria. + + Shared machinery for `CubMLQMC` and `CubMLQMCCont`: replication-based + level statistics (`update_data`, `_update_bias_estimate`), level growth + (`_add_level`), and resume checkpoint validation/replay used across MLQMC + stopping criteria. + """ @staticmethod def _append_level_replication_sums(data, level, rep_sums, n_increment): - """Append replayable per-replication sums for one MLQMC level update.""" + """Append replayable per-replication sums for one MLQMC level update. + """ if (not hasattr(data, "level_rep_sums")) or (not hasattr(data, "level_n_increments")): return while len(data.level_rep_sums) <= level: @@ -98,14 +108,40 @@ def _resume_match_from_snapshots(snapshots, checkpoint): return resume_iter_count, snapshots[i:] return None, None - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rel_tol is None, "rel_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance, converted to an RMSE + tolerance via `self.alpha`. Ignored if `rmse_tol` is supplied. + rel_tol (Union[None, float]): Unsupported; must be `None`. + rmse_tol (Union[None, float]): Root mean squared error tolerance. Takes + precedence over `abs_tol` if both are supplied. + + Raises: + AssertionError: If `rel_tol` is supplied. + """ + if not (rel_tol is None): + raise AssertionError("rel_tol not supported by this stopping criterion.") if rmse_tol != None: self.rmse_tol = float(rmse_tol) elif abs_tol != None: self.rmse_tol = float(abs_tol) / norm.ppf(1 - self.alpha / 2.0) - def update_data(self, data): + def update_data(self, data: Data): + """Double the sample count on every active level and refresh statistics. + + For each level with `data.eval_level[l]` set, doubles its replicated + sample count (or draws `self.n_init` if the level is new), evaluates + the paired coarse/fine integrand there, and folds the new + per-replication sums into `data.mean_level_reps`, `data.mean_level`, + `data.var_level`, and `data.var_cost_ratio_level`. Then refreshes the + bias estimate, `data.n_total`, and `data.solution`, and clears + `data.eval_level`. + + Args: + data (Data): Integration state to update in place. + """ # update sample sums for l in range(data.levels): if not data.eval_level[l]: diff --git a/qmcpy/stopping_criterion/abstract_cub_qmc_ld_g.py b/qmcpy/stopping_criterion/abstract_cub_qmc_ld_g.py index 1e7d48e49..4d924810b 100644 --- a/qmcpy/stopping_criterion/abstract_cub_qmc_ld_g.py +++ b/qmcpy/stopping_criterion/abstract_cub_qmc_ld_g.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_stopping_criterion import AbstractStoppingCriterion from ..util.data import Data @@ -18,6 +19,16 @@ def _lstsq_pyfunc(x, y): class AbstractCubQMCLDG(AbstractStoppingCriterion): + """Abstract base class for guaranteed low-discrepancy QMC stopping criteria. + + Implements the fast-transform (FFT/FWT) cubature error bound shared by + `CubQMCLatticeG`, `CubQMCNetG`, and similar guaranteed lattice/digital-net + stopping criteria: doubling sample counts each iteration, maintaining the + running transform coefficients (`_ytildefull`, `_kappanumap`), optional + control-variate correction, and the cone-condition check that certifies + the error bound. + """ + _RESUME_REQUIRED_FIELDS = ( "solution", "comb_bound_low", "comb_bound_high", "comb_bound_diff", "comb_flags", "n", "n_max", "xfull", "yfull" ) @@ -41,7 +52,7 @@ def __init__( allowed_distribs, cast_complex, error_fun, - ): + ) -> None: self.parameters = ["abs_tol", "rel_tol", "n_init", "n_limit"] # Input Checks if np.log2(n_init) % 1 != 0 or n_init < 2**8: @@ -74,7 +85,8 @@ def __init__( ParameterWarning, ) self.n_limit = dd_n_limit - assert isinstance(error_fun, str) or callable(error_fun) + if not (isinstance(error_fun, str) or callable(error_fun)): + raise AssertionError # _error_fun_key stores a simple, serializable string and ensures correct state saving # in __getstate__(), bypassing serialization of complex lambda functions, which often fails. self.error_fun, self._error_fun_key = self._resolve_error_fun(error_fun) @@ -97,20 +109,23 @@ def __init__( super(AbstractCubQMCLDG, self).__init__( allowed_distribs=allowed_distribs, allow_vectorized_integrals=True ) - assert ( + if not ( self.integrand.discrete_distrib.no_replications == True - ), "Require the discrete distribution has replications=None" - assert ( + ): + raise AssertionError("Require the discrete distribution has replications=None") + if not ( self.integrand.discrete_distrib.randomize != "FALSE" - ), "Require discrete distribution is randomized" + ): + raise AssertionError("Require discrete distribution is randomized") self.set_tolerance(abs_tol, rel_tol) # control variates self._init_control_variates(control_variates, control_variate_means) self.update_beta = update_beta if self.ncv > 0: - assert self.cv_mu.shape == ( + if not (self.cv_mu.shape == ( (self.ncv,) + self.integrand.d_indv - ), "Control variate means should have shape (len(control variates),d_indv)." + )): + raise AssertionError("Control variate means should have shape (len(control variates),d_indv).") self.parameters += ["cv", "cv_mu", "update_beta"] else: self.update_beta = False @@ -217,7 +232,24 @@ def _validate_resume(self, data): "beta", data.beta, self.integrand.d_indv + (self.ncv,) ) - def integrate(self, resume=None): + def integrate(self, resume: Union[None, Data] = None) -> tuple: + """Determine the samples needed to satisfy the target tolerance. + + Doubles the sample count each iteration, updates the running fast + transform (`_ytildefull`) and its permutation (`_kappanumap`), + optionally corrects for control variates, and (if `self.check_cone`) + checks the cone condition that certifies the low-discrepancy error + bound. Stops once every combined output is within tolerance or + `self.n_limit` would be exceeded. + + Args: + resume (Union[None, Data]): Existing integration state to resume from, if + supported. Defaults to None. + + Returns: + tuple: Approximation to the integral with shape ``integrand.d_comb`` + and the corresponding data object. + """ t_start = time() resume_provenance = self._capture_resume_provenance(resume) first_resume_iter = False @@ -475,8 +507,21 @@ def integrate(self, resume=None): trace.finalize() return data.solution, data - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rmse_tol is None, "rmse_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance, broadcast to + `self.abs_tols` with shape `integrand.d_comb`. + rel_tol (Union[None, float]): Relative error tolerance, broadcast to + `self.rel_tols` with shape `integrand.d_comb`. + rmse_tol (Union[None, float]): Unsupported; must be `None`. + + Raises: + AssertionError: If `rmse_tol` is supplied. + """ + if not (rmse_tol is None): + raise AssertionError("rmse_tol not supported by this stopping criterion.") if abs_tol is not None: self.abs_tol = abs_tol self.abs_tols = np.full(self.integrand.d_comb, self.abs_tol) diff --git a/qmcpy/stopping_criterion/abstract_stopping_criterion.py b/qmcpy/stopping_criterion/abstract_stopping_criterion.py index 8c1a56588..b784d7b04 100644 --- a/qmcpy/stopping_criterion/abstract_stopping_criterion.py +++ b/qmcpy/stopping_criterion/abstract_stopping_criterion.py @@ -1,9 +1,11 @@ +from ..util.data import Data +from typing import TYPE_CHECKING, Any, Callable, TextIO, Union import copy import sys -from typing import TYPE_CHECKING from .diagnostics import ( # noqa: F401 _IterationTraceLogger, + _IterationHistoryTable, _format_iteration_log, _get_iteration_log_frame, _print_iteration_log, @@ -22,10 +24,20 @@ class AbstractStoppingCriterion(object): + """Abstract base class for QMCPy stopping criteria. + + A stopping criterion drives adaptive sampling for a given `integrand` + until an error tolerance is met, via `integrate`. Concrete stopping + criteria (e.g. `CubQMCNetG`, `CubMCCLT`) implement `integrate` and + `set_tolerance`; this base class handles shared bookkeeping: checkpoint + resume/save, iteration logging, and validating the integrand/true + measure/discrete distribution combination. + """ + _RESUME_FORMAT_VERSION = 1 # Increment when checkpoint format changes in a non-backwards-compatible way _ITERATION_LOG_VIEWS = ("all", "current", "without_resume", "stage_last") - def __init__(self, allowed_distribs, allow_vectorized_integrals): + def __init__(self, allowed_distribs: list, allow_vectorized_integrals: bool) -> None: """Initialize a stopping criterion base class. Args: @@ -68,12 +80,12 @@ def __init__(self, allowed_distribs, allow_vectorized_integrals): self.parameters = [] self.elapsed_time = float(getattr(self, "elapsed_time", 0.0)) - def integrate(self, resume=None) -> tuple: + def integrate(self, resume: Union[None, Data] = None) -> tuple: """Determine the samples needed to satisfy the target tolerance. Args: - resume (Data, optional): Existing integration state to resume from, - if supported. A valid resume checkpoint must continue the same + resume (Union[None, Data]): Existing integration state to resume from, if + supported. A valid resume checkpoint must continue the same numerical experiment without duplicating samples, losing accumulated statistics, or weakening the requested tolerance guarantee. Supported resume implementations validate and copy @@ -81,9 +93,8 @@ def integrate(self, resume=None) -> tuple: object is preserved. Defaults to None. Returns: - tuple[Union[float, np.ndarray], Data]: Approximation to the integral - with shape ``integrand.d_comb`` and the corresponding data - object. + tuple: Approximation to the integral with shape ``integrand.d_comb`` and + the corresponding data object. """ raise MethodImplementationError(self, "integrate") @@ -92,7 +103,7 @@ def _make_trace_logger(self) -> _IterationTraceLogger: Returns: _IterationTraceLogger: Trace logger configured from the stopping - criterion's optional trace attributes. + criterion's optional trace attributes. """ requested_trace_iterations = bool(getattr(self, "trace_iterations", False)) trace_verbose = bool(getattr(self, "verbose", False)) @@ -123,20 +134,21 @@ def _make_trace_logger(self) -> _IterationTraceLogger: def get_iteration_log( self, - history=None, - printed_only=True, - drop_empty_columns=True, - formatted=True, - view="all", + history: Union[None, list] = None, + printed_only: bool = True, + drop_empty_columns: bool = True, + formatted: bool = True, + view: str = "all", ) -> "pandas.DataFrame": """Return the latest iteration log as a pandas DataFrame. Args: - history (list[dict] | None): Iteration history to format. If ``None``, - uses ``self.iteration_history`` when available. + history (Union[None, list]): Iteration history to format. If + ``None``, uses ``self.iteration_history`` when available. printed_only (bool): If ``True``, include only rows that were selected for printed output. - drop_empty_columns (bool): If ``True``, drop columns with no values. + drop_empty_columns (bool): If ``True``, drop columns with no + values. formatted (bool): If ``True``, return formatted display values when available. view (str): Which log view to return. ``"all"`` and ``"current"`` @@ -204,17 +216,17 @@ def _apply_iteration_log_view(log_df, view): positions = positions[positions >= 0] return log_df.loc[non_resume_indices[positions]].reset_index(drop=True) - def format_iteration_log(self, history=None, printed_only=True, include_header=True) -> str: + def format_iteration_log(self, history: "Union[None, _IterationHistoryTable]" = None, printed_only: bool = True, include_header: bool = True) -> str: """Return the iteration log as formatted text. Args: - history (IterationHistoryTable | None, optional): Iteration history - to format. If ``None``, uses ``self.iteration_history`` when + history (Union[None, _IterationHistoryTable]): Iteration history to + format. If ``None``, uses ``self.iteration_history`` when available. Defaults to None. - printed_only (bool, optional): If ``True``, include only rows that - were selected for printed output. Defaults to True. - include_header (bool, optional): If ``True``, include the trace - label header before the table. Defaults to True. + printed_only (bool): If ``True``, include only rows that were + selected for printed output. Defaults to True. + include_header (bool): If ``True``, include the trace label header + before the table. Defaults to True. Returns: str: Formatted iteration log text. @@ -228,18 +240,18 @@ def format_iteration_log(self, history=None, printed_only=True, include_header=T include_header=include_header, ) - def print_iteration_log(self, history=None, printed_only=True, include_header=True, file=None) -> None: + def print_iteration_log(self, history: "Union[None, _IterationHistoryTable]" = None, printed_only: bool = True, include_header: bool = True, file: Union[None, TextIO] = None) -> None: """Print the iteration log for the latest run or supplied history. Args: - history (IterationHistoryTable | None, optional): Iteration history - to print. If ``None``, uses ``self.iteration_history`` when - available. Defaults to None. - printed_only (bool, optional): If ``True``, print only rows that - were selected for printed output. Defaults to True. - include_header (bool, optional): If ``True``, include the trace - label header before the table. Defaults to True. - file (typing.TextIO | None, optional): Output stream. Defaults to + history (Union[None, _IterationHistoryTable]): Iteration history to print. + If ``None``, uses ``self.iteration_history`` when available. + Defaults to None. + printed_only (bool): If ``True``, print only rows that were + selected for printed output. Defaults to True. + include_header (bool): If ``True``, include the trace label header + before the table. Defaults to True. + file (Union[None, TextIO]): Output stream. Defaults to ``sys.stdout`` when None. Returns: @@ -255,18 +267,18 @@ def print_iteration_log(self, history=None, printed_only=True, include_header=Tr file=file, ) - def _prepare_resume_data(self, resume, validate_resume, restore_resume): + def _prepare_resume_data(self, resume: Union[None, Data], validate_resume: Callable, restore_resume: Callable): """Validate and restore a resume checkpoint before integration. Args: - resume (Data or None): Resume checkpoint passed to ``integrate``. - validate_resume (callable): Validator taking ``resume`` and raising + resume (Union[None, Data]): Resume checkpoint passed to ``integrate``. + validate_resume (Callable): Validator taking ``resume`` and raising on incompatible state. - restore_resume (callable): Restorer taking ``resume`` and mutating + restore_resume (Callable): Restorer taking ``resume`` and mutating the current stopping criterion into a compatible resumed state. Returns: - Data or None: A validated deep copy of the supplied checkpoint, or + Union[None, Data]: A validated deep copy of the supplied checkpoint, or None when no checkpoint was supplied. """ if resume is None: @@ -279,7 +291,9 @@ def _prepare_resume_data(self, resume, validate_resume, restore_resume): @staticmethod def _detach_resume_stopping_criterion_history(data): - """Detach copied solver-owned history caches while preserving checkpoint history.""" + """Detach copied solver-owned history caches while preserving + checkpoint history. + """ stopping_crit = getattr(data, "stopping_crit", None) if stopping_crit is None: return @@ -288,20 +302,25 @@ def _detach_resume_stopping_criterion_history(data): if hasattr(stopping_crit, "history_df"): stopping_crit.history_df = None - def _restore_resume_state(self, data): + def _restore_resume_state(self, data: Data): """Optional hook for subclasses to align state before resuming. Subclasses that need to restore RNG state or rewrite checkpoint fields - may override this method. The default implementation contains no operation. + may override this method. The default implementation contains no + operation. Args: data (Data): Deep-copied resume checkpoint that will be mutated by the resumed integration run. + + Returns: + None """ return None - def _capture_resume_provenance(self, resume): - """Capture resume bookkeeping before the live ``Data`` object is mutated. + def _capture_resume_provenance(self, resume: Data or None): + """Capture resume bookkeeping before the live ``Data`` object is + mutated. Args: resume (Data or None): Resume checkpoint passed to ``integrate``. @@ -357,14 +376,14 @@ def _annotate_checkpoint_metadata(self, data): self.discrete_distrib ) - def _finalize_integration_data(self, data, elapsed, resume_provenance=None): + def _finalize_integration_data(self, data: Data, elapsed: float, resume_provenance: Union[None, dict] = None): """Attach shared integration metadata before returning ``Data``. Args: data (Data): Integration state to finalize. elapsed (float): Wall-clock time spent in the current ``integrate`` call. - resume_provenance (dict or None, optional): Output of + resume_provenance (Union[None, dict]): Output of :meth:`_capture_resume_provenance`. Defaults to None. """ data.stopping_crit = self @@ -388,12 +407,13 @@ def _finalize_integration_data(self, data, elapsed, resume_provenance=None): data.history_df = getattr(self, "history_df", None) self._annotate_checkpoint_metadata(data) - def _resume_value_equal(self, current, saved): - """Deep equality check tolerant of arrays, lists, dicts, and QMCPy objects. + def _resume_value_equal(self, current: Any, saved: Any): + """Deep equality check tolerant of arrays, lists, dicts, and QMCPy + objects. Args: - current: Value from the live stopping criterion. - saved: Value from the resume checkpoint. + current (Any): Value from the live stopping criterion. + saved (Any): Value from the resume checkpoint. Returns: bool: True when the two values are considered equal. @@ -441,8 +461,9 @@ def _resume_value_equal(self, current, saved): def _is_sparse(value): return hasattr(value, "nnz") and hasattr(value, "shape") - def _require_resume_attrs(self, data, attrs): - """Raise ParameterError if any attribute in *attrs* is absent from *data*. + def _require_resume_attrs(self, data: Data, attrs: tuple[str, ...]): + """Raise ParameterError if any attribute in *attrs* is absent from + *data*. Args: data (Data): Resume checkpoint. @@ -458,16 +479,17 @@ def _require_resume_attrs(self, data, attrs): % ", ".join(sorted(missing)) ) - def _validate_resume_object(self, label, current, saved, attrs): - """Validate that a saved sub-object is compatible with the current one. + def _validate_resume_object(self, label: str, current: object, saved: object, attrs: tuple[str, ...]): + """Validate that a saved sub-object is compatible with the current + one. Checks type equality and then compares each attribute listed in *attrs* using :meth:`_resume_value_equal`. Args: label (str): Human-readable name used in error messages. - current: Live object (integrand, true measure, etc.). - saved: Saved object from the resume checkpoint. + current (object): Live object (integrand, true measure, etc.). + saved (object): Saved object from the resume checkpoint. attrs (tuple[str, ...]): Attribute names to compare. Raises: @@ -494,7 +516,7 @@ def _validate_resume_object(self, label, current, saved, attrs): "resume data has incompatible %s.%s." % (label, attr) ) - def _validate_resume_data(self, data, required_fields=()): + def _validate_resume_data(self, data: Data, required_fields: tuple[str, ...] =()): """Run standard cross-cutting resume compatibility checks. Validates stopping criterion type, integrand, true measure, discrete @@ -503,8 +525,8 @@ def _validate_resume_data(self, data, required_fields=()): Args: data (Data): Resume checkpoint to validate. - required_fields (tuple[str, ...], optional): Additional attribute - names that must be present on *data*. Defaults to ``()``. + required_fields (tuple[str, ...]): Additional attribute names that + must be present on *data*. Defaults to ``()``. Raises: ParameterError: If any compatibility check fails. @@ -536,7 +558,7 @@ def _validate_resume_data(self, data, required_fields=()): if int(data.n_total) > int(self.n_limit): raise ParameterError("resume data n_total=%d exceeds current n_limit=%d." % (int(data.n_total), int(self.n_limit))) - def _validate_resume_with_state(self, data, required_fields=(), state_fields=()): + def _validate_resume_with_state(self, data: Data, required_fields: tuple[str, ...] =(), state_fields: tuple[str, ...] =()): """Validate resume data including algorithm-specific state fields. Calls :meth:`_validate_resume_data` and additionally checks that all @@ -544,10 +566,10 @@ def _validate_resume_with_state(self, data, required_fields=(), state_fields=()) Args: data (Data): Resume checkpoint to validate. - required_fields (tuple[str, ...], optional): Extra data attributes - required beyond the standard set. Defaults to ``()``. - state_fields (tuple[str, ...], optional): Algorithm-state attributes - that must also be present. Defaults to ``()``. + required_fields (tuple[str, ...]): Extra data attributes required + beyond the standard set. Defaults to ``()``. + state_fields (tuple[str, ...]): Algorithm-state attributes that + must also be present. Defaults to ``()``. Raises: ParameterError: If any compatibility check fails. @@ -573,20 +595,20 @@ def _is_power_of_two(n): return n > 0 and (n & (n - 1)) == 0 @staticmethod - def _resolve_error_fun(error_fun): + def _resolve_error_fun(error_fun: Union[str, Callable]): """Resolve an *error_fun* argument from a string keyword or callable. Args: - error_fun (Union[str, callable]): ``'EITHER'`` or ``'BOTH'`` or a + error_fun (Union[str, Callable]): ``'EITHER'`` or ``'BOTH'`` or a callable with signature ``(sv, abs_tol, rel_tol) -> tol``. Returns: - tuple[callable, str or None]: The resolved callable and its canonical - string key (``'EITHER'`` or ``'BOTH'``), or ``None`` when the - input was already a callable. + tuple[callable, str or None]: The resolved callable and its canonical string key (``'EITHER'`` or + ``'BOTH'``), or ``None`` when the input was already a callable. Raises: - ParameterError: If a string argument is not ``'EITHER'`` or ``'BOTH'``. + ParameterError: If a string argument is not ``'EITHER'`` or + ``'BOTH'``. """ _error_fun_key = None if isinstance(error_fun, str): @@ -600,7 +622,7 @@ def _resolve_error_fun(error_fun): return error_fun, _error_fun_key @staticmethod - def _checkpoint_rmse_tol(data): + def _checkpoint_rmse_tol(data: Data): """Extract the RMSE tolerance stored in a resume checkpoint. Args: @@ -617,18 +639,18 @@ def _checkpoint_rmse_tol(data): pass return None - def _init_control_variates(self, control_variates, control_variate_means): + def _init_control_variates(self, control_variates: Union[list, AbstractIntegrand], control_variate_means: np.ndarray): """Validate and store control variates and their means. - Sets ``self.cv``, ``self.cv_mu``, and ``self.ncv`` after validating that - every entry in *control_variates* is an ``AbstractIntegrand`` instance - that shares the same discrete distribution and ``d_indv`` as the main - integrand. + Sets ``self.cv``, ``self.cv_mu``, and ``self.ncv`` after validating + that every entry in *control_variates* is an ``AbstractIntegrand`` + instance that shares the same discrete distribution and ``d_indv`` as + the main integrand. Args: - control_variates (list or AbstractIntegrand): Control variate + control_variates (Union[list, AbstractIntegrand]): Control variate integrand(s). - control_variate_means (array-like): Known means of each control + control_variate_means (np.ndarray): Known means of each control variate. Returns: @@ -642,7 +664,8 @@ def _init_control_variates(self, control_variates, control_variate_means): if isinstance(self.cv, AbstractIntegrand): self.cv = [self.cv] self.cv_mu = self.cv_mu[None, ...] - assert isinstance(self.cv, list), "cv must be a list of AbstractIntegrand objects" + if not (isinstance(self.cv, list)): + raise AssertionError("cv must be a list of AbstractIntegrand objects") for cv in self.cv: if ( (not isinstance(cv, AbstractIntegrand)) @@ -658,7 +681,7 @@ def _init_control_variates(self, control_variates, control_variate_means): self.ncv = len(self.cv) return self.ncv - def _restore_resume_rng_state(self, data): + def _restore_resume_rng_state(self, data: Data): """Deep-copy the saved RNG state into the live discrete distribution. Ensures that samples drawn after resuming are independent of those @@ -678,21 +701,22 @@ def _restore_resume_rng_state(self, data): saved_distrib.rng.bit_generator.state ) - def _compute_indv_alphas(self, alphas_comb): - """Distribute combined confidence levels to individual integrand dimensions. + def _compute_indv_alphas(self, alphas_comb: np.ndarray): + """Distribute combined confidence levels to individual integrand + dimensions. Uses the integrand dependency map to allocate the per-combined-output alpha budget down to each individual output dimension. Args: - alphas_comb (np.ndarray): Per-combined-output confidence levels with - shape ``integrand.d_comb``. + alphas_comb (np.ndarray): Per-combined-output confidence levels + with shape ``integrand.d_comb``. Returns: - tuple[np.ndarray, bool]: ``(alphas_indv, identity_dependency)`` - where *alphas_indv* has shape ``integrand.d_indv`` and - *identity_dependency* is True when each combined output depends - on exactly its matching individual output. + tuple[np.ndarray, bool]: ``(alphas_indv, identity_dependency)`` where *alphas_indv* has + shape ``integrand.d_indv`` and *identity_dependency* is True when + each combined output depends on exactly its matching individual + output. """ alphas_indv = np.tile(1, self.integrand.d_indv) identity_dependency = True @@ -712,15 +736,15 @@ def _compute_indv_alphas(self, alphas_comb): alphas_indv = np.where(alpha_k_mat == 0, alphas_indv, np.minimum(alpha_k_mat, alphas_indv)) return alphas_indv, identity_dependency - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: """Reset the tolerances. Args: - abs_tol (float): Absolute tolerance, when supported. If supplied, + abs_tol (Union[None, float]): Absolute tolerance, when supported. If supplied, reset it; otherwise ignore it. - rel_tol (float): Relative tolerance, when supported. If supplied, + rel_tol (Union[None, float]): Relative tolerance, when supported. If supplied, reset it; otherwise ignore it. - rmse_tol (float): RMSE tolerance, when supported. If supplied, + rmse_tol (Union[None, float]): RMSE tolerance, when supported. If supplied, reset it; otherwise ignore it. If ``rmse_tol`` is not supplied but ``abs_tol`` is, then ``rmse_tol = abs_tol / norm.ppf(1 - alpha / 2)``. diff --git a/qmcpy/stopping_criterion/cub_mc_clt.py b/qmcpy/stopping_criterion/cub_mc_clt.py index c2e72870b..cd35f5bf4 100644 --- a/qmcpy/stopping_criterion/cub_mc_clt.py +++ b/qmcpy/stopping_criterion/cub_mc_clt.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_stopping_criterion import AbstractStoppingCriterion from ..util.data import Data @@ -14,8 +15,8 @@ class CubMCCLT(AbstractStoppingCriterion): - r""" - IID Monte Carlo stopping criterion based on the Central Limit Theorem in a two step method. + r"""IID Monte Carlo stopping criterion based on the Central Limit Theorem + in a two step method. Examples: >>> ao = FinancialOption(IIDStdUniform(52,seed=7)) @@ -127,27 +128,30 @@ class CubMCCLT(AbstractStoppingCriterion): def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0.0, - n_init=1024, - n_limit=2**30, - inflate=1.2, - alpha=0.01, - control_variates=None, - control_variate_means=None, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0.0, + n_init: int = 1024, + n_limit: int = 2**30, + inflate: float = 1.2, + alpha: Union[float, np.ndarray] = 0.01, + control_variates: Union[None, list] = None, + control_variate_means: Union[None, np.ndarray] = None, + ) -> None: + r"""Initialize a CubMCCLT stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - inflate (float): Inflation factor $\geq 1$ to multiply by the variance estimate to make it more conservative. - alpha (np.ndarray): Uncertainty level in $(0,1)$. - control_variates (list): Integrands to use as control variates, each with the same underlying discrete distribution instance. - control_variate_means (np.ndarray): Means of each control variate. + inflate (float): Inflation factor $\geq 1$ to multiply by the + variance estimate to make it more conservative. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. + control_variates (Union[None, list]): Integrands to use as control variates, + each with the same underlying discrete distribution instance. + control_variate_means (Union[None, np.ndarray]): Means of each control variate. """ if control_variates is None: control_variates = [] @@ -166,13 +170,16 @@ def __init__( self.rel_tol = rel_tol self.n_init = n_init self.n_limit = n_limit - assert self.n_limit > ( + if not (self.n_limit > ( 2 * self.n_init - ), "require n_limit is at least twic as much as n_init" + )): + raise AssertionError("require n_limit is at least twic as much as n_init") self.alpha = alpha self.inflate = inflate - assert self.inflate >= 1 - assert 0 < self.alpha < 1 + if not (self.inflate >= 1): + raise AssertionError + if not (0 < self.alpha < 1): + raise AssertionError # QMCPy Objs self.integrand = integrand self.true_measure = self.integrand.true_measure @@ -181,13 +188,15 @@ def __init__( allowed_distribs=[AbstractIIDDiscreteDistribution], allow_vectorized_integrals=True, ) - assert self.integrand.d_indv == () + if not (self.integrand.d_indv == ()): + raise AssertionError # control variates self._init_control_variates(control_variates, control_variate_means) if self.ncv > 0: - assert self.cv_mu.shape == ( + if not (self.cv_mu.shape == ( (self.ncv,) + self.integrand.d_indv - ), "Control variate means should have shape (len(control variates),d_indv)." + )): + raise AssertionError("Control variate means should have shape (len(control variates),d_indv).") self.parameters += ["cv", "cv_mu"] self.z_star = -norm.ppf(self.alpha / 2.0) @@ -198,7 +207,24 @@ def _get_main_stage_samples(self, data): ycv = np.array(data.ycvfull[:, self.n_init :], copy=False) return y - ((ycv - self.cv_mu[:, None]) * self.beta[:, None]).sum(0) - def integrate(self, resume=None): + def integrate(self, resume: Union[None, Data] = None) -> tuple: + """Determine the samples needed to satisfy the target tolerance. + + Draws an initial `self.n_init` samples to estimate the standard + deviation, then uses the CLT-based normal quantile (`self.z_star`, + inflated by `self.inflate`) to size and draw a second, final batch, + producing a symmetric confidence-interval bound on the integral. + + Args: + resume (Union[None, Data]): Unsupported; must be `None`, as `CubMCCLT` cannot + resume a prior checkpoint. + + Returns: + tuple: Approximation to the integral and the corresponding data object. + + Raises: + ParameterError: If `resume` is not `None`. + """ t_start = time() trace = self._make_trace_logger() if resume is not None: @@ -273,8 +299,19 @@ def integrate(self, resume=None): trace.finalize() return data.solution, data - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rmse_tol is None, "rmse_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance. + rel_tol (Union[None, float]): Relative error tolerance. + rmse_tol (Union[None, float]): Unsupported; must be `None`. + + Raises: + AssertionError: If `rmse_tol` is supplied. + """ + if not (rmse_tol is None): + raise AssertionError("rmse_tol not supported by this stopping criterion.") if abs_tol is not None: self.abs_tol = abs_tol if rel_tol is not None: diff --git a/qmcpy/stopping_criterion/cub_mc_clt_vec.py b/qmcpy/stopping_criterion/cub_mc_clt_vec.py index 66ecebf6e..b955f1991 100644 --- a/qmcpy/stopping_criterion/cub_mc_clt_vec.py +++ b/qmcpy/stopping_criterion/cub_mc_clt_vec.py @@ -1,3 +1,5 @@ +from ..integrand.abstract_integrand import AbstractIntegrand +from typing import Union, Callable from .abstract_stopping_criterion import AbstractStoppingCriterion from ..util.data import Data @@ -14,8 +16,8 @@ class CubMCCLTVec(AbstractStoppingCriterion): - r""" - IID Monte Carlo stopping criterion stopping criterion based on the Central Limit Theorem with doubling sample sizes. + r"""IID Monte Carlo stopping criterion stopping criterion based on the + Central Limit Theorem with doubling sample sizes. Examples: >>> k = Keister(IIDStdUniform(seed=7)) @@ -159,36 +161,40 @@ class CubMCCLTVec(AbstractStoppingCriterion): def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0.0, - n_init=256.0, - n_limit=2**30, - error_fun="EITHER", - inflate=1, - alpha=0.01, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0.0, + n_init: int = 256, + n_limit: int = 2**30, + error_fun: Union[str, Callable] = "EITHER", + inflate: float = 1, + alpha: Union[float, np.ndarray] = 0.01, + ) -> None: + r"""Initialize a CubMCCLTVec stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - error_fun (Union[str, callable]): Function mapping the approximate solution, absolute error tolerance, and relative error tolerance to the current error bound. + error_fun (Union[str, Callable]): Function mapping the approximate + solution, absolute error tolerance, and relative error + tolerance to the current error bound. - - `'EITHER'`, the default, requires the approximation error must be below either the absolue *or* relative tolerance. + - `'EITHER'`, the default, requires the approximation error to be below either the absolute *or* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.maximum(abs_tol,abs(sv)*rel_tol) ``` - - `'BOTH'` requires the approximation error to be below both the absolue *and* relative tolerance. + - `'BOTH'` requires the approximation error to be below both the absolute *and* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.minimum(abs_tol,abs(sv)*rel_tol) ``` - inflate (float): Inflation factor $\geq 1$ to multiply by the variance estimate to make it more conservative. - alpha (np.ndarray): Uncertainty level in $(0,1)$. + inflate (float): Inflation factor $\geq 1$ to multiply by the + variance estimate to make it more conservative. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. """ self.parameters = [ "inflate", @@ -212,11 +218,13 @@ def __init__( # Set Attributes self.n_init = int(n_init) self.n_limit = int(n_limit) - assert isinstance(error_fun, str) or callable(error_fun) + if not (isinstance(error_fun, str) or callable(error_fun)): + raise AssertionError self.error_fun, _ = self._resolve_error_fun(error_fun) self.alpha = alpha self.inflate = float(inflate) - assert self.inflate >= 1 + if not (self.inflate >= 1): + raise AssertionError # QMCPy Objs self.integrand = integrand self.true_measure = self.integrand.true_measure @@ -225,9 +233,10 @@ def __init__( allowed_distribs=[AbstractIIDDiscreteDistribution], allow_vectorized_integrals=True, ) - assert ( + if not ( self.integrand.discrete_distrib.no_replications == True - ), "Require the discrete distribution has replications=None" + ): + raise AssertionError("Require the discrete distribution has replications=None") self.alphas_indv, _ = self._compute_indv_alphas( np.full(self.integrand.d_comb, self.alpha) ) @@ -257,7 +266,22 @@ def _restore_resume_state(self, data): self.integrand.discrete_distrib = self.discrete_distrib self.integrand.true_measure.discrete_distrib = self.discrete_distrib - def integrate(self, resume=None): + def integrate(self, resume: Union[None, Data] = None) -> tuple: + """Determine the samples needed to satisfy the target tolerance. + + Doubles the sample count each iteration and forms a CLT-based + confidence interval (`self.z_star`, inflated by `self.inflate`) on + each not-yet-converged output. Stops once every combined output is + within tolerance or `self.n_limit` would be exceeded. + + Args: + resume (Union[None, Data]): Existing integration state to resume from, if + supported. Defaults to None. + + Returns: + tuple: Approximation to the integral with shape ``integrand.d_comb`` + and the corresponding data object. + """ t_start = time() resume_provenance = self._capture_resume_provenance(resume) trace = self._make_trace_logger() @@ -352,8 +376,21 @@ def integrate(self, resume=None): trace.finalize() return data.solution, data - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rmse_tol is None, "rmse_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance, broadcast to + `self.abs_tols` with shape `integrand.d_comb`. + rel_tol (Union[None, float]): Relative error tolerance, broadcast to + `self.rel_tols` with shape `integrand.d_comb`. + rmse_tol (Union[None, float]): Unsupported; must be `None`. + + Raises: + AssertionError: If `rmse_tol` is supplied. + """ + if not (rmse_tol is None): + raise AssertionError("rmse_tol not supported by this stopping criterion.") if abs_tol is not None: self.abs_tol = abs_tol self.abs_tols = np.full(self.integrand.d_comb, self.abs_tol) diff --git a/qmcpy/stopping_criterion/cub_mc_g.py b/qmcpy/stopping_criterion/cub_mc_g.py index 99d151f45..d2978fada 100644 --- a/qmcpy/stopping_criterion/cub_mc_g.py +++ b/qmcpy/stopping_criterion/cub_mc_g.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_stopping_criterion import AbstractStoppingCriterion from ..util.data import Data @@ -15,8 +16,8 @@ class CubMCG(AbstractStoppingCriterion): - r""" - IID Monte Carlo stopping criterion using Berry-Esseen inequalities in a two step method with guarantees for functions with bounded kurtosis. + r"""IID Monte Carlo stopping criterion using Berry-Esseen inequalities in + a two step method with guarantees for functions with bounded kurtosis. Examples: >>> ao = FinancialOption(IIDStdUniform(52,seed=7)) @@ -253,27 +254,30 @@ class CubMCG(AbstractStoppingCriterion): def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0.0, - n_init=1024, - n_limit=2**30, - inflate=1.2, - alpha=0.01, - control_variates=None, - control_variate_means=None, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0.0, + n_init: int = 1024, + n_limit: int = 2**30, + inflate: float = 1.2, + alpha: Union[float, np.ndarray] = 0.01, + control_variates: Union[None, list] = None, + control_variate_means: Union[None, np.ndarray] = None, + ) -> None: + r"""Initialize a CubMCG stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - inflate (float): Inflation factor $\geq 1$ to multiply by the variance estimate to make it more conservative. - alpha (np.ndarray): Uncertainty level in $(0,1)$. - control_variates (list): Integrands to use as control variates, each with the same underlying discrete distribution instance. - control_variate_means (np.ndarray): Means of each control variate. + inflate (float): Inflation factor $\geq 1$ to multiply by the + variance estimate to make it more conservative. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. + control_variates (Union[None, list]): Integrands to use as control variates, + each with the same underlying discrete distribution instance. + control_variate_means (Union[None, np.ndarray]): Means of each control variate. """ if control_variates is None: control_variates = [] @@ -301,13 +305,15 @@ def __init__( allowed_distribs=[AbstractIIDDiscreteDistribution], allow_vectorized_integrals=False, ) - assert self.integrand.d_indv == () + if not (self.integrand.d_indv == ()): + raise AssertionError # control variates self._init_control_variates(control_variates, control_variate_means) if self.ncv > 0: - assert self.cv_mu.shape == ( + if not (self.cv_mu.shape == ( (self.ncv,) + self.integrand.d_indv - ), "Control variate means should have shape (len(control variates),d_indv)." + )): + raise AssertionError("Control variate means should have shape (len(control variates),d_indv).") self.parameters += ["cv", "cv_mu"] def _get_main_stage_samples(self, data): @@ -323,7 +329,26 @@ def _update_main_stage_solution(self, data): data.solution = y_main.mean() data.n_total = data.yfull.shape[-1] - def integrate(self, resume=None): + def integrate(self, resume: Union[None, Data] = None) -> tuple: + """Determine the samples needed to satisfy the target tolerance. + + Draws an initial `self.n_init` samples to estimate the standard + deviation and kurtosis. If `self.rel_tol` is 0, sizes and draws one + additional batch via a Chebyshev/Berry-Esseen bound (`_nchebe`). + Otherwise, iteratively grows the sample size (`_ncbinv`) until the + Berry-Esseen confidence bound meets both the absolute and relative + tolerance or `self.n_limit` would be exceeded. + + Args: + resume (Union[None, Data]): Unsupported; must be `None`, as `CubMCG` cannot + resume a prior checkpoint. + + Returns: + tuple: Approximation to the integral and the corresponding data object. + + Raises: + ParameterError: If `resume` is not `None`. + """ t_start = time() trace = self._make_trace_logger() if resume is not None: @@ -529,31 +554,47 @@ def _ncbinv(self, n1, alpha1, kurtmax): # take the min of Chebyshev and Berry Esseen tolerance return eps - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rmse_tol is None, "rmse_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance. + rel_tol (Union[None, float]): Relative error tolerance. + rmse_tol (Union[None, float]): Unsupported; must be `None`. + + Raises: + AssertionError: If `rmse_tol` is supplied. + """ + if not (rmse_tol is None): + raise AssertionError("rmse_tol not supported by this stopping criterion.") if abs_tol != None: self.abs_tol = abs_tol if rel_tol != None: self.rel_tol = rel_tol -def _tol_fun(abs_tol, rel_tol, theta, mu, toltype): - # """ - # Generalized error tolerance function. +def _tol_fun(abs_tol: float, rel_tol: float, theta: float, mu: float, toltype: str): + """Generalized error tolerance function. - # Args: - # abs_tol (float): absolute error tolerance - # rel_tol (float): relative error tolerance - # theta (float): parameter in 'theta' case - # mu (float): true mean - # toltype (str): different options of tolerance function + Args: + abs_tol (float): Absolute error tolerance. + rel_tol (float): Relative error tolerance. + theta (float): Weight in `"combine"` case; 0 gives pure relative + tolerance, 1 gives pure absolute tolerance. + mu (float): True mean. + toltype (str): `"combine"` for a weighted sum of the two tolerances, + or `"max"` for their max. - # Returns: - # float: tolerance as weighted sum of absolute and relative tolerance - # """ + Returns: + float: Tolerance as a combination of absolute and relative tolerance. + """ if toltype == "combine": # the linear combination of two tolerances # theta == 0 --> relative error tolerance # theta == 1 --> absolute error tolerance return theta * abs_tol + (1 - theta) * rel_tol * abs(mu) elif toltype == "max": # the max case return max(abs_tol, rel_tol * abs(mu)) + else: + raise ParameterError( + f"unknown toltype {toltype!r}; expected 'combine' or 'max'." + ) diff --git a/qmcpy/stopping_criterion/cub_mlmc.py b/qmcpy/stopping_criterion/cub_mlmc.py index 5fba7672f..1643058a4 100644 --- a/qmcpy/stopping_criterion/cub_mlmc.py +++ b/qmcpy/stopping_criterion/cub_mlmc.py @@ -1,4 +1,7 @@ +from typing import Union from .abstract_cub_mlmc import AbstractCubMLMC +from ..integrand.abstract_integrand import AbstractIntegrand +from ..util.data import Data import copy from ..discrete_distribution import IIDStdUniform from ..discrete_distribution.abstract_discrete_distribution import ( @@ -13,10 +16,6 @@ class CubMLMC(AbstractCubMLMC): - _RESUME_REQUIRED_FIELDS = ( - "levels", "n_level", "sum_level", "diff_n_level", "cost_level", "level_integrands" - ) - """ Multilevel IID Monte Carlo stopping criterion. @@ -27,7 +26,7 @@ class CubMLMC(AbstractCubMLMC): >>> data Data (Data) solution 1.785 - n_total 3577556 + n_total 3199033 levels 2^(2) n_level [2438191 490331 207606 62905] mean_level [1.715 0.053 0.013 0.003] @@ -35,7 +34,7 @@ class CubMLMC(AbstractCubMLMC): cost_per_sample [ 2. 4. 8. 16.] alpha 2.008 beta 1.997 - gamma 1.000 + gamma ... time_integrate ... CubMLMC (AbstractStoppingCriterion) rmse_tol 0.006 @@ -73,34 +72,45 @@ class CubMLMC(AbstractCubMLMC): 2. [http://people.maths.ox.ac.uk/~gilesm/mlmc/#MATLAB](http://people.maths.ox.ac.uk/~gilesm/mlmc/#MATLAB). """ + _RESUME_REQUIRED_FIELDS = ( + "levels", "n_level", "sum_level", "diff_n_level", "cost_level", "level_integrands" + ) + def __init__( self, - integrand, - abs_tol=0.05, - rmse_tol=None, - n_init=256, - n_limit=1e10, - alpha=0.01, - levels_min=2, - levels_max=10, - alpha0=-1.0, - beta0=-1.0, - gamma0=-1.0, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 0.05, + rmse_tol: Union[None, np.ndarray] = None, + n_init: int = 256, + n_limit: int = 10**10, + alpha: Union[float, np.ndarray] = 0.01, + levels_min: int = 2, + levels_max: int = 10, + alpha0: float = -1.0, + beta0: float = -1.0, + gamma0: float = -1.0, + ) -> None: + r"""Initialize a CubMLMC stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rmse_tol (np.ndarray): Root mean squared error tolerance. - If supplied, then absolute tolerance and alpha are ignored in favor of the rmse tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rmse_tol (Union[None, np.ndarray]): Root mean squared error tolerance. If + supplied, then absolute tolerance and alpha are ignored in + favor of the rmse tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - alpha (np.ndarray): Uncertainty level in $(0,1)$. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. levels_min (int): Minimum level of refinement $\geq 2$. levels_max (int): Maximum level of refinement $\geq$ `levels_min`. - alpha0 (float): Weak error is $\mathcal{O}(2^{-\alpha_0\ell})$ in the level $\ell$. If `alpha0`$\leq 0$ then it will be estimated. - beta0 (float): Variance is $\mathcal{O}(2^{-\beta_0\ell})$ in the level $\ell$. If `beta0`$\leq 0$ then it will be estimated. - gamma0 (float): Sample cost is $\mathcal{O}(2^{\gamma_0\ell})$ in the level $\ell$. If `gamma0`$\leq 0$ then it will be estimated. + alpha0 (float): Weak error is $\mathcal{O}(2^{-\alpha_0\ell})$ in + the level $\ell$. If `alpha0`$\leq 0$ then it will be + estimated. + beta0 (float): Variance is $\mathcal{O}(2^{-\beta_0\ell})$ in the + level $\ell$. If `beta0`$\leq 0$ then it will be estimated. + gamma0 (float): Sample cost is $\mathcal{O}(2^{\gamma_0\ell})$ in + the level $\ell$. If `gamma0`$\leq 0$ then it will be + estimated. """ self.parameters = ["rmse_tol", "n_init", "levels_min", "levels_max", "theta"] if levels_min < 2: @@ -115,7 +125,8 @@ def __init__( else: # use absolute tolerance self.rmse_tol = float(abs_tol) / norm.ppf(1 - alpha / 2) self.alpha = alpha - assert 0 < self.alpha < 1 + if not (0 < self.alpha < 1): + raise AssertionError self.n_init = n_init self.n_limit = n_limit self.levels_min = levels_min @@ -222,7 +233,9 @@ def _run_integrate_loop( return snapshots def _replay_resume_exactly(self, checkpoint, t_start=None, resume_provenance=None): - """Replay cached per-level diffs to reconstruct checkpoint state and trace rows.""" + """Replay cached per-level diffs to reconstruct checkpoint state and + trace rows. + """ shadow = self._construct_data() shadow.level_integrands = list(checkpoint.level_integrands) shadow.cached_level_diffs = [ @@ -262,17 +275,17 @@ def _replay_resume_exactly(self, checkpoint, t_start=None, resume_provenance=Non delattr(shadow, attr) return shadow, snapshots[absorb_index:], replay_iter_count - def integrate(self, resume=None) -> tuple: + def integrate(self, resume: Union[None, Data] = None) -> tuple: """Run (or continue) the MLMC integration. Args: - resume (Data, optional): Checkpoint returned by a previous - ``integrate()`` call. The new tolerance may be tighter *or* - looser than the one used when the checkpoint was created. - With a tighter tolerance the algorithm draws additional samples - from where it left off. With a looser tolerance the existing - samples already satisfy the requirement and the method returns - immediately with no new sampling. + resume (Union[None, Data]): Checkpoint returned by a previous ``integrate()`` + call. The new tolerance may be tighter *or* looser than the + one used when the checkpoint was created. With a tighter + tolerance the algorithm draws additional samples from where it + left off. With a looser tolerance the existing samples already + satisfy the requirement and the method returns immediately with + no new sampling. Returns: tuple: ``(solution, data)``. diff --git a/qmcpy/stopping_criterion/cub_mlmc_cont.py b/qmcpy/stopping_criterion/cub_mlmc_cont.py index 5f47411df..9617eff10 100644 --- a/qmcpy/stopping_criterion/cub_mlmc_cont.py +++ b/qmcpy/stopping_criterion/cub_mlmc_cont.py @@ -1,4 +1,7 @@ +from typing import Union from .abstract_cub_mlmc import AbstractCubMLMC +from ..integrand.abstract_integrand import AbstractIntegrand +from ..util.data import Data import copy from ..discrete_distribution import IIDStdUniform from ..discrete_distribution.abstract_discrete_distribution import ( @@ -13,10 +16,6 @@ class CubMLMCCont(AbstractCubMLMC): - _RESUME_REQUIRED_FIELDS = ( - "levels", "n_level", "sum_level", "diff_n_level", "cost_level", "level_integrands" - ) - r""" Multilevel IID Monte Carlo stopping criterion with continuation. @@ -27,15 +26,15 @@ class CubMLMCCont(AbstractCubMLMC): >>> data Data (Data) solution 1.771 - n_total 2291120 + n_total 1480870 levels 3 - n_level [1094715 222428 79666 912 256] - mean_level [1.71 0.048 0.012] - var_level [21.826 1.768 0.453] + n_level [1145480 230538 104852] + mean_level [1.71 0.048 0.013] + var_level [21.819 1.766 0.451] cost_per_sample [2. 4. 8.] - alpha 1.970 - beta 1.965 - gamma 1.000 + alpha 1.868 + beta 1.969 + gamma ... time_integrate ... CubMLMCCont (AbstractStoppingCriterion) rmse_tol 0.006 @@ -45,7 +44,7 @@ class CubMLMCCont(AbstractCubMLMC): n_tols 10 inflate 1.668 theta_init 2^(-1) - theta 0.010 + theta 0.051 FinancialOption (AbstractIntegrand) option ASIAN call_put CALL @@ -72,30 +71,36 @@ class CubMLMCCont(AbstractCubMLMC): 1. [https://github.com/PieterjanRobbe/MultilevelEstimators.jl](https://github.com/PieterjanRobbe/MultilevelEstimators.jl). """ + _RESUME_REQUIRED_FIELDS = ( + "levels", "n_level", "sum_level", "diff_n_level", "cost_level", "level_integrands" + ) + def __init__( self, - integrand, - abs_tol=0.05, - rmse_tol=None, - n_init=256, - n_limit=1e10, - inflate=100 ** (1 / 9), - alpha=0.01, - levels_min=2, - levels_max=10, - n_tols=10, - theta_init=0.5, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 0.05, + rmse_tol: Union[None, np.ndarray] = None, + n_init: int = 256, + n_limit: int = 10**10, + inflate: float = 100 ** (1 / 9), + alpha: Union[float, np.ndarray] = 0.01, + levels_min: int = 2, + levels_max: int = 10, + n_tols: int = 10, + theta_init: float = 0.5, + ) -> None: + r"""Initialize a CubMLMCCont stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rmse_tol (np.ndarray): Root mean squared error tolerance. - If supplied, then absolute tolerance and alpha are ignored in favor of the rmse tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rmse_tol (Union[None, np.ndarray]): Root mean squared error tolerance. If + supplied, then absolute tolerance and alpha are ignored in + favor of the rmse tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. inflate (float): Coarser tolerance multiplication factor $\geq 1$. - alpha (np.ndarray): Uncertainty level in $(0,1)$. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. levels_min (int): Minimum level of refinement $\geq 2$. levels_max (int): Maximum level of refinement $\geq$ `levels_min`. n_tols (int): Number of coarser tolerances to run. @@ -140,8 +145,10 @@ def __init__( self._active_trace = None self.alpha = alpha self.inflate = inflate - assert self.inflate >= 1 - assert 0 < self.alpha < 1 + if not (self.inflate >= 1): + raise AssertionError + if not (0 < self.alpha < 1): + raise AssertionError super(CubMLMCCont, self).__init__( allowed_distribs=[AbstractIIDDiscreteDistribution], allow_vectorized_integrals=False, @@ -161,18 +168,18 @@ def _can_replay_resume_exactly(self, data): return False return hasattr(data, "level_diffs") and len(data.level_diffs) == len(data.n_level) - def integrate(self, resume=None) -> tuple: + def integrate(self, resume: Union[None, Data] = None) -> tuple: """Run (or continue) the continuation-MLMC integration. Args: - resume (Data, optional): Checkpoint returned by a previous - ``integrate()`` call. The new tolerance may be tighter *or* - looser than the one used when the checkpoint was created. - With a tighter tolerance the algorithm picks up the tolerance - ladder from ``max(checkpoint_rmse_tol, target_rmse_tol)`` and - continues down to ``target_rmse_tol``. With a looser tolerance - the first step immediately converges on the existing samples - and no additional ladder steps are needed. + resume (Union[None, Data]): Checkpoint returned by a previous ``integrate()`` + call. The new tolerance may be tighter *or* looser than the + one used when the checkpoint was created. With a tighter + tolerance the algorithm picks up the tolerance ladder from + ``max(checkpoint_rmse_tol, target_rmse_tol)`` and continues + down to ``target_rmse_tol``. With a looser tolerance the first + step immediately converges on the existing samples and no + additional ladder steps are needed. Returns: tuple: ``(solution, data)``. @@ -261,8 +268,10 @@ def _update_trace_solution(data): ).sum() def _replay_resume_exactly(self, checkpoint, t_start=None, resume_provenance=None): - """Ensure iteration number in `replay_iter_count` same in LOOSE-last and RESUMED-first iterations, - by simply saving `level_rep_sums` and `level_n_increments`.""" + """Ensure iteration number in `replay_iter_count` same in LOOSE-last + and RESUMED-first iterations, by simply saving `level_rep_sums` and + `level_n_increments`. + """ shadow_trace = self._active_trace = None try: shadow = self._construct_data() diff --git a/qmcpy/stopping_criterion/cub_mlqmc.py b/qmcpy/stopping_criterion/cub_mlqmc.py index 7df24ac15..177848db0 100644 --- a/qmcpy/stopping_criterion/cub_mlqmc.py +++ b/qmcpy/stopping_criterion/cub_mlqmc.py @@ -1,4 +1,6 @@ +from typing import Union from .abstract_cub_mlqmc import AbstractCubMLQMC +from ..integrand.abstract_integrand import AbstractIntegrand from ..util.data import Data import copy from ..discrete_distribution import DigitalNetB2, Lattice, Halton @@ -14,11 +16,6 @@ class CubMLQMC(AbstractCubMLQMC): - _RESUME_REQUIRED_FIELDS = ( - "levels", "n_level", "eval_level", "mean_level_reps", "mean_level", - "var_level", "cost_level", "var_cost_ratio_level", "bias_estimate", "level_integrands" - ) - """ Multilevel Quasi-Monte Carlo stopping criterion. @@ -76,26 +73,33 @@ class CubMLQMC(AbstractCubMLQMC): [http://people.maths.ox.ac.uk/~gilesm/files/radon.pdf](http://people.maths.ox.ac.uk/~gilesm/files/radon.pdf). """ + _RESUME_REQUIRED_FIELDS = ( + "levels", "n_level", "eval_level", "mean_level_reps", "mean_level", + "var_level", "cost_level", "var_cost_ratio_level", "bias_estimate", "level_integrands" + ) + def __init__( self, - integrand, - abs_tol=0.05, - rmse_tol=None, - n_init=256, - n_limit=1e10, - alpha=0.01, - levels_min=2, - levels_max=10, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 0.05, + rmse_tol: Union[None, np.ndarray] = None, + n_init: int = 256, + n_limit: int = 10**10, + alpha: Union[float, np.ndarray] = 0.01, + levels_min: int = 2, + levels_max: int = 10, + ) -> None: + r"""Initialize a CubMLQMC stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rmse_tol (np.ndarray): Root mean squared error tolerance. - If supplied, then absolute tolerance and alpha are ignored in favor of the rmse tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rmse_tol (Union[None, np.ndarray]): Root mean squared error tolerance. If + supplied, then absolute tolerance and alpha are ignored in + favor of the rmse tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - alpha (np.ndarray): Uncertainty level in $(0,1)$. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. levels_min (int): Minimum level of refinement $\geq 2$. levels_max (int): Maximum level of refinement $\geq$ `levels_min`. """ @@ -106,7 +110,8 @@ def __init__( else: # use absolute tolerance self.rmse_tol = float(abs_tol) / norm.ppf(1 - alpha / 2) self.alpha = alpha - assert 0 < self.alpha < 1 + if not (0 < self.alpha < 1): + raise AssertionError self.n_init = n_init self.n_limit = n_limit self.levels_min = levels_min @@ -120,7 +125,8 @@ def __init__( allow_vectorized_integrals=False, ) self.replications = self.discrete_distrib.replications - assert self.replications >= 4, "require at least 4 replications" + if not (self.replications >= 4): + raise AssertionError("require at least 4 replications") def _validate_resume(self, data): self._validate_resume_data(data, required_fields=self._RESUME_REQUIRED_FIELDS) @@ -218,17 +224,17 @@ def _run_integrate_loop( break return snapshots - def integrate(self, resume=None) -> tuple: + def integrate(self, resume: Union[None, Data] = None) -> tuple: """Run (or continue) the MLQMC integration. Args: - resume (Data, optional): Checkpoint returned by a previous - ``integrate()`` call. The new tolerance may be tighter *or* - looser than the one used when the checkpoint was created. - With a tighter tolerance the algorithm draws additional samples - from where it left off. With a looser tolerance the existing - samples already satisfy the requirement and the method returns - immediately with no new sampling. + resume (Union[None, Data]): Checkpoint returned by a previous ``integrate()`` + call. The new tolerance may be tighter *or* looser than the + one used when the checkpoint was created. With a tighter + tolerance the algorithm draws additional samples from where it + left off. With a looser tolerance the existing samples already + satisfy the requirement and the method returns immediately with + no new sampling. Returns: tuple: ``(solution, data)``. diff --git a/qmcpy/stopping_criterion/cub_mlqmc_cont.py b/qmcpy/stopping_criterion/cub_mlqmc_cont.py index a6393a871..3a1f71cef 100644 --- a/qmcpy/stopping_criterion/cub_mlqmc_cont.py +++ b/qmcpy/stopping_criterion/cub_mlqmc_cont.py @@ -1,4 +1,6 @@ +from typing import Union from .abstract_cub_mlqmc import AbstractCubMLQMC +from ..integrand.abstract_integrand import AbstractIntegrand from ..util.data import Data import copy from ..discrete_distribution import DigitalNetB2, Lattice, Halton @@ -14,11 +16,6 @@ class CubMLQMCCont(AbstractCubMLQMC): - _RESUME_REQUIRED_FIELDS = ( - "levels", "n_level", "eval_level", "mean_level_reps", "mean_level", - "var_level", "cost_level", "var_cost_ratio_level", "bias_estimate", "level_integrands" - ) - """ Multilevel Quasi-Monte Carlo stopping criterion with continuation. @@ -79,30 +76,37 @@ class CubMLQMCCont(AbstractCubMLQMC): 1. [https://github.com/PieterjanRobbe/MultilevelEstimators.jl](https://github.com/PieterjanRobbe/MultilevelEstimators.jl). """ + _RESUME_REQUIRED_FIELDS = ( + "levels", "n_level", "eval_level", "mean_level_reps", "mean_level", + "var_level", "cost_level", "var_cost_ratio_level", "bias_estimate", "level_integrands" + ) + def __init__( self, - integrand, - abs_tol=0.05, - rmse_tol=None, - n_init=256, - n_limit=1e10, - inflate=100 ** (1 / 9), - alpha=0.01, - levels_min=2, - levels_max=10, - n_tols=10, - theta_init=0.5, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 0.05, + rmse_tol: Union[None, np.ndarray] = None, + n_init: int = 256, + n_limit: int = 10**10, + inflate: float = 100 ** (1 / 9), + alpha: Union[float, np.ndarray] = 0.01, + levels_min: int = 2, + levels_max: int = 10, + n_tols: int = 10, + theta_init: float = 0.5, + ) -> None: + r"""Initialize a CubMLQMCCont stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rmse_tol (np.ndarray): Root mean squared error tolerance. - If supplied, then absolute tolerance and alpha are ignored in favor of the rmse tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rmse_tol (Union[None, np.ndarray]): Root mean squared error tolerance. If + supplied, then absolute tolerance and alpha are ignored in + favor of the rmse tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. inflate (float): Coarser tolerance multiplication factor $\geq 1$. - alpha (np.ndarray): Uncertainty level in $(0,1)$. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. levels_min (int): Minimum level of refinement $\geq 2$. levels_max (int): Maximum level of refinement $\geq$ `levels_min`. n_tols (int): Number of coarser tolerances to run. @@ -136,8 +140,10 @@ def __init__( self._active_trace = None self.alpha = alpha self.inflate = inflate - assert self.inflate >= 1 - assert 0 < self.alpha < 1 + if not (self.inflate >= 1): + raise AssertionError + if not (0 < self.alpha < 1): + raise AssertionError # QMCPy Objs self.integrand = integrand self.true_measure = self.integrand.true_measure @@ -147,7 +153,8 @@ def __init__( allow_vectorized_integrals=False, ) self.replications = self.discrete_distrib.replications - assert self.replications >= 4, "require at least 4 replications" + if not (self.replications >= 4): + raise AssertionError("require at least 4 replications") def _validate_resume(self, data): self._validate_resume_data(data, required_fields=self._RESUME_REQUIRED_FIELDS) @@ -168,18 +175,18 @@ def _can_replay_resume_exactly(self, data): return False return hasattr(data, "level_rep_sums") and hasattr(data, "level_n_increments") - def integrate(self, resume=None) -> tuple: + def integrate(self, resume: Union[None, Data] = None) -> tuple: """Run (or continue) the continuation-MLQMC integration. Args: - resume (Data, optional): Checkpoint returned by a previous - ``integrate()`` call. The new tolerance may be tighter *or* - looser than the one used when the checkpoint was created. - With a tighter tolerance the algorithm picks up the tolerance - ladder from ``max(checkpoint_rmse_tol, target_rmse_tol)`` and - continues down to ``target_rmse_tol``. With a looser tolerance - the first step immediately converges on the existing samples - and no additional ladder steps are needed. + resume (Union[None, Data]): Checkpoint returned by a previous ``integrate()`` + call. The new tolerance may be tighter *or* looser than the + one used when the checkpoint was created. With a tighter + tolerance the algorithm picks up the tolerance ladder from + ``max(checkpoint_rmse_tol, target_rmse_tol)`` and continues + down to ``target_rmse_tol``. With a looser tolerance the first + step immediately converges on the existing samples and no + additional ladder steps are needed. Returns: tuple: ``(solution, data)``. diff --git a/qmcpy/stopping_criterion/cub_qmc_bayes_lattice_g.py b/qmcpy/stopping_criterion/cub_qmc_bayes_lattice_g.py index ff54774e8..26b9a7282 100644 --- a/qmcpy/stopping_criterion/cub_qmc_bayes_lattice_g.py +++ b/qmcpy/stopping_criterion/cub_qmc_bayes_lattice_g.py @@ -1,20 +1,18 @@ +from ..integrand.abstract_integrand import AbstractIntegrand +from typing import Union, Callable from .abstract_cub_bayes_ld_g import AbstractCubBayesLDG from ..discrete_distribution import Lattice from ..integrand import Keister, BoxIntegral, Genz, SensitivityIndices from ..fast_transform import fftbr, omega_fftbr -from ..util import ParameterError # , ParameterWarning #MaxSamplesWarning, +from ..util import ParameterError -# from math import factorial import numpy as np -# from time import time -# import warnings - class CubQMCBayesLatticeG(AbstractCubBayesLDG): - r""" - Quasi-Monte Carlo stopping criterion using fast Bayesian cubature and rank-1 lattices - with guarantees for Gaussian processes having certain shift invariant kernels. + r"""Quasi-Monte Carlo stopping criterion using fast Bayesian cubature and + rank-1 lattices with guarantees for Gaussian processes having certain shift + invariant kernels. Examples: >>> k = Keister(Lattice(2, seed=123456789)) @@ -182,44 +180,49 @@ class CubQMCBayesLatticeG(AbstractCubBayesLDG): def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0, - n_init=2**8, - n_limit=2**22, - error_fun="EITHER", - alpha=0.01, - ptransform="C1SIN", - errbd_type="MLE", - order=2, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0, + n_init: int = 2**8, + n_limit: int = 2**22, + error_fun: Union[str, Callable] = "EITHER", + alpha: Union[float, np.ndarray] = 0.01, + ptransform: str = "C1SIN", + errbd_type: str = "MLE", + order: int = 2, + ) -> None: + r"""Initialize a CubQMCBayesLatticeG stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - error_fun (Union[str, callable]): Function mapping the approximate solution, absolute error tolerance, and relative error tolerance to the current error bound. + error_fun (Union[str, Callable]): Function mapping the approximate + solution, absolute error tolerance, and relative error + tolerance to the current error bound. - - `'EITHER'`, the default, requires the approximation error must be below either the absolue *or* relative tolerance. + - `'EITHER'`, the default, requires the approximation error to be below either the absolute *or* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.maximum(abs_tol,abs(sv)*rel_tol) ``` - - `'BOTH'` requires the approximation error to be below both the absolue *and* relative tolerance. + - `'BOTH'` requires the approximation error to be below both the absolute *and* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.minimum(abs_tol,abs(sv)*rel_tol) ``` - alpha (np.ndarray): Uncertainty level in $(0,1)$. - ptransform (str): Periodization transform, see the options in `AbstractIntegrand.f`. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. + ptransform (str): Periodization transform, see the options in + `AbstractIntegrand.f`. errbd_type (str): Options are - `'MLE'`: Marginal Log Likelihood. - `'GCV'`: Generalized Cross Validation. - `'FULL'`: Full Bayes. - order (int): Bernoulli kernel's order. If zero, choose order automatically + order (int): Bernoulli kernel's order. If zero, choose order + automatically """ super(CubQMCBayesLatticeG, self).__init__( integrand, diff --git a/qmcpy/stopping_criterion/cub_qmc_bayes_net_g.py b/qmcpy/stopping_criterion/cub_qmc_bayes_net_g.py index 9aa96fa5b..780e6184e 100644 --- a/qmcpy/stopping_criterion/cub_qmc_bayes_net_g.py +++ b/qmcpy/stopping_criterion/cub_qmc_bayes_net_g.py @@ -1,3 +1,5 @@ +from ..integrand.abstract_integrand import AbstractIntegrand +from typing import Union, Callable from .abstract_cub_bayes_ld_g import AbstractCubBayesLDG from ..discrete_distribution import DigitalNetB2 from ..integrand import Keister, BoxIntegral, Genz, SensitivityIndices @@ -15,9 +17,9 @@ class CubQMCBayesNetG(AbstractCubBayesLDG): - r""" - Quasi-Monte Carlo stopping criterion using fast Bayesian cubature and digital nets - with guarantees for Gaussian processes having certain digitally shift invariant kernels. + r"""Quasi-Monte Carlo stopping criterion using fast Bayesian cubature and + digital nets with guarantees for Gaussian processes having certain + digitally shift invariant kernels. Examples: >>> k = Keister(DigitalNetB2(2, seed=123456789)) @@ -191,35 +193,38 @@ class CubQMCBayesNetG(AbstractCubBayesLDG): def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0, - n_init=2**8, - n_limit=2**22, - error_fun="EITHER", - alpha=0.01, - errbd_type="MLE", - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0, + n_init: int = 2**8, + n_limit: int = 2**22, + error_fun: Union[str, Callable] = "EITHER", + alpha: Union[float, np.ndarray] = 0.01, + errbd_type: str = "MLE", + ) -> None: + r"""Initialize a CubQMCBayesNetG stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - error_fun (Union[str, callable]): Function mapping the approximate solution, absolute error tolerance, and relative error tolerance to the current error bound. + error_fun (Union[str, Callable]): Function mapping the approximate + solution, absolute error tolerance, and relative error + tolerance to the current error bound. - - `'EITHER'`, the default, requires the approximation error must be below either the absolue *or* relative tolerance. + - `'EITHER'`, the default, requires the approximation error to be below either the absolute *or* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.maximum(abs_tol,abs(sv)*rel_tol) ``` - - `'BOTH'` requires the approximation error to be below both the absolue *and* relative tolerance. + - `'BOTH'` requires the approximation error to be below both the absolute *and* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.minimum(abs_tol,abs(sv)*rel_tol) ``` - alpha (np.ndarray): Uncertainty level in $(0,1)$. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. errbd_type (str): Options are - `'MLE'`: Marginal Log Likelihood. @@ -296,9 +301,21 @@ def _shift_inv_kernel_digital( return vec_lambda, vec_lambda_ring, lambda_factor - # Builds High order walsh kernel function @staticmethod - def BuildKernelFunc(order): + def BuildKernelFunc(order: int) -> Callable: + """Build a 1-D high-order Walsh kernel function. + + Args: + order (int): Smoothness order of the digital-net Walsh kernel; + 1, 2, or 3. + + Returns: + Callable: Function mapping an array of 1-D coordinates to the + corresponding Walsh kernel values. + + Raises: + NotYetImplemented: If `order` is not 1, 2, or 3. + """ # a1 = @(x)(-np.floor(np.log2(x))) def a1(x): out = -np.floor(np.log2(x + np.finfo(float).eps)) diff --git a/qmcpy/stopping_criterion/cub_qmc_lattice_g.py b/qmcpy/stopping_criterion/cub_qmc_lattice_g.py index d888ccc3a..073ec5156 100644 --- a/qmcpy/stopping_criterion/cub_qmc_lattice_g.py +++ b/qmcpy/stopping_criterion/cub_qmc_lattice_g.py @@ -1,3 +1,5 @@ +from ..integrand.abstract_integrand import AbstractIntegrand +from typing import Union, Callable from .abstract_cub_qmc_ld_g import AbstractCubQMCLDG, _default_fudge from ..discrete_distribution import Lattice from ..true_measure import Gaussian, Uniform @@ -10,9 +12,9 @@ class CubQMCLatticeG(AbstractCubQMCLDG): - r""" - Quasi-Monte Carlo stopping criterion using rank-1 lattice cubature - with guarantees for cones of functions with a predictable decay in the Fourier coefficients. + r"""Quasi-Monte Carlo stopping criterion using rank-1 lattice cubature + with guarantees for cones of functions with a predictable decay in the + Fourier coefficients. Examples: >>> k = Keister(Lattice(seed=7)) @@ -175,38 +177,44 @@ class CubQMCLatticeG(AbstractCubQMCLDG): def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0.0, - n_init=2**10, - n_limit=2**30, - error_fun="EITHER", - fudge=_default_fudge, - check_cone=False, - ptransform="BAKER", - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0.0, + n_init: int = 2**10, + n_limit: int = 2**30, + error_fun: Union[str, Callable] = "EITHER", + fudge: Callable = _default_fudge, + check_cone: bool = False, + ptransform: str = "BAKER", + ) -> None: + r"""Initialize a CubQMCLatticeG stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - error_fun (Union[str, callable]): Function mapping the approximate solution, absolute error tolerance, and relative error tolerance to the current error bound. + error_fun (Union[str, Callable]): Function mapping the approximate + solution, absolute error tolerance, and relative error + tolerance to the current error bound. - - `'EITHER'`, the default, requires the approximation error must be below either the absolue *or* relative tolerance. + - `'EITHER'`, the default, requires the approximation error to be below either the absolute *or* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.maximum(abs_tol,abs(sv)*rel_tol) ``` - - `'BOTH'` requires the approximation error to be below both the absolue *and* relative tolerance. + - `'BOTH'` requires the approximation error to be below both the absolute *and* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.minimum(abs_tol,abs(sv)*rel_tol) ``` - fudge (function): Positive function multiplying the finite sum of the Fourier coefficients specified in the cone of functions. - check_cone (bool): Whether or not to check if the function falls in the cone. - ptransform (str): Periodization transform, see the options in `AbstractIntegrand.f`. + fudge (Callable): Positive function multiplying the finite sum of + the Fourier coefficients specified in the cone of functions. + check_cone (bool): Whether or not to check if the function falls in + the cone. + ptransform (str): Periodization transform, see the options in + `AbstractIntegrand.f`. """ super(CubQMCLatticeG, self).__init__( integrand, diff --git a/qmcpy/stopping_criterion/cub_qmc_net_g.py b/qmcpy/stopping_criterion/cub_qmc_net_g.py index 6a215b0e8..0ca1579bc 100644 --- a/qmcpy/stopping_criterion/cub_qmc_net_g.py +++ b/qmcpy/stopping_criterion/cub_qmc_net_g.py @@ -1,3 +1,5 @@ +from ..integrand.abstract_integrand import AbstractIntegrand +from typing import Union, Callable from .abstract_cub_qmc_ld_g import AbstractCubQMCLDG, _default_fudge from ..fast_transform import fwht, omega_fwht from ..util import ParameterError @@ -10,9 +12,9 @@ class CubQMCNetG(AbstractCubQMCLDG): - r""" - Quasi-Monte Carlo stopping criterion using digital net cubature - with guarantees for cones of functions with a predictable decay in the Walsh coefficients. + r"""Quasi-Monte Carlo stopping criterion using digital net cubature with + guarantees for cones of functions with a predictable decay in the Walsh + coefficients. Examples: >>> k = Keister(DigitalNetB2(seed=7)) @@ -217,43 +219,50 @@ class CubQMCNetG(AbstractCubQMCLDG): def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0.0, - n_init=2**10, - n_limit=2**35, - error_fun="EITHER", - fudge=_default_fudge, - check_cone=False, - control_variates=None, - control_variate_means=None, - update_cv_coeffs=False, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0.0, + n_init: int = 2**10, + n_limit: int = 2**35, + error_fun: Union[str, Callable] = "EITHER", + fudge: Callable = _default_fudge, + check_cone: bool = False, + control_variates: Union[None, list] = None, + control_variate_means: Union[None, np.ndarray] = None, + update_cv_coeffs: bool = False, + ) -> None: + r"""Initialize a CubQMCNetG stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - error_fun (Union[str, callable]): Function mapping the approximate solution, absolute error tolerance, and relative error tolerance to the current error bound. + error_fun (Union[str, Callable]): Function mapping the approximate + solution, absolute error tolerance, and relative error + tolerance to the current error bound. - - `'EITHER'`, the default, requires the approximation error must be below either the absolue *or* relative tolerance. + - `'EITHER'`, the default, requires the approximation error to be below either the absolute *or* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.maximum(abs_tol,abs(sv)*rel_tol) ``` - - `'BOTH'` requires the approximation error to be below both the absolue *and* relative tolerance. + - `'BOTH'` requires the approximation error to be below both the absolute *and* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.minimum(abs_tol,abs(sv)*rel_tol) ``` - fudge (function): Positive function multiplying the finite sum of the Fourier coefficients specified in the cone of functions. - check_cone (bool): Whether or not to check if the function falls in the cone. - control_variates (list): Integrands to use as control variates, each with the same underlying discrete distribution instance. - control_variate_means (np.ndarray): Means of each control variate. - update_cv_coeffs (bool): If set to true, the control variate coefficients are recomputed at each iteration. - Otherwise they are estimated once after the initial sampling and then fixed. + fudge (Callable): Positive function multiplying the finite sum of + the Fourier coefficients specified in the cone of functions. + check_cone (bool): Whether or not to check if the function falls in + the cone. + control_variates (Union[None, list]): Integrands to use as control variates, + each with the same underlying discrete distribution instance. + control_variate_means (Union[None, np.ndarray]): Means of each control variate. + update_cv_coeffs (bool): If set to true, the control variate + coefficients are recomputed at each iteration. Otherwise they + are estimated once after the initial sampling and then fixed. """ if control_variates is None: control_variates = [] diff --git a/qmcpy/stopping_criterion/cub_qmc_rep_student_t.py b/qmcpy/stopping_criterion/cub_qmc_rep_student_t.py index 4461a5add..12d242f50 100644 --- a/qmcpy/stopping_criterion/cub_qmc_rep_student_t.py +++ b/qmcpy/stopping_criterion/cub_qmc_rep_student_t.py @@ -1,3 +1,5 @@ +from ..integrand.abstract_integrand import AbstractIntegrand +from typing import Union, Callable from .abstract_stopping_criterion import AbstractStoppingCriterion from ..util.data import Data from ..discrete_distribution import DigitalNetB2 @@ -20,8 +22,6 @@ class CubQMCRepStudentT(AbstractStoppingCriterion): - _RESUME_REQUIRED_FIELDS = ("xfull", "yfull", "n", "n_rep", "_ysums", "n_max") - r""" Quasi-Monte Carlo stopping criterion based on Student's $t$-distribution for multiple replications. @@ -202,38 +202,44 @@ class CubQMCRepStudentT(AbstractStoppingCriterion): [https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=10408613](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=10408613). """ + _RESUME_REQUIRED_FIELDS = ("xfull", "yfull", "n", "n_rep", "_ysums", "n_max") + def __init__( self, - integrand, - abs_tol=1e-2, - rel_tol=0.0, - n_init=256.0, - n_limit=2**30, - error_fun="EITHER", - inflate=1, - alpha=0.01, - ): - r""" + integrand: AbstractIntegrand, + abs_tol: Union[float, np.ndarray] = 1e-2, + rel_tol: Union[float, np.ndarray] = 0.0, + n_init: int = 256, + n_limit: int = 2**30, + error_fun: Union[str, Callable] = "EITHER", + inflate: float = 1, + alpha: Union[float, np.ndarray] = 0.01, + ) -> None: + r"""Initialize a CubQMCRepStudentT stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. - abs_tol (np.ndarray): Absolute error tolerance. - rel_tol (np.ndarray): Relative error tolerance. + abs_tol (Union[float, np.ndarray]): Absolute error tolerance. + rel_tol (Union[float, np.ndarray]): Relative error tolerance. n_init (int): Initial number of samples. n_limit (int): Maximum number of samples. - error_fun (Union[str, callable]): Function mapping the approximate solution, absolute error tolerance, and relative error tolerance to the current error bound. + error_fun (Union[str, Callable]): Function mapping the approximate + solution, absolute error tolerance, and relative error + tolerance to the current error bound. - - `'EITHER'`, the default, requires the approximation error must be below either the absolue *or* relative tolerance. + - `'EITHER'`, the default, requires the approximation error to be below either the absolute *or* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.maximum(abs_tol,abs(sv)*rel_tol) ``` - - `'BOTH'` requires the approximation error to be below both the absolue *and* relative tolerance. + - `'BOTH'` requires the approximation error to be below both the absolute *and* relative tolerance. Equivalent to setting ```python error_fun = lambda sv,abs_tol,rel_tol: np.minimum(abs_tol,abs(sv)*rel_tol) ``` - inflate (float): Inflation factor $\geq 1$ to multiply by the variance estimate to make it more conservative. - alpha (np.ndarray): Uncertainty level in $(0,1)$. + inflate (float): Inflation factor $\geq 1$ to multiply by the + variance estimate to make it more conservative. + alpha (Union[float, np.ndarray]): Uncertainty level in $(0,1)$. """ self.parameters = ["inflate", "alpha", "abs_tol", "rel_tol", "n_init", "n_limit"] # Input Checks @@ -250,7 +256,8 @@ def __init__( # Set Attributes self.n_init = int(n_init) self.n_limit = int(n_limit) - assert isinstance(error_fun, str) or callable(error_fun) + if not (isinstance(error_fun, str) or callable(error_fun)): + raise AssertionError if isinstance(error_fun, str): if error_fun.upper() == "EITHER": error_fun = lambda sv, abs_tol, rel_tol: np.maximum( @@ -265,8 +272,10 @@ def __init__( self.error_fun = error_fun self.alpha = alpha self.inflate = float(inflate) - assert self.inflate >= 1 - assert 0 < self.alpha < 1 + if not (self.inflate >= 1): + raise AssertionError + if not (0 < self.alpha < 1): + raise AssertionError # QMCPy Objs self.integrand = integrand self.true_measure = self.integrand.true_measure @@ -275,12 +284,14 @@ def __init__( allowed_distribs=[AbstractLDDiscreteDistribution], allow_vectorized_integrals=True, ) - assert ( + if not ( self.integrand.discrete_distrib.replications > 1 - ), "Require the discrete distribution has replications>1" - assert ( + ): + raise AssertionError("Require the discrete distribution has replications>1") + if not ( self.integrand.discrete_distrib.randomize != "FALSE" - ), "Require discrete distribution is randomized" + ): + raise AssertionError("Require discrete distribution is randomized") self.alphas_indv, _ = self._compute_indv_alphas( np.full(self.integrand.d_comb, self.alpha) ) @@ -289,7 +300,24 @@ def __init__( self.alphas_indv / 2, df=self.integrand.discrete_distrib.replications - 1 ) - def integrate(self, resume=None): + def integrate(self, resume: Union[None, Data] = None) -> tuple: + """Determine the samples needed to satisfy the target tolerance. + + Doubles the per-replication sample count each iteration and forms a + Student's $t$ confidence interval (`self.t_star`, inflated by + `self.inflate`, with degrees of freedom set by + `self.discrete_distrib.replications`) on each not-yet-converged + output. Stops once every combined output is within tolerance or + `self.n_limit` would be exceeded. + + Args: + resume (Union[None, Data]): Existing integration state to resume from, if + supported. Defaults to None. + + Returns: + tuple: Approximation to the integral with shape ``integrand.d_comb`` + and the corresponding data object. + """ t_start = time() resume_provenance = self._capture_resume_provenance(resume) trace = self._make_trace_logger() @@ -425,8 +453,21 @@ def _restore_resume_state(self, data): self.integrand.discrete_distrib = self.discrete_distrib self.integrand.true_measure.discrete_distrib = self.discrete_distrib - def set_tolerance(self, abs_tol=None, rel_tol=None, rmse_tol=None): - assert rmse_tol is None, "rmse_tol not supported by this stopping criterion." + def set_tolerance(self, abs_tol: Union[None, float] = None, rel_tol: Union[None, float] = None, rmse_tol: Union[None, float] = None) -> None: + """Update the stopping criterion's target tolerance. + + Args: + abs_tol (Union[None, float]): Absolute error tolerance, broadcast to + `self.abs_tols` with shape `integrand.d_comb`. + rel_tol (Union[None, float]): Relative error tolerance, broadcast to + `self.rel_tols` with shape `integrand.d_comb`. + rmse_tol (Union[None, float]): Unsupported; must be `None`. + + Raises: + AssertionError: If `rmse_tol` is supplied. + """ + if not (rmse_tol is None): + raise AssertionError("rmse_tol not supported by this stopping criterion.") if abs_tol is not None: self.abs_tol = abs_tol self.abs_tols = np.full(self.integrand.d_comb, self.abs_tol) diff --git a/qmcpy/stopping_criterion/diagnostics.py b/qmcpy/stopping_criterion/diagnostics.py index c58ec5a1b..ac498ab0c 100644 --- a/qmcpy/stopping_criterion/diagnostics.py +++ b/qmcpy/stopping_criterion/diagnostics.py @@ -1,11 +1,16 @@ """Diagnostics helpers for stopping-criterion iteration tracing.""" +from __future__ import annotations +from typing import TYPE_CHECKING, Union import io import numpy as np import sys from contextlib import redirect_stdout from math import log10 +if TYPE_CHECKING: + from .abstract_stopping_criterion import AbstractStoppingCriterion + # ITER rows up to this count are always printed; above it the log-scale throttle applies. _THROTTLE_ITER_THRESHOLD = 30 @@ -41,7 +46,7 @@ def __init__(self): @property def _column_names(self): - """Return the ordered column names stored in the table. """ + """Return the ordered column names stored in the table.""" return _ITERATION_HISTORY_COLUMNS @property @@ -86,15 +91,15 @@ def _append(self, stage, row, visible_columns=None, printed=True): return len(self) - 1 def _mark_printed(self, index, printed=True): - """Update whether a stored row is marked as printed. """ + """Update whether a stored row is marked as printed.""" self._columns["printed"][index] = bool(printed) def _row(self, index): - """Return one row as a dictionary. """ + """Return one row as a dictionary.""" return {column: self._columns[column][index] for column in self._column_names} def _rows(self): - """Return all rows as dictionaries. """ + """Return all rows as dictionaries.""" return [self._row(index) for index in range(len(self))] def _to_dict(self): @@ -450,15 +455,16 @@ def _format_iteration_log( class _IterationTraceLogger(object): - def __init__(self, stopping_criterion): + def __init__(self, stopping_criterion: AbstractStoppingCriterion): """Create a trace logger bound to the given stopping criterion. Args: - stopping_criterion: Stopping criterion instance. The logger reads + stopping_criterion (AbstractStoppingCriterion): Stopping + criterion instance. The logger reads the optional attributes ``trace_iterations`` (bool), ``trace_label`` (str), ``verbose`` (bool), ``trace_print`` - (bool), and the internal ``_trace_store_*`` flags to - configure storage and live printing. + (bool), and the internal ``_trace_store_*`` flags to configure + storage and live printing. """ requested_trace_iterations = bool( getattr(stopping_criterion, "trace_iterations", False) @@ -511,7 +517,9 @@ def __init__(self, stopping_criterion): ) def _would_be_throttled(self, iter_count): - """Return True if an ITER row with this count would be suppressed by throttling.""" + """Return True if an ITER row with this count would be suppressed by + throttling. + """ if self.verbose: return False if iter_count is None or iter_count <= _THROTTLE_ITER_THRESHOLD: @@ -520,8 +528,9 @@ def _would_be_throttled(self, iter_count): return iter_count % step != 0 @staticmethod - def _state_signature(data): - """Return a hashable snapshot of the data fields used to detect duplicate rows. + def _state_signature(data: object): + """Return a hashable snapshot of the data fields used to detect + duplicate rows. Args: data (object): Integration state object. @@ -543,8 +552,9 @@ def _print_header_once(self): print(f"=== {self.label} iteration log ===") self.header_printed = True - def _get_visible_columns(self, data, row=None): - """Return the ordered list of column names to display, inferred from data. + def _get_visible_columns(self, data: object, row: Union[None, dict] = None): + """Return the ordered list of column names to display, inferred from + data. The result is cached after the first call so all rows share the same columns. @@ -552,12 +562,13 @@ def _get_visible_columns(self, data, row=None): Args: data (object): Integration state object used to determine which optional columns are present. + row (Union[None, dict]): Pre-extracted diagnostic row to infer columns from, + if already available. Computed from `data` when `None`. Returns: - tuple[str, ...]: Column names from the set ``{'stage', 'iter', - 'solution', 'bound_diff', 'comb_bound_diff', - 'bound_half_width', 'bias_estimate', 'n_min', 'n_total', - 'm', 'xfull.shape'}``. + tuple[str, ...]: Column names from the set ``{'stage', 'iter', 'solution', + 'bound_diff', 'comb_bound_diff', 'bound_half_width', + 'bias_estimate', 'n_min', 'n_total', 'm', 'xfull.shape'}``. """ if self.visible_columns is not None: return self.visible_columns @@ -565,19 +576,22 @@ def _get_visible_columns(self, data, row=None): self.visible_columns = _visible_columns_from_row(row) return self.visible_columns - def emit(self, stage, data, step_value=None, increment=False, iter_value=None): + def emit(self, stage: str, data: object, step_value: Union[None, int] = None, increment: bool = False, iter_value: Union[None, int] = None): """Print one diagnostic row for the given stage label. Args: stage (str): Row label, e.g. ``"ITER"`` or ``"RESUME"``. data (object): Integration state object. - step_value (int | None, optional): Value to assign to ``data.m`` - before printing. Defaults to None. - increment (bool, optional): If True, advance the internal iteration - counter and assign the new value to ``data._iter_count``. - Defaults to False. - iter_value (int | None, optional): Explicit iteration count to - display (overrides ``increment``). Defaults to None. + step_value (Union[None, int]): Value to assign to ``data.m`` before + printing. Defaults to None. + increment (bool): If True, advance the internal iteration counter + and assign the new value to ``data._iter_count``. Defaults to + False. + iter_value (Union[None, int]): Explicit iteration count to display + (overrides ``increment``). Defaults to None. + + Returns: + None """ if not self.enabled: return @@ -622,8 +636,9 @@ def emit(self, stage, data, step_value=None, increment=False, iter_value=None): ) self.table_header_printed = True - def resume(self, data, step_value=None): - """Emit a RESUME row and snapshot the current state for duplicate suppression. + def resume(self, data: object, step_value: Union[None, int] = None): + """Emit a RESUME row and snapshot the current state for duplicate + suppression. Reads ``data._iter_count`` to restore the iteration counter so that the next :meth:`iteration` call continues counting from the right number. @@ -632,8 +647,8 @@ def resume(self, data, step_value=None): Args: data (object): Integration state object from the resume checkpoint. - step_value (int | None, optional): Value to assign to ``data.m`` - before printing. Defaults to None. + step_value (Union[None, int]): Value to assign to ``data.m`` before + printing. Defaults to None. """ self._seed_history_from_resume(data) previous_iter_count = getattr(data, "_iter_count", None) @@ -666,7 +681,7 @@ def _seed_history_from_resume(self, data): self.stopping_criterion.iteration_history = self.history self._resume_seeded = True - def iteration(self, data, step_value=None): + def iteration(self, data: object, step_value: Union[None, int] = None): """Emit an ITER row, unless state is unchanged since the last resume. If :meth:`resume` was just called and the data state has not changed @@ -675,8 +690,11 @@ def iteration(self, data, step_value=None): Args: data (object): Current integration state object. - step_value (int | None, optional): Value to assign to ``data.m`` - before printing. Defaults to None. + step_value (Union[None, int]): Value to assign to ``data.m`` before + printing. Defaults to None. + + Returns: + None """ current_signature = self._state_signature(data) if ( @@ -720,7 +738,9 @@ def _flush_last_if_suppressed(self): self._last_printed_iter_count = self._last_iter_count def finalize(self): - """Force-print the last ITER row if throttling suppressed it, then clear snapshot.""" + """Force-print the last ITER row if throttling suppressed it, then + clear snapshot. + """ self._flush_last_if_suppressed() # DataFrame is built lazily in get_iteration_log() to avoid pandas # construction overhead on every integrate() call when tracing is off. @@ -736,11 +756,11 @@ def finalize(self): def _print_diagnostic( - label, - data, - table_header=False, - verbose=True, - visible_columns=None, + label: str, + data: object, + table_header: bool = False, + verbose: bool = True, + visible_columns: Union[None, tuple, list] = None, ): """Print diagnostic information for an integration state. @@ -748,13 +768,15 @@ def _print_diagnostic( label (str): Stage label shown in the first column. data (object): Integration state carrying fields such as ``solution``, ``n_total``, ``n_min``, ``m``, and ``xfull``. - table_header (bool, optional): Whether to print the compact table - header before the row. Defaults to False. - verbose (bool, optional): Whether to print every ``ITER`` row. - Defaults to True. If False, the current iteration-log throttling - rules are applied. - visible_columns (tuple[str, ...] | list[str] | None, optional): Ordered - columns to print. Defaults to all supported columns. + table_header (bool): Whether to print the compact table header before + the row. Defaults to False. + verbose (bool): Whether to print every ``ITER`` row. Defaults to True. + If False, the current iteration-log throttling rules are applied. + visible_columns (Union[None, tuple, list]): Ordered columns + to print. Defaults to all supported columns. + + Returns: + None """ row = _extract_diagnostic_row(data) iter_display = row["iter"] @@ -776,4 +798,4 @@ def _print_diagnostic( header_line = " ".join(header_values[column] for column in visible_columns) print(header_line) print("-" * len(header_line)) - print(" ".join(aligned_row_values[column] for column in visible_columns)) \ No newline at end of file + print(" ".join(aligned_row_values[column] for column in visible_columns)) diff --git a/qmcpy/stopping_criterion/pf_gp_ci.py b/qmcpy/stopping_criterion/pf_gp_ci.py index 5da66259f..76e2e83de 100644 --- a/qmcpy/stopping_criterion/pf_gp_ci.py +++ b/qmcpy/stopping_criterion/pf_gp_ci.py @@ -1,3 +1,7 @@ +from __future__ import annotations + +from ..integrand.abstract_integrand import AbstractIntegrand +from typing import TYPE_CHECKING, Union, Callable from .abstract_stopping_criterion import AbstractStoppingCriterion from ..discrete_distribution import DigitalNetB2 from ..integrand.ishigami import Ishigami @@ -17,17 +21,50 @@ import torch import gpytorch +if TYPE_CHECKING: + import matplotlib.figure + class Suggester(object): + """Base class for future-sample suggestion schemes used by `PFGPCI`. + + Subclasses implement `suggest` to propose the next batch of sample + locations, typically concentrated near the estimated failure boundary. + """ + pass class PFSampleErrorDensityAR(Suggester): - def __init__(self, verbose=False): + """Suggest new samples via acceptance-rejection on the GP error density. + + Draws uniform candidates and accepts them with probability proportional + to the current GP's misclassification-error density, concentrating new + samples near the estimated failure boundary. + """ + + def __init__(self, verbose=False) -> None: self.verbose = verbose super(PFSampleErrorDensityAR, self).__init__() - def suggest(self, n, d, gp, rng, efficiency, pct=0.5): + def suggest(self, n: int, d: int, gp: ExactGPyTorchRegressionModel, rng: np.random.Generator, efficiency: float, pct: float = 0.5) -> np.ndarray: + """Draw `n` new sample locations via acceptance-rejection. + + Args: + n (int): Number of samples to return. + d (int): Dimension of the sampling domain. + gp (ExactGPyTorchRegressionModel): Current GP surrogate, used to + evaluate the error density at candidate points. + rng (np.random.Generator): Random number generator for + candidate draws. + efficiency (float): Estimated acceptance rate, used to size each + batch of candidate draws. + pct (float): Target probability of accepting at least `n` points + within one candidate batch. + + Returns: + np.ndarray: `n` accepted sample locations, shape `(n, d)`. + """ if self.verbose: print( "\tAR sampling with efficiency %.1e, expect %d draws: " @@ -55,16 +92,42 @@ def suggest(self, n, d, gp, rng, efficiency, pct=0.5): class SuggesterSimple(Suggester): - def __init__(self, sampler): + """Suggest new samples by drawing the next block from a fixed sampler. + + Wraps an `AbstractTrueMeasure`/`AbstractDiscreteDistribution` (or any + callable with the same interface) and advances through it sequentially, + ignoring the current GP state. + """ + + def __init__(self, sampler) -> None: self.sampler = sampler if isinstance(self.sampler, AbstractTrueMeasure): - assert (self.sampler.range == [0, 1]).all() + if not ((self.sampler.range == [0, 1]).all()): + raise AssertionError self.n_min = 0 super(SuggesterSimple, self).__init__() - def suggest(self, n, d, gp, rng, **kwargs): + def suggest(self, n: int, d: int, gp: ExactGPyTorchRegressionModel, rng: np.random.Generator, **kwargs) -> np.ndarray: + """Draw the next `n` sample locations from `self.sampler`. + + Args: + n (int): Number of samples to return. + d (int): Dimension of the sampling domain; must match + `self.sampler.d`. + gp (ExactGPyTorchRegressionModel): Unused; accepted for + interface compatibility with other `Suggester` + implementations. + rng (np.random.Generator): Unused; accepted for interface + compatibility with other `Suggester` implementations. + **kwargs: Unused; accepted for interface compatibility with + other `Suggester` implementations. + + Returns: + np.ndarray: `n` sample locations, shape `(n, d)`. + """ n_max = self.n_min + n - assert d == self.sampler.d + if not (d == self.sampler.d): + raise AssertionError try: x = self.sampler(n_min=self.n_min, n_max=n_max) except TypeError: @@ -74,8 +137,8 @@ def suggest(self, n, d, gp, rng, **kwargs): class PFGPCI(AbstractStoppingCriterion): - """ - Probability of failure estimation using adaptive Gaussian process construction and resulting credible intervals. + """Probability of failure estimation using adaptive Gaussian process + construction and resulting credible intervals. Examples: >>> pfgpci = PFGPCI( @@ -163,60 +226,85 @@ class PFGPCI(AbstractStoppingCriterion): def __init__( self, - integrand, - failure_threshold, - failure_above_threshold, - abs_tol=5e-3, - n_init=64, - n_limit=1000, - alpha=1e-2, - init_samples=None, - batch_sampler=PFSampleErrorDensityAR(), - n_batch=4, - n_approx=2**20, - gpytorch_prior_mean=gpytorch.means.ZeroMean(), - gpytorch_prior_cov=gpytorch.kernels.ScaleKernel( + integrand: AbstractIntegrand, + failure_threshold: float, + failure_above_threshold: bool, + abs_tol: float = 5e-3, + n_init: float = 64, + n_limit: int = 1000, + alpha: float = 1e-2, + init_samples: Union[None, float] = None, + batch_sampler: Union[Suggester, AbstractDiscreteDistribution] = PFSampleErrorDensityAR(), + n_batch: int = 4, + n_approx: int = 2**20, + gpytorch_prior_mean: gpytorch.means = gpytorch.means.ZeroMean(), + gpytorch_prior_cov: gpytorch.kernels = gpytorch.kernels.ScaleKernel( gpytorch.kernels.MaternKernel(nu=2.5) ), - gpytorch_likelihood=gpytorch.likelihoods.GaussianLikelihood( + gpytorch_likelihood: gpytorch.likelihoods = gpytorch.likelihoods.GaussianLikelihood( noise_constraint=gpytorch.constraints.Interval(1e-12, 1e-8) ), - gpytorch_marginal_log_likelihood_func=lambda likelihood, gpyt_model: gpytorch.mlls.ExactMarginalLogLikelihood( + gpytorch_marginal_log_likelihood_func: Callable = lambda likelihood, gpyt_model: gpytorch.mlls.ExactMarginalLogLikelihood( likelihood, gpyt_model ), - torch_optimizer_func=lambda gpyt_model: torch.optim.Adam( + torch_optimizer_func: Callable = lambda gpyt_model: torch.optim.Adam( gpyt_model.parameters(), lr=0.1 ), - gpytorch_train_iter=100, - gpytorch_use_gpu=False, - verbose=False, - n_ref_approx=2**22, - seed_ref_approx=None, - ): - """ + gpytorch_train_iter: int = 100, + gpytorch_use_gpu: bool = False, + verbose: Union[bool, int] = False, + n_ref_approx: int = 2**22, + seed_ref_approx: Union[None, int] = None, + ) -> None: + """Initialize a PFGPCI stopping criterion. + Args: integrand (AbstractIntegrand): The integrand. failure_threshold (float): Thresholds for failure. - failure_above_threshold (bool): Set to `True` if failure occurs when the simulation exceeds `failure_threshold` and False otherwise. - abs_tol (float): The desired maximum distance from the estimate to either end of the credible interval. - n_init (float): Initial number of samples from integrand.discrete_distrib from which to build the first surrogate GP + failure_above_threshold (bool): Set to `True` if failure occurs + when the simulation exceeds `failure_threshold` and False + otherwise. + abs_tol (float): The desired maximum distance from the estimate to + either end of the credible interval. + n_init (float): Initial number of samples from + integrand.discrete_distrib from which to build the first + surrogate GP n_limit (int): Budget of simulations. - n_batch (int): The number of samples per batch to draw from batch_sampler. - alpha (float): The credible interval is constructed to hold with probability at least 1 - alpha - init_samples (float): If the simulation has already been run, pass in (x,y) where x are past samples from the discrete distribution and y are corresponding simulation evaluations. - batch_sampler (Suggester or AbstractDiscreteDistribution): A suggestion scheme for future samples. - n_approx (int): Number of points from integrand.discrete_distrib used to approximate estimate and credible interval bounds + alpha (float): The credible interval is constructed to hold with + probability at least 1 - alpha + init_samples (Union[None, float]): If the simulation has already been run, pass + in (x,y) where x are past samples from the discrete + distribution and y are corresponding simulation evaluations. + batch_sampler (Union[Suggester, AbstractDiscreteDistribution]): + A suggestion scheme for future samples. + n_batch (int): The number of samples per batch to draw from + batch_sampler. + n_approx (int): Number of points from integrand.discrete_distrib + used to approximate estimate and credible interval bounds gpytorch_prior_mean (gpytorch.means): prior mean function of the GP - gpytorch_prior_cov (gpytorch.kernels): Prior covariance kernel of the GP - gpytorch_likelihood (gpytorch.likelihoods): GP likelihood, require one of gpytorch.likelihoods.{GaussianLikelihood, GaussianLikelihoodWithMissingObs, FixedNoiseGaussianLikelihood} - gpytorch_marginal_log_likelihood_func (callable): Function taking in the likelihood and gpytorch model and returning a marginal log likelihood from gpytorch.mlls - torch_optimizer_func (callable): Function taking in the gpytorch model and returning an optimizer from torch.optim - gpytorch_train_iter (int): Training iterations for the GP in gpytorch - gpytorch_use_gpu (bool): If True, have gpytorch use a GPU for fitting and trining the GP - verbose (int): If verbose > 0, print information through the call to integrate() - n_ref_approx (int): If n_ref_approx > 0, use n_ref_approx points to get a reference QMC approximation of the true solution. - Caution: If n_ref_approx > 0, it should be a large int e.g. 2**22, in which case it is only helpful for cheap to evaluate simulations - seed_ref_approx (int): Seed for the reference approximation. Only applies when n_ref_approx>0 + gpytorch_prior_cov (gpytorch.kernels): Prior covariance kernel of + the GP + gpytorch_likelihood (gpytorch.likelihoods): GP likelihood, require + one of gpytorch.likelihoods.{GaussianLikelihood, + GaussianLikelihoodWithMissingObs, FixedNoiseGaussianLikelihood} + gpytorch_marginal_log_likelihood_func (Callable): Function taking + in the likelihood and gpytorch model and returning a marginal + log likelihood from gpytorch.mlls + torch_optimizer_func (Callable): Function taking in the gpytorch + model and returning an optimizer from torch.optim + gpytorch_train_iter (int): Training iterations for the GP in + gpytorch + gpytorch_use_gpu (bool): If True, have gpytorch use a GPU for + fitting and training the GP + verbose (Union[bool, int]): If verbose > 0, print information through the call + to integrate() + n_ref_approx (int): If n_ref_approx > 0, use n_ref_approx points to + get a reference QMC approximation of the true solution. + Caution: If n_ref_approx > 0, it should be a large int e.g. + 2**22, in which case it is only helpful for cheap to evaluate + simulations + seed_ref_approx (Union[None, int]): Seed for the reference approximation. Only + applies when n_ref_approx>0 """ self.parameters = ["abs_tol", "n_init", "n_limit", "n_batch"] self.integrand = integrand @@ -228,23 +316,29 @@ def __init__( self.failure_above_threshold = failure_above_threshold self.abs_tol = abs_tol self.alpha = alpha - assert 0 < self.alpha < 1 + if not (0 < self.alpha < 1): + raise AssertionError self.n_init = n_init self.init_samples = init_samples is not None if self.init_samples: self.x_init, self.y_init = init_samples - assert self.x_init.ndim == 2 and self.y_init.ndim == 1 - assert self.x_init.shape[1] == self.d and len(self.y_init) == len( + if not (self.x_init.ndim == 2 and self.y_init.ndim == 1): + raise AssertionError + if not (self.x_init.shape[1] == self.d and len(self.y_init) == len( self.x_init - ) - assert self.n_init == len(self.x_init) + )): + raise AssertionError + if not (self.n_init == len(self.x_init)): + raise AssertionError self.ytf_init = self._affine_tf(self.y_init) self.batch_sampler = batch_sampler self.n_batch = n_batch self.n_limit = n_limit - assert self.n_limit >= self.n_init + if not (self.n_limit >= self.n_init): + raise AssertionError self.n_approx = n_approx - assert (self.n_approx + self.n_init) <= 2**32 + if not ((self.n_approx + self.n_init) <= 2**32): + raise AssertionError self.gpytorch_prior_mean = gpytorch_prior_mean self.gpytorch_prior_cov = gpytorch_prior_cov self.gpytorch_likelihood = gpytorch_likelihood @@ -277,7 +371,32 @@ def _affine_tf(self, y): else self.failure_threshold - y ) - def integrate(self, seed=None, refit=False, resume=None): + def integrate(self, seed: Union[None, int] = None, refit: bool = False, resume: Union[None, Data] = None) -> tuple: + """Determine the samples needed to satisfy the target tolerance. + + Draws an initial batch (`self.n_init` points, or `init_samples` if + supplied), fits a GP surrogate, then repeatedly draws + `self.n_batch` more points via `self.batch_sampler`, updates the GP, + and refines the credible-interval bound on the probability of + failure. Stops once the bound is within `self.abs_tol` or + `self.n_limit` would be exceeded. + + Args: + seed (Union[None, int]): Seed for the internal `DigitalNetB2` sampler used to + approximate the solution and (if `init_samples` was not + supplied) draw the initial batch. + refit (bool): If `True`, refit the GP hyperparameters from + scratch every batch rather than only on the first batch. + resume (Union[None, Data]): Unsupported; must be `None`, as `PFGPCI` cannot + resume a prior checkpoint. + + Returns: + tuple: Approximation to the probability of failure + and the corresponding data object. + + Raises: + ParameterError: If `resume` is not `None`. + """ t0 = time.time() trace = self._make_trace_logger() if resume is not None: @@ -395,7 +514,7 @@ def __init__( gpytorch_use_gpu, verbose, approx_true_solution, - ): + ) -> None: self.stopping_crit = stopping_crit self.integrand = integrand self.true_measure = true_measure @@ -437,7 +556,20 @@ def __init__( parameters=["solution", "error_bound", "bound_low", "bound_high", "n_total", "time_integrate"] ) - def update_data(self, batch_count, xdraw, ydrawtf): + def update_data(self, batch_count: int, xdraw: np.ndarray, ydrawtf: np.ndarray): + """Fold one new batch of samples into the GP surrogate and credible interval. + + Refits the GP from scratch (on the first batch, or every batch if + `self.refit`), otherwise incrementally adds the new data to the + existing GP. Recomputes the probability-of-failure estimate and its + credible interval from the updated surrogate. + + Args: + batch_count (int): Index of this batch (0 for the initial batch). + xdraw (np.ndarray): New sample locations, shape `(n_new, d)`. + ydrawtf (np.ndarray): Affine-transformed integrand values at + `xdraw` (positive indicates failure), shape `(n_new,)`. + """ self.n_batch.append(len(xdraw)) self.x, self.y = np.vstack([self.x, xdraw]), np.hstack([self.y, ydrawtf]) if batch_count == 0 or self.refit: @@ -496,7 +628,15 @@ def update_data(self, batch_count, xdraw, ydrawtf): ) ) - def get_results_dict(self): + def get_results_dict(self) -> dict: + """Collect the per-iteration history as arrays. + + Returns: + dict: Per-iteration `"iter"`, `"n_sum"` (cumulative sample + count), `"n_batch"`, `"error_bounds"`, `"ci_low"`, `"ci_high"`, + and `"solutions"` arrays; plus `"solutions_ref"`, `"error_ref"`, + and `"in_ci"` if `self.approx_true_solution`. + """ df = { "iter": np.arange(len(self.n_sum)), "n_sum": self.n_sum, @@ -514,7 +654,18 @@ def get_results_dict(self): ) return df - def plot(self, trace_only=False, **kwargs): + def plot(self, trace_only: bool = False, **kwargs) -> matplotlib.figure.Figure: + """Plot the convergence trace, plus a per-batch GP diagnostic panel if `d` is 1 or 2. + + Args: + trace_only (bool): If `True` (or if `d` is not 1 or 2, or no GP + has been fit yet), plot only the convergence trace. + **kwargs: Passed through to `plot_1d`/`plot_2d` when a + per-batch diagnostic panel is drawn. + + Returns: + matplotlib.figure.Figure: The assembled figure. + """ from matplotlib import pyplot if self.d == 1 and not trace_only and self.saved_gps != []: @@ -572,7 +723,20 @@ def plot(self, trace_only=False, **kwargs): ) return fig - def plot_1d(self, meshticks=1025, ci_percentage=0.95, **kwargs): + def plot_1d(self, meshticks: int = 1025, ci_percentage: float = 0.95, **kwargs) -> matplotlib.figure.Figure: + """Plot, for each batch, the 1-D error density and GP fit with a credible band. + + Args: + meshticks (int): Number of points in the `[0,1]` plotting mesh. + ci_percentage (float): Credible level for the plotted GP + prediction band. + **kwargs: Unused; accepted for interface compatibility with + `plot`. + + Returns: + matplotlib.figure.Figure: The assembled figure. + matplotlib.gridspec.GridSpec: The figure's grid layout. + """ from matplotlib import pyplot, gridspec beta = norm.ppf(np.mean([ci_percentage, 1])) @@ -637,7 +801,20 @@ def plot_1d(self, meshticks=1025, ci_percentage=0.95, **kwargs): ax.xaxis.set_visible(False) return fig, gs - def plot_2d(self, meshticks=257, clevels=32, **kwargs): + def plot_2d(self, meshticks: int = 257, clevels: int = 32, **kwargs) -> matplotlib.figure.Figure: + """Plot, for each batch, 2-D contours of the true function, error density, and GP mean. + + Args: + meshticks (int): Number of points per axis in the `[0,1]^2` + plotting mesh. + clevels (int): Number of contour levels. + **kwargs: Unused; accepted for interface compatibility with + `plot`. + + Returns: + matplotlib.figure.Figure: The assembled figure. + matplotlib.gridspec.GridSpec: The figure's grid layout. + """ from matplotlib import pyplot, gridspec, colormaps n_batches = len(self.n_batch) diff --git a/qmcpy/true_measure/abstract_true_measure.py b/qmcpy/true_measure/abstract_true_measure.py index b7611b1f3..68e934647 100644 --- a/qmcpy/true_measure/abstract_true_measure.py +++ b/qmcpy/true_measure/abstract_true_measure.py @@ -1,3 +1,4 @@ +from typing import Tuple, Union from ..util import MethodImplementationError, _univ_repr, ParameterError from ..discrete_distribution.abstract_discrete_distribution import ( AbstractDiscreteDistribution, @@ -7,8 +8,17 @@ class AbstractTrueMeasure(object): + """Abstract base class for QMCPy true measures. - def __init__(self): + A true measure composes a transform (`self.transform`) on top of a + sampler (an `AbstractDiscreteDistribution`, or another + `AbstractTrueMeasure` for recursive composition), mapping unit-cube + samples to samples from the target measure. Concrete measures (e.g. + `Gaussian`, `Uniform`) set `self.domain`, `self.range`, and implement the + transform/weight/moment logic this base class exposes. + """ + + def __init__(self) -> None: prefix = "A concrete implementation of TrueMeasure must have " if not hasattr(self, "domain"): raise ParameterError( @@ -42,14 +52,18 @@ def _set_moments(self, mean, variance, standard_deviation, covariance): @staticmethod def _read_only_view(value): - """Return a view which cannot be made writeable while its base is read only.""" + """Return a view which cannot be made writeable while its base is + read only. + """ view = value.view() view.setflags(write=False) return view def _scalar_if_univariate(self, value): - """For univariate (``d == 1``) measures, return a Python ``float`` scalar - (via :func:`numpy.squeeze`); otherwise return a read only array view.""" + """For univariate (``d == 1``) measures, return a Python ``float`` + scalar (via :func:`numpy.squeeze`); otherwise return a read only array + view. + """ if getattr(self, "d", None) == 1: return float(np.squeeze(value)) return self._read_only_view(value) @@ -60,18 +74,34 @@ def _scalar_if_univariate(self, value): @property def mean(self): + """Union[float, np.ndarray]: The measure's mean, set via `_set_moments`. + A Python `float` for univariate (`d == 1`) measures, otherwise a + read-only array. + """ return self._scalar_if_univariate(self._mean) @property def variance(self): + """Union[float, np.ndarray]: The measure's variance, set via + `_set_moments`. A Python `float` for univariate (`d == 1`) measures, + otherwise a read-only array. + """ return self._scalar_if_univariate(self._variance) @property def standard_deviation(self): + """Union[float, np.ndarray]: The measure's standard deviation, set + via `_set_moments`. A Python `float` for univariate (`d == 1`) + measures, otherwise a read-only array. + """ return self._scalar_if_univariate(self._standard_deviation) @property def covariance(self): + """Union[np.ndarray, scipy.sparse.spmatrix]: The measure's + covariance, set via `_set_moments`. Read-only; sparse covariances + are returned as-is, dense ones as a read-only view. + """ covariance = self._covariance if sparse.issparse(covariance): return covariance @@ -111,7 +141,7 @@ def _parse_sampler(self, sampler): "sampler input should either be a AbstractDiscreteDistribution or AbstractTrueMeasure" ) - def __call__(self, n=None, n_min=None, n_max=None, return_weights=False, warn=True): + def __call__(self, n: Union[None, int] = None, n_min: Union[None, int] = None, n_max: Union[None, int] = None, return_weights: bool = False, warn: bool = True) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: r""" - If just `n` is supplied, generate samples from the sequence at indices 0,...,`n`-1. - If `n_min` and `n_max` are supplied, generate samples from the sequence at indices `n_min`,...,`n_max`-1. @@ -125,11 +155,13 @@ def __call__(self, n=None, n_min=None, n_max=None, return_weights=False, warn=Tr warn (bool): If `False`, disable warnings when generating samples. Returns: - t (np.ndarray): Samples from the sequence. + Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: Samples from the + sequence when `return_weights=False`, otherwise the pair + `(samples, jacobian_weights)`. - - If `replications` is `None` then this will be of size (`n_max`-`n_min`) $\times$ `dimension` - - If `replications` is a positive int, then `t` will be of size `replications` $\times$ (`n_max`-`n_min`) $\times$ `dimension` - weights (np.ndarray): Only returned when `return_weights=True`. The Jacobian weights for the transformation + - If `replications` is `None` the samples are of size (`n_max`-`n_min`) $\times$ `dimension` + - If `replications` is a positive int, they are of size `replications` $\times$ (`n_max`-`n_min`) $\times$ `dimension` + - The jacobian weights, when returned, drop the trailing `dimension` axis """ return self.gen_samples( n=n, n_min=n_min, n_max=n_max, return_weights=return_weights, warn=warn @@ -138,8 +170,12 @@ def __call__(self, n=None, n_min=None, n_max=None, return_weights=False, warn=Tr def gen_samples( self, n=None, n_min=None, n_max=None, return_weights=False, warn=True ): + r"""Generate samples from the measure. Called by `__call__`; see its + docstring for the full `Args:`/`Returns:` description. + """ x = self.discrete_distrib(n=n, n_min=n_min, n_max=n_max, warn=warn) - assert isinstance(return_weights, bool) + if not (isinstance(return_weights, bool)): + raise AssertionError return self._jacobian_transform_r(x=x, return_weights=return_weights) def _jacobian_transform_r(self, x, return_weights): @@ -173,17 +209,17 @@ def _jacobian_transform_r(self, x, return_weights): return t def _transform(self, x): - r"""Transformation from the standard uniform to the true measure distribution.""" + r"""Transformation from the standard uniform to the true measure + distribution. + """ raise MethodImplementationError( self, "_transform. Try setting sampler to be in a PDF AbstractTrueMeasure to importance sample by.", ) - def _weight(self, x): - r""" - Non-negative weight function. - This is often a PDF, but is not required to be - e.g., Lebesgue weight is always 1, but is not a PDF. + def _weight(self, x: np.ndarray) -> np.ndarray: + r"""Non-negative weight function. This is often a PDF, but is not + required to be e.g., Lebesgue weight is always 1, but is not a PDF. Args: x (np.ndarray): n x d matrix of samples @@ -195,20 +231,22 @@ def _weight(self, x): self, "weight. Try a different true measure with a _weight method." ) - def spawn(self, s=1, dimensions=None): - r""" - Spawn new instances of the current true measure but with new seeds and dimensions. - Used by multi-level QMC algorithms which require different seeds and dimensions on each level. + def spawn(self, s: int = 1, dimensions: Union[None, np.ndarray] = None) -> list: + r"""Spawn new instances of the current true measure but with new seeds + and dimensions. Used by multi-level QMC algorithms which require + different seeds and dimensions on each level. - Note: - Use `replications` instead of using `spawn` when possible, e.g., when spawning copies which all have the same dimension. + Notes: + Use `replications` instead of using `spawn` when possible, e.g., + when spawning copies which all have the same dimension. Args: s (int): Number of copies to spawn - dimensions (np.ndarray): Length `s` array of dimensions for each copy. Defaults to the current dimension. + dimensions (Union[None, np.ndarray]): Length `s` array of dimensions for each + copy. Defaults to the current dimension. Returns: - spawned_true_measures (list): True measure with new seeds and dimensions. + list: True measure with new seeds and dimensions. """ sampler = self.discrete_distrib if self.transform == self else self.transform sampler_spawns = sampler.spawn(s=s, dimensions=dimensions) diff --git a/qmcpy/true_measure/acceptance_rejection.py b/qmcpy/true_measure/acceptance_rejection.py index 13eb05ece..98a837028 100644 --- a/qmcpy/true_measure/acceptance_rejection.py +++ b/qmcpy/true_measure/acceptance_rejection.py @@ -1,4 +1,8 @@ +from typing import Callable, List, Union from .abstract_true_measure import AbstractTrueMeasure +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) from ..util import MethodImplementationError, ParameterError import numpy as np import warnings @@ -13,34 +17,31 @@ def _next_pow2(n): class AcceptanceRejection(AbstractTrueMeasure): - """ - Deterministic Acceptance-Rejection (DAR) sampler on the unit cube. + """Deterministic Acceptance-Rejection (DAR) sampler on the unit cube. - Implements Algorithm 2 from Zhu & Dick (2014). A (t,m,s)-net in - dimension s = d+1 is used as the driver, where the first d coordinates - form the candidate point and the last coordinate is the acceptance - threshold. This gives a star discrepancy bound of O(N^{-1/s}) on the - accepted samples, compared to O(N^{-1/2}) for standard random - acceptance-rejection. + Implements Algorithm 2 from Zhu & Dick (2014). A (t,m,s)-net in dimension s + = d+1 is used as the driver, where the first d coordinates form the + candidate point and the last coordinate is the acceptance threshold. This + gives a star discrepancy bound of O(N^{-1/s}) on the accepted samples, + compared to O(N^{-1/2}) for standard random acceptance-rejection. - The sampler dimension must be d+1 where d is the target dimension. - The number of driver points is always a power of 2 (required for the + The sampler dimension must be d+1 where d is the target dimension. The + number of driver points is always a power of 2 (required for the (t,m,s)-net property of Theorem 1). Args: - sampler (AbstractDiscreteDistribution): A QMCPy discrete - distribution of dimension s = target_dim + 1. Must mimic - StdUniform. The last coordinate is used as the acceptance - threshold. - target_density (callable): Unnormalised target density psi(x) - where x has shape (N, d). Must return shape (N,) and be - non-negative on [0,1]^d. - upper_bound (float): L = sup_{x in [0,1]^d} psi(x). Every - evaluation of psi must be <= L. - density_integral (float): C = integral_{[0,1]^d} psi(x) dx. - The acceptance rate is C/L. - max_retries (int): Number of times gen_samples will double the - driver size if not enough points are accepted. Default 4. + sampler (AbstractDiscreteDistribution): A QMCPy discrete distribution + of dimension s = target_dim + 1. Must mimic StdUniform. The last + coordinate is used as the acceptance threshold. + target_density (Callable): Unnormalised target density psi(x) where x + has shape (N, d). Must return shape (N,) and be non-negative on + [0,1]^d. + upper_bound (float): L = sup_{x in [0,1]^d} psi(x). Every evaluation of + psi must be <= L. + density_integral (float): C = integral_{[0,1]^d} psi(x) dx. The + acceptance rate is C/L. + max_retries (int): Number of times gen_samples will double the driver + size if not enough points are accepted. Default 4. Examples: >>> import numpy as np @@ -78,7 +79,7 @@ class AcceptanceRejection(AbstractTrueMeasure): qmcpy.util.exceptions_warnings.ParameterError: n_min > 0 but no prior call was made. Call gen_samples with n_min=0 first. """ - def __init__(self, sampler, target_density, upper_bound, density_integral, max_retries=4): + def __init__(self, sampler: AbstractDiscreteDistribution, target_density: Callable, upper_bound: float, density_integral: float, max_retries: int = 4) -> None: self.parameters = ['target_dim', 'upper_bound', 'density_integral', 'acceptance_rate'] self.domain = np.array([[0, 1]]) self._parse_sampler(sampler) @@ -107,36 +108,36 @@ def __init__(self, sampler, target_density, upper_bound, density_integral, max_r self._driver_offset = None super(AcceptanceRejection, self).__init__() - def gen_samples(self, n=None, n_min=None, n_max=None, return_weights=False, warn=True): - """ - Generate accepted samples from the target density. + def gen_samples(self, n: Union[None, int] = None, n_min: Union[None, int] = None, n_max: Union[None, int] = None, return_weights: bool = False, warn: bool = True) -> np.ndarray: + """Generate accepted samples from the target density. - Unlike other TrueMeasures, this method cannot be decomposed into - a fixed 1-to-1 _transform because acceptance-rejection produces - a variable number of outputs from a fixed driver batch. gen_samples - is therefore overridden directly. + Unlike other TrueMeasures, this method cannot be decomposed into a + fixed 1-to-1 _transform because acceptance-rejection produces a + variable number of outputs from a fixed driver batch. gen_samples is + therefore overridden directly. - Supports continued sampling: calling with n_min=0 starts fresh, - and subsequent calls with n_min>0 continue from the same driver - sequence position. + Supports continued sampling: calling with n_min=0 starts fresh, and + subsequent calls with n_min>0 continue from the same driver sequence + position. Args: - n (int): Number of accepted samples to return. Treated as - n_min=0, n_max=n (always resets the driver sequence). - n_min (int): Starting accepted-sample index. Use 0 to reset - and start fresh. Use a positive value to continue from - the previous call. - n_max (int): Ending accepted-sample index (exclusive). - Number of samples returned is n_max - n_min. + n (Union[None, int]): Number of accepted samples to return. Treated as n_min=0, + n_max=n (always resets the driver sequence). + n_min (Union[None, int]): Starting accepted-sample index. Use 0 to reset and + start fresh. Use a positive value to continue from the previous + call. + n_max (Union[None, int]): Ending accepted-sample index (exclusive). Number of + samples returned is n_max - n_min. return_weights (bool): If True, also return importance weights psi(x)/C for each accepted sample. - warn (bool): If True, warn when fewer than n samples are - returned after all retries. + warn (bool): If True, warn when fewer than n samples are returned + after all retries. Returns: - samples (np.ndarray): Shape (n, target_dim). - weights (np.ndarray): Shape (n,). Only returned when - return_weights=True. + np.ndarray: Accepted samples of shape (n, target_dim). When + return_weights=True, a tuple (samples, weights) is returned + instead, where weights has shape (n,) and holds the importance + weights psi(x)/C. """ if n_max is not None: if n_min is None: @@ -210,49 +211,47 @@ def _spawn(self, sampler, dimension): class AcceptanceRejectionReal(AbstractTrueMeasure): - """ - Deterministic Acceptance-Rejection (DAR) sampler on real space R^d. + """Deterministic Acceptance-Rejection (DAR) sampler on real space R^d. - Implements Algorithm 3 from Zhu & Dick (2014). Extends Algorithm 2 - to densities on R^d by mapping the unit-cube driver through marginal - quantile functions (inverse Rosenblatt transform, Lemma 4) before - applying the acceptance test. + Implements Algorithm 3 from Zhu & Dick (2014). Extends Algorithm 2 to + densities on R^d by mapping the unit-cube driver through marginal quantile + functions (inverse Rosenblatt transform, Lemma 4) before applying the + acceptance test. The driver point (u_1, ..., u_d, u_{d+1}) is transformed as: - z_j = F_j^{-1}(u_j) for j = 1, ..., d - u = u_{d+1} threshold coordinate (unchanged) + z_j = F_j^{-1}(u_j) for j = 1, ..., d u = u_{d+1} threshold + coordinate (unchanged) Acceptance condition: psi(z) >= L * H(z) * u - where H is the auxiliary bound function satisfying psi(z) <= L * H(z) - for all z in R^d. This gives the same discrepancy bound O(N^{-1/s}) - as Algorithm 2. + where H is the auxiliary bound function satisfying psi(z) <= L * H(z) for + all z in R^d. This gives the same discrepancy bound O(N^{-1/s}) as + Algorithm 2. - Note: - inv_cdfs applies each quantile function independently per - dimension. This is exact when H factors as a product of - independent marginals (e.g. a product of univariate distributions). + Notes: + inv_cdfs applies each quantile function independently per dimension. + This is exact when H factors as a product of independent marginals + (e.g. a product of univariate distributions). Args: - sampler (AbstractDiscreteDistribution): A QMCPy discrete - distribution of dimension s = target_dim + 1. Must mimic - StdUniform. - target_density (callable): Unnormalised target density psi(z) - where z has shape (N, d). Must return shape (N,). - Must satisfy psi(z) <= L * H(z) for all z. - inv_cdfs (list of callable): List of d quantile functions - [F_1^{-1}, ..., F_d^{-1}], one per dimension. Each maps - a 1-D array of uniforms in [0,1] to R. - Example: [scipy.stats.norm.ppf] for a 1-D standard Gaussian. - H_func (callable): Auxiliary bound function H(z) where z has - shape (N, d). Must return shape (N,) and satisfy - psi(z) <= L * H(z) for all z in R^d. + sampler (AbstractDiscreteDistribution): A QMCPy discrete distribution + of dimension s = target_dim + 1. Must mimic StdUniform. + target_density (Callable): Unnormalised target density psi(z) where z + has shape (N, d). Must return shape (N,). Must satisfy psi(z) <= L + * H(z) for all z. + inv_cdfs (List[Callable]): List of d quantile functions [F_1^{-1}, + ..., F_d^{-1}], one per dimension. Each maps a 1-D array of + uniforms in [0,1] to R. Example: [scipy.stats.norm.ppf] for a 1-D + standard Gaussian. + H_func (Callable): Auxiliary bound function H(z) where z has shape (N, + d). Must return shape (N,) and satisfy psi(z) <= L * H(z) for all z + in R^d. upper_bound (float): L satisfying psi(z) <= L * H(z) for all z. - density_integral (float): C = integral_{R^d} psi(z) dz. - The acceptance rate is C/L. - max_retries (int): Number of times gen_samples will double the - driver size if not enough points are accepted. Default 4. + density_integral (float): C = integral_{R^d} psi(z) dz. The acceptance + rate is C/L. + max_retries (int): Number of times gen_samples will double the driver + size if not enough points are accepted. Default 4. Examples: >>> import numpy as np @@ -276,7 +275,8 @@ class AcceptanceRejectionReal(AbstractTrueMeasure): density_integral 1 acceptance_rate 2^(-1) - Continued sampling: batches resume the driver sequence without restarting. + Continued sampling: batches resume the driver sequence without + restarting. >>> inv_cdfs = [lambda u: norm.ppf(u, loc=0, scale=2)] >>> m1 = AcceptanceRejectionReal(DigitalNetB2(dimension=2, seed=7), psi, inv_cdfs=inv_cdfs, H_func=H, upper_bound=2., density_integral=1.) @@ -294,8 +294,8 @@ class AcceptanceRejectionReal(AbstractTrueMeasure): qmcpy.util.exceptions_warnings.ParameterError: n_min > 0 but no prior call was made. Call gen_samples with n_min=0 first. """ - def __init__(self, sampler, target_density, inv_cdfs, H_func, - upper_bound, density_integral, max_retries=4): + def __init__(self, sampler: AbstractDiscreteDistribution, target_density: Callable, inv_cdfs: List[Callable], H_func: Callable, + upper_bound: float, density_integral: float, max_retries: int = 4) -> None: self.parameters = ['target_dim', 'upper_bound', 'density_integral', 'acceptance_rate'] self.domain = np.array([[0, 1]]) self._parse_sampler(sampler) @@ -325,36 +325,36 @@ def __init__(self, sampler, target_density, inv_cdfs, H_func, self._driver_offset = None super(AcceptanceRejectionReal, self).__init__() - def gen_samples(self, n=None, n_min=None, n_max=None, return_weights=False, warn=True): - """ - Generate accepted samples from the target density on R^d. + def gen_samples(self, n: Union[None, int] = None, n_min: Union[None, int] = None, n_max: Union[None, int] = None, return_weights: bool = False, warn: bool = True) -> np.ndarray: + """Generate accepted samples from the target density on R^d. - Unlike other TrueMeasures, this method cannot be decomposed into - a fixed 1-to-1 _transform because acceptance-rejection produces - a variable number of outputs from a fixed driver batch. gen_samples - is therefore overridden directly. + Unlike other TrueMeasures, this method cannot be decomposed into a + fixed 1-to-1 _transform because acceptance-rejection produces a + variable number of outputs from a fixed driver batch. gen_samples is + therefore overridden directly. - Supports continued sampling: calling with n_min=0 starts fresh, - and subsequent calls with n_min>0 continue from the same driver - sequence position. + Supports continued sampling: calling with n_min=0 starts fresh, and + subsequent calls with n_min>0 continue from the same driver sequence + position. Args: - n (int): Number of accepted samples to return. Treated as - n_min=0, n_max=n (always resets the driver sequence). - n_min (int): Starting accepted-sample index. Use 0 to reset - and start fresh. Use a positive value to continue from - the previous call. - n_max (int): Ending accepted-sample index (exclusive). - Number of samples returned is n_max - n_min. + n (Union[None, int]): Number of accepted samples to return. Treated as n_min=0, + n_max=n (always resets the driver sequence). + n_min (Union[None, int]): Starting accepted-sample index. Use 0 to reset and + start fresh. Use a positive value to continue from the previous + call. + n_max (Union[None, int]): Ending accepted-sample index (exclusive). Number of + samples returned is n_max - n_min. return_weights (bool): If True, also return importance weights psi(z)/C for each accepted sample. - warn (bool): If True, warn when fewer than n samples are - returned after all retries. + warn (bool): If True, warn when fewer than n samples are returned + after all retries. Returns: - samples (np.ndarray): Shape (n, target_dim). - weights (np.ndarray): Shape (n,). Only returned when - return_weights=True. + np.ndarray: Accepted samples of shape (n, target_dim). When + return_weights=True, a tuple (samples, weights) is returned + instead, where weights has shape (n,) and holds the importance + weights psi(z)/C. """ if n_max is not None: if n_min is None: diff --git a/qmcpy/true_measure/bernoulli_cont.py b/qmcpy/true_measure/bernoulli_cont.py index cadf9f727..4784ebca4 100644 --- a/qmcpy/true_measure/bernoulli_cont.py +++ b/qmcpy/true_measure/bernoulli_cont.py @@ -1,3 +1,7 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from typing import Union from .abstract_true_measure import AbstractTrueMeasure from ..util import DimensionError from ..discrete_distribution import DigitalNetB2 @@ -5,8 +9,9 @@ class BernoulliCont(AbstractTrueMeasure): - r""" - Continuous Bernoulli distribution with independent marginals as described in [https://en.wikipedia.org/wiki/Continuous_Bernoulli_distribution](https://en.wikipedia.org/wiki/Continuous_Bernoulli_distribution). + r"""Continuous Bernoulli distribution with independent marginals as + described in + [https://en.wikipedia.org/wiki/Continuous_Bernoulli_distribution](https://en.wikipedia.org/wiki/Continuous_Bernoulli_distribution). Examples: >>> true_measure = BernoulliCont(DigitalNetB2(2,seed=7),lam=.2) @@ -36,14 +41,17 @@ class BernoulliCont(AbstractTrueMeasure): [0.6345258 , 0.60241448, 0.84822692]]]) """ - def __init__(self, sampler, lam=1 / 2): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], lam: Union[float, np.ndarray] = 1 / 2) -> None: + r"""Initialize a BernoulliCont true measure. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. - lam (Union[float, np.ndarray]): Vector of shape parameters, each in $(0,1)$. + lam (Union[float, np.ndarray]): Vector of shape parameters, each in + $(0,1)$. """ self.parameters = ["lam"] self.domain = np.array([[0, 1]]) diff --git a/qmcpy/true_measure/brownian_motion.py b/qmcpy/true_measure/brownian_motion.py index 564223a00..dc2e593da 100644 --- a/qmcpy/true_measure/brownian_motion.py +++ b/qmcpy/true_measure/brownian_motion.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .gaussian import Gaussian from ..discrete_distribution import DigitalNetB2 from ..util import ParameterError, ParameterWarning @@ -7,9 +12,10 @@ class BrownianMotion(Gaussian): - r""" - Brownian Motion as described in [https://en.wikipedia.org/wiki/Brownian_motion](https://en.wikipedia.org/wiki/Brownian_motion). - For a standard Brownian Motion $W$ we define the Brownian Motion $B$ with initial value $B_0$, drift $\gamma$, and diffusion $\sigma^2$ to be + r"""Brownian Motion as described in + [https://en.wikipedia.org/wiki/Brownian_motion](https://en.wikipedia.org/wiki/Brownian_motion). + For a standard Brownian Motion $W$ we define the Brownian Motion $B$ with + initial value $B_0$, drift $\gamma$, and diffusion $\sigma^2$ to be $$B(t) = B_0 + \gamma t + \sigma W(t).$$ @@ -70,7 +76,8 @@ class BrownianMotion(Gaussian): bridge_construction_times [1. 0.5 0.75 0.25] bridge_output_times [0.25 0.5 0.75 1. ] - Example 4: With Brownian Bridge construction and independent replications + Example 4: With Brownian Bridge construction and independent + replications >>> x = BrownianMotion(DigitalNetB2(4,seed=7,replications=3),decomp_type='BrownianBridge')(2) >>> x.shape @@ -85,9 +92,10 @@ class BrownianMotion(Gaussian): [[ 0.59845146, 1.10849282, 1.34022073, 1.02092441], [-0.20298903, -0.23324496, -0.3026512 , -0.35202342]]]) - Example 5: With custom monitoring times and passing bridge_vdc_gray_ordering=False (reaches all four cases) + Example 5: With custom monitoring times and passing + bridge_vdc_gray_ordering=False (reaches all four cases) - >>> true_measure = BrownianMotion(DigitalNetB2(4,seed=7),decomp_type='BrownianBridge',monitoring_times=[0.6,1.0,0.3,0.8],bridge_vdc_gray_ordering=False) + >>> true_measure = BrownianMotion(DigitalNetB2(4,seed=7),decomp_type='BrownianBridge',monitoring_times=[0.6,1.0,0.3,0.8],bridge_vdc_gray_ordering=False) >>> true_measure.time_vec array([0.3, 0.6, 0.8, 1. ]) >>> true_measure(2) @@ -98,7 +106,8 @@ class BrownianMotion(Gaussian): >>> true_measure.bridge_output_times array([0.3, 0.6, 0.8, 1. ]) - Example 6: With custom monitoring times. By default the times are sorted and inserted in van der Corput order + Example 6: With custom monitoring times. By default the times are + sorted and inserted in van der Corput order >>> true_measure = BrownianMotion(DigitalNetB2(4,seed=7),decomp_type='BrownianBridge',monitoring_times=[0.6,1.0,0.3,0.8]) >>> true_measure.time_vec @@ -126,7 +135,7 @@ class BrownianMotion(Gaussian): **References:** - 1. Art B. Owen. + 1. Art B. Owen. Monte Carlo theory, methods and examples. Section 6.4, Detailed Simulation of Brownian Motion, 2013 [https://artowen.su.domains/mc/](https://artowen.su.domains/mc/) @@ -134,20 +143,22 @@ class BrownianMotion(Gaussian): def __init__( self, - sampler, - t_final=1, - initial_value=0, - drift=0, - diffusion=1, - decomp_type="PCA", - lazy_decomp=True, - monitoring_times=None, - bridge_vdc_gray_ordering=True, - bridge_output_order='increasing', - ): - r""" + sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], + t_final: float = 1, + initial_value: float = 0, + drift: int = 0, + diffusion: int = 1, + decomp_type: str = "PCA", + lazy_decomp: bool = True, + monitoring_times: Union[None, np.ndarray, list] = None, + bridge_vdc_gray_ordering: bool = True, + bridge_output_order: str = 'increasing', + ) -> None: + r"""Initialize a BrownianMotion true measure. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -155,19 +166,25 @@ def __init__( initial_value (float): Initial value $B_0$. drift (int): Drift $\gamma$. diffusion (int): Diffusion $\sigma^2$. - decomp_type (str): Method for decomposition for covariance matrix. Options include + decomp_type (str): Method for decomposition for covariance matrix. + Options include - `'PCA'` for principal component analysis, - `'Cholesky'` for cholesky decomposition, or - `'BrownianBridge'` or `'Bridge'` for brownian bridge construction. - lazy_decomp (bool): If True, defer expensive matrix decomposition until needed. - monitoring_times (Union[np.ndarray, list]): Optional custom sampling times for `'BrownianBridge'` - with length d. The given order is the insertion order if `'bridge_vdc_gray_ordering'` is False. - bridge_vdc_gray_ordering (bool): For `'BrownianBridge'` when monitoring_times is specified. If True, - monitoring_times is sorted to match van der Corput ordering. - bridge_output_order (str): If `'increasing'`, output is returned in increasing order. If `'input'`, - output matches the order given in `'monitoring_times'`. If a custom monitoring times is not given, - the output is given in increasing order. + lazy_decomp (bool): If True, defer expensive matrix decomposition + until needed. + monitoring_times (Union[None, np.ndarray, list]): Optional custom + sampling times for `'BrownianBridge'` with length d. The given + order is the insertion order if `'bridge_vdc_gray_ordering'` is + False. + bridge_vdc_gray_ordering (bool): For `'BrownianBridge'` when + monitoring_times is specified. If True, monitoring_times is + sorted to match van der Corput ordering. + bridge_output_order (str): If `'increasing'`, output is returned in + increasing order. If `'input'`, output matches the order given + in `'monitoring_times'`. If a custom monitoring times is not + given, the output is given in increasing order. """ if str(decomp_type).upper() == "BRIDGE": decomp_type = "BrownianBridge" diff --git a/qmcpy/true_measure/clayton_copula.py b/qmcpy/true_measure/clayton_copula.py index fb25b2648..6b28cf8d1 100644 --- a/qmcpy/true_measure/clayton_copula.py +++ b/qmcpy/true_measure/clayton_copula.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .copula import ( AbstractCopula, _clip_unit_interval, @@ -11,24 +16,22 @@ class ClaytonCopula(AbstractCopula): - r""" - Clayton copula transform with user supplied marginals. + r"""Clayton copula transform with user supplied marginals. This implementation supports general dimension for ``theta > 0``. It maps independent uniforms to Clayton-dependent uniforms using the conditional inverse / inverse Rosenblatt transform. For coordinate ``j`` after - observing the previous ``m = j - 1`` coordinates, the conditional inverse is + observing the previous ``m = j - 1`` coordinates, the conditional inverse + is - $$ - v = \left(1 + A - \left(w^{-\theta/(1 + m \theta)} - 1\right)\right)^{-1/\theta}, - $$ + $$ v = \left(1 + A \left(w^{-\theta/(1 + m \theta)} - + 1\right)\right)^{-1/\theta}, $$ - where ``A = 1 + sum(phi(u_i))`` over previous coordinates and - ``phi(u) = u^{-theta} - 1``. + where ``A = 1 + sum(phi(u_i))`` over previous coordinates and ``phi(u) = + u^{-theta} - 1``. - The base ``AbstractCopula`` class then applies each marginal quantile function. - SciPy calls the quantile function ``ppf``. + The base ``AbstractCopula`` class then applies each marginal quantile + function. SciPy calls the quantile function ``ppf``. Clayton copulas have positive lower-tail dependence for ``theta > 0``. @@ -81,8 +84,9 @@ class ClaytonCopula(AbstractCopula): [doi:10.1016/j.jmva.2012.02.019](https://doi.org/10.1016/j.jmva.2012.02.019). """ - def __init__(self, sampler, marginals, theta): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], marginals: list, theta: float) -> None: + r"""Initialize a ClaytonCopula true measure. + Args: sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A sampler or transform whose range is the unit cube. diff --git a/qmcpy/true_measure/copula.py b/qmcpy/true_measure/copula.py index 67ade5612..ddd4ee09c 100644 --- a/qmcpy/true_measure/copula.py +++ b/qmcpy/true_measure/copula.py @@ -7,28 +7,24 @@ class AbstractCopula(AbstractTrueMeasure): - r""" - Abstract base class for copula TrueMeasures. + r"""Abstract base class for copula TrueMeasures. A copula layer maps independent uniform input points to dependent uniform points on the unit cube: - $$ - U \in [0,1]^d \mapsto V = T(U) \in [0,1]^d. - $$ + $$ U \in [0,1]^d \mapsto V = T(U) \in [0,1]^d. $$ The base class then applies marginal quantile functions to obtain final target samples, - $$ - X_j = F_j^{-1}(V_j). - $$ + $$ X_j = F_j^{-1}(V_j). $$ SciPy calls the quantile function ``ppf``. Concrete subclasses implement - ``_transform_to_uniform`` for the family-specific copula sampling transform. + ``_transform_to_uniform`` for the family-specific copula sampling + transform. """ - def __init__(self, sampler, marginals): + def __init__(self, sampler, marginals) -> None: self.domain = np.array([[0, 1]]) self._parse_sampler(sampler) @@ -40,14 +36,13 @@ def __init__(self, sampler, marginals): super(AbstractCopula, self).__init__() def _transform_to_uniform(self, x) -> np.ndarray: - r""" - Transform independent uniforms ``U`` into dependent copula uniforms ``V``. + r"""Transform independent uniforms ``U`` into dependent copula + uniforms ``V``. """ raise MethodImplementationError(self, "_transform_to_uniform") - def copula_transform(self, u) -> np.ndarray: - r""" - Apply only the copula layer ``U -> V``. + def copula_transform(self, u: np.ndarray) -> np.ndarray: + r"""Apply only the copula layer ``U -> V``. Args: u (np.ndarray): Independent uniform points on ``[0,1]^d``. @@ -60,8 +55,8 @@ def copula_transform(self, u) -> np.ndarray: def gen_copula_samples( self, n=None, n_min=None, n_max=None, warn=True ) -> np.ndarray: - r""" - Generate dependent copula uniforms without applying marginal quantiles. + r"""Generate dependent copula uniforms without applying marginal + quantiles. This is the copula-only workflow ``U -> V``. Calling the object itself keeps the ordinary TrueMeasure workflow ``U -> V -> X``. @@ -72,8 +67,7 @@ def gen_copula_samples( return self._transform_to_uniform(u) def _apply_marginal_quantiles(self, v) -> np.ndarray: - r""" - Apply marginal quantile functions to dependent uniforms. + r"""Apply marginal quantile functions to dependent uniforms. SciPy frozen distributions expose the quantile function as ``ppf``. """ diff --git a/qmcpy/true_measure/frank_copula.py b/qmcpy/true_measure/frank_copula.py index 3f17d38c0..ac591f12c 100644 --- a/qmcpy/true_measure/frank_copula.py +++ b/qmcpy/true_measure/frank_copula.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .copula import ( AbstractCopula, _clip_unit_interval, @@ -11,11 +16,9 @@ def _eulerian_coefficients(n): - """ - Return Eulerian coefficients for Li_{-n}(z). + """Return Eulerian coefficients for Li_{-n}(z). - For nonnegative integer n, - Li_{-n}(z) = z * A_n(z) / (1 - z) ** (n + 1), + For nonnegative integer n, Li_{-n}(z) = z * A_n(z) / (1 - z) ** (n + 1), where A_n is the Eulerian polynomial. """ if n == 0: @@ -33,8 +36,7 @@ def _eulerian_coefficients(n): class FrankCopula(AbstractCopula): - r""" - Frank copula transform with user supplied univariate marginals. + r"""Frank copula transform with user supplied univariate marginals. This implementation supports general dimension for ``theta > 0``. Negative ``theta`` is supported only for the bivariate case, where the negative @@ -43,9 +45,9 @@ class FrankCopula(AbstractCopula): The transform uses the inverse Rosenblatt construction for the Frank Archimedean copula. It maps independent uniforms to dependent uniforms by - recursively inverting conditional CDFs. The base ``AbstractCopula`` class then - applies each marginal quantile function. SciPy calls the quantile function - ``ppf``. + recursively inverting conditional CDFs. The base ``AbstractCopula`` class + then applies each marginal quantile function. SciPy calls the quantile + function ``ppf``. Examples: >>> import numpy as np @@ -105,16 +107,17 @@ class FrankCopula(AbstractCopula): [doi:10.1016/j.jmva.2012.02.019](https://doi.org/10.1016/j.jmva.2012.02.019). """ - def __init__(self, sampler, marginals, theta): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], marginals: list, theta: float) -> None: + r"""Initialize a FrankCopula true measure. + Args: sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A sampler or transform whose range is the unit cube. marginals (list): Length d list of SciPy-like univariate distributions implementing a quantile function, called ``ppf`` in SciPy. - theta (float): Frank dependence parameter. Must be nonzero. Negative - values are currently supported only for ``d=2``. + theta (float): Frank dependence parameter. Must be nonzero. + Negative values are currently supported only for ``d=2``. """ self.parameters = ["marginals", "theta"] super(FrankCopula, self).__init__(sampler=sampler, marginals=marginals) diff --git a/qmcpy/true_measure/gaussian.py b/qmcpy/true_measure/gaussian.py index 37d58982b..3d8896d5e 100644 --- a/qmcpy/true_measure/gaussian.py +++ b/qmcpy/true_measure/gaussian.py @@ -1,6 +1,9 @@ from .abstract_true_measure import AbstractTrueMeasure from ..util import DimensionError, ParameterError from ..discrete_distribution import DigitalNetB2 +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) import numpy as np from numpy.linalg import cholesky, slogdet from scipy.stats import norm, multivariate_normal @@ -9,10 +12,10 @@ class Gaussian(AbstractTrueMeasure): - """ - Gaussian (Normal) distribution as described in [https://en.wikipedia.org/wiki/Multivariate_normal_distribution](https://en.wikipedia.org/wiki/Multivariate_normal_distribution). + """Gaussian (Normal) distribution as described in + [https://en.wikipedia.org/wiki/Multivariate_normal_distribution](https://en.wikipedia.org/wiki/Multivariate_normal_distribution). - Note: + Notes: - `Normal` is an alias for `Gaussian` Examples: @@ -48,16 +51,20 @@ class Gaussian(AbstractTrueMeasure): [ 1.1844196 , 0.44964332, 1.27760936]]]) """ - def __init__(self, sampler, mean=0.0, covariance=1.0, decomp_type="PCA"): - """ + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], mean: Union[float, np.ndarray] = 0.0, covariance: Union[float, np.ndarray] = 1.0, decomp_type: str = "PCA") -> None: + """Initialize a Gaussian true measure. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. mean (Union[float, np.ndarray]): Mean vector. - covariance (Union[float, np.ndarray]): Covariance matrix. A float or vector will be expanded into a diagonal matrix. - decomp_type (str): Method for decomposition for covariance matrix. Options include + covariance (Union[float, np.ndarray]): Covariance matrix. A float + or vector will be expanded into a diagonal matrix. + decomp_type (str): Method for decomposition for covariance matrix. + Options include - `'PCA'` for principal component analysis, or - `'Cholesky'` for cholesky decomposition. @@ -69,7 +76,8 @@ def __init__(self, sampler, mean=0.0, covariance=1.0, decomp_type="PCA"): self._parse_gaussian_params(mean, covariance, decomp_type) self.range = np.array([[-np.inf, np.inf]]) super(Gaussian, self).__init__() - assert self.mu.shape == (self.d,) and self.a.shape == (self.d, self.d) + if not (self.mu.shape == (self.d,) and self.a.shape == (self.d, self.d)): + raise AssertionError def _parse_gaussian_params(self, mean, covariance, decomp_type, lazy_decomp=False): self.decomp_type = decomp_type.upper() @@ -108,7 +116,9 @@ def _parse_gaussian_params(self, mean, covariance, decomp_type, lazy_decomp=Fals self._setup_scipy_mvn() def _compute_decomposition(self): - """Compute matrix decomposition (PCA or Cholesky). Raises ParameterError for BrownianBridge.""" + """Compute matrix decomposition (PCA or Cholesky). Raises + ParameterError for BrownianBridge. + """ if self._a_cache is not None: return self._a_cache diff --git a/qmcpy/true_measure/gaussian_copula.py b/qmcpy/true_measure/gaussian_copula.py index 276106157..4908ec3e0 100644 --- a/qmcpy/true_measure/gaussian_copula.py +++ b/qmcpy/true_measure/gaussian_copula.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .copula import ( AbstractCopula, _clip_unit_interval, @@ -14,8 +19,7 @@ class GaussianCopula(AbstractCopula): - r""" - Gaussian copula transform with user supplied univariate marginals. + r"""Gaussian copula transform with user supplied univariate marginals. This TrueMeasure separates the dependence model from the marginal distributions: @@ -26,9 +30,9 @@ class GaussianCopula(AbstractCopula): 4. apply each marginal quantile function. SciPy calls the quantile function ``ppf``. The marginal objects must expose - this method. If they also expose - ``cdf`` and ``pdf`` or ``logpdf``, then ``_weight`` computes the Gaussian - copula joint density. Otherwise weights are treated as one with a warning. + this method. If they also expose ``cdf`` and ``pdf`` or ``logpdf``, then + ``_weight`` computes the Gaussian copula joint density. Otherwise weights + are treated as one with a warning. Examples: >>> import numpy as np @@ -76,15 +80,17 @@ class GaussianCopula(AbstractCopula): [arXiv:1508.03483](https://arxiv.org/abs/1508.03483). """ - def __init__(self, sampler, marginals, correlation): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], marginals: list, correlation: np.ndarray) -> None: + r"""Initialize a GaussianCopula true measure. + Args: sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A sampler or transform whose range is the unit cube. marginals (list): Length d list of SciPy-like univariate distributions implementing a quantile function, called ``ppf`` in SciPy. - correlation (np.ndarray): d x d positive definite correlation matrix. + correlation (np.ndarray): d x d positive definite correlation + matrix. """ self.parameters = ["marginals", "correlation"] super(GaussianCopula, self).__init__(sampler=sampler, marginals=marginals) diff --git a/qmcpy/true_measure/geometric_brownian_motion.py b/qmcpy/true_measure/geometric_brownian_motion.py index 9daa76c24..0943cda25 100644 --- a/qmcpy/true_measure/geometric_brownian_motion.py +++ b/qmcpy/true_measure/geometric_brownian_motion.py @@ -1,5 +1,9 @@ from .brownian_motion import BrownianMotion +from .abstract_true_measure import AbstractTrueMeasure from ..discrete_distribution import DigitalNetB2 +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) from ..util import ParameterError from typing import Union, Tuple from numpy import ( @@ -21,10 +25,11 @@ class GeometricBrownianMotion(BrownianMotion): - r""" - A Geometric Brownian Motion (GBM) with initial value $S_0$, drift $\gamma$, and diffusion $\sigma^2$ is + r"""A Geometric Brownian Motion (GBM) with initial value $S_0$, drift + $\gamma$, and diffusion $\sigma^2$ is - $$\mathrm{GBM}(t) = S_0 \exp[(\gamma - \sigma^2/2) t + \sigma \mathrm{BM}(t)]$$ + $$\mathrm{GBM}(t) = S_0 \exp[(\gamma - \sigma^2/2) t + \sigma + \mathrm{BM}(t)]$$ where BM is a Brownian Motion drift $\gamma$ and diffusion $\sigma^2$. @@ -48,25 +53,33 @@ class GeometricBrownianMotion(BrownianMotion): def __init__( self, - sampler, - t_final=1, - initial_value=1, - drift=0, - diffusion=1, - decomp_type="PCA", - lazy_load=True, - lazy_decomp=True, - ): - r""" + sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], + t_final: float = 1, + initial_value: float = 1, + drift: float = 0, + diffusion: float = 1, + decomp_type: str = "PCA", + lazy_load: bool = True, + lazy_decomp: bool = True, + ) -> None: + r"""Initialize a GeometricBrownianMotion true measure. + Args: - sampler (DiscreteDistribution/TrueMeasure): A discrete distribution or true measure. - t_final (float): End time for the geometric Brownian motion, non-negative. - initial_value (float): Positive initial value of the process, $S_0$. + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A discrete distribution + or true measure. + t_final (float): End time for the geometric Brownian motion, + non-negative. + initial_value (float): Positive initial value of the process, + $S_0$. drift (float): Drift coefficient $\gamma$. - diffusion (float): Positive diffusion coefficient $\sigma^2$, where $\sigma$ is volatility. - decomp_type (str): Method of decomposition, either "PCA", "Cholesky", or "BrownianBridge". - lazy_load (bool): If True, defer GBM-specific computations until needed. - lazy_decomp (bool): If True, defer expensive matrix decomposition until needed. + diffusion (float): Positive diffusion coefficient $\sigma^2$, where + $\sigma$ is volatility. + decomp_type (str): Method of decomposition, either "PCA", + "Cholesky", or "BrownianBridge". + lazy_load (bool): If True, defer GBM-specific computations until + needed. + lazy_decomp (bool): If True, defer expensive matrix decomposition + until needed. """ super().__init__( sampler, @@ -196,14 +209,16 @@ def _spawn(self, sampler, dimension): ) def _validate_input(self): - """ - Validates the input parameters of the GeometricBrownianMotion class. + """Validates the input parameters of the GeometricBrownianMotion + class. Raises: ValueError: If the end time `t_final' is negative. - ValueError: If the diffusion coefficient is less than or equal to zero. + ValueError: If the diffusion coefficient is less than or equal to + zero. ValueError: If the initial value is less than or equal to zero. - ParameterError: If the decomposition type is not 'PCA', 'Cholesky', or 'BrownianBridge'. + ParameterError: If the decomposition type is not 'PCA', 'Cholesky', + or 'BrownianBridge'. """ if self.t < 0: raise ValueError( @@ -223,8 +238,8 @@ def _validate_input(self): ) def _validate_samples(self, samples, strict=False): - """ - Validate that generated GBM samples meet mathematical requirements. + """Validate that generated GBM samples meet mathematical + requirements. """ min_val = samples.min() max_val = samples.max() @@ -260,7 +275,8 @@ def _validate_samples(self, samples, strict=False): return validation_results def _setup_lognormal_distribution(self): - """Setup scipy multivariate normal for the log-transformed variables.""" + """Setup scipy multivariate normal for the log-transformed variables. + """ # Mean of log(S(t)/S0): (drift - 0.5*diffusion) * t log_mean = (self.drift - 0.5 * self.diffusion) * self.time_vec @@ -272,10 +288,10 @@ def _setup_lognormal_distribution(self): mean=log_mean, cov=log_cov, allow_singular=True ) - def _weight(self, x): - """ - Compute PDF of multivariate log-normal distribution. - For log-normal: f(x) = (1/∏x_i) * φ(log(x/S0)) where φ is multivariate normal PDF. + def _weight(self, x: ndarray): + """Compute PDF of multivariate log-normal distribution. For + log-normal: f(x) = (1/∏x_i) * φ(log(x/S0)) where φ is multivariate + normal PDF. Args: x (ndarray): GBM sample paths of shape (n_samples, n_timepoints) @@ -298,19 +314,18 @@ def _weight(self, x): return normal_pdf * jacobian def gen_samples( - self, n=None, n_min=None, n_max=None, return_weights=False, warn=True + self, n: Union[None, int] = None, n_min: Union[None, int] = None, n_max: Union[None, int] = None, return_weights: bool = False, warn: bool = True ) -> Union[ndarray, Tuple[ndarray, ndarray]]: - """ - Generate GBM samples using the parent's transform pipeline. - + """Generate GBM samples using the parent's transform pipeline. + Args: - n (int): number of samples to generate - n_min (int): minimum index of sequence - n_max (int): maximum index of sequence + n (Union[None, int]): number of samples to generate + n_min (Union[None, int]): minimum index of sequence + n_max (Union[None, int]): maximum index of sequence return_weights (bool): whether to return Jacobian weights warn (bool): whether to warn about sample generation - + Returns: - samples (Union[ndarray,tuple]): GBM samples, optionally with weights if return_weights=True + Union[ndarray, Tuple[ndarray, ndarray]]: GBM samples, optionally with weights if return_weights=True """ return super().gen_samples(n=n, n_min=n_min, n_max=n_max, return_weights=return_weights, warn=warn) diff --git a/qmcpy/true_measure/gumbel_copula.py b/qmcpy/true_measure/gumbel_copula.py index 0ecbdf7c8..96bff94d0 100644 --- a/qmcpy/true_measure/gumbel_copula.py +++ b/qmcpy/true_measure/gumbel_copula.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .copula import ( AbstractCopula, _clip_unit_interval, @@ -11,17 +16,16 @@ class GumbelCopula(AbstractCopula): - r""" - Gumbel copula transform with user supplied marginals. + r"""Gumbel copula transform with user supplied marginals. - This implementation supports general dimension for ``theta >= 1``. It - maps independent uniforms to Gumbel-dependent uniforms by numerically - inverting the conditional CDFs from the inverse Rosenblatt construction. - The base ``AbstractCopula`` class then applies marginal quantile functions. - SciPy calls the quantile function ``ppf``. + This implementation supports general dimension for ``theta >= 1``. It maps + independent uniforms to Gumbel-dependent uniforms by numerically inverting + the conditional CDFs from the inverse Rosenblatt construction. The base + ``AbstractCopula`` class then applies marginal quantile functions. SciPy + calls the quantile function ``ppf``. - Gumbel copulas have positive upper-tail dependence for ``theta > 1``. - The boundary case ``theta = 1`` is the independent copula. + Gumbel copulas have positive upper-tail dependence for ``theta > 1``. The + boundary case ``theta = 1`` is the independent copula. Examples: >>> import numpy as np @@ -76,15 +80,17 @@ class GumbelCopula(AbstractCopula): [doi:10.1016/j.jmva.2012.02.019](https://doi.org/10.1016/j.jmva.2012.02.019). """ - def __init__(self, sampler, marginals, theta): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], marginals: list, theta: float) -> None: + r"""Initialize a GumbelCopula true measure. + Args: sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A sampler or transform whose range is the unit cube. marginals (list): Length d list of SciPy-like univariate distributions implementing a quantile function, called ``ppf`` in SciPy. - theta (float): Gumbel dependence parameter, requiring ``theta >= 1``. + theta (float): Gumbel dependence parameter, requiring ``theta >= + 1``. """ self.parameters = ["marginals", "theta"] super(GumbelCopula, self).__init__(sampler=sampler, marginals=marginals) diff --git a/qmcpy/true_measure/johnsons_su.py b/qmcpy/true_measure/johnsons_su.py index 1d0cc54e2..e6c270b01 100644 --- a/qmcpy/true_measure/johnsons_su.py +++ b/qmcpy/true_measure/johnsons_su.py @@ -1,3 +1,7 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from typing import Union from .abstract_true_measure import AbstractTrueMeasure from ..util import DimensionError, ParameterError from ..discrete_distribution import DigitalNetB2 @@ -6,8 +10,9 @@ class JohnsonsSU(AbstractTrueMeasure): - r""" - Johnson's $S_U$-distribution with independent marginals as described in [https://en.wikipedia.org/wiki/Johnson%27s_SU-distribution](https://en.wikipedia.org/wiki/Johnson%27s_SU-distribution). + r"""Johnson's $S_U$-distribution with independent marginals as described + in + [https://en.wikipedia.org/wiki/Johnson%27s_SU-distribution](https://en.wikipedia.org/wiki/Johnson%27s_SU-distribution). Examples: >>> true_measure = JohnsonsSU(DigitalNetB2(2,seed=7),gamma=1,xi=2,delta=3,lam=4) @@ -40,10 +45,12 @@ class JohnsonsSU(AbstractTrueMeasure): [ 1.57765245, 1.00275 , 1.64972468]]]) """ - def __init__(self, sampler, gamma=1, xi=1, delta=2, lam=2): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], gamma: Union[float, np.ndarray] = 1, xi: Union[float, np.ndarray] = 1, delta: Union[float, np.ndarray] = 2, lam: Union[float, np.ndarray] = 2) -> None: + r"""Initialize a JohnsonsSU true measure. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -84,12 +91,13 @@ def __init__(self, sampler, gamma=1, xi=1, delta=2, lam=2): if not ((self._delta > 0).all() and (self._lam > 0).all()): raise ParameterError("delta and lam must be all be positive") super(JohnsonsSU, self).__init__() - assert ( + if not ( self._gamma.shape == (self.d,) and self._xi.shape == (self.d,) and self._delta.shape == (self.d,) and self._lam.shape == (self.d,) - ) + ): + raise AssertionError def _transform(self, x): return self._lam * np.sinh((norm.ppf(x) - self._gamma) / self._delta) + self._xi diff --git a/qmcpy/true_measure/kumaraswamy.py b/qmcpy/true_measure/kumaraswamy.py index 47d4bdd9d..a5765473b 100644 --- a/qmcpy/true_measure/kumaraswamy.py +++ b/qmcpy/true_measure/kumaraswamy.py @@ -1,3 +1,7 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from typing import Union from .abstract_true_measure import AbstractTrueMeasure from ..util import DimensionError, ParameterError from ..discrete_distribution import DigitalNetB2 @@ -7,8 +11,8 @@ class Kumaraswamy(AbstractTrueMeasure): - r""" - Kumaraswamy distribution as described in [https://en.wikipedia.org/wiki/Kumaraswamy_distribution](https://en.wikipedia.org/wiki/Kumaraswamy_distribution). + r"""Kumaraswamy distribution as described in + [https://en.wikipedia.org/wiki/Kumaraswamy_distribution](https://en.wikipedia.org/wiki/Kumaraswamy_distribution). Examples: >>> true_measure = Kumaraswamy(DigitalNetB2(2,seed=7),a=[1,2],b=[3,4]) @@ -50,10 +54,12 @@ class Kumaraswamy(AbstractTrueMeasure): [0.37253319, 0.45379743, 0.63366422]]]) """ - def __init__(self, sampler, a=2, b=2): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], a: Union[float, np.ndarray] = 2, b: Union[float, np.ndarray] = 2) -> None: + r"""Initialize a Kumaraswamy true measure. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -92,22 +98,21 @@ def __init__(self, sampler, a=2, b=2): covariance=diags(variance, format="dia"), ) super(Kumaraswamy, self).__init__() - assert self.alpha.shape == (self.d,) and self.beta.shape == (self.d,) + if not (self.alpha.shape == (self.d,) and self.beta.shape == (self.d,)): + raise AssertionError def _compute_moments(self): - r""" - Compute the marginal mean and variance of each coordinate. + r"""Compute the marginal mean and variance of each coordinate. The Kumaraswamy raw moments are $M_n = b\,B(1 + n/a, b)$ [1], so the - mean is $M_1$ and the variance is $M_2 - M_1^2$. Forming that difference - directly causes cancellation error once the variance is small relative - to $M_1^2$ (e.g. large $a$). + mean is $M_1$ and the variance is $M_2 - M_1^2$. Forming that + difference directly causes cancellation error once the variance is + small relative to $M_1^2$ (e.g. large $a$). Instead, with the log-moment function $K(r) = \log M_r$, - $$\text{mean} = e^{K(1)}, \qquad - \operatorname{Var}[X] = \text{mean}^2\,(e^{q} - 1), \qquad - q = K(2) - 2K(1).$$ + $$\text{mean} = e^{K(1)}, \qquad \operatorname{Var}[X] = + \text{mean}^2\,(e^{q} - 1), \qquad q = K(2) - 2K(1).$$ Each log-moment is available in closed form via the log-Beta function [2], $K(r) = \log b + \ln B(1 + r/a, b)$, so ``mean`` and $q$ are diff --git a/qmcpy/true_measure/lebesgue.py b/qmcpy/true_measure/lebesgue.py index ec507bb00..efba73152 100644 --- a/qmcpy/true_measure/lebesgue.py +++ b/qmcpy/true_measure/lebesgue.py @@ -7,8 +7,8 @@ class Lebesgue(AbstractTrueMeasure): - r""" - Lebesgue measure as described in [https://en.wikipedia.org/wiki/Lebesgue_measure](https://en.wikipedia.org/wiki/Lebesgue_measure). + r"""Lebesgue measure as described in + [https://en.wikipedia.org/wiki/Lebesgue_measure](https://en.wikipedia.org/wiki/Lebesgue_measure). Examples: >>> Lebesgue(Gaussian(DigitalNetB2(2,seed=7))) @@ -35,10 +35,12 @@ class Lebesgue(AbstractTrueMeasure): (1, 1) 0.08333333333333333 """ - def __init__(self, sampler): - r""" + def __init__(self, sampler: AbstractTrueMeasure) -> None: + r"""Initialize a Lebesgue true measure. + Args: - sampler (AbstractTrueMeasure): A true measure by which to compose a transform. + sampler (AbstractTrueMeasure): A true measure by which to compose a + transform. """ self.parameters = [] if not isinstance(sampler, AbstractTrueMeasure): diff --git a/qmcpy/true_measure/matern_gp.py b/qmcpy/true_measure/matern_gp.py index df389bcc3..93f5d1212 100644 --- a/qmcpy/true_measure/matern_gp.py +++ b/qmcpy/true_measure/matern_gp.py @@ -12,8 +12,7 @@ class MaternGP(Gaussian): - r""" - A Gaussian process with Matérn covariance kernel. + r"""A Gaussian process with Matérn covariance kernel. Examples: >>> true_measure = MaternGP(DigitalNetB2(dimension=3,seed=7),points=np.linspace(0,1,3)[:,None],nu=3/2,length_scale=[3,4,5],variance=0.01,mean=np.array([.3,.4,.5])) @@ -67,22 +66,28 @@ class MaternGP(Gaussian): def __init__( self, - sampler, - points, - length_scale=1.0, - nu=1.5, - variance=1.0, - mean=0.0, - nugget=1e-6, - decomp_type="PCA", - ): - r""" + sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], + points: np.ndarray, + length_scale: Union[float, np.ndarray] = 1.0, + nu: float = 1.5, + variance: float = 1.0, + mean: Union[float, np.ndarray] = 0.0, + nugget: float = 1e-6, + decomp_type: str = "PCA", + ) -> None: + r"""Initialize a MaternGP true measure. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. - points (np.ndarray): The positions of points on a metric space. The array should have shape $(d,k)$ where $d$ is the dimension of the sampler and $k$ is the latent dimension. + points (np.ndarray): The positions of points on a metric space. The + array should have shape $(d,k)$ where $d$ is the dimension of + the sampler and $k$ is the latent dimension. + length_scale (Union[float, np.ndarray]): Determines "peakiness", or + how correlated two points are based on their distance. nu (float): The "smoothness" of the MaternGP function, e.g., - $\nu = 1/2$ is equivalent to the absolute exponential kernel, @@ -90,15 +95,17 @@ def __init__( - $\nu = 5/2$ implies twice differentiability. - as $\nu \to \infty$ the kernel becomes equivalent to the RBF kernel, see [`sklearn.gaussian_process.kernels.RBF`](https://scikit-learn.org/stable/modules/generated/sklearn.gaussian_process.kernels.RBF.html#sklearn.gaussian_process.kernels.RBF). - Note that when $\nu \notin \{1/2, 3/2, 5/2, \infty \}$ the kernel is around $10$ times slower to evaluate. - length_scale (Union[float, np.ndarray]): Determines "peakiness", or how correlated two points are based on their distance. + Note that when $\nu \notin \{1/2, 3/2, 5/2, \infty \}$ the + kernel is around $10$ times slower to evaluate. variance (float): Global scaling factor of the kernel. Retrievable after construction via the `kernel_variance` property. (The inherited `variance` attribute is the vector of marginal - variances, i.e. the diagonal of `covariance`.) - mean (Union[float, np.ndarray]): Mean vector for multivariate `Gaussian`. + variances, i.e., the diagonal of `covariance`.) + mean (Union[float, np.ndarray]): Mean vector for multivariate + `Gaussian`. nugget (float): Positive nugget to add to diagonal. - decomp_type (str): Method for decomposition for covariance matrix. Options include + decomp_type (str): Method for decomposition for covariance matrix. + Options include - `'PCA'` for principal component analysis, or - `'Cholesky'` for cholesky decomposition. @@ -116,24 +123,30 @@ def __init__( raise ParameterError("points must be a one or two dimensional np.ndarray.") if points.ndim == 1: points = points[:, None] - assert ( + if not ( points.ndim == 2 and points.shape[0] == sampler.d - ), "points should be a two dimension array with the number of points equal to the dimension of the sampler" + ): + raise AssertionError("points should be a two dimension array with the number of points equal to the dimension of the sampler") mean = np.array(mean) if mean.size == 1: mean = mean.item() * np.ones(sampler.d) - assert mean.shape == (sampler.d,), "mean should be a length d vector" - assert np.isscalar(nu) and nu > 0, "nu should be a positive scalar" + if not (mean.shape == (sampler.d,)): + raise AssertionError("mean should be a length d vector") + if not (np.isscalar(nu) and nu > 0): + raise AssertionError("nu should be a positive scalar") length_scale = np.array(length_scale) if length_scale.size == 1: length_scale = length_scale.item() * np.ones(sampler.d) - assert ( + if not ( length_scale.shape == (sampler.d,) and (length_scale > 0).all() - ), "length_scale should be a vector with length equal to the dimension of the sampler" - assert ( + ): + raise AssertionError("length_scale should be a vector with length equal to the dimension of the sampler") + if not ( np.isscalar(variance) and variance > 0 - ), "variance should be a positive scalar" - assert np.isscalar(nugget) and nugget > 0, "nugget should be a positive scalar" + ): + raise AssertionError("variance should be a positive scalar") + if not (np.isscalar(nugget) and nugget > 0): + raise AssertionError("nugget should be a positive scalar") self.points = points self.length_scale = length_scale self.nu = nu diff --git a/qmcpy/true_measure/product_measure.py b/qmcpy/true_measure/product_measure.py index 3f860e84b..dc1809441 100644 --- a/qmcpy/true_measure/product_measure.py +++ b/qmcpy/true_measure/product_measure.py @@ -1,3 +1,4 @@ +from typing import Union import numpy as np from scipy import sparse @@ -9,8 +10,8 @@ class ProductMeasure(AbstractTrueMeasure): - r""" - Product true measure for independent composition of marginal true measures. + r"""Product true measure for independent composition of marginal true + measures. ``ProductMeasure`` represents an independent product of smaller true measures. Each marginal may be one-dimensional or multidimensional. If the @@ -29,12 +30,11 @@ class ProductMeasure(AbstractTrueMeasure): For example, if the marginals are - marginal 1: 2D Gaussian - marginal 2: 1D zero-inflated exponential + marginal 1: 2D Gaussian marginal 2: 1D zero-inflated exponential then ``ProductMeasure`` uses a 3D sampler and returns samples with three - coordinates. The first two coordinates come from the Gaussian marginal, - and the third coordinate comes from the zero-inflated exponential marginal. + coordinates. The first two coordinates come from the Gaussian marginal, and + the third coordinate comes from the zero-inflated exponential marginal. The marginal true measures still have their own samplers because QMCPy's current ``AbstractTrueMeasure`` API requires every true measure to be @@ -45,8 +45,7 @@ class ProductMeasure(AbstractTrueMeasure): samplerless/template true-measure mode may be useful, but that is separate from this class. - Notes - ----- + Notes: For independent marginal blocks, means, variances, and standard deviations are concatenated in marginal order, while covariance is block diagonal. @@ -55,8 +54,7 @@ class ProductMeasure(AbstractTrueMeasure): QMCPy's recursive transform helper, but exact final-space product weights are not currently implemented here. - Examples - -------- + Examples: Combine two one-dimensional uniform true measures: >>> from qmcpy import DigitalNetB2, DummySampler, ProductMeasure, Uniform @@ -101,29 +99,27 @@ class ProductMeasure(AbstractTrueMeasure): (4, 3) """ - def __init__(self, sampler, marginals): - """ - Initialize a product measure from one sampler and several marginals. - - Parameters - ---------- - sampler : AbstractDiscreteDistribution - The sampler for the whole product measure. Its dimension must - equal the sum of the marginal dimensions. + def __init__(self, sampler: AbstractDiscreteDistribution, marginals: Union[list, tuple]) -> None: + """Initialize a product measure from one sampler and several + marginals. - marginals : list or tuple of AbstractTrueMeasure - Independent true measures to place side by side. A marginal may - itself be multidimensional. + Args: + sampler (AbstractDiscreteDistribution): The sampler for the whole + product measure. Its dimension must equal the sum of the + marginal dimensions. + marginals (Union[list, tuple]): Independent true + measures to place side by side. A marginal may itself be + multidimensional. - Why one sampler? - ---------------- - The product measure should be driven by one total-dimensional QMC - point set. We do not generate separate QMC samples from each marginal. - Instead, one sample u in [0,1]^d is split into blocks: + Notes: + Why one sampler? The product measure should be driven by one + total-dimensional QMC point set. We do not generate separate QMC + samples from each marginal. Instead, one sample u in [0,1]^d is + split into blocks: - u = (u_marginal_1, u_marginal_2, ..., u_marginal_k). + u = (u_marginal_1, u_marginal_2, ..., u_marginal_k). - This preserves the intended total-dimensional QMC construction. + This preserves the intended total-dimensional QMC construction. """ if not isinstance(marginals, (list, tuple)) or len(marginals) == 0: raise ParameterError("ProductMeasure requires a nonempty list of marginals.") @@ -196,7 +192,9 @@ def __init__(self, sampler, marginals): self.parameters.append(statistic) def _marginal_statistic(self, marginal, marginal_index, statistic): - """Return a statistic or identify the marginal that does not provide it.""" + """Return a statistic or identify the marginal that does not provide + it. + """ try: return getattr(marginal, statistic) except AttributeError as error: @@ -226,18 +224,29 @@ def _concatenate_marginal_statistic(self, statistic): @property def mean(self): + """np.ndarray: The measure's mean, concatenated from each marginal's + `mean` in marginal order and cached after first access. + """ if self._mean_cache is None: self._mean_cache = self._concatenate_marginal_statistic("mean") return self._mean_cache @property def variance(self): + """np.ndarray: The measure's variance, concatenated from each + marginal's `variance` in marginal order and cached after first + access. + """ if self._variance_cache is None: self._variance_cache = self._concatenate_marginal_statistic("variance") return self._variance_cache @property def standard_deviation(self): + """np.ndarray: The measure's standard deviation, concatenated from + each marginal's `standard_deviation` in marginal order and cached + after first access. + """ if self._standard_deviation_cache is None: self._standard_deviation_cache = self._concatenate_marginal_statistic( "standard_deviation" @@ -286,12 +295,19 @@ def _compute_covariance(self): @property def covariance(self): + """Union[np.ndarray, scipy.sparse.spmatrix]: The measure's + block-diagonal covariance, built from each marginal's `covariance` + and cached after first access. Sparse if any marginal's covariance + is sparse, dense otherwise. + """ if self._covariance_cache is None: self._covariance_cache = self._compute_covariance() return self._covariance_cache def __repr__(self): - """Represent ProductMeasure without expanding marginal sparse matrices.""" + """Represent ProductMeasure without expanding marginal sparse + matrices. + """ lines = [f"{type(self).__name__} (AbstractTrueMeasure)"] for parameter in dict.fromkeys(self.parameters): if parameter == "marginals": @@ -314,12 +330,12 @@ def __repr__(self): @staticmethod def _expand_bounds(bounds, dimension, name): - """ - Expand a marginal's bounds so they have one row per output coordinate. + """Expand a marginal's bounds so they have one row per output + coordinate. - Some true measures store bounds as shape (1, 2), meaning the same - bound applies to all coordinates. Others store bounds as shape - (dimension, 2), meaning each coordinate has its own bound. + Some true measures store bounds as shape (1, 2), meaning the same bound + applies to all coordinates. Others store bounds as shape (dimension, + 2), meaning each coordinate has its own bound. ProductMeasure needs all marginal ranges stacked together, so every marginal range must be represented as shape (dimension, 2). @@ -338,20 +354,18 @@ def _expand_bounds(bounds, dimension, name): @property def _has_recursive_marginal(self): - """ - Check whether any marginal is itself recursively composed. + """Check whether any marginal is itself recursively composed. In QMCPy, a true measure can sometimes be built on top of another true measure. Sampling can still be handled by the recursive transform helper, but exact product weights in the final transformed space are - more delicate. For now, ProductMeasure only computes exact weights - when all marginals are direct true measures. + more delicate. For now, ProductMeasure only computes exact weights when + all marginals are direct true measures. """ return any(marginal.transform != marginal for marginal in self.marginals) def _split_blocks(self, x): - """ - Split an input array into marginal coordinate blocks. + """Split an input array into marginal coordinate blocks. The split always happens along the final axis, so this works for both ordinary samples with shape (n, d) and replicated samples with shape @@ -367,18 +381,16 @@ def _split_blocks(self, x): return np.split(x, self._split_indices, axis=-1) def _transform(self, x): - """ - Transform unit-cube samples into product-measure samples. + """Transform unit-cube samples into product-measure samples. - Steps - ----- + Steps ----- 1. Split the full unit-cube sample into marginal blocks. 2. Send each block to the matching marginal true measure. 3. Concatenate the transformed marginal outputs. This implements - T(u) = (T_1(u_1), T_2(u_2), ..., T_k(u_k)), + T(u) = (T_1(u_1), T_2(u_2), ..., T_k(u_k)), where each marginal T_j acts only on its own coordinate block. """ @@ -392,17 +404,16 @@ def _transform(self, x): return np.concatenate(transformed_blocks, axis=-1) def _weight(self, x): - """ - Compute the product density/weight for independent marginals. + """Compute the product density/weight for independent marginals. For independent components, the joint weight is the product of the marginal weights: w(x) = w_1(x_1) * w_2(x_2) * ... * w_k(x_k). - This method supports direct marginal true measures. Recursive - marginals are blocked for now because their final-space weights need - more careful handling. + This method supports direct marginal true measures. Recursive marginals + are blocked for now because their final-space weights need more careful + handling. """ if self._has_recursive_marginal: raise ParameterError( @@ -419,8 +430,7 @@ def _weight(self, x): return weight def _spawn(self, sampler, dimension): - """ - Spawn a new ProductMeasure with a new outer sampler. + """Spawn a new ProductMeasure with a new outer sampler. QMCPy's spawn mechanism creates new randomized copies of a sampler or true measure. ProductMeasure preserves the same marginal structure and diff --git a/qmcpy/true_measure/scipy_wrapper.py b/qmcpy/true_measure/scipy_wrapper.py index 9f528b175..67303c71b 100644 --- a/qmcpy/true_measure/scipy_wrapper.py +++ b/qmcpy/true_measure/scipy_wrapper.py @@ -1,3 +1,4 @@ +from typing import Union from .abstract_true_measure import AbstractTrueMeasure from ..util import DimensionError, ParameterError from ..discrete_distribution.abstract_discrete_distribution import ( @@ -60,8 +61,7 @@ def _custom_univariate_sanity_issues(dist, n_grid=64): class _MVNAdapter: - """ - Small adapter that turns a SciPy multivariate normal like object into + """Small adapter that turns a SciPy multivariate normal like object into something with a simple ``transform(u)`` interface. Idea: @@ -97,8 +97,7 @@ def __init__(self, mvn_like): self._chol = np.linalg.cholesky(cov) def transform(self, u): - """ - Take u in (0,1)^d and turn it into correlated normal samples. + """Take u in (0,1)^d and turn it into correlated normal samples. """ u = np.asarray(u, dtype=float) if u.shape[-1] != self.dim: @@ -119,8 +118,7 @@ def transform(self, u): return x_flat.reshape(z.shape) def logpdf(self, x): - """ - Forward to the SciPy logpdf, keeping shapes tidy. + """Forward to the SciPy logpdf, keeping shapes tidy. """ x = np.asarray(x, dtype=float) if x.shape[-1] != self.dim: @@ -134,12 +132,10 @@ def logpdf(self, x): class SciPyWrapper(AbstractTrueMeasure): - r""" - True measure that wraps SciPy style distributions. + r"""True measure that wraps SciPy style distributions. - This class keeps the original behavior of SciPyWrapper with - independent 1D marginals and adds an optional "joint" mode for - dependent distributions. + This class keeps the original behavior of SciPyWrapper with independent 1D + marginals and adds an optional "joint" mode for dependent distributions. Examples: Independent marginals from ``scipy.stats``: @@ -180,23 +176,23 @@ class SciPyWrapper(AbstractTrueMeasure): (4, 2) """ - def __init__(self, sampler, scipy_distribs): - """ - Parameters - ---------- - sampler : AbstractDiscreteDistribution - Low discrepancy or iid sampler in dimension d, living on [0,1)^d. - scipy_distribs : - One of the following: - - - A single SciPy 1D continuous frozen distribution. - - A list of such frozen distributions (independent marginals). - - A custom 1D distribution object with ``ppf`` and ``pdf`` or - ``logpdf`` methods. - - A joint object with: - * ``transform(u)`` method - * optional ``logpdf(x)`` method - * ``dim`` or ``dimension`` attribute (otherwise ``sampler.d``). + def __init__(self, sampler: AbstractDiscreteDistribution, scipy_distribs: Union[scipy.stats._distn_infrastructure.rv_continuous_frozen, list, object]) -> None: + """Wrap one or more SciPy distributions as a QMCPy true measure. + + Args: + sampler (AbstractDiscreteDistribution): Low discrepancy or iid + sampler in dimension d, living on [0,1)^d. + scipy_distribs (Union[scipy.stats._distn_infrastructure.rv_continuous_frozen, list, object]): One + of the following: + + - A single SciPy 1D continuous frozen distribution. + - A list of such frozen distributions (independent marginals). + - A custom 1D distribution object with ``ppf`` and ``pdf`` or + ``logpdf`` methods. + - A joint object with: + * ``transform(u)`` method + * optional ``logpdf(x)`` method + * ``dim`` or ``dimension`` attribute (otherwise ``sampler.d``). """ self.domain = np.array([[0.0, 1.0]]) @@ -234,8 +230,7 @@ def __init__(self, sampler, scipy_distribs): # ------------------------------------------------------------------ def _looks_like_joint(self, obj): - """ - Heuristic check to decide if the user passed a joint distribution. + """Heuristic check to decide if the user passed a joint distribution. We treat it as "joint" if: - it already has a ``transform(u)`` method, or @@ -257,8 +252,7 @@ def _looks_like_joint(self, obj): return False def _setup_joint(self, joint_obj): - """ - Configure the wrapper in "joint" mode. + """Configure the wrapper in "joint" mode. Either: - wrap a SciPy style multivariate normal in _MVNAdapter, or @@ -308,11 +302,10 @@ def _setup_joint(self, joint_obj): self.range = np.tile(np.array([-np.inf, np.inf]), (self.d, 1)) def _setup_marginals(self, scipy_distribs): - """ - Configure the wrapper in "independent marginals" mode. + """Configure the wrapper in "independent marginals" mode. - We accept a single frozen dist or a list, and we also allow - user defined 1D distributions that have the right methods. + We accept a single frozen dist or a list, and we also allow user + defined 1D distributions that have the right methods. """ rv_cont = scipy.stats._distn_infrastructure.rv_continuous_frozen @@ -373,14 +366,14 @@ def _setup_marginals(self, scipy_distribs): self.range = np.asarray(ranges) self._is_joint = False - assert len(self.sds) == self.d + if not (len(self.sds) == self.d): + raise AssertionError def _sanity_check_univariate(self, dist): - """ - Light sanity check for a custom 1D distribution. + """Light sanity check for a custom 1D distribution. - The goal is not to be perfect, just to catch obvious mistakes and - warn the user. We never raise here, only emit warnings. + The goal is not to be perfect, just to catch obvious mistakes and warn + the user. We never raise here, only emit warnings. We check on a grid 0.01..0.99 that: - ppf is finite and roughly increasing, @@ -444,11 +437,10 @@ def _sanity_check_univariate(self, dist): # ------------------------------------------------------------------ def _transform(self, x): - """ - Map unit cube samples to the physical space. + """Map unit cube samples to the physical space. - For joint mode we delegate to the joint object. - For marginal mode we call ``ppf`` dimension wise. + For joint mode we delegate to the joint object. For marginal mode we + call ``ppf`` dimension wise. """ x = np.asarray(x, dtype=float) @@ -461,8 +453,7 @@ def _transform(self, x): return t def _weight(self, x): - """ - Compute unnormalised density weights. + """Compute unnormalised density weights. - For joint distributions with logpdf we simply exp(logpdf). - For joint distributions with no density we return 1. @@ -501,8 +492,7 @@ def _weight(self, x): return rho def _spawn(self, sampler, dimension): - """ - Create a child true measure that shares the same distribution + """Create a child true measure that shares the same distribution configuration but uses a new sampler. We simply reuse the original ``scipy_distribs`` argument so the diff --git a/qmcpy/true_measure/student_t.py b/qmcpy/true_measure/student_t.py index 510b11388..dcf91a7f3 100644 --- a/qmcpy/true_measure/student_t.py +++ b/qmcpy/true_measure/student_t.py @@ -6,8 +6,7 @@ class _StudentTAdapter: - """ - Multivariate Student t adapter for SciPyWrapper. + """Multivariate Student t adapter for SciPyWrapper. - transform(u): sequential conditioning using univariate t conditionals - logpdf(x): forwarded to scipy.stats.multivariate_t (if available) @@ -103,11 +102,10 @@ def logpdf(self, x): class StudentT(SciPyWrapper): - """ - Convenience true measure: multivariate Student t. + """Convenience true measure: multivariate Student t. """ - def __init__(self, sampler, loc, shape, df): + def __init__(self, sampler, loc, shape, df) -> None: super().__init__( sampler=sampler, scipy_distribs=_StudentTAdapter(loc=loc, shape=shape, df=df), diff --git a/qmcpy/true_measure/student_t_copula.py b/qmcpy/true_measure/student_t_copula.py index 998831ced..4316c1f47 100644 --- a/qmcpy/true_measure/student_t_copula.py +++ b/qmcpy/true_measure/student_t_copula.py @@ -1,3 +1,8 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure +from typing import Union from .copula import ( AbstractCopula, _clip_unit_interval, @@ -13,19 +18,18 @@ class StudentTCopula(AbstractCopula): - r""" - Student-t copula transform with user supplied univariate marginals. + r"""Student-t copula transform with user supplied univariate marginals. - This TrueMeasure uses the same marginal workflow as ``GaussianCopula``, - but builds dependent uniforms through a multivariate Student-t copula with + This TrueMeasure uses the same marginal workflow as ``GaussianCopula``, but + builds dependent uniforms through a multivariate Student-t copula with correlation matrix ``correlation`` and degrees of freedom ``df``. - The transform uses the inverse Rosenblatt construction for the - multivariate Student-t distribution. This is equivalent in distribution to - the standard correlated-normal plus shared chi-square scaling construction, - but it only needs d deterministic uniforms from the base QMCPy sampler. - It is not the incorrect shortcut of applying univariate ``t.ppf``, a - Cholesky factor, and then univariate ``t.cdf``. + The transform uses the inverse Rosenblatt construction for the multivariate + Student-t distribution. This is equivalent in distribution to the standard + correlated-normal plus shared chi-square scaling construction, but it only + needs d deterministic uniforms from the base QMCPy sampler. It is not the + incorrect shortcut of applying univariate ``t.ppf``, a Cholesky factor, and + then univariate ``t.cdf``. Examples: >>> import numpy as np @@ -87,15 +91,17 @@ class StudentTCopula(AbstractCopula): "Weights will be treated as 1." ) - def __init__(self, sampler, marginals, correlation, df): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], marginals: list, correlation: np.ndarray, df: float) -> None: + r"""Initialize a StudentTCopula true measure. + Args: sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A sampler or transform whose range is the unit cube. marginals (list): Length d list of SciPy-like univariate distributions implementing a quantile function, called ``ppf`` in SciPy. - correlation (np.ndarray): d x d positive definite correlation matrix. + correlation (np.ndarray): d x d positive definite correlation + matrix. df (float): Positive Student-t degrees of freedom. """ self.parameters = ["marginals", "correlation", "df"] @@ -120,8 +126,7 @@ def _parse_df(self, df): return df def _dependent_t_samples(self, u): - """ - Map independent uniforms to a multivariate Student-t sample. + """Map independent uniforms to a multivariate Student-t sample. A direct scale-mixture construction would need d normal uniforms plus one extra chi-square uniform for the shared radial scale. Since diff --git a/qmcpy/true_measure/triangular.py b/qmcpy/true_measure/triangular.py index aa035ab9b..534a29eb6 100644 --- a/qmcpy/true_measure/triangular.py +++ b/qmcpy/true_measure/triangular.py @@ -5,15 +5,13 @@ class TriangularDistribution: - """ - Triangular distribution matching scipy.stats.triang behavior. + """Triangular distribution matching scipy.stats.triang behavior. - Support: [loc, loc + scale] - Mode: loc + c*scale, with 0 < c < 1 - Provides ppf and pdf for SciPyWrapper custom-marginal usage. + Support: [loc, loc + scale] Mode: loc + c*scale, with 0 < c < 1 Provides + ppf and pdf for SciPyWrapper custom-marginal usage. """ - def __init__(self, c=0.5, loc=0.0, scale=1.0): + def __init__(self, c=0.5, loc=0.0, scale=1.0) -> None: c = float(c) loc = float(loc) scale = float(scale) @@ -31,7 +29,15 @@ def __init__(self, c=0.5, loc=0.0, scale=1.0): self._b = loc + scale self._m = loc + c * scale - def pdf(self, x): + def pdf(self, x: np.ndarray) -> np.ndarray: + """Probability density function of the triangular distribution. + + Args: + x (np.ndarray): Points at which to evaluate the density. + + Returns: + np.ndarray: Density values, same shape as `x`. + """ x = np.asarray(x, dtype=float) a, m, b = self._a, self._m, self._b out = np.zeros_like(x, dtype=float) @@ -43,7 +49,15 @@ def pdf(self, x): out[right] = 2.0 * (b - x[right]) / ((b - a) * (b - m)) return out - def ppf(self, u): + def ppf(self, u: np.ndarray) -> np.ndarray: + """Percent point function (inverse CDF) of the triangular distribution. + + Args: + u (np.ndarray): Probabilities in `[0,1]` at which to evaluate the inverse CDF. + + Returns: + np.ndarray: Quantile values, same shape as `u`. + """ u = np.asarray(u, dtype=float) a, m, b = self._a, self._m, self._b Fm = (m - a) / (b - a) @@ -60,7 +74,7 @@ def ppf(self, u): class Triangular(SciPyWrapper): """Convenience TrueMeasure wrapper around TriangularDistribution.""" - def __init__(self, sampler, c=0.5, loc=0.0, scale=1.0): + def __init__(self, sampler, c=0.5, loc=0.0, scale=1.0) -> None: super().__init__( sampler=sampler, scipy_distribs=TriangularDistribution(c=c, loc=loc, scale=scale), diff --git a/qmcpy/true_measure/uniform.py b/qmcpy/true_measure/uniform.py index 0dc8c496e..1bd7f8e0d 100644 --- a/qmcpy/true_measure/uniform.py +++ b/qmcpy/true_measure/uniform.py @@ -1,3 +1,7 @@ +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from typing import Union from .abstract_true_measure import AbstractTrueMeasure from ..util import DimensionError, ParameterError from ..discrete_distribution import DigitalNetB2 @@ -6,8 +10,8 @@ class Uniform(AbstractTrueMeasure): - r""" - Uniform distribution, see [https://en.wikipedia.org/wiki/Continuous_uniform_distribution](https://en.wikipedia.org/wiki/Continuous_uniform_distribution). + r"""Uniform distribution, see + [https://en.wikipedia.org/wiki/Continuous_uniform_distribution](https://en.wikipedia.org/wiki/Continuous_uniform_distribution). Examples: >>> true_measure = Uniform(DigitalNetB2(2,seed=7),lower_bound=[0,.5],upper_bound=[2,3]) @@ -49,10 +53,12 @@ class Uniform(AbstractTrueMeasure): [1.37943573, 1.10241448, 1.13481488]]]) """ - def __init__(self, sampler, lower_bound=0, upper_bound=1): - r""" + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], lower_bound: Union[float, np.ndarray] = 0, upper_bound: Union[float, np.ndarray] = 1) -> None: + r"""Initialize a Uniform true measure. + Args: - sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): Either + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): + Either - a discrete distribution from which to transform samples, or - a true measure by which to compose a transform. @@ -94,7 +100,8 @@ def __init__(self, sampler, lower_bound=0, upper_bound=1): (self.a.reshape((self.d, 1)), self.b.reshape((self.d, 1))) ) super(Uniform, self).__init__() - assert self.a.shape == (self.d,) and self.b.shape == (self.d,) + if not (self.a.shape == (self.d,) and self.b.shape == (self.d,)): + raise AssertionError def _transform(self, x): return x * self.delta + self.a diff --git a/qmcpy/true_measure/uniform_triangle.py b/qmcpy/true_measure/uniform_triangle.py index 624b72229..0fd22a28b 100644 --- a/qmcpy/true_measure/uniform_triangle.py +++ b/qmcpy/true_measure/uniform_triangle.py @@ -1,18 +1,20 @@ +from typing import Union import numpy as np from ..util import DimensionError +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure from .scipy_wrapper import SciPyWrapper from ..discrete_distribution import DigitalNetB2 class _UniformTriangleAdapter: - """ - Uniform on triangle T = {(x, y): 0 <= y <= x <= 1} + """Uniform on triangle T = {(x, y): 0 <= y <= x <= 1} Exact transform: - u1, u2 ~ U(0, 1) - x = sqrt(u1) - y = u2 * x + u1, u2 ~ U(0, 1) x = sqrt(u1) y = u2 * x """ def __init__(self): @@ -50,10 +52,9 @@ def logpdf(self, x): class UniformTriangle(SciPyWrapper): - """ - Uniform distribution on the triangle {(x, y): 0 <= y <= x <= 1}. + """Uniform distribution on the triangle {(x, y): 0 <= y <= x <= 1}. - Example: + Examples: >>> tm = UniformTriangle(sampler=DigitalNetB2(2, seed=7)) >>> x = tm(4) >>> x.shape @@ -62,5 +63,12 @@ class UniformTriangle(SciPyWrapper): True """ - def __init__(self, sampler): + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure]) -> None: + """Initialize a UniformTriangle true measure. + + Args: + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A + 2-dimensional sampler generating unit-cube samples to be + transformed to the triangle. + """ super().__init__(sampler=sampler, scipy_distribs=_UniformTriangleAdapter()) diff --git a/qmcpy/true_measure/zero_inflated_exp_uniform.py b/qmcpy/true_measure/zero_inflated_exp_uniform.py index 11526556f..f1f5707c6 100644 --- a/qmcpy/true_measure/zero_inflated_exp_uniform.py +++ b/qmcpy/true_measure/zero_inflated_exp_uniform.py @@ -1,20 +1,24 @@ +from typing import Union import warnings import numpy as np from ..util import DimensionError, ParameterError +from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, +) +from ..true_measure.abstract_true_measure import AbstractTrueMeasure from .scipy_wrapper import SciPyWrapper class _ZeroInflatedExponential: - """ - One-dimensional zero-inflated exponential distribution. + """One-dimensional zero-inflated exponential distribution. This distribution has probability mass ``p_zero`` at zero and an exponential distribution with rate ``lam`` on positive values. - It implements ``ppf`` so it can be passed to ``SciPyWrapper`` as a - custom univariate marginal. + It implements ``ppf`` so it can be passed to ``SciPyWrapper`` as a custom + univariate marginal. """ def __init__(self, p_zero=0.4, lam=1.5): @@ -27,13 +31,11 @@ def __init__(self, p_zero=0.4, lam=1.5): self.lam = float(lam) def ppf(self, u): - """ - Generalized inverse CDF of the zero-inflated exponential. + """Generalized inverse CDF of the zero-inflated exponential. SciPyWrapper supplies one coordinate at a time. For example: - sampler output: (n, 1) - ppf input: (n,) + sampler output: (n, 1) ppf input: (n,) """ u = np.asarray(u, dtype=float) @@ -58,8 +60,7 @@ def ppf(self, u): class _DeprecatedZeroInflatedExpUniform2D: - """ - Adapter for the deprecated two-dimensional ``y_split`` construction. + """Adapter for the deprecated two-dimensional ``y_split`` construction. """ dim = 2 @@ -108,14 +109,12 @@ def logpdf(self, x): class ZeroInflatedExpUniform(SciPyWrapper): - """ - One-dimensional zero-inflated exponential true measure. + """One-dimensional zero-inflated exponential true measure. - The ``y_split`` keyword is retained temporarily for backward - compatibility with the deprecated two-dimensional construction. + The ``y_split`` keyword is retained temporarily for backward compatibility + with the deprecated two-dimensional construction. - Examples - -------- + Examples: Without replications: >>> from qmcpy import DigitalNetB2, ZeroInflatedExpUniform @@ -186,7 +185,20 @@ class ZeroInflatedExpUniform(SciPyWrapper): True """ - def __init__(self, sampler, p_zero=0.4, lam=1.5, y_split=None): + def __init__(self, sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], p_zero: float = 0.4, lam: float = 1.5, y_split: Union[None, float] = None) -> None: + """Initialize a ZeroInflatedExpUniform true measure. + + Args: + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): A + 1-dimensional sampler generating unit-cube samples to be + transformed. If `y_split` is set, a 2-dimensional sampler is + required instead (deprecated construction). + p_zero (float): Probability mass at zero. + lam (float): Rate parameter of the exponential component. + y_split (Union[None, float]): Deprecated. If set, uses the legacy + 2-dimensional zero-inflated exponential-uniform construction + instead of the 1-dimensional interface. + """ if y_split is not None: warnings.warn( "`y_split` is deprecated. The 2D zero-inflated " @@ -255,8 +267,7 @@ def __init__(self, sampler, p_zero=0.4, lam=1.5, y_split=None): ] def _compute_moments(self): - r""" - Closed-form mean and variance of the zero-inflated exponential. + r"""Closed-form mean and variance of the zero-inflated exponential. The distribution is a two component mixture that places probability mass $p = $ ``p_zero`` at $X = 0$ and, with probability $1 - p$, draws @@ -268,15 +279,14 @@ def _compute_moments(self): component raw moments [2]. Because the point mass sits exactly at zero, that component adds nothing to either moment, leaving - $$\mathbb{E}[X] = (1 - p)\,\frac{1}{\lambda}, \qquad - \mathbb{E}[X^2] = (1 - p)\,\frac{2}{\lambda^2}.$$ + $$\mathbb{E}[X] = (1 - p)\,\frac{1}{\lambda}, \qquad \mathbb{E}[X^2] = + (1 - p)\,\frac{2}{\lambda^2}.$$ The variance then follows from $\operatorname{Var}[X] = \mathbb{E}[X^2] - \mathbb{E}[X]^2$ (equivalently, the law of total variance [3]): - $$\operatorname{Var}[X] - = \frac{(1 - p)(1 + p)}{\lambda^2} - = \frac{1 - p^2}{\lambda^2}.$$ + $$\operatorname{Var}[X] = \frac{(1 - p)(1 + p)}{\lambda^2} = \frac{1 - + p^2}{\lambda^2}.$$ The measure is one dimensional, so ``mean`` and ``variance`` are returned as length-1 arrays for consistency with the other true diff --git a/qmcpy/util/abstraction_functions.py b/qmcpy/util/abstraction_functions.py index e2258b1cb..caf315090 100644 --- a/qmcpy/util/abstraction_functions.py +++ b/qmcpy/util/abstraction_functions.py @@ -2,9 +2,8 @@ from copy import copy -def _univ_repr(qmc_object, abc_class_name, attributes): - """ - Clean way to represent qmc_object data. +def _univ_repr(qmc_object: object, abc_class_name: str, attributes: list): + """Clean way to represent qmc_object data. Args: qmc_object (object): an qmc_object instance @@ -12,11 +11,11 @@ def _univ_repr(qmc_object, abc_class_name, attributes): attributes (list): list of attributes to include Returns: - s (str): string representation of this qmcpy object + str: string representation of this qmcpy object - Note: - print(qmc_object) is equivalent to print(qmc_object.__repr__()). - See an abstract classes __repr__ method for example call to this method. + Notes: + print(qmc_object) is equivalent to print(qmc_object.__repr__()). See an + abstract classes __repr__ method for example call to this method. """ with np.printoptions(precision=3, threshold=10): unique_attributes = [] diff --git a/qmcpy/util/data.py b/qmcpy/util/data.py index 7ed43f8ea..b5fa7d915 100644 --- a/qmcpy/util/data.py +++ b/qmcpy/util/data.py @@ -1,15 +1,22 @@ import gzip import pickle +from pathlib import Path +from typing import Union from ..util import _univ_repr class Data(object): + """Container for the state a stopping criterion accumulates while integrating. - def __init__(self, parameters): + Holds the parameters reported in the integration results and supports saving + to and loading from disk so a run can be resumed. + """ + + def __init__(self, parameters) -> None: self.parameters = parameters - def save(self, path, compress=False, overwrite=False): + def save(self, path: Union[str, Path], compress: bool = False, overwrite: bool = False) -> str: """Save this Data object to disk using pickle. Warning: @@ -18,18 +25,17 @@ def save(self, path, compress=False, overwrite=False): come from a trusted source. Args: - path (str or pathlib.Path): File path to save to. If + path (Union[str, Path]): File path to save to. If ``compress=True``, a ``.gz`` suffix is appended automatically when not already present. - compress (bool, optional): Gzip-compress the saved file. Defaults - to False. - overwrite (bool, optional): If False (default), raise - ``FileExistsError`` when the file already exists. If True, - overwrite any existing file. + compress (bool): Gzip-compress the saved file. Defaults to False. + overwrite (bool): If False (default), raise ``FileExistsError`` + when the file already exists. If True, overwrite any existing + file. Returns: - str: The final path the file was written to (may differ from - *path* when ``compress=True`` appends ``.gz``). + str: The final path the file was written to (may differ from *path* when + ``compress=True`` appends ``.gz``). Raises: FileExistsError: If the target path already exists and @@ -49,7 +55,7 @@ def save(self, path, compress=False, overwrite=False): return path @classmethod - def load(cls, path): + def load(cls, path: Union[str, Path]) -> "Data": """Load a Data object from disk. Warning: @@ -58,8 +64,8 @@ def load(cls, path): trusted source. Args: - path (str or pathlib.Path): Path to the saved file. Files ending - in ``.gz`` are decompressed automatically. + path (Union[str, Path]): Path to the saved file. Files ending in + ``.gz`` are decompressed automatically. Returns: Data: The loaded Data object. diff --git a/qmcpy/util/dig_shift_invar_ops.py b/qmcpy/util/dig_shift_invar_ops.py index 566176a87..d7459893e 100644 --- a/qmcpy/util/dig_shift_invar_ops.py +++ b/qmcpy/util/dig_shift_invar_ops.py @@ -1,14 +1,19 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union +if TYPE_CHECKING: + import torch + import numpy as np from .exceptions_warnings import ParameterError from .torch_numpy_ops import get_npt -def k4sumterm(x, t, cutoff=1e-8): - r""" - $$K_4(x) = \sum_{a=0}^{t-1} \frac{x_a}{2^{3a}}$$ +def k4sumterm(x: Union[np.ndarray, torch.Tensor], t: int, cutoff: float = 1e-8) -> Union[np.ndarray, torch.Tensor]: + r"""$$K_4(x) = \sum_{a=0}^{t-1} \frac{x_a}{2^{3a}}$$ - where $x_a$ is the bit at index $a$ in the binary expansion of $x$ - e.g. $x = 6$ with $t=3$ has $(x_0,x_1,x_2) = (1,1,0)$ + where $x_a$ is the bit at index $a$ in the binary expansion of $x$ e.g. $x + = 6$ with $t=3$ has $(x_0,x_1,x_2) = (1,1,0)$ Examples: >>> t = 3 @@ -31,11 +36,13 @@ def k4sumterm(x, t, cutoff=1e-8): [-1.14, -0.89, -0.89, -0.86]]) Args: - x (Union[np.ndarray torch.Tensor]): Integer arrays. + x (Union[np.ndarray, torch.Tensor]): Integer arrays. t (int): Number of bits in each integer. + cutoff (float): Stop accumulating terms once `1/2**(3*a)` falls + below this threshold. Returns: - y (Union[np.ndarray torch.Tensor]): The $K_4$ sum term. + Union[np.ndarray, torch.Tensor]: The $K_4$ sum term. """ total = 0.0 for a in range(0, t): @@ -65,17 +72,18 @@ def k4sumterm(x, t, cutoff=1e-8): } -def weighted_walsh_funcs(alpha, xb, t): - r""" - Weighted walsh functions +def weighted_walsh_funcs(alpha: int, xb: Union[np.ndarray, torch.Tensor], t: int) -> Union[np.ndarray, torch.Tensor]: + r"""Weighted walsh functions $$\sum_{k=0}^\infty \mathrm{wal}_k(x) 2^{-\mu_\alpha(k)}$$ - where $\mathrm{wal}_k$ is the $k^\text{th}$ Walsh function - and $\mu_\alpha$ is the Dick weight function which sums the first $\alpha$ largest indices of $1$ bits in the binary expansion of $k$ - e.g. $k=13=1101_2$ has 1-bit indexes $(4,3,1)$ so + where $\mathrm{wal}_k$ is the $k^\text{th}$ Walsh function and $\mu_\alpha$ + is the Dick weight function which sums the first $\alpha$ largest indices + of $1$ bits in the binary expansion of $k$ e.g. $k=13=1101_2$ has 1-bit + indexes $(4,3,1)$ so - $$\mu_1(k) = 4, \mu_2(k) = 4+3, \mu_3(k) = 4+3+1 = \mu_4(k) = \mu_5(k) = \dots$$ + $$\mu_1(k) = 4, \mu_2(k) = 4+3, \mu_3(k) = 4+3+1 = \mu_4(k) = \mu_5(k) = + \dots$$ Examples: >>> t = 3 @@ -114,29 +122,33 @@ def weighted_walsh_funcs(alpha, xb, t): Args: alpha (int): Weighted walsh functions order. - xb (Union[np.ndarray, torch.Tensor]): Integer points at which to evaluate the weighted Walsh function. + xb (Union[np.ndarray, torch.Tensor]): Integer points at which to + evaluate the weighted Walsh function. t (int): Number of bits in each integer in xb. - returns: - y (Union[np.ndarray, torch.Tensor]): Weighted Walsh function values. + Returns: + Union[np.ndarray, torch.Tensor]: Weighted Walsh function values. **References:** - 1. Dick, Josef. - "Walsh spaces containing smooth functions and quasi–Monte Carlo rules of arbitrary high order." - SIAM Journal on Numerical Analysis 46.3 (2008): 1519-1553. + 1. Dick, Josef. + "Walsh spaces containing smooth functions and quasi–Monte Carlo rules of arbitrary high order." + SIAM Journal on Numerical Analysis 46.3 (2008): 1519-1553. - 2. Dick, Josef. - "The decay of the Walsh coefficients of smooth functions." - Bulletin of the Australian Mathematical Society 80.3 (2009): 430-453. + 2. Dick, Josef. + "The decay of the Walsh coefficients of smooth functions." + Bulletin of the Australian Mathematical Society 80.3 (2009): 430-453. """ - assert isinstance(alpha, int) - assert alpha in WEIGHTEDWALSHFUNCSPOS, ( - "alpha = %d not in WEIGHTEDWALSHFUNCSPOS" % alpha - ) - assert alpha in WEIGHTEDWALSHFUNCSZEROS, ( - "alpha = %d not in WEIGHTEDWALSHFUNCSZEROS" % alpha - ) + if not (isinstance(alpha, int)): + raise AssertionError + if not (alpha in WEIGHTEDWALSHFUNCSPOS): + raise AssertionError( + "alpha = %d not in WEIGHTEDWALSHFUNCSPOS" % alpha + ) + if not (alpha in WEIGHTEDWALSHFUNCSZEROS): + raise AssertionError( + "alpha = %d not in WEIGHTEDWALSHFUNCSZEROS" % alpha + ) if isinstance(xb, np.ndarray): np_or_torch = np y = np.ones(xb.shape) @@ -153,9 +165,9 @@ def weighted_walsh_funcs(alpha, xb, t): return y -def to_bin(x, t): - r""" - Convert floating point representations of digital net samples in base $b=2$ to binary representations. +def to_bin(x: Union[np.ndarray, torch.Tensor], t: int) -> Union[np.ndarray, torch.Tensor]: + r"""Convert floating point representations of digital net samples in base + $b=2$ to binary representations. Examples: >>> xf = np.random.Generator(np.random.PCG64(7)).uniform(low=0,high=1,size=(5)) @@ -178,11 +190,14 @@ def to_bin(x, t): Args: - x (Union[np.ndarray, torch.Tensor]): floating point representation of samples. - t (int): number of bits in binary represtnations. Typically `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. + x (Union[np.ndarray, torch.Tensor]): floating point representation of + samples. + t (int): number of bits in binary representations. Typically `dnb2.t` + where `isinstance(dnb2,DigitalNetB2)`. Returns: - xb (Unioin[np.ndarray,torch.Tensor]): binary representation of samples with `dtype` either `np.uint64` or `torch.int64`. + Union[np.ndarray, torch.Tensor]: binary representation of samples with `dtype` either `np.uint64` or + `torch.int64`. """ npt = get_npt(x) if npt == np: @@ -201,9 +216,9 @@ def to_bin(x, t): raise ParameterError("x.dtype must be float or int, got %s" % str(x.dtype)) -def to_float(x, t): - r""" - Convert binary representations of digital net samples in base $b=2$ to floating point representations. +def to_float(x: Union[np.ndarray, torch.Tensor], t: int) -> Union[np.ndarray, torch.Tensor]: + r"""Convert binary representations of digital net samples in base $b=2$ to + floating point representations. Examples: >>> xb = np.arange(8,dtype=np.uint64) @@ -218,11 +233,13 @@ def to_float(x, t): tensor([0.0000, 0.1250, 0.2500, 0.3750, 0.5000, 0.6250, 0.7500, 0.8750]) Args: - x (Union[np.ndarray, torch.Tensor]): binary representation of samples with `dtype` either `np.uint64` or `torch.int64`. - t (int): number of bits in binary represtnations. Typically `dnb2.t` where `isinstance(dnb2,DigitalNetB2)`. + x (Union[np.ndarray, torch.Tensor]): binary representation of samples + with `dtype` either `np.uint64` or `torch.int64`. + t (int): number of bits in binary representations. Typically `dnb2.t` + where `isinstance(dnb2,DigitalNetB2)`. Returns: - xf (Unioin[np.ndarray,torch.Tensor]): floating point representation of samples. + Union[np.ndarray, torch.Tensor]: floating point representation of samples. """ npt = get_npt(x) if npt == np: # npt==torch @@ -241,9 +258,9 @@ def to_float(x, t): raise ParameterError("x.dtype must be torch.int64, got %s" % str(x.dtype)) -def bin_from_numpy_to_torch(xb): - r""" - Convert `numpy.uint64` to `torch.int64`, useful for converting binary samples from `DigitalNetB2` to torch representations. +def bin_from_numpy_to_torch(xb: Union[np.ndarray]) -> Union[torch.Tensor]: + r"""Convert `numpy.uint64` to `torch.int64`, useful for converting binary + samples from `DigitalNetB2` to torch representations. Examples: >>> xb = np.arange(8,dtype=np.uint64) @@ -253,13 +270,16 @@ def bin_from_numpy_to_torch(xb): tensor([0, 1, 2, 3, 4, 5, 6, 7]) Args: - xb (Union[np.ndarray]): binary representation of samples with `dtype=np.uint64` + xb (Union[np.ndarray]): binary representation of samples with + `dtype=np.uint64` Returns: - xbtorch (Unioin[torch.Tensor]): binary representation of samples with `dtype=torch.int64`. + Union[torch.Tensor]: binary representation of samples with `dtype=torch.int64`. """ - assert xb.dtype == np.uint64 - assert xb.max() <= (2**63 - 1), "require all xb < 2^63" + if not (xb.dtype == np.uint64): + raise AssertionError + if not (xb.max() <= (2**63 - 1)): + raise AssertionError("require all xb < 2^63") import torch return torch.from_numpy(xb.astype(np.int64)) diff --git a/qmcpy/util/exact_gpytorch_regression_model.py b/qmcpy/util/exact_gpytorch_regression_model.py index f3dc73864..d5104de8d 100644 --- a/qmcpy/util/exact_gpytorch_regression_model.py +++ b/qmcpy/util/exact_gpytorch_regression_model.py @@ -1,39 +1,65 @@ +from typing import Union import numpy as np import torch import gpytorch class ExactGPyTorchRegressionModel(gpytorch.models.ExactGP): + """Exact Gaussian process regression model backed by GPyTorch. + + Wraps ``gpytorch.models.ExactGP`` with fitting, chunked prediction, and + incremental data addition, optionally on the GPU. + """ + allowed_likelihood_types = ( gpytorch.likelihoods.GaussianLikelihood, gpytorch.likelihoods.GaussianLikelihoodWithMissingObs, gpytorch.likelihoods.FixedNoiseGaussianLikelihood, ) - def __init__(self, x_t, y_t, prior_mean, prior_cov, likelihood, use_gpu=False): + def __init__(self, x_t, y_t, prior_mean, prior_cov, likelihood, use_gpu=False) -> None: if isinstance(x_t, np.ndarray): x_t = torch.from_numpy(x_t) if isinstance(y_t, np.ndarray): y_t = torch.from_numpy(y_t) - assert x_t.ndim == 2 and y_t.ndim == 1 and len(x_t) == len(y_t) + if not (x_t.ndim == 2 and y_t.ndim == 1 and len(x_t) == len(y_t)): + raise AssertionError super(ExactGPyTorchRegressionModel, self).__init__(x_t, y_t, likelihood) - assert isinstance( + if not (isinstance( self.likelihood, ExactGPyTorchRegressionModel.allowed_likelihood_types - ) + )): + raise AssertionError self.mean_module, self.covar_module = prior_mean, prior_cov self.d = x_t.shape[1] self.use_gpu = use_gpu if self.use_gpu: - assert torch.cuda.is_available() + if not (torch.cuda.is_available()): + raise AssertionError self = self.cuda() self.likelihood = self.likelihood.cuda() - def forward(self, x): + def forward(self, x: torch.Tensor) -> gpytorch.distributions.MultivariateNormal: + """Evaluate the GP prior at the given inputs. + + Args: + x (torch.Tensor): Inputs of shape ``(n, d)``. + + Returns: + gpytorch.distributions.MultivariateNormal: Prior distribution at ``x``. + """ mean_x = self.mean_module(x) covar_x = self.covar_module(x) return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) - def fit(self, optimizer, mll, training_iter, verbose=0): + def fit(self, optimizer: torch.optim.Optimizer, mll: gpytorch.mlls.MarginalLogLikelihood, training_iter: int, verbose: int = 0): + """Fit the model hyperparameters by maximizing the marginal log likelihood. + + Args: + optimizer (torch.optim.Optimizer): Optimizer over the model parameters. + mll (gpytorch.mlls.MarginalLogLikelihood): Objective to maximize. + training_iter (int): Number of optimizer steps. + verbose (int): Print progress every ``verbose`` iterations; ``0`` is silent. + """ self.train() self.likelihood.train() if verbose: @@ -49,10 +75,23 @@ def fit(self, optimizer, mll, training_iter, verbose=0): print("\t\t\t%s %.2e" % (name.ljust(50, "."), val)) optimizer.step() - def predict(self, x, noise_const=0, chunk_size=2**15): + def predict(self, x: Union[np.ndarray, torch.Tensor], noise_const: float = 0, chunk_size: int = 2**15) -> tuple: + """Predict the posterior mean and standard deviation at new inputs. + + Inputs are processed in chunks so large batches do not exhaust memory. + + Args: + x (Union[np.ndarray, torch.Tensor]): Inputs of shape ``(n, d)``. + noise_const (float): Observation noise assumed at each new input. + chunk_size (int): Number of inputs evaluated per batch. + + Returns: + tuple: Posterior mean and standard deviation, each of length ``n``. + """ if isinstance(x, np.ndarray): x = torch.from_numpy(x) - assert x.ndim == 2 and x.shape[1] == self.d + if not (x.ndim == 2 and x.shape[1] == self.d): + raise AssertionError self.eval() self.likelihood.eval() n = len(x) @@ -81,17 +120,28 @@ def _predict_batch(self, x, noise): torch.cuda.empty_cache() return mean_post.numpy(), std_post.numpy() - def add_data(self, x_t_new, y_t_new): + def add_data(self, x_t_new: Union[np.ndarray, torch.Tensor], y_t_new: Union[np.ndarray, torch.Tensor]) -> "ExactGPyTorchRegressionModel": + """Add observations to the training set and condition the model on them. + + Args: + x_t_new (Union[np.ndarray, torch.Tensor]): New inputs of shape ``(n, d)``. + y_t_new (Union[np.ndarray, torch.Tensor]): New responses of length ``n``. + + Returns: + ExactGPyTorchRegressionModel: Fantasy model conditioned on the combined + training set. The receiver is left unchanged. + """ if isinstance(x_t_new, np.ndarray): x_t_new = torch.from_numpy(x_t_new) if isinstance(y_t_new, np.ndarray): y_t_new = torch.from_numpy(y_t_new) - assert ( + if not ( x_t_new.ndim == 2 and x_t_new.shape[1] == self.d and y_t_new.ndim == 1 and len(x_t_new) == len(y_t_new) - ) + ): + raise AssertionError if self.use_gpu: x_t_new, y_t_new = x_t_new.cuda(), y_t_new.cuda() fantasy_model = self.get_fantasy_model(x_t_new, y_t_new) diff --git a/qmcpy/util/exceptions_warnings.py b/qmcpy/util/exceptions_warnings.py index 663791b12..d0ab7edd1 100644 --- a/qmcpy/util/exceptions_warnings.py +++ b/qmcpy/util/exceptions_warnings.py @@ -7,30 +7,26 @@ class DimensionError(Exception): - """ - Class for raising error about dimension + """Class for raising error about dimension """ class DistributionCompatibilityError(Exception): - """ - Class for raising error about incompatible distribution + """Class for raising error about incompatible distribution """ class NotYetImplemented(Exception): - """ - Class for raising error when a component has been implemented yet + """Class for raising error when a component has been implemented yet """ class MethodImplementationError(Exception): - """ - Class for raising error when an abstract method has not been implemented - in the child class. + """Class for raising error when an abstract method has not been + implemented in the child class. """ - def __init__(self, subclass, method_name): + def __init__(self, subclass, method_name) -> None: s_f = ( "%s does not have an implementation of the %s method. " + "See superclass for method description." @@ -41,30 +37,25 @@ def __init__(self, subclass, method_name): class ParameterError(Exception): - """ - Class for raising error about input parameters + """Class for raising error about input parameters """ class ParameterWarning(Warning): - """ - Class for issuing warnings about unacceptable parameters + """Class for issuing warnings about unacceptable parameters """ class MaxSamplesWarning(Warning): - """ - Class for issuing warning about using maximum number of data samples + """Class for issuing warning about using maximum number of data samples """ class MaxLevelsWarning(Warning): - """ - Class for issuing warning about using maximum number of data samples + """Class for issuing warning about using maximum number of data samples """ class CubatureWarning(Warning): - """ - Class for issuing warnings throughout cubature algorithms + """Class for issuing warnings throughout cubature algorithms """ diff --git a/qmcpy/util/latnetbuilder_linker.py b/qmcpy/util/latnetbuilder_linker.py index a37536682..e35cac16e 100644 --- a/qmcpy/util/latnetbuilder_linker.py +++ b/qmcpy/util/latnetbuilder_linker.py @@ -2,20 +2,20 @@ import numpy as np -def latnetbuilder_linker(lnb_dir="./", out_dir="./", fout_prefix="lnb4qmcpy"): - """ +def latnetbuilder_linker(lnb_dir: str = "./", out_dir: str = "./", fout_prefix: str = "lnb4qmcpy") -> str: + """Convert a LatNet Builder output directory into a QMCPy generating vector or matrix. + Args: - lnb_dir (str): relative path to directory where `outputMachine.txt` is stored - e.g. 'my_lnb/poly_lat/' + lnb_dir (str): relative path to directory where `outputMachine.txt` is + stored e.g. 'my_lnb/poly_lat/' out_dir (str): relative path to directory where output should be stored e.g. 'my_lnb/poly_lat_qmcpy/' - fout_prefix (str): start of output file name. - e.g. 'my_poly_lat_vec' + fout_prefix (str): start of output file name. e.g. 'my_poly_lat_vec' Returns: - str: path to file which can be passed into QMCPy's Lattice or Sobol' in order to use - the linked latnetbuilder generating vector/matrix - e.g. 'my_poly_lat_vec.10.16.npy' + str: path to file which can be passed into QMCPy's Lattice or Sobol' in + order to use the linked latnetbuilder generating vector/matrix e.g. + 'my_poly_lat_vec.10.16.npy' Adapted from latnetbuilder parser: https://github.com/umontreal-simul/latnetbuilder/blob/master/python-wrapper/latnetbuilder/parse_output.py#L74 diff --git a/qmcpy/util/mlmc_test.py b/qmcpy/util/mlmc_test.py index 5e421f31b..ec417a7f5 100644 --- a/qmcpy/util/mlmc_test.py +++ b/qmcpy/util/mlmc_test.py @@ -1,27 +1,27 @@ import qmcpy as qp +from ..integrand.abstract_integrand import AbstractIntegrand import numpy as np def mlmc_test( - integrand, - n = 20000, - l = 8, - n_init = 200, - rmse_tols = np.array([.005, 0.01, 0.02, 0.05, 0.1]), - levels_min = 2, - levels_max = 10, + integrand: AbstractIntegrand, + n: int = 20000, + l: int = 8, + n_init: int = 200, + rmse_tols: np.ndarray = np.array([.005, 0.01, 0.02, 0.05, 0.1]), + levels_min: int = 2, + levels_max: int = 10, ): - r""" - Multilevel Monte Carlo test routine. + r"""Multilevel Monte Carlo test routine. Examples: >>> fo = qp.FinancialOption( ... sampler=qp.IIDStdUniform(seed=7), ... option = "ASIAN", ... asian_mean = "GEOMETRIC", - ... volatility = 0.2, - ... start_price = 100, - ... strike_price = 100, - ... interest_rate = 0.05, + ... volatility = 0.2, + ... start_price = 100, + ... strike_price = 100, + ... interest_rate = 0.05, ... t_final = 1) >>> print('Exact Value: %s'%fo.get_exact_value_inf_dim()) Exact Value: 5.546818633789201 @@ -43,12 +43,12 @@ def mlmc_test( gamma = 1.000000 (exponent for MLMC cost) MLMC complexity tests rmse_tol value mlmc_cost std_cost savings N_l - 5.000e-03 5.545e+00 3.339e+07 1.038e+08 3.11 8605392 1566846 559701 198886 70359 - 1.000e-02 5.539e+00 7.272e+06 1.243e+07 1.71 2009192 365451 130781 46623 - 2.000e-02 5.549e+00 1.827e+06 3.108e+06 1.70 503397 91821 33196 11736 - 5.000e-02 5.474e+00 2.324e+05 2.556e+05 1.10 71432 13143 4617 - 1.000e-01 5.466e+00 6.220e+04 6.389e+04 1.03 19477 3361 1225 - + 5.000e-03 5.545e+00 3.339e+07 1.038e+08 3.11 8605392 1566846 559701 198886 70359 + 1.000e-02 5.539e+00 7.272e+06 1.243e+07 1.71 2009192 365451 130781 46623 + 2.000e-02 5.549e+00 1.827e+06 3.108e+06 1.70 503397 91821 33196 11736 + 5.000e-02 5.474e+00 2.324e+05 2.556e+05 1.10 71432 13143 4617 + 1.000e-01 5.466e+00 6.220e+04 6.389e+04 1.03 19477 3361 1225 + Args: integrand (AbstractIntegrand): multilevel integrand n (int): number of samples for convergence tests @@ -76,7 +76,7 @@ def mlmc_test( cst = 0 integrand_spawn = integrand_spawns[ll] for j in range(1,101): - # evaluate integral at sampleing points samples + # Evaluate the integral at sampled points. samples = integrand_spawn.discrete_distrib.gen_samples(n=n/100) Pc,Pf = integrand_spawn.f(samples) dP = Pf-Pc @@ -160,5 +160,7 @@ def mlmc_test( mlmc_cost = sum(nl*cl) idx = np.minimum(len(var2),len(nl))-1 std_cost = var2[idx]*cl[-1] / ((1.-theta)*rmse_tols[i]**2) - print(' %-15.3e%-15.3e%-15.3e%-15.3e%-15.2f%s'\ - %(rmse_tols[i], p, mlmc_cost, std_cost, std_cost/mlmc_cost,''.join('%-13d'%nli for nli in nl))) + output = ' %-15.3e%-15.3e%-15.3e%-15.3e%-15.2f%s' \ + % (rmse_tols[i], p, mlmc_cost, std_cost, std_cost/mlmc_cost, + ''.join('%-13d' % nli for nli in nl)) + print(output.rstrip()) diff --git a/qmcpy/util/plot_functions.py b/qmcpy/util/plot_functions.py index 6c7c15ff4..4e722cfc2 100644 --- a/qmcpy/util/plot_functions.py +++ b/qmcpy/util/plot_functions.py @@ -1,38 +1,62 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union import numpy as np import os import qmcpy as qp +if TYPE_CHECKING: + from ..discrete_distribution.abstract_discrete_distribution import ( + AbstractDiscreteDistribution, + ) + from ..true_measure.abstract_true_measure import AbstractTrueMeasure + import matplotlib.figure + def plot_proj( - sampler, - n=64, - d_horizontal=1, - d_vertical=2, - math_ind=True, - marker_size=5, - figfac=5, - fig_title="Projection of Samples", - axis_pad=0, - want_grid=True, - font_family="sans-serif", - where_title=1, - **kwargs -): - """ + sampler: Union[AbstractDiscreteDistribution, AbstractTrueMeasure], + n: Union[int, list] = 64, + d_horizontal: Union[int, list] = 1, + d_vertical: Union[int, list] = 2, + math_ind: bool = True, + marker_size: float = 5, + figfac: float = 5, + fig_title: str = "Projection of Samples", + axis_pad: float = 0, + want_grid: bool = True, + font_family: str = "sans-serif", + where_title: float = 1, + **kwargs: dict +) -> matplotlib.figure.Figure: + """Plot two-dimensional projections of a point set. + Args: - sampler (DiscreteDistribution,TrueMeasure): The generator of samples to be plotted. - n (Union[int, list]): The number of samples or a list of samples(used for extensibility) to be plotted. - d_horizontal (Union[int, list]): The dimension or list of dimensions to be plotted on the horizontal axes. - d_vertical (Union[int, list]): The dimension or list of dimensions to be plotted on the vertical axes. - math_ind (bool): Setting to `True` will enable user to pass in math indices. + sampler (Union[AbstractDiscreteDistribution, AbstractTrueMeasure]): The generator of samples + to be plotted. + n (Union[int, list]): The number of samples or a list of samples(used + for extensibility) to be plotted. + d_horizontal (Union[int, list]): The dimension or list of dimensions to + be plotted on the horizontal axes. + d_vertical (Union[int, list]): The dimension or list of dimensions to + be plotted on the vertical axes. + math_ind (bool): Setting to `True` will enable user to pass in math + indices. marker_size (float): The marker size (typographic points are 1/72 in.). figfac (float): The figure size factor. fig_title (str): The title of the figure. - axis_pad (float): The padding of the axis so that points on the boundaries can be seen. + axis_pad (float): The padding of the axis so that points on the + boundaries can be seen. want_grid (bool): Setting to `True` will enable grid on the plot. font_family (str): The font family of the plot. - where_title (float): the position of the title on the plot. Default value is 1. - **kwargs (dict): Additional keyword arguments passed to `matplotlib.pyplot.scatter`. + where_title (float): the position of the title on the plot. Default + value is 1. + **kwargs (dict): Additional keyword arguments passed to + `matplotlib.pyplot.scatter`. + + Returns: + matplotlib.figure.Figure: The created figure. + matplotlib.axes.Axes: Array of subplot axes, one per + (`d_horizontal`, `d_vertical`) pair. """ try: import matplotlib.pyplot as plt diff --git a/qmcpy/util/shift_invar_ops.py b/qmcpy/util/shift_invar_ops.py index 35ba5cac8..9aa061cc9 100644 --- a/qmcpy/util/shift_invar_ops.py +++ b/qmcpy/util/shift_invar_ops.py @@ -1,3 +1,9 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union +if TYPE_CHECKING: + import torch + import numpy as np @@ -10,17 +16,19 @@ class Polynomial: >>> assert np.allclose(y,y_true,atol=1e-12) """ - def __init__(self, coeffs): - """ - Polynomial evaluation with Horner's rule + def __init__(self, coeffs: Union[list, np.ndarray, torch.Tensor]) -> None: + """Polynomial evaluation with Horner's rule Args: - coeffs (list or np.ndarray or torch.Tensor): vector of coefficients - e.g. coeffs = [a, b, c] corresponds to the quadratic polynomial a*x**2 + b*x + c + coeffs (Union[list, np.ndarray, torch.Tensor]): Vector of + coefficients, e.g., `coeffs = [a, b, c]` corresponds to the + quadratic polynomial `a*x**2 + b*x + c`. """ - assert isinstance(coeffs, list) + if not (isinstance(coeffs, list)): + raise AssertionError self.order = len(coeffs) - assert self.order >= 1 + if not (self.order >= 1): + raise AssertionError self.coeffs = coeffs def __call__(self, x): @@ -52,9 +60,8 @@ def __call__(self, x): } -def bernoulli_poly(n, x): - r""" - $n^\text{th}$ Bernoulli polynomial +def bernoulli_poly(n: int, x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + r"""$n^\text{th}$ Bernoulli polynomial Examples: >>> x = np.arange(6).reshape((2,3))/6 @@ -103,13 +110,16 @@ def bernoulli_poly(n, x): Args: n (int): Polynomial order. - x (Union[np.ndarray, torch.Tensor]): Points at which to evaluate the Bernoulli polynomial. + x (Union[np.ndarray, torch.Tensor]): Points at which to evaluate the + Bernoulli polynomial. Returns: - y (Union[np.ndarray, torch.Tensor]): Bernoulli polynomial values. + Union[np.ndarray, torch.Tensor]: Bernoulli polynomial values. """ - assert isinstance(n, int) - assert n in BERNOULLIPOLYSDICT, "n = %d not in BERNOULLIPOLYSDICT" % n + if not (isinstance(n, int)): + raise AssertionError + if not (n in BERNOULLIPOLYSDICT): + raise AssertionError("n = %d not in BERNOULLIPOLYSDICT" % n) bpoly = BERNOULLIPOLYSDICT[n] y = bpoly(x) return y diff --git a/qmcpy/util/stop_notebook.py b/qmcpy/util/stop_notebook.py index 1794c835f..3b1e08dc9 100644 --- a/qmcpy/util/stop_notebook.py +++ b/qmcpy/util/stop_notebook.py @@ -1,5 +1,17 @@ -def stop_notebook(query="Type 'yes' to continue running notebook"): - # This is a function to be able to stop a notebook when you run all cells +def stop_notebook(query: str = "Type 'yes' to continue running notebook"): + """Prompt at a notebook checkpoint and halt execution unless the user confirms. + + Placed between cells so that "Run All" pauses instead of running an + expensive section unattended. Any answer other than ``yes`` (case + insensitive) calls :func:`sys.exit`, which the notebook kernel reports as a + stopped cell rather than a traceback. + + Args: + query (str): Prompt shown to the user. + + Raises: + SystemExit: If the answer is not ``yes``. + """ keep_running = input(query) if keep_running.casefold() != "yes": import sys diff --git a/qmcpy/util/torch_numpy_ops.py b/qmcpy/util/torch_numpy_ops.py index 144e5719b..53618ef67 100644 --- a/qmcpy/util/torch_numpy_ops.py +++ b/qmcpy/util/torch_numpy_ops.py @@ -1,11 +1,30 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union +import types import numpy as np +if TYPE_CHECKING: + import torch + + +def get_npt(x: Union[np.ndarray, torch.Tensor]) -> types.ModuleType: + """Return the array backend module matching the input. + + Args: + x (Union[np.ndarray, torch.Tensor]): Array whose backend is wanted. + + Returns: + types.ModuleType: ``numpy`` for an ``np.ndarray``, otherwise ``torch``. -def get_npt(x): + Raises: + AssertionError: If ``x`` is neither an ``np.ndarray`` nor a ``torch.Tensor``. + """ if isinstance(x, np.ndarray): return np else: import torch - assert isinstance(x, torch.Tensor) + if not (isinstance(x, torch.Tensor)): + raise AssertionError return torch diff --git a/qmcpy/util/transforms.py b/qmcpy/util/transforms.py index 2c1ae0b22..1700ec991 100644 --- a/qmcpy/util/transforms.py +++ b/qmcpy/util/transforms.py @@ -1,51 +1,144 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union +import types import numpy as np import scipy.special from .torch_numpy_ops import get_npt +if TYPE_CHECKING: + import torch + EPS64 = float(np.finfo(np.float64).eps) -def insert_batch_dims(param, ndims, k): +def insert_batch_dims(param: Union[np.ndarray, torch.Tensor], ndims: int, k: int) -> Union[np.ndarray, torch.Tensor]: + """Insert singleton dimensions into a parameter so it broadcasts against batched inputs. + + Args: + param (Union[np.ndarray, torch.Tensor]): Parameter to reshape. + ndims (int): Number of singleton dimensions to insert. + k (int): Position at which to insert them. + + Returns: + Union[np.ndarray, torch.Tensor]: ``param`` with ``ndims`` singleton axes + inserted after its first ``k`` axes. + """ ones = [1] * ndims return param.reshape(list(param.shape[:k]) + ones + list(param.shape[k:])) -def tf_exp(x): +def tf_exp(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Exponential transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``exp(x)``. + """ npt = get_npt(x) return npt.exp(x) -def tf_exp_inv(x): +def tf_exp_inv(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Inverse of the exponential transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``log(x)``. + """ npt = get_npt(x) return npt.log(x) -def tf_exp_eps(x): +def tf_exp_eps(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Exponential transform offset by machine epsilon. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``exp(x) + eps``, kept strictly positive. + """ return tf_exp(x) + EPS64 -def tf_exp_eps_inv(x): +def tf_exp_eps_inv(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Inverse of the epsilon-offset exponential transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``log(x - eps)``. + """ return tf_exp_inv(x - EPS64) -def tf_square(x): +def tf_square(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Square transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``x**2``. + """ return x**2 -def tf_square_inv(x): +def tf_square_inv(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Inverse of the square transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``sqrt(x)``. + """ npt = get_npt(x) return npt.sqrt(x) -def tf_square_eps(x): +def tf_square_eps(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Square transform offset by machine epsilon. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``x**2 + eps``, kept strictly positive. + """ return tf_square(x) + EPS64 -def tf_square_eps_inv(x): +def tf_square_eps_inv(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Inverse of the epsilon-offset square transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``sqrt(x - eps)``. + """ return tf_square_inv(x - EPS64) -def tf_explinear(x): +def tf_explinear(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Exponential-linear (softplus) transform. + + Behaves like ``exp(x)`` for small ``x`` and like ``x`` for large ``x``, so it + maps the real line to the positive reals without overflowing. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``log(1 + exp(x))``, computed stably. + """ npt = get_npt(x) if npt == np: return -scipy.special.log_expit(-x) @@ -53,35 +146,88 @@ def tf_explinear(x): return -npt.nn.functional.logsigmoid(-x) -def tf_explinear_inv(x): +def tf_explinear_inv(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Inverse of the exponential-linear transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``log(expm1(x))``, falling back to ``x`` once ``x >= 34`` + where the two agree to machine precision. + """ npt = get_npt(x) return npt.where(x < 34, npt.log(npt.expm1(x)), x) -def tf_explinear_eps(x): +def tf_explinear_eps(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Exponential-linear transform offset by machine epsilon. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``tf_explinear(x) + eps``, kept strictly positive. + """ return tf_explinear(x) + EPS64 -def tf_explinear_eps_inv(x): +def tf_explinear_eps_inv(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Inverse of the epsilon-offset exponential-linear transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``tf_explinear_inv(x - eps)``. + """ return tf_explinear_inv(x - EPS64) -def tf_identity(x): +def tf_identity(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]: + """Identity transform. + + Args: + x (Union[np.ndarray, torch.Tensor]): Input values. + + Returns: + Union[np.ndarray, torch.Tensor]: ``x`` unchanged. + """ return x def parse_assign_param( - pname, - param, - shape_param, - requires_grad_param, - tfs_param, - endsize_ops, - constraints, - torchify, - npt, - nptkwargs, -): + pname: str, + param: Union[float, np.ndarray, torch.Tensor], + shape_param: list, + requires_grad_param: bool, + tfs_param: tuple, + endsize_ops: list, + constraints: list, + torchify: bool, + npt: types.ModuleType, + nptkwargs: dict, +) -> tuple: + """Validate and normalize one kernel parameter, returning it in array form. + + A scalar is broadcast to ``shape_param``; an array-like is converted to the + backend array type and checked against the supplied constraints. + + Args: + pname (str): Parameter name, used in error messages. + param (Union[float, np.ndarray, torch.Tensor]): Value to normalize. + shape_param (list): Target shape used when ``param`` is a scalar. + requires_grad_param (bool): Whether the torch parameter requires a gradient. + tfs_param (tuple): Pair of forward and inverse transforms for this parameter. + endsize_ops (list): Permitted sizes for the trailing dimension. + constraints (list): Constraints the parameter must satisfy. + torchify (bool): Return a ``torch.Tensor`` rather than an ``np.ndarray``. + npt (types.ModuleType): Array backend, either ``numpy`` or ``torch``. + nptkwargs (dict): Backend keyword arguments such as ``dtype`` and ``device``. + + Returns: + tuple: The normalized parameter, its shape, and its transformed value. + """ if np.isscalar(param): param = param * npt.ones(shape_param, **nptkwargs) else: @@ -89,35 +235,46 @@ def parse_assign_param( if not isinstance(param, npt.Tensor): param = npt.tensor(param) param = npt.atleast_1d(param) - assert isinstance(param, npt.Tensor), ( - "%s must be a scalar or torch.Tensor" % pname - ) + if not (isinstance(param, npt.Tensor)): + raise AssertionError( + "%s must be a scalar or torch.Tensor" % pname + ) else: if not isinstance(param, npt.ndarray): param = npt.array(param) param = npt.atleast_1d(param) - assert isinstance(param, npt.ndarray), ( - "%s must be a scalar or np.ndarray" % pname - ) + if not (isinstance(param, npt.ndarray)): + raise AssertionError( + "%s must be a scalar or np.ndarray" % pname + ) shape_param = list(param.shape) - assert len(shape_param) >= 1, "invalid shape_%s = %s" % (pname, str(shape_param)) - assert len(tfs_param) == 2, "tfs_scale should be a tuple of length 2" - assert callable(tfs_param[0]), "tfs_scale[0] should be a callable e.g. torch.log" - assert callable(tfs_param[1]), "tfs_scale[1] should be a callable e.g. torch.exp" + if not (len(shape_param) >= 1): + raise AssertionError("invalid shape_%s = %s" % (pname, str(shape_param))) + if not (len(tfs_param) == 2): + raise AssertionError("tfs_scale should be a tuple of length 2") + if not (callable(tfs_param[0])): + raise AssertionError("tfs_scale[0] should be a callable e.g. torch.log") + if not (callable(tfs_param[1])): + raise AssertionError("tfs_scale[1] should be a callable e.g. torch.exp") raw_param = tfs_param[0](param) if torchify: - assert isinstance(requires_grad_param, bool) + if not (isinstance(requires_grad_param, bool)): + raise AssertionError if requires_grad_param: raw_param = 1.0 * raw_param raw_param = npt.nn.Parameter(raw_param, requires_grad=requires_grad_param) - assert shape_param[-1] in endsize_ops, "%s not in %s" % ( - str(shape_param[-1]), - str(endsize_ops), - ) + if not (shape_param[-1] in endsize_ops): + raise AssertionError("%s not in %s" % ( + str(shape_param[-1]), + str(endsize_ops), + )) if "POSITIVE" in constraints: - assert (param > 0).all(), "%s must be positive" % pname + if not ((param > 0).all()): + raise AssertionError("%s must be positive" % pname) if "NON-NEGATIVE" in constraints: - assert (param >= 0).all(), "%s must be non-negative" % pname + if not ((param >= 0).all()): + raise AssertionError("%s must be non-negative" % pname) if "INTEGER" in constraints: - assert (param % 1 == 0).all(), "%s must be integers" % pname + if not ((param % 1 == 0).all()): + raise AssertionError("%s must be integers" % pname) return raw_param diff --git a/scripts/add_docstring_arg_types.py b/scripts/add_docstring_arg_types.py new file mode 100644 index 000000000..06f2f85ec --- /dev/null +++ b/scripts/add_docstring_arg_types.py @@ -0,0 +1,666 @@ +#!/usr/bin/env python3 +"""Synchronize Google-style docstring types from Python annotations. + +This helper is intentionally conservative: it rewrites existing ``Args:`` +entries for public functions and methods only when the corresponding argument +has an explicit annotation in the signature. With ``--include-outputs``, it +also updates existing ``Returns:`` and ``Yields:`` descriptions from return +annotations. It does not infer types from implementation code and it does not +invent missing descriptions or sections. +""" +from __future__ import annotations + +import argparse +import ast +import re +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + + +SECTION_HEADER = re.compile(r"^\s*[A-Z][A-Za-z]*(?: [A-Z][A-Za-z]*)*:\s*$") +ARG_ENTRY = re.compile( + r"^(?P\s*)" + r"(?P\*{0,2}[A-Za-z_][A-Za-z0-9_]*)" + r"\s*" + r"(?:\((?P[^)]*)\))?" + r"\s*:\s*" + r"(?P.*)$" +) +OUTPUT_ENTRY = re.compile( + r"^(?P\s*)(?P[^:]+):\s*(?P.*)$" +) +YIELD_CONTAINER_NAMES = { + "AsyncGenerator", + "AsyncIterator", + "Generator", + "Iterable", + "Iterator", +} + + +@dataclass +class Update: + path: Path + line: int + function: str + argument: str + annotation: str + previous_type: str | None + section: str = "Args" + + +@dataclass +class Skip: + path: Path + line: int + function: str + reason: str + + +@dataclass +class FileResult: + path: Path + updates: list[Update] + skips: list[Skip] + changed: bool + + +def doc_node(node: ast.AST) -> ast.Constant | None: + """Return the string-literal node holding ``node``'s docstring, if any. + + Args: + node (ast.AST): Node whose docstring literal is wanted. + + Returns: + ast.Constant | None: The docstring node, or ``None`` when absent. + """ + body = getattr(node, "body", None) + if ( + body + and isinstance(body[0], ast.Expr) + and isinstance(body[0].value, ast.Constant) + and isinstance(body[0].value.value, str) + ): + return body[0].value + return None + + +def iter_public_functions(tree: ast.Module): + """Yield public module functions and methods from public classes. + + Args: + tree (ast.Module): Parsed module to walk. + + Yields: + ast.FunctionDef | ast.AsyncFunctionDef: Each public function or method. + """ + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if not node.name.startswith("_"): + yield node, node.name + elif isinstance(node, ast.ClassDef) and not node.name.startswith("_"): + for sub in node.body: + if not isinstance(sub, (ast.FunctionDef, ast.AsyncFunctionDef)): + continue + if sub.name == "__init__" or not sub.name.startswith("_"): + yield sub, f"{node.name}.{sub.name}" + + +def _annotation_text(source: str, annotation: ast.AST | None) -> str | None: + """Return the source spelling of a type annotation.""" + if annotation is None: + return None + if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + return annotation.value + text = ast.get_source_segment(source, annotation) + if text is not None: + text = text.strip() + if ( + len(text) >= 2 + and text[0] in {"'", '"'} + and text[-1] == text[0] + ): + try: + value = ast.literal_eval(text) + except (SyntaxError, ValueError): + return text + if isinstance(value, str): + return value + # ``ast.unparse`` turns a multiline annotation into a safe, single-line + # representation for a Google-style argument entry. + return ast.unparse(annotation) + + +def _argument_annotations(node: ast.FunctionDef | ast.AsyncFunctionDef, source: str): + """Map argument names to explicit annotation text.""" + annotations = {} + args = ( + list(node.args.posonlyargs) + + list(node.args.args) + + list(node.args.kwonlyargs) + ) + for arg in args: + if arg.arg in {"self", "cls"}: + continue + annotation = _annotation_text(source, arg.annotation) + if annotation is not None: + annotations[arg.arg] = annotation + if node.args.vararg is not None: + annotation = _annotation_text(source, node.args.vararg.annotation) + if annotation is not None: + annotations[node.args.vararg.arg] = annotation + if node.args.kwarg is not None: + annotation = _annotation_text(source, node.args.kwarg.annotation) + if annotation is not None: + annotations[node.args.kwarg.arg] = annotation + return annotations + + +def line_without_ending(line: str) -> tuple[str, str]: + """Split a line into content and original line ending. + + Args: + line (str): Source line, with or without a line ending. + + Returns: + tuple[str, str]: The content and the line ending that was removed. + """ + if line.endswith("\r\n"): + return line[:-2], "\r\n" + if line.endswith("\n"): + return line[:-1], "\n" + return line, "" + + +def find_section( + lines: list[str], start: int, end: int, name: str +) -> tuple[int, int] | None: + """Return the header and end indexes for a Google-style section. + + Args: + lines (list[str]): Docstring lines to search. + start (int): First index to consider. + end (int): Index one past the last to consider. + name (str): Section header to look for, such as ``"Args"``. + + Returns: + tuple[int, int] | None: Header and end indexes, or ``None`` when the + section is absent. + """ + header_line = None + header_indent = None + for i in range(start, end + 1): + content, _ = line_without_ending(lines[i]) + if content.strip() == f"{name}:": + header_line = i + header_indent = len(content) - len(content.lstrip()) + break + if header_line is None or header_indent is None: + return None + + section_end = end + for i in range(header_line + 1, end + 1): + content, _ = line_without_ending(lines[i]) + stripped = content.strip() + if not stripped: + continue + indent = len(content) - len(content.lstrip()) + if indent <= header_indent and SECTION_HEADER.match(content): + section_end = i - 1 + break + return header_line, section_end + + +def find_args_section( + lines: list[str], start: int, end: int +) -> tuple[int, int] | None: + """Return ``(args_line, section_end)`` indexes for a Google Args section. + + Args: + lines (list[str]): Docstring lines to search. + start (int): First index to consider. + end (int): Index one past the last to consider. + + Returns: + tuple[int, int] | None: Header and end indexes, or ``None`` when there + is no Args section. + """ + return find_section(lines, start, end, "Args") + + +def _yield_annotation_text(source: str, annotation: ast.AST) -> str | None: + """Extract the yielded item type from a standard iterator annotation.""" + if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + try: + annotation = ast.parse(annotation.value, mode="eval").body + except SyntaxError: + return None + if not isinstance(annotation, ast.Subscript): + return None + value = annotation.value + if isinstance(value, ast.Name): + container = value.id + elif isinstance(value, ast.Attribute): + container = value.attr + else: + return None + if container not in YIELD_CONTAINER_NAMES: + return None + + item = annotation.slice + if container in {"Generator", "AsyncGenerator"} and isinstance(item, ast.Tuple): + if not item.elts: + return None + item = item.elts[0] + return _annotation_text(source, item) + + +def looks_like_type(text: str) -> bool: + """Return whether text is syntactically usable as a type expression. + + Args: + text (str): Candidate type expression. + + Returns: + bool: Whether ``text`` parses as a Python expression. + """ + try: + ast.parse(text, mode="eval") + except SyntaxError: + return False + return True + + +def _update_output_section( + path: Path, + lines: list[str], + node: ast.FunctionDef | ast.AsyncFunctionDef, + function: str, + section_name: str, + annotation: str, + overwrite_existing: bool, +) -> tuple[list[Update], list[Skip]]: + """Add a signature-derived type to an existing output description.""" + dnode = doc_node(node) + if dnode is None or dnode.end_lineno is None: + return [], [Skip(path, node.lineno, function, "missing docstring")] + + section = find_section( + lines, + dnode.lineno - 1, + dnode.end_lineno - 1, + section_name, + ) + if section is None: + return [], [ + Skip( + path, + dnode.lineno, + function, + f"missing {section_name} section for annotated output", + ) + ] + + header_content, _ = line_without_ending(lines[section[0]]) + header_indent = len(header_content) - len(header_content.lstrip()) + for i in range(section[0] + 1, section[1] + 1): + content, ending = line_without_ending(lines[i]) + if not content.strip(): + continue + indent = len(content) - len(content.lstrip()) + if indent <= header_indent: + continue + + match = OUTPUT_ENTRY.match(content) + previous_type = None + description = content.strip() + entry_indent = content[:indent] + if match is not None and looks_like_type(match.group("type").strip()): + previous_type = match.group("type").strip() + if not overwrite_existing: + return [], [] + description = match.group("description").lstrip() + entry_indent = match.group("indent") + + suffix = f" {description}" if description else "" + replacement = f"{entry_indent}{annotation}:{suffix}{ending}" + if replacement == lines[i]: + return [], [] + lines[i] = replacement + slot = "yield" if section_name == "Yields" else "return" + return [ + Update( + path=path, + line=i + 1, + function=function, + argument=slot, + annotation=annotation, + previous_type=previous_type, + section=section_name, + ) + ], [] + + return [], [ + Skip(path, node.lineno, function, f"empty {section_name} section") + ] + + +def _update_args_section( + path: Path, + lines: list[str], + node: ast.FunctionDef | ast.AsyncFunctionDef, + function: str, + annotations: dict[str, str], + overwrite_existing: bool, +) -> tuple[list[Update], list[Skip]]: + """Add annotation text to matching ``Args:`` entries.""" + dnode = doc_node(node) + if dnode is None or dnode.end_lineno is None: + return [], [Skip(path, node.lineno, function, "missing docstring")] + + section = find_args_section(lines, dnode.lineno - 1, dnode.end_lineno - 1) + if section is None: + return [], [Skip(path, dnode.lineno, function, "missing Args section")] + + updates = [] + seen = set() + _, section_end = section + for i in range(section[0] + 1, section_end + 1): + content, ending = line_without_ending(lines[i]) + match = ARG_ENTRY.match(content) + if match is None: + continue + display_name = match.group("name") + argument = display_name.lstrip("*") + if argument not in annotations: + continue + seen.add(argument) + previous_type = match.group("type") + if previous_type is not None and not overwrite_existing: + continue + annotation = annotations[argument] + description = match.group("description").lstrip() + suffix = f" {description}" if description else "" + replacement = ( + f"{match.group('indent')}{display_name} ({annotation}):{suffix}{ending}" + ) + if replacement == lines[i]: + continue + lines[i] = replacement + updates.append( + Update( + path=path, + line=i + 1, + function=function, + argument=argument, + annotation=annotation, + previous_type=previous_type, + ) + ) + + skips = [ + Skip( + path, + node.lineno, + function, + f"missing Args entry for annotated argument `{name}`", + ) + for name in sorted(set(annotations) - seen) + ] + return updates, skips + + +def update_file( + path: Path, + check: bool = False, + overwrite_existing: bool = False, + include_outputs: bool = False, +) -> FileResult: + """Update Google-style types in one Python file. + + Args: + path (Path): Python file to update. + check (bool): Report what would change without writing. + overwrite_existing (bool): Replace types already present rather than only + filling in missing ones. + include_outputs (bool): Also update the Returns section. + + Returns: + FileResult: Counts of updates made and entries skipped. + """ + source = path.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(path)) + lines = source.splitlines(keepends=True) + updates = [] + skips = [] + + for node, function in iter_public_functions(tree): + annotations = _argument_annotations(node, source) + if annotations: + node_updates, node_skips = _update_args_section( + path=path, + lines=lines, + node=node, + function=function, + annotations=annotations, + overwrite_existing=overwrite_existing, + ) + updates.extend(node_updates) + skips.extend(node_skips) + + if not include_outputs or node.name == "__init__" or node.returns is None: + continue + return_annotation = _annotation_text(source, node.returns) + if return_annotation in {None, "None", "NoneType"}: + continue + + dnode = doc_node(node) + if dnode is None or dnode.end_lineno is None: + skips.append(Skip(path, node.lineno, function, "missing docstring")) + continue + doc_start = dnode.lineno - 1 + doc_end = dnode.end_lineno - 1 + yields_section = find_section(lines, doc_start, doc_end, "Yields") + section_name = "Yields" if yields_section is not None else "Returns" + output_annotation = return_annotation + if section_name == "Yields": + output_annotation = _yield_annotation_text(source, node.returns) + if output_annotation is None: + skips.append( + Skip( + path, + node.lineno, + function, + "cannot derive yielded item type from return annotation", + ) + ) + continue + node_updates, node_skips = _update_output_section( + path=path, + lines=lines, + node=node, + function=function, + section_name=section_name, + annotation=output_annotation, + overwrite_existing=overwrite_existing, + ) + updates.extend(node_updates) + skips.extend(node_skips) + + changed = bool(updates) + if changed and not check: + path.write_text("".join(lines), encoding="utf-8") + return FileResult(path=path, updates=updates, skips=skips, changed=changed) + + +def _changed_files(ref: str) -> list[Path]: + """Return Python files changed relative to ``ref`` using ``git diff``.""" + result = subprocess.run( + ["git", "diff", "--name-only", "--diff-filter=ACMR", ref, "--", "*.py"], + capture_output=True, + text=True, + check=True, + ) + return [Path(name) for name in result.stdout.splitlines()] + + +def _is_under(path: Path, root: Path) -> bool: + """Return whether a relative or absolute path is under root.""" + try: + path.resolve().relative_to(root.resolve()) + except ValueError: + return False + return True + + +def python_files( + paths: list[str], diff_ref: str | None, root: str | None = None +) -> list[Path]: + """Collect Python files from paths, or from ``git diff`` when requested. + + Args: + paths (list[str]): Files or directories to collect from. + diff_ref (str | None): Git ref to diff against instead of using ``paths``. + root (str | None): Repository root for the diff; defaults to the cwd. + + Returns: + list[Path]: Python files to process, in sorted order. + """ + if diff_ref is not None: + candidates = _changed_files(diff_ref) + else: + candidates = [Path(p) for p in (paths or ["qmcpy"])] + + if root is not None and diff_ref is not None: + root_path = Path(root) + candidates = [path for path in candidates if _is_under(path, root_path)] + + files = [] + for path in candidates: + if path.is_dir(): + files.extend(sorted(path.rglob("*.py"))) + elif path.suffix == ".py" and path.exists(): + files.append(path) + return sorted(dict.fromkeys(files)) + + +def _parse_args(argv: list[str]) -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "paths", + nargs="*", + help="Python files or directories to update. Defaults to qmcpy.", + ) + parser.add_argument( + "--diff", + metavar="REF", + help="Update Python files reported by `git diff --name-only REF -- '*.py'`.", + ) + parser.add_argument( + "--check", + action="store_true", + help="Report files that would change without writing them.", + ) + parser.add_argument( + "--overwrite-existing", + action="store_true", + help="Replace existing Google Args types with signature annotations.", + ) + parser.add_argument( + "--include-outputs", + action="store_true", + help="Also update existing Returns and Yields descriptions.", + ) + parser.add_argument( + "--root", + help="Restrict files selected by --diff to this directory.", + ) + parser.add_argument( + "--quiet", + action="store_true", + help="Only print the final summary.", + ) + return parser.parse_args(argv) + + +def main(argv: list[str]) -> int: + """Run the command-line interface. + + Args: + argv (list[str]): Command-line arguments, excluding the program name. + + Returns: + int: Process exit status; ``0`` on success. + """ + args = _parse_args(argv) + try: + files = python_files(args.paths, args.diff, root=args.root) + except subprocess.CalledProcessError as exc: + print(f"git diff failed: {exc}", file=sys.stderr) + return 2 + + if not files: + print("No Python files to inspect.") + return 0 + + results = [] + had_parse_error = False + for path in files: + try: + result = update_file( + path, + check=args.check, + overwrite_existing=args.overwrite_existing, + include_outputs=args.include_outputs, + ) + except SyntaxError as exc: + had_parse_error = True + print(f"{path}: skipped syntax error: {exc}", file=sys.stderr) + continue + results.append(result) + + updates = [update for result in results for update in result.updates] + skips = [skip for result in results for skip in result.skips] + if not args.quiet and (updates or skips): + print() + for update in updates: + action = "would update" if args.check else "updated" + old = ( + "" + if update.previous_type is None + else f" replacing `{update.previous_type}`" + ) + print( + f" - {update.path}:{update.line}: {action} " + f"{update.function}.{update.argument} ({update.annotation}){old}" + ) + for skip in skips: + print(f" - {skip.path}:{skip.line}: skipped {skip.function}: {skip.reason}") + + args_updates = [update for update in updates if update.section == "Args"] + output_updates = [update for update in updates if update.section != "Args"] + changed_files = sum(1 for result in results if result.changed) + verb = "would change" if args.check else "changed" + print( + f" - {len(files)} file(s) inspected; {len(args_updates)} Args type update(s); " + f"{len(output_updates)} output type update(s); " + f"{changed_files} file(s) {verb}." + ) + + if changed_files == 0: + print(f"clean (0 of {len(files)} files)") + elif args.check: + print(f"ERROR: {changed_files} would change ({changed_files} of {len(files)} files)") + else: + print(f"{changed_files} changed ({changed_files} of {len(files)} files)") + + if args.check and updates: + return 1 + return 2 if had_parse_error else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/annotate_public_api_types.py b/scripts/annotate_public_api_types.py new file mode 100644 index 000000000..f88847bbd --- /dev/null +++ b/scripts/annotate_public_api_types.py @@ -0,0 +1,810 @@ +#!/usr/bin/env python3 +"""Add conservative public-API annotations from Google-style docstrings. + +Only public module functions, public methods, and constructors of public +classes are considered. A docstring type is applied only when it is valid +Python annotation syntax and every referenced name is already bound by the +module or is a built-in type. Existing annotations are never overwritten; +conflicts are reported for review. +""" +from __future__ import annotations + +import argparse +import ast +import re +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path + +import libcst as cst +from libcst.metadata import MetadataWrapper, PositionProvider + +from scripts import add_docstring_arg_types as docstrings + + +BUILTIN_TYPE_NAMES = { + "bool", + "bytearray", + "bytes", + "complex", + "dict", + "float", + "frozenset", + "int", + "list", + "memoryview", + "object", + "range", + "set", + "slice", + "str", + "tuple", + "type", +} + + +@dataclass(frozen=True) +class FunctionSpec: + """Docstring-derived annotations for one public callable.""" + + function: str + line: int + arguments: dict[str, str] + rejected_arguments: dict[str, tuple[str, str]] + return_type: str | None + + +@dataclass(frozen=True) +class Update: + """One annotation inserted into a function signature.""" + + path: Path + line: int + function: str + slot: str + annotation: str + + +@dataclass(frozen=True) +class Conflict: + """A disagreement between an annotation and its docstring type.""" + + path: Path + line: int + function: str + slot: str + signature_type: str + docstring_type: str + + +@dataclass(frozen=True) +class Skip: + """A docstring type that is unsafe to place in a signature.""" + + path: Path + line: int + function: str + slot: str + docstring_type: str + reason: str + + +@dataclass(frozen=True) +class UnsafeExistingAnnotation: + """An existing annotation that repeats a rejected docstring type.""" + + path: Path + line: int + function: str + slot: str + annotation: str + reason: str + + +@dataclass(frozen=True) +class SourceResult: + """Result of analyzing and transforming one source string.""" + + source: str + updates: tuple[Update, ...] + conflicts: tuple[Conflict, ...] + skips: tuple[Skip, ...] + unsafe_existing: tuple[UnsafeExistingAnnotation, ...] + + +@dataclass(frozen=True) +class FileResult: + """Result of inspecting one Python file.""" + + path: Path + updates: tuple[Update, ...] + conflicts: tuple[Conflict, ...] + skips: tuple[Skip, ...] + unsafe_existing: tuple[UnsafeExistingAnnotation, ...] + changed: bool + + +def _is_type_checking_test(test: ast.expr) -> bool: + """True for an `if` test of `TYPE_CHECKING` or `typing.TYPE_CHECKING`. + + Names bound only under this guard are never executed at runtime, so an + annotation that references them is safe exactly when the module also + enables postponed evaluation (`from __future__ import annotations`) -- + the annotation is then stored as an unevaluated string. + """ + if isinstance(test, ast.Name): + return test.id == "TYPE_CHECKING" + return isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING" + + +def _collect_import_names(node: ast.stmt, names: set[str]) -> None: + """Add the names one import-like or binding statement introduces.""" + if isinstance(node, ast.Import): + for alias in node.names: + names.add(alias.asname or alias.name.split(".")[0]) + elif isinstance(node, ast.ImportFrom): + for alias in node.names: + if alias.name != "*": + names.add(alias.asname or alias.name) + elif isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): + names.add(node.name) + elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.NamedExpr)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for target in targets: + if isinstance(target, ast.Name): + names.add(target.id) + + +def _has_future_annotations(tree: ast.Module) -> bool: + """True if the module enables postponed evaluation of annotations.""" + return any( + isinstance(node, ast.ImportFrom) + and node.module == "__future__" + and any(alias.name == "annotations" for alias in node.names) + for node in tree.body + ) + + +def _module_names(tree: ast.Module, before_line: int) -> set[str]: + """Collect module names bound before a callable's definition. + + Names introduced only inside an `if TYPE_CHECKING:` guard are included + too, but only when the module also has `from __future__ import + annotations`: such names are unavailable at runtime, and without + postponed evaluation a signature referencing one would raise NameError + the moment the function is defined. + """ + names = set(BUILTIN_TYPE_NAMES) + include_type_checking = _has_future_annotations(tree) + for node in tree.body: + if getattr(node, "lineno", before_line) >= before_line: + continue + _collect_import_names(node, names) + if ( + include_type_checking + and isinstance(node, ast.If) + and _is_type_checking_test(node.test) + ): + for sub in node.body: + _collect_import_names(sub, names) + return names + + +def _annotation_expression( + text: str, + available_names: set[str], +) -> tuple[str | None, str | None]: + """Validate and normalize a docstring type for runtime-safe insertion.""" + text = text.strip() + try: + expression = ast.parse(text, mode="eval").body + except SyntaxError: + return None, "not valid Python annotation syntax" + + if isinstance(expression, ast.Constant) and isinstance(expression.value, str): + return None, "string descriptions are not inserted as annotations" + unsafe = ( + ast.BoolOp, + ast.Call, + ast.Compare, + ast.Dict, + ast.DictComp, + ast.GeneratorExp, + ast.IfExp, + ast.Lambda, + ast.ListComp, + ast.Set, + ast.SetComp, + ) + if any(isinstance(node, unsafe) for node in ast.walk(expression)): + return None, "contains an expression that is unsafe in an annotation" + if any(isinstance(node, ast.BinOp) for node in ast.walk(expression)): + return None, "uses an operator that is not safe for Python 3.9 annotations" + + referenced_names = { + node.id for node in ast.walk(expression) if isinstance(node, ast.Name) + } + missing = sorted(referenced_names - available_names) + if missing: + return None, f"name(s) not available in module: {', '.join(missing)}" + + normalized = ast.unparse(expression) + try: + cst.parse_expression(normalized) + except cst.ParserSyntaxError: + return None, "cannot be represented by LibCST" + return normalized, None + + +def _argument_defaults( + node: ast.FunctionDef | ast.AsyncFunctionDef, +) -> dict[str, ast.expr]: + """Map parameter names to defaults present in the signature.""" + positional = list(node.args.posonlyargs) + list(node.args.args) + defaults = {} + if node.args.defaults: + defaults.update( + { + argument.arg: default + for argument, default in zip( + positional[-len(node.args.defaults):], + node.args.defaults, + ) + } + ) + defaults.update( + { + argument.arg: default + for argument, default in zip( + node.args.kwonlyargs, + node.args.kw_defaults, + ) + if default is not None + } + ) + return defaults + + +def _default_compatibility_reason( + annotation: str, + default: ast.expr | None, +) -> str | None: + """Reject obvious contradictions between an annotation and a default.""" + if default is None: + return None + expression = ast.parse(annotation, mode="eval").body + identifiers = { + node.id for node in ast.walk(expression) if isinstance(node, ast.Name) + } + identifiers.update( + node.attr for node in ast.walk(expression) if isinstance(node, ast.Attribute) + ) + if identifiers & {"Any", "object"}: + return None + + try: + value = ast.literal_eval(default) + except (ValueError, TypeError): + return None + + if value is None: + permits_none = "Optional" in identifiers or any( + isinstance(node, ast.Constant) and node.value is None + for node in ast.walk(expression) + ) + if not permits_none: + return "default is None but the documented type is not optional" + return None + + if isinstance(value, bool): + compatible = {"bool"} + elif isinstance(value, int): + compatible = {"complex", "float", "int", "Integral", "Number", "Real"} + elif isinstance(value, float): + compatible = {"complex", "float", "Number", "Real"} + elif isinstance(value, str): + compatible = {"str"} + elif isinstance(value, bytes): + compatible = {"bytes"} + elif isinstance(value, list): + compatible = { + "Collection", + "Iterable", + "List", + "MutableSequence", + "Sequence", + "list", + } + elif isinstance(value, tuple): + compatible = {"Collection", "Iterable", "Sequence", "Tuple", "tuple"} + elif isinstance(value, dict): + compatible = {"Dict", "Mapping", "MutableMapping", "dict"} + elif isinstance(value, (set, frozenset)): + compatible = { + "AbstractSet", + "Collection", + "FrozenSet", + "Iterable", + "Set", + "frozenset", + "set", + } + else: + return None + if identifiers & compatible: + return None + return ( + f"default value of type {type(value).__name__} conflicts with " + "the documented type" + ) + + +def _docstring_argument_types( + node: ast.FunctionDef | ast.AsyncFunctionDef, + lines: list[str], +) -> dict[str, tuple[str, int]]: + """Extract explicit Google-style argument types and their source lines.""" + dnode = docstrings.doc_node(node) + if dnode is None or dnode.end_lineno is None: + return {} + section = docstrings.find_args_section( + lines, + dnode.lineno - 1, + dnode.end_lineno - 1, + ) + if section is None: + return {} + + types = {} + for i in range(section[0] + 1, section[1] + 1): + content, _ = docstrings.line_without_ending(lines[i]) + match = docstrings.ARG_ENTRY.match(content) + if match is None or match.group("type") is None: + continue + name = match.group("name").lstrip("*") + types[name] = (match.group("type").strip(), i + 1) + return types + + +def _docstring_output_type( + node: ast.FunctionDef | ast.AsyncFunctionDef, + lines: list[str], +) -> tuple[str, int] | None: + """Extract an explicit aggregate type from an existing Returns section.""" + dnode = docstrings.doc_node(node) + if dnode is None or dnode.end_lineno is None: + return None + section = docstrings.find_section( + lines, + dnode.lineno - 1, + dnode.end_lineno - 1, + "Returns", + ) + if section is None: + return None + + header, _ = docstrings.line_without_ending(lines[section[0]]) + header_indent = len(header) - len(header.lstrip()) + for i in range(section[0] + 1, section[1] + 1): + content, _ = docstrings.line_without_ending(lines[i]) + if not content.strip(): + continue + indent = len(content) - len(content.lstrip()) + if indent <= header_indent: + continue + match = docstrings.OUTPUT_ENTRY.match(content) + if match is None: + return None + candidate = match.group("type").strip() + if not docstrings.looks_like_type(candidate): + return None + return candidate, i + 1 + return None + + +def _collect_specs( + source: str, + path: Path, +) -> tuple[dict[tuple[int, str], FunctionSpec], list[Skip]]: + """Collect validated docstring types for public functions and methods.""" + tree = ast.parse(source, filename=str(path)) + lines = source.splitlines(keepends=True) + specs = {} + skips = [] + + for node, function in docstrings.iter_public_functions(tree): + available_names = _module_names(tree, node.lineno) + if "." in function: + # A class is not bound to its module name until its body finishes. + available_names.discard(function.split(".", maxsplit=1)[0]) + arguments = {} + rejected_arguments = {} + defaults = _argument_defaults(node) + for name, (text, line) in _docstring_argument_types(node, lines).items(): + annotation, reason = _annotation_expression(text, available_names) + if annotation is not None: + reason = _default_compatibility_reason( + annotation, + defaults.get(name), + ) + if reason is not None: + annotation = None + if annotation is None: + rejection_reason = reason or "unsafe" + rejected_arguments[name] = (text, rejection_reason) + skips.append( + Skip(path, line, function, name, text, rejection_reason) + ) + else: + arguments[name] = annotation + + return_type = "None" if node.name == "__init__" else None + output = _docstring_output_type(node, lines) + if output is not None and node.name != "__init__": + text, line = output + annotation, reason = _annotation_expression(text, available_names) + if annotation is None: + skips.append( + Skip(path, line, function, "return", text, reason or "unsafe") + ) + else: + return_type = annotation + + specs[(node.lineno, node.name)] = FunctionSpec( + function=function, + line=node.lineno, + arguments=arguments, + rejected_arguments=rejected_arguments, + return_type=return_type, + ) + return specs, skips + + +def _annotation_code(annotation: cst.Annotation) -> str: + """Render one LibCST annotation expression without surrounding syntax.""" + return cst.Module(body=[]).code_for_node(annotation.annotation) + + +def _normalized_annotation(text: str) -> str: + """Normalize annotations for conflict comparison.""" + try: + expression = ast.parse(text, mode="eval").body + except SyntaxError: + return re.sub(r"\s+", "", text) + if isinstance(expression, ast.Constant) and isinstance(expression.value, str): + try: + expression = ast.parse(expression.value, mode="eval").body + except SyntaxError: + return expression.value + return ast.dump(expression, include_attributes=False) + + +class PublicAPIAnnotationTransformer(cst.CSTTransformer): + """Insert validated docstring types into matching public signatures.""" + + METADATA_DEPENDENCIES = (PositionProvider,) + + def __init__(self, path: Path, specs: dict[tuple[int, str], FunctionSpec]): + """Record the file and the annotations to apply. + + Args: + path (Path): File being transformed, used in diagnostics. + specs (dict[tuple[int, str], FunctionSpec]): Annotation specification keyed + by ``(line number, function name)``. + """ + self.path = path + self.specs = specs + self.updates: list[Update] = [] + self.conflicts: list[Conflict] = [] + self.unsafe_existing: list[UnsafeExistingAnnotation] = [] + + def _update_param( + self, + original: cst.Param, + updated: cst.Param, + spec: FunctionSpec, + ) -> cst.Param: + """Annotate one parameter or report a signature/docstring conflict.""" + name = original.name.value + desired = spec.arguments.get(name) + line = self.get_metadata(PositionProvider, original.name).start.line + rejected = spec.rejected_arguments.get(name) + if desired is None: + if original.annotation is not None and rejected is not None: + existing = _annotation_code(original.annotation) + rejected_type, reason = rejected + if _normalized_annotation(existing) == _normalized_annotation( + rejected_type + ): + self.unsafe_existing.append( + UnsafeExistingAnnotation( + self.path, + line, + spec.function, + name, + existing, + reason, + ) + ) + return updated + if original.annotation is not None: + existing = _annotation_code(original.annotation) + if _normalized_annotation(existing) != _normalized_annotation(desired): + self.conflicts.append( + Conflict( + self.path, + line, + spec.function, + name, + existing, + desired, + ) + ) + return updated + + self.updates.append(Update(self.path, line, spec.function, name, desired)) + changes = {"annotation": cst.Annotation(cst.parse_expression(desired))} + if updated.default is not None and isinstance(updated.equal, cst.AssignEqual): + changes["equal"] = updated.equal.with_changes( + whitespace_before=cst.SimpleWhitespace(" "), + whitespace_after=cst.SimpleWhitespace(" "), + ) + return updated.with_changes(**changes) + + def leave_FunctionDef( + self, + original_node: cst.FunctionDef, + updated_node: cst.FunctionDef, + ) -> cst.FunctionDef: + """Update an eligible function or method signature. + + Args: + original_node (cst.FunctionDef): Node before any child updates. + updated_node (cst.FunctionDef): Node with child updates already applied. + + Returns: + cst.FunctionDef: The annotated node, or ``updated_node`` unchanged when + the function is not eligible. + """ + line = self.get_metadata(PositionProvider, original_node.name).start.line + spec = self.specs.get((line, original_node.name.value)) + if spec is None: + return updated_node + + original_params = original_node.params + updated_params = updated_node.params + posonly_params = tuple( + self._update_param(original, updated, spec) + for original, updated in zip( + original_params.posonly_params, + updated_params.posonly_params, + ) + ) + params = tuple( + self._update_param(original, updated, spec) + for original, updated in zip(original_params.params, updated_params.params) + ) + kwonly_params = tuple( + self._update_param(original, updated, spec) + for original, updated in zip( + original_params.kwonly_params, + updated_params.kwonly_params, + ) + ) + star_arg = updated_params.star_arg + if isinstance(original_params.star_arg, cst.Param) and isinstance( + updated_params.star_arg, cst.Param + ): + star_arg = self._update_param( + original_params.star_arg, + updated_params.star_arg, + spec, + ) + star_kwarg = updated_params.star_kwarg + if original_params.star_kwarg is not None and star_kwarg is not None: + star_kwarg = self._update_param( + original_params.star_kwarg, + star_kwarg, + spec, + ) + + returns = updated_node.returns + if spec.return_type is not None: + if original_node.returns is None: + self.updates.append( + Update( + self.path, + line, + spec.function, + "return", + spec.return_type, + ) + ) + returns = cst.Annotation(cst.parse_expression(spec.return_type)) + else: + existing = _annotation_code(original_node.returns) + if _normalized_annotation(existing) != _normalized_annotation( + spec.return_type + ): + self.conflicts.append( + Conflict( + self.path, + line, + spec.function, + "return", + existing, + spec.return_type, + ) + ) + + return updated_node.with_changes( + params=updated_params.with_changes( + posonly_params=posonly_params, + params=params, + kwonly_params=kwonly_params, + star_arg=star_arg, + star_kwarg=star_kwarg, + ), + returns=returns, + ) + + +def transform_source(source: str, path: Path = Path("")) -> SourceResult: + """Annotate one source string without writing it. + + Args: + source (str): Python source to annotate. + path (Path): Path reported in diagnostics. + + Returns: + SourceResult: Annotated source together with updates and conflicts. + """ + specs, skips = _collect_specs(source, path) + module = cst.parse_module(source) + transformer = PublicAPIAnnotationTransformer(path, specs) + transformed = MetadataWrapper(module).visit(transformer) + return SourceResult( + source=transformed.code, + updates=tuple(transformer.updates), + conflicts=tuple(transformer.conflicts), + skips=tuple(skips), + unsafe_existing=tuple(transformer.unsafe_existing), + ) + + +def update_file(path: Path, check: bool = False) -> FileResult: + """Annotate one Python file. + + Args: + path (Path): Python file to annotate. + check (bool): Report what would change without writing. + + Returns: + FileResult: Whether the file changed, and the updates and conflicts found. + """ + source = path.read_text(encoding="utf-8") + result = transform_source(source, path=path) + changed = result.source != source + if changed and not check: + path.write_text(result.source, encoding="utf-8") + return FileResult( + path=path, + updates=result.updates, + conflicts=result.conflicts, + skips=result.skips, + unsafe_existing=result.unsafe_existing, + changed=changed, + ) + + +def _parse_args(argv: list[str]) -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "paths", + nargs="*", + help="Python files or directories to update. Defaults to qmcpy.", + ) + parser.add_argument( + "--diff", + metavar="REF", + help="Use Python files reported by git diff REF.", + ) + parser.add_argument( + "--root", + default="qmcpy", + help="Restrict files selected by --diff. Defaults to qmcpy.", + ) + parser.add_argument( + "--check", + action="store_true", + help="Report annotations without writing files.", + ) + parser.add_argument( + "--quiet", + action="store_true", + help="Only print the final summary.", + ) + return parser.parse_args(argv) + + +def main(argv: list[str]) -> int: + """Run the command-line interface. + + Args: + argv (list[str]): Command-line arguments, excluding the program name. + + Returns: + int: Process exit status; ``0`` on success. + """ + args = _parse_args(argv) + try: + files = docstrings.python_files(args.paths, args.diff, root=args.root) + except subprocess.CalledProcessError as exc: + print(f"git diff failed: {exc}", file=sys.stderr) + return 2 + + if not files: + print("No Python files to inspect.") + return 0 + + results = [] + had_parse_error = False + for path in files: + try: + results.append(update_file(path, check=args.check)) + except (SyntaxError, cst.ParserSyntaxError) as exc: + had_parse_error = True + print(f"{path}: skipped syntax error: {exc}", file=sys.stderr) + + updates = [update for result in results for update in result.updates] + conflicts = [conflict for result in results for conflict in result.conflicts] + skips = [skip for result in results for skip in result.skips] + unsafe_existing = [ + issue for result in results for issue in result.unsafe_existing + ] + if not args.quiet: + action = "would annotate" if args.check else "annotated" + for update in updates: + print( + f"{update.path}:{update.line}: {action} " + f"{update.function}.{update.slot} as {update.annotation}" + ) + for conflict in conflicts: + print( + f"{conflict.path}:{conflict.line}: conflict " + f"{conflict.function}.{conflict.slot}: signature " + f"`{conflict.signature_type}` != docstring " + f"`{conflict.docstring_type}`" + ) + for skip in skips: + print( + f"{skip.path}:{skip.line}: skipped {skip.function}.{skip.slot} " + f"`{skip.docstring_type}`: {skip.reason}" + ) + for issue in unsafe_existing: + print( + f"{issue.path}:{issue.line}: unsafe existing annotation " + f"{issue.function}.{issue.slot} `{issue.annotation}`: " + f"{issue.reason}" + ) + + changed_files = sum(result.changed for result in results) + verb = "would change" if args.check else "changed" + print( + f"{len(files)} file(s) inspected; {len(updates)} signature update(s); " + f"{len(conflicts)} conflict(s); {len(skips)} unsafe type(s) skipped; " + f"{len(unsafe_existing)} unsafe existing annotation(s); " + f"{changed_files} file(s) {verb}." + ) + + if had_parse_error: + return 2 + if conflicts or unsafe_existing or (args.check and updates): + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/baseline_counts.json b/scripts/baseline_counts.json new file mode 100644 index 000000000..cc7ab45b8 --- /dev/null +++ b/scripts/baseline_counts.json @@ -0,0 +1,5 @@ +{ + "check_docstring": 0, + "pydoclint": 0, + "unsafe_annotations": 0 +} diff --git a/scripts/check_baseline.py b/scripts/check_baseline.py new file mode 100644 index 000000000..f015e1d25 --- /dev/null +++ b/scripts/check_baseline.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""Ratchet gate for the informational docstring/annotation checks. + +`check_docstring`, `pydoclint`, and `annotate_public_api_types` are +informational today (see F9/F10 in the PR #613 review) because fixing every +existing violation before enabling them as hard gates is a large, separate +undertaking. This script tracks each check's full-tree violation count in +`scripts/baseline_counts.json` and fails only if a count *increases* -- +new violations are blocked; the existing backlog is not required to be +cleared just to land an unrelated change. + +Usage: + python scripts/check_baseline.py # compare against the baseline + python scripts/check_baseline.py --update # write current counts as the new baseline + +`--update` is for a change that intentionally reduces (or, with justification +in the PR description, increases) one of these counts. + +A mid-migration branch (e.g. adding type hints to signatures across many +files) can make a count spike well above the committed baseline before it +comes back down -- pydoclint's DOC105/106/107 cross-check every arg's +signature type against its docstring type, so partially-applied hints +surface more mismatches than having no hints at all. That is expected, not +a bug in this script: `make check` will keep reporting "REGRESSED" for that +check until the migration is complete and `--update` is run to record the +new, lower count. +""" +import json +import re +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +BASELINE_PATH = Path(__file__).resolve().parent / "baseline_counts.json" + +CHECKS = { + "check_docstring": { + "cmd": [sys.executable, "scripts/check_docstring.py", "qmcpy"], + # check_docstring.py's summary line reads either "N issue(s) across + # M file(s)" or, once N reaches zero, "no issues in M file(s)" -- + # match both so the ratchet keeps working after a check is fully fixed. + "pattern": re.compile( + r"^\s*(?:- )?\d+ file\(s\) scanned: " + r"(?:(\d+) issue\(s\) across|no issues in) \d+ file\(s\)", + re.M, + ), + }, + "pydoclint": { + "cmd": ["pydoclint", "-q", "qmcpy"], + "line_pattern": re.compile(r"^\s*\d+: DOC\d+:", re.M), + }, + "unsafe_annotations": { + "cmd": [sys.executable, "-m", "scripts.annotate_public_api_types", "--check", "--root", "qmcpy"], + "pattern": re.compile(r"(\d+) unsafe existing annotation\(s\)"), + }, +} + + +def run_check(spec): + result = subprocess.run(spec["cmd"], capture_output=True, text=True, cwd=REPO_ROOT) + output = result.stdout + result.stderr + if "line_pattern" in spec: + return len(spec["line_pattern"].findall(output)) + match = spec["pattern"].search(output) + if match is None: + raise RuntimeError(f"could not parse a count from output of {spec['cmd']}") + return int(match.group(1) or 0) + + +def main(argv): + update = "--update" in argv + baseline = json.loads(BASELINE_PATH.read_text()) if BASELINE_PATH.exists() else {} + + current = {} + regressed = [] + for name, spec in CHECKS.items(): + count = run_check(spec) + current[name] = count + base = baseline.get(name) + if base is None: + status = "no baseline yet" + elif count > base: + status = f"REGRESSED from {base}" + regressed.append(name) + elif count < base: + status = f"improved from {base}" + else: + status = "unchanged" + print(f" - {name}: {count} ({status})") + + if not regressed: + print(f"clean (0 of {len(CHECKS)} counts)") + else: + print(f"ERROR: {len(regressed)} regressed ({len(regressed)} of {len(CHECKS)} counts)") + + if update: + BASELINE_PATH.write_text(json.dumps(current, indent=2, sort_keys=True) + "\n") + print(f"\nWrote new baseline to {BASELINE_PATH.relative_to(REPO_ROOT)}") + return 0 + + if regressed: + print( + f"\nRegression in: {', '.join(regressed)}. Fix the new violations, " + "or if the increase is intentional and justified in the PR " + "description, run `python scripts/check_baseline.py --update` " + "and commit the updated baseline file.", + file=sys.stderr, + ) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/check_docstring.py b/scripts/check_docstring.py new file mode 100644 index 000000000..d8f399890 --- /dev/null +++ b/scripts/check_docstring.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +"""Check that public docstrings under ``qmcpy/`` follow Google style. + +A *public* object is a module, a class / function / method whose name does not +start with ``_``, or the ``__init__`` of a public class (QMCPy documents +constructor arguments in ``__init__``'s own docstring). Only module-level +functions/classes and the methods of public classes are inspected (helpers +nested inside functions are skipped). For every such docstring this script +flags: + +* ``missing`` -- public class / function / method has no + docstring (suppressed by ``--skip-missing``) +* ``missing-summary`` -- the docstring opens straight with a section + header (``Args:``, ``Returns:``, ...) instead + of a one-line summary. This is the common + cause of pydoclint's opaque ``DOC001``. +* ``numpy-section`` -- a section written NumPy-style (``Returns`` + followed by a ``-----`` underline) instead of + Google style (``Returns:``) +* ``no-blank-before-section`` -- a Google section header (``Args:``, + ``Returns:``, ``Raises:``, ...) is not preceded + by a blank line +* ``malformed-section-header`` -- a line that names a known section but is not + the canonical ``Name:`` form: a missing colon + (``Examples``), wrong casing (``EXAMPLES:``, + ``examples:``), or stray characters around the + colon (``Args :``) + +Usage: + python scripts/check_docstring.py [PATH ...] [--strict] [--quiet] + [--skip-missing] [--diff [REF]] + +PATH defaults to ``qmcpy``. Informational by default (exit 0); ``--strict`` +makes the exit code non-zero when anything is flagged, so it can gate CI +(``STRICT=--strict make check_docstring``). ``--diff [REF]`` (REF defaults to +``develop``) prints a second summary restricted to the scanned files that +changed relative to REF -- committed on the branch, modified in the working +tree, or untracked. +""" +from __future__ import annotations + +import ast +import re +import subprocess +import sys +from pathlib import Path + +# Canonical Google section headers, written as ``Name:`` on their own line. +GOOGLE_SECTIONS = { + "Args", "Arguments", "Attributes", "Example", "Examples", "Keyword Args", + "Note", "Notes", "Raises", "References", "Return", "Returns", "See Also", + "Todo", "Warning", "Warnings", "Warns", "Yield", "Yields", +} +# Section words that, followed by a dashed underline, mean the docstring is +# using NumPy style rather than Google style. +NUMPY_SECTIONS = { + "Parameters", "Other Parameters", "Returns", "Raises", "Yields", + "Attributes", "Notes", "Examples", "See Also", "References", "Warns", + "Warnings", "Methods", +} +_DASHES = re.compile(r"^-{3,}$") +# Canonical header: capitalised word(s), a single colon, nothing else. +_HEADER = re.compile(r"^([A-Z][A-Za-z]*(?: [A-Z][A-Za-z]*)*):$") +# Case-insensitive lookup from any known section label to its canonical spelling. +_CANON = {name.lower(): name for name in GOOGLE_SECTIONS | NUMPY_SECTIONS} + + +def _is_property_setter_or_deleter(node): + """True if ``node`` is decorated ``@.setter`` or ``@.deleter``. + + Such methods share their contract with the ``@property`` getter of the + same name (which is separately checked), so requiring their own + docstring would be a false positive -- no Python convention expects one. + """ + for decorator in node.decorator_list: + if (isinstance(decorator, ast.Attribute) + and decorator.attr in ("setter", "deleter")): + return True + return False + + +def _iter_public(tree): + """Yield ``(node, kind)`` for the module plus its public API objects.""" + yield tree, "module" + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if not node.name.startswith("_"): + yield node, "function" + elif isinstance(node, ast.ClassDef) and not node.name.startswith("_"): + yield node, "class" + for sub in node.body: + if not isinstance(sub, (ast.FunctionDef, ast.AsyncFunctionDef)): + continue + if _is_property_setter_or_deleter(sub): + continue + if not sub.name.startswith("_"): + yield sub, "method" + elif sub.name == "__init__": + yield sub, "constructor" + + +def _doc_node(node): + """Return the string-literal node holding ``node``'s docstring, or None.""" + body = getattr(node, "body", None) + if (body and isinstance(body[0], ast.Expr) + and isinstance(body[0].value, ast.Constant) + and isinstance(body[0].value.value, str)): + return body[0].value + return None + + +def _is_section_word(s): + """Return the section word if ``s`` is a lone Google/NumPy section header.""" + word = s[:-1].strip() if s.endswith(":") else s + if word in GOOGLE_SECTIONS or word in NUMPY_SECTIONS: + return word + return None + + +def check_file(path, skip_missing=False): + """Return a list of ``(lineno, category, detail)`` findings for one file.""" + findings = [] + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node, kind in _iter_public(tree): + dnode = _doc_node(node) + if dnode is None: + # A bare "no docstring" finding is noise for __init__ (pydoclint + # owns constructor-argument coverage) and meaningless for a module. + if kind not in ("module", "constructor") and not skip_missing: + findings.append(( + getattr(node, "lineno", 1), "missing", + f"public {kind} `{getattr(node, 'name', path.stem)}` has no docstring", + )) + continue + lines = dnode.value.split("\n") + first_nonblank = next((j for j, ln in enumerate(lines) if ln.strip()), None) + for i, raw in enumerate(lines): + s = raw.strip() + if not s: + continue + word = s[:-1].strip() if s.endswith(":") else s + nxt = lines[i + 1].strip() if i + 1 < len(lines) else "" + if word in NUMPY_SECTIONS and _DASHES.match(nxt): + findings.append(( + dnode.lineno + i, "numpy-section", + f"`{word}` written NumPy-style; use Google `{word}:`", + )) + continue + m = _HEADER.match(s) + if m and m.group(1) in GOOGLE_SECTIONS: + if i == first_nonblank and kind != "module": + findings.append(( + dnode.lineno + i, "missing-summary", + f"docstring opens with `{s}`; add a one-line summary first", + )) + elif i > 0 and lines[i - 1].strip() != "": + findings.append(( + dnode.lineno + i, "no-blank-before-section", + f"add a blank line before `{s}`", + )) + elif not _DASHES.match(nxt): + canon = _CANON.get(re.sub(r"\s+", " ", word).strip().lower()) + if canon is not None and s != f"{canon}:": + if not s.rstrip().endswith(":"): + why = "missing colon" + elif word != canon: + why = ( + f"label must be `{canon}` " + "(first letter capitalised, the rest lower-case)" + ) + else: + why = "stray characters around the colon" + findings.append(( + dnode.lineno + i, "malformed-section-header", + f"`{s}` should be `{canon}:` ({why})", + )) + return findings + + +def _changed_files(ref): + """Return resolved paths of *.py files that changed relative to ``ref``. + + Union of files committed on the branch (``ref...HEAD``), files modified in + the working tree, and untracked files. Raises ``RuntimeError`` if git is + unavailable or ``ref`` cannot be resolved. + """ + commands = ( + ["git", "diff", "--name-only", "--diff-filter=ACMR", f"{ref}...HEAD"], + ["git", "diff", "--name-only", "--diff-filter=ACMR", "HEAD"], + ["git", "ls-files", "--others", "--exclude-standard"], + ) + names = set() + for cmd in commands: + try: + out = subprocess.run( + cmd, capture_output=True, text=True, check=True, + ).stdout + except (OSError, subprocess.CalledProcessError) as exc: + raise RuntimeError(f"`{' '.join(cmd)}` failed: {exc}") from exc + names.update(n for n in out.splitlines() if n.endswith(".py")) + return {Path(n).resolve() for n in names} + + +def _summary(total, n_files, by_cat, label): + """Format one summary line.""" + if total == 0: + return f"{label}: no issues in {n_files} file(s)" + breakdown = ", ".join(f"{v} {k}" for k, v in sorted(by_cat.items())) + return f"{label}: {total} issue(s) across {n_files} file(s): {breakdown}" + + +def _parse_diff_flag(argv): + """Pull ``--diff [REF]`` out of ``argv``; return (remaining_argv, ref|None).""" + args, ref, i = [], None, 0 + while i < len(argv): + a = argv[i] + if a == "--diff": + nxt = argv[i + 1] if i + 1 < len(argv) else "" + if nxt and not nxt.startswith("-"): + ref, i = nxt, i + 2 + else: + ref, i = "develop", i + 1 + continue + if a.startswith("--diff="): + ref = a.split("=", 1)[1] or "develop" + i += 1 + continue + args.append(a) + i += 1 + return args, ref + + +def main(argv): + argv, diff_ref = _parse_diff_flag(list(argv)) + strict = "--strict" in argv + quiet = "--quiet" in argv + skip_missing = "--skip-missing" in argv + paths = [a for a in argv if not a.startswith("-")] or ["qmcpy"] + + files = [] + for p in map(Path, paths): + files.extend(sorted(p.rglob("*.py")) if p.is_dir() else [p]) + if not files: + print(f"no *.py files under {', '.join(paths)}", file=sys.stderr) + return 1 + + total = 0 + by_cat = {} + per_file = {} + for f in files: + try: + findings = check_file(f, skip_missing=skip_missing) + except SyntaxError as exc: + print(f"{f.as_posix()}: skipped (syntax error: {exc})", file=sys.stderr) + continue + per_file[f] = findings + for _, cat, _ in findings: + by_cat[cat] = by_cat.get(cat, 0) + 1 + total += 1 + + if total and not quiet: + print() + for f, findings in per_file.items(): + for lineno, cat, detail in findings: + print(f" - {f.as_posix()}:{lineno}: {cat}: {detail}") + print(" - " + _summary(total, len(files), by_cat, f"{len(files)} file(s) scanned")) + + if diff_ref is not None: + try: + changed = _changed_files(diff_ref) + except RuntimeError as exc: + print(f"--diff {diff_ref}: skipped ({exc})", file=sys.stderr) + else: + sub_cat, sub_total, sub_files = {}, 0, 0 + for f, findings in per_file.items(): + if f.resolve() not in changed: + continue + sub_files += 1 + for _, cat, _ in findings: + sub_cat[cat] = sub_cat.get(cat, 0) + 1 + sub_total += 1 + print(" - " + _summary(sub_total, sub_files, sub_cat, f"changed vs {diff_ref}")) + + files_with_issues = sum(1 for findings in per_file.values() if findings) + if files_with_issues == 0: + print(f"clean (0 of {len(files)} files)") + else: + prefix = "ERROR" if (strict and total) else "WARNING" + print(f"{prefix}: {files_with_issues} problem(s) ({files_with_issues} of {len(files)} files)") + return 1 if (strict and total) else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/check_links.py b/scripts/check_links.py index 8fc1405a1..4468f5be5 100644 --- a/scripts/check_links.py +++ b/scripts/check_links.py @@ -217,23 +217,30 @@ def main() -> int: site_dir = Path(args.site_dir) site_url = read_site_url() + pages = sum(1 for _ in site_dir.rglob("*.html")) problems = check_internal(site_dir, site_url=site_url) - print(f"Checked internal links under {site_dir}: {len(problems)} problem(s).") - for p in problems: - print(f" {p}") + print(f" - Checked internal links under {site_dir}: {len(problems)} problem(s).") + if problems: + print() + for p in problems: + print(f" - {p}") if args.external: ext_broken, ext_warnings = check_external(site_dir, site_url=site_url) print( - f"\nChecked external links: {len(ext_broken)} broken link(s), " + f"\n - Checked external links: {len(ext_broken)} broken link(s), " f"{len(ext_warnings)} warning(s)." ) for p in ext_broken: - print(f" {p}") + print(f" - {p}") for p in ext_warnings: - print(f" [warning] {p}") + print(f" - [warning] {p}") problems += ext_broken + if problems: + print(f"ERROR: {len(problems)} problem(s) ({len(problems)} of {pages} pages)") + else: + print(f"clean (0 of {pages} pages)") return 1 if problems else 0 diff --git a/scripts/check_test_style.py b/scripts/check_test_style.py new file mode 100755 index 000000000..68c95b329 --- /dev/null +++ b/scripts/check_test_style.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +"""Check ``test/test_*.py`` files against two suite conventions. + +1. **Object class.** A test file should be written as a ``unittest.TestCase`` + subclass, not as bare ``def test_*`` pytest functions. The numeric-correctness + backbone (``test_tm_true_measures.py``, ``test_sc_stopping_criteria.py``, + ``test_dd_discrete_distribs.py``, ...) already follows this; newer per-measure + and tooling files do not. The split is listed so it stays visible in review. + +2. **Area prefix.** A test file should be named ``test__.py`` where + ```` marks the ``qmcpy`` subpackage under test (or a cross-cutting + bucket). Recognized areas: + + dd discrete_distribution tm true_measure + ft fast_transform ut util + ig integrand ee end-to-end / cross-cutting pipeline + kn kernel sr scripts/ tooling, packaging, docs checks + sc stopping_criterion + + A test that spans two areas goes under the component actually under test, + with the other named in ```` (e.g. ``test_sc_cubbayes_kernels.py``); + ``ee`` is only for tests where neither side is the clear subject. Only the + codes above are accepted -- new two-letter codes are a ``--strict`` failure. + +Usage: + python scripts/check_test_style.py [TEST_DIR] [--strict] [--quiet] + +TEST_DIR defaults to ``test``. With ``--strict`` the exit code is non-zero when +any file violates either convention (so it can gate CI); otherwise it is always +0 and the output is informational. +""" +import ast +import re +import sys +from pathlib import Path + +AREA_PREFIXES = { + "dd": "discrete_distribution", + "ft": "fast_transform", + "ig": "integrand", + "kn": "kernel", + "sc": "stopping_criterion", + "tm": "true_measure", + "ut": "util", + "ee": "end-to-end / cross-cutting pipeline", + "sr": "scripts/ tooling, packaging, docs checks", +} +AREA_RE = re.compile(r"^test_(?:" + "|".join(sorted(AREA_PREFIXES)) + r")_.+\.py$") + + +def _area_ok(path): + """True if the filename starts with a recognized ``test__`` prefix.""" + return bool(AREA_RE.match(path.name)) + + +def _subclasses_testcase(node): + """True if a ClassDef lists ``TestCase`` / ``unittest.TestCase`` as a base.""" + for base in node.bases: + if isinstance(base, ast.Attribute) and base.attr == "TestCase": + return True + if isinstance(base, ast.Name) and base.id == "TestCase": + return True + return False + + +def classify(path): + """Return (has_testcase_class, has_bare_top_level_test, has_any_test). + + ``has_bare_top_level_test`` only looks at module-level functions, so a + file with a proper TestCase class that *also* has a stray top-level + ``def test_*():`` still flags the violation instead of being masked by + the class. ``has_any_test`` still walks the whole tree, to distinguish a + file with no tests at all from one whose tests just aren't bare/top-level + (e.g. methods on a non-TestCase class). + """ + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + has_class = any( + isinstance(n, ast.ClassDef) and _subclasses_testcase(n) + for n in ast.walk(tree) + ) + has_bare_top_level_test = any( + isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) + and n.name.startswith("test_") + for n in tree.body + ) + has_any_test = any( + isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) + and n.name.startswith("test_") + for n in ast.walk(tree) + ) + return has_class, has_bare_top_level_test, has_any_test + + +def main(argv): + strict = "--strict" in argv + quiet = "--quiet" in argv + positional = [a for a in argv if not a.startswith("-")] + test_dir = Path(positional[0]) if positional else Path("test") + + # Recursive: a misnamed/bare-function test file placed in a subdirectory + # should still be caught. test/booktests/ is excluded -- it has its own + # separate, documented naming convention (tb_*.py, generated from + # demos/) and isn't meant to comply with the test__*.py convention + # this script enforces; test/booktests/test_runtimes.py in particular + # isn't a test at all, just a runtime-estimates data module that happens + # to start with "test_". + files = sorted( + f for f in test_dir.rglob("test_*.py") if "booktests" not in f.parts + ) + if not files: + print(f"no test_*.py files under {test_dir}/", file=sys.stderr) + return 1 + + class_based, function_based, no_tests = [], [], [] + for f in files: + has_class, has_bare_top_level_test, has_any_test = classify(f) + if has_bare_top_level_test: + # Flagged regardless of has_class: a stray top-level `def + # test_*():` violates the convention even in a file that also + # has a proper TestCase class. + function_based.append(f) + elif has_class: + class_based.append(f) + elif has_any_test: + function_based.append(f) + else: + no_tests.append(f) + + misnamed = [f for f in files if not _area_ok(f)] + + if function_based or no_tests or misnamed: + print() + + if not quiet: + print(f" - {len(class_based)} of {len(files)} files use a unittest.TestCase class") + if function_based: + print(f" - {len(function_based)} file(s) use bare pytest functions (no unittest.TestCase class):") + for f in function_based: + print(f" - {f.as_posix()}") + if no_tests and not quiet: + print(f" - {len(no_tests)} file(s) define no test_* callables:") + for f in no_tests: + print(f" - {f.as_posix()}") + + if not quiet: + print( + f" - {len(files) - len(misnamed)} of {len(files)} files use a " + f"test__ prefix ({', '.join(sorted(AREA_PREFIXES))})" + ) + if misnamed: + print(f" - {len(misnamed)} file(s) have no recognized test__ prefix:") + for f in misnamed: + print(f" - {f.as_posix()}") + + bad = {*function_based, *misnamed, *no_tests} + if not bad: + print(f"clean (0 of {len(files)} files)") + else: + prefix = "ERROR" if strict else "WARNING" + print(f"{prefix}: {len(bad)} problem(s) ({len(bad)} of {len(files)} files)") + return 1 if (strict and (function_based or misnamed or no_tests)) else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/colab_notebooks_manifest.json b/scripts/colab_notebooks_manifest.json index ed62ba8be..e683f8291 100644 --- a/scripts/colab_notebooks_manifest.json +++ b/scripts/colab_notebooks_manifest.json @@ -21,6 +21,7 @@ "demos/lattice_random_generator.ipynb", "demos/lebesgue_integration.ipynb", "demos/linear-scrambled-halton.ipynb", + "demos/makefile_dev_tools.ipynb", "demos/nei_demo.ipynb", "demos/plot_proj_function.ipynb", "demos/pricing_options.ipynb", diff --git a/scripts/convert_asserts.py b/scripts/convert_asserts.py new file mode 100644 index 000000000..3ebba6c67 --- /dev/null +++ b/scripts/convert_asserts.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +"""Convert Python assertions to explicit exception raises. + +The codemod preserves formatting and comments with LibCST. By default it +converts ``assert condition, message`` to an explicit ``AssertionError`` so +the validation is not removed by ``python -O``. A developer may select a +different exception that is already in scope, but the tool deliberately does +not guess domain-specific exception classes. +""" +from __future__ import annotations + +import argparse +import re +import subprocess +import sys +from collections import Counter +from dataclasses import dataclass +from pathlib import Path + +import libcst as cst +from libcst.metadata import MetadataWrapper, PositionProvider + + +EXCEPTION_NAME = re.compile( + r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*$" +) + + +@dataclass(frozen=True) +class SourceResult: + """Result of transforming one source string.""" + + source: str + converted_lines: tuple[int, ...] + skipped_lines: tuple[int, ...] + + +@dataclass(frozen=True) +class FileResult: + """Result of inspecting one Python file.""" + + path: Path + converted_lines: tuple[int, ...] + skipped_lines: tuple[int, ...] + changed: bool + + +def _parenthesize(expression: cst.BaseExpression) -> cst.BaseExpression: + """Parenthesize an expression unless it is already parenthesized.""" + if expression.lpar: + return expression + return expression.with_changes( + lpar=(cst.LeftParen(),), + rpar=(cst.RightParen(),), + ) + + +def _exception_call( + exception: cst.BaseExpression, + message: cst.BaseExpression, +) -> cst.Call: + """Build an exception call while reusing message-parenthesis whitespace.""" + if ( + not message.lpar + or not message.rpar + or isinstance(message, (cst.Tuple, cst.Yield)) + ): + return cst.Call(func=exception, args=[cst.Arg(message)]) + + opening = message.lpar[0] + closing = message.rpar[-1] + unwrapped_message = message.with_changes( + lpar=message.lpar[1:], + rpar=message.rpar[:-1], + ) + return cst.Call( + func=exception, + args=[ + cst.Arg( + unwrapped_message, + whitespace_after_arg=closing.whitespace_before, + ) + ], + whitespace_before_args=opening.whitespace_after, + ) + + +class ConvertAssertTransformer(cst.CSTTransformer): + """Rewrite standalone assertion statements as explicit conditional raises.""" + + METADATA_DEPENDENCIES = (PositionProvider,) + + def __init__(self, exception: str): + """Record the exception to raise in place of each assertion. + + Args: + exception (str): Exception expression to raise, such as ``"AssertionError"``. + """ + self.exception = cst.parse_expression(exception) + self.seen_lines = [] + self.converted_lines = [] + + def visit_Assert(self, node: cst.Assert) -> None: + """Record every assertion, including forms that cannot be rewritten. + + Args: + node (cst.Assert): Assertion encountered in the tree. + """ + position = self.get_metadata(PositionProvider, node) + self.seen_lines.append(position.start.line) + + def leave_SimpleStatementLine( + self, + original_node: cst.SimpleStatementLine, + updated_node: cst.SimpleStatementLine, + ) -> cst.BaseStatement: + """Rewrite an assert when it is the line's only small statement. + + Args: + original_node (cst.SimpleStatementLine): Node before any child updates. + updated_node (cst.SimpleStatementLine): Node with child updates applied. + + Returns: + cst.BaseStatement: The rewritten statement, or ``updated_node`` unchanged + when the line holds more than the assertion. + """ + if len(updated_node.body) != 1: + return updated_node + assertion = updated_node.body[0] + if not isinstance(assertion, cst.Assert): + return updated_node + + condition = cst.UnaryOperation( + operator=cst.Not(whitespace_after=cst.SimpleWhitespace(" ")), + expression=_parenthesize(assertion.test), + ) + exception = self.exception.deep_clone() + if assertion.msg is None: + raised_exception = exception + else: + raised_exception = _exception_call(exception, assertion.msg) + + position = self.get_metadata(PositionProvider, original_node) + self.converted_lines.append(position.start.line) + return cst.If( + test=condition, + body=cst.IndentedBlock( + header=updated_node.trailing_whitespace, + body=[ + cst.SimpleStatementLine( + body=[cst.Raise(exc=raised_exception)] + ) + ], + ), + leading_lines=updated_node.leading_lines, + ) + + +def transform_source(source: str, exception: str = "AssertionError") -> SourceResult: + """Transform standalone assertions in a Python source string. + + Args: + source (str): Python source to transform. + exception (str): Exception expression to raise in place of each assertion. + + Returns: + SourceResult: Transformed source with the lines seen and converted. + """ + _validate_exception(exception) + module = cst.parse_module(source) + transformer = ConvertAssertTransformer(exception) + updated = MetadataWrapper(module).visit(transformer) + + skipped = Counter(transformer.seen_lines) + skipped.subtract(transformer.converted_lines) + skipped_lines = tuple( + line + for line, count in sorted(skipped.items()) + for _ in range(max(count, 0)) + ) + return SourceResult( + source=updated.code, + converted_lines=tuple(transformer.converted_lines), + skipped_lines=skipped_lines, + ) + + +def convert_file( + path: Path, + exception: str = "AssertionError", + check: bool = False, +) -> FileResult: + """Convert assertions in one Python file. + + Args: + path (Path): Python file to convert. + exception (str): Exception expression to raise in place of each assertion. + check (bool): Report what would change without writing. + + Returns: + FileResult: Whether the file changed, and the conversion counts. + """ + source = path.read_text(encoding="utf-8") + result = transform_source(source, exception=exception) + changed = result.source != source + if changed and not check: + path.write_text(result.source, encoding="utf-8") + return FileResult( + path=path, + converted_lines=result.converted_lines, + skipped_lines=result.skipped_lines, + changed=changed, + ) + + +def _validate_exception(exception: str) -> None: + """Require a simple or dotted exception name, not arbitrary code.""" + if not EXCEPTION_NAME.fullmatch(exception): + raise ValueError( + "exception must be a name already in scope, such as " + "AssertionError, ValueError, or qmcpy.util.ParameterError" + ) + + +def _changed_files(ref: str) -> list[Path]: + """Return changed production Python files relative to ``ref``.""" + result = subprocess.run( + [ + "git", + "diff", + "--name-only", + "--diff-filter=ACMR", + ref, + "--", + "qmcpy/*.py", + ], + capture_output=True, + text=True, + check=True, + ) + return [Path(name) for name in result.stdout.splitlines()] + + +def _python_files(paths: list[str], diff_ref: str | None) -> list[Path]: + """Collect Python files from paths or a production-code diff.""" + if diff_ref is not None: + candidates = _changed_files(diff_ref) + else: + candidates = [Path(path) for path in (paths or ["qmcpy"])] + + files = [] + for path in candidates: + if path.is_dir(): + files.extend(sorted(path.rglob("*.py"))) + elif path.suffix == ".py" and path.exists(): + files.append(path) + return sorted(dict.fromkeys(files)) + + +def _parse_args(argv: list[str]) -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "paths", + nargs="*", + help="Python files or directories to update. Defaults to qmcpy.", + ) + parser.add_argument( + "--diff", + metavar="REF", + help="Use changed qmcpy/*.py files reported by git diff REF.", + ) + parser.add_argument( + "--exception", + default="AssertionError", + help=( + "Exception name already in scope for every selected file. " + "Defaults to AssertionError." + ), + ) + parser.add_argument( + "--check", + action="store_true", + help="Report convertible assertions without writing files.", + ) + parser.add_argument( + "--quiet", + action="store_true", + help="Only print the final summary.", + ) + return parser.parse_args(argv) + + +def main(argv: list[str]) -> int: + """Run the command-line interface. + + Args: + argv (list[str]): Command-line arguments, excluding the program name. + + Returns: + int: Process exit status; ``0`` on success. + """ + args = _parse_args(argv) + try: + _validate_exception(args.exception) + files = _python_files(args.paths, args.diff) + except (ValueError, subprocess.CalledProcessError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + if not files: + print("No Python files to inspect.") + return 0 + + results = [] + had_parse_error = False + for path in files: + try: + result = convert_file( + path, + exception=args.exception, + check=args.check, + ) + except (cst.ParserSyntaxError, UnicodeError) as error: + had_parse_error = True + print(f"{path}: skipped parse error: {error}", file=sys.stderr) + continue + results.append(result) + + if not args.quiet and any(r.converted_lines or r.skipped_lines for r in results): + action = "would convert" if args.check else "converted" + print() + for result in results: + for line in result.converted_lines: + print( + f" - {result.path}:{line}: {action} assert to " + f"explicit {args.exception}" + ) + for line in result.skipped_lines: + print( + f" - {result.path}:{line}: skipped assert in a compound " + "one-line statement" + ) + + converted = sum(len(result.converted_lines) for result in results) + skipped = sum(len(result.skipped_lines) for result in results) + changed_files = sum(result.changed for result in results) + verb = "would change" if args.check else "changed" + print( + f" - {len(files)} file(s) inspected; {converted} assert(s) converted; " + f"{skipped} assert(s) skipped; {changed_files} file(s) {verb}." + ) + + if changed_files == 0: + print(f"clean (0 of {len(files)} files)") + elif args.check: + print(f"ERROR: {changed_files} would change ({changed_files} of {len(files)} files)") + else: + print(f"{changed_files} changed ({changed_files} of {len(files)} files)") + + if args.check and converted: + return 1 + return 2 if had_parse_error else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/flatten_qmcpy_imports.py b/scripts/flatten_qmcpy_imports.py index 033952768..9ee50ce0b 100644 --- a/scripts/flatten_qmcpy_imports.py +++ b/scripts/flatten_qmcpy_imports.py @@ -747,10 +747,19 @@ def flatten_imports( ) -> tuple[bytes, int]: """Flatten, combine, alphabetize, and deduplicate public imports. - `public_names` is qmcpy's public API surface (see `_load_qmcpy_public_names`). - When it's None, nested imports are left unchanged and existing top-level - star imports are deduplicated but left unexpanded. `protect_python` should - be true for Python files so strings and comments are never rewritten. + `public_names` is qmcpy's public API surface (see `_load_qmcpy_public_names`). + When it's None, nested imports are left unchanged and existing top-level + star imports are deduplicated but left unexpanded. `protect_python` should + be true for Python files so strings and comments are never rewritten. + + Args: + content (bytes): File contents to rewrite. + public_names (frozenset[str] | None): Names treated as public; defaults to + the package's own public API. + protect_python (bool): Leave imports inside Python code blocks untouched. + + Returns: + bytes: The rewritten contents, unchanged when nothing needed flattening. """ change_count = 0 @@ -832,7 +841,14 @@ def _is_supported(path: Path) -> bool: def iter_target_files(paths: Iterable[Path]) -> Iterator[Path]: - """Yield supported files under paths, pruning generated and cache directories.""" + """Yield supported files under paths, pruning generated and cache directories. + + Args: + paths (Iterable[Path]): Files or directories to walk. + + Yields: + Path: Each supported file, skipping generated and cache directories. + """ seen: set[Path] = set() for path in paths: @@ -876,6 +892,15 @@ def _display_path(path: Path, base: Path) -> Path: def main(argv: list[str] | None = None) -> int: + """Run the command-line interface. + + Args: + argv (list[str] | None): Command-line arguments, excluding the program + name; defaults to ``sys.argv[1:]``. + + Returns: + int: Process exit status; ``0`` on success. + """ parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--check", @@ -907,7 +932,7 @@ def main(argv: list[str] | None = None) -> int: file=sys.stderr, ) - changed_files = 0 + changed = [] # list of (display_path, import_count) changed_imports = 0 for path in targets: original = path.read_bytes() @@ -919,29 +944,31 @@ def main(argv: list[str] | None = None) -> int: if not count: continue - changed_files += 1 + changed.append((_display_path(path, repository_root), count)) changed_imports += count if not args.check: path.write_bytes(updated) - action = "Would update" if args.check else "Updated" - import_label = "import" if count == 1 else "imports" - print( - f"{action}: {_display_path(path, repository_root)} " - f"({count} {import_label})" - ) - if changed_files: - action = "need updates" if args.check else "updated" + action = "would update" if args.check else "updated" + if changed: + file_label = "file" if len(changed) == 1 else "files" import_label = "import" if changed_imports == 1 else "imports" - file_label = "file" if changed_files == 1 else "files" + print() print( - f"{changed_imports} {import_label} in " - f"{changed_files} {file_label} {action}." + f" - qmcpy imports {action}: {len(changed)} {file_label}, " + f"{changed_imports} {import_label}:" ) + for display_path, count in sorted(changed): + per = "import" if count == 1 else "imports" + print(f" - {display_path} ({count} {per})") + + if not changed: + print(f"clean (0 of {len(targets)} files)") + elif args.check: + print(f"ERROR: {len(changed)} would change ({len(changed)} of {len(targets)} files)") else: - print("All eligible QMCPy imports already use the top-level package.") - - return int(args.check and changed_files > 0) + print(f"{len(changed)} changed ({len(changed)} of {len(targets)} files)") + return int(args.check and bool(changed)) if __name__ == "__main__": diff --git a/scripts/harden_colab_notebook.py b/scripts/harden_colab_notebook.py index 196c80aa8..d0cd83711 100644 --- a/scripts/harden_colab_notebook.py +++ b/scripts/harden_colab_notebook.py @@ -486,16 +486,20 @@ def main() -> int: manifest_path, ) for notebook_rel in successes: - print(f"Hardened {notebook_rel} for Colab.") + print(f" - Hardened {notebook_rel} for Colab.") if failures: - print("") - print("Not yet hardened:") + print() + print(" - Not yet hardened:") for notebook_rel, error in failures: - print(f"- {notebook_rel}: {error}") - print("") + print(f" - {notebook_rel}: {error}") print( - f"Hardened {len(successes)} notebook(s); {len(failures)} notebook(s) still need manual follow-up." + f" - Hardened {len(successes)} notebook(s); {len(failures)} notebook(s) still need manual follow-up." ) + total = len(successes) + len(failures) + if not failures: + print(f"clean (0 of {total} notebooks)") + else: + print(f"ERROR: {len(failures)} need follow-up ({len(failures)} of {total} notebooks)") return 0 if not failures else 1 diff --git a/scripts/remove_trailing_whitespace.py b/scripts/remove_trailing_whitespace.py index e8def8d17..307cfe2c1 100644 --- a/scripts/remove_trailing_whitespace.py +++ b/scripts/remove_trailing_whitespace.py @@ -38,6 +38,14 @@ def iter_source_files(paths: list[str]) -> list[Path]: + """Collect the tracked and untracked files eligible for whitespace cleanup. + + Args: + paths (list[str]): Files or directories to restrict the search to. + + Returns: + list[Path]: Sorted regular files with a supported name or suffix. + """ command = [ "git", "ls-files", @@ -105,6 +113,15 @@ def _strip_python_source(original: bytes) -> bytes: def remove_trailing_whitespace(path: Path, check: bool) -> bool: + """Strip trailing whitespace from one file. + + Args: + path (Path): File to process; binary files are left untouched. + check (bool): Report whether the file would change without writing. + + Returns: + bool: Whether the file changed, or would change under ``check``. + """ original = path.read_bytes() if b"\0" in original: return False @@ -125,11 +142,24 @@ def main() -> int: parser.add_argument("paths", nargs="+", help="tracked files or directories to process") args = parser.parse_args() - changed = [ - path for path in iter_source_files(args.paths) if remove_trailing_whitespace(path, args.check) - ] + scanned = list(iter_source_files(args.paths)) + changed = sorted( + path for path in scanned + if remove_trailing_whitespace(path, args.check) + ) action = "would update" if args.check else "updated" - print(f"trailing whitespace {action}: {len(changed)} file(s)") + if changed: + print() + print(f" - trailing whitespace {action}: {len(changed)} file(s):") + for path in changed: + print(f" - {path}") + + if not changed: + print(f"clean (0 of {len(scanned)} files)") + elif args.check: + print(f"ERROR: {len(changed)} would change ({len(changed)} of {len(scanned)} files)") + else: + print(f"{len(changed)} changed ({len(changed)} of {len(scanned)} files)") return int(args.check and bool(changed)) diff --git a/scripts/unwrap_markdown.py b/scripts/unwrap_markdown.py index 7841d8cc7..ad273bf48 100755 --- a/scripts/unwrap_markdown.py +++ b/scripts/unwrap_markdown.py @@ -24,6 +24,15 @@ def iter_targets(paths: list[str]) -> tuple[list[Path], list[str]]: + """Collect the Markdown and notebook files to process. + + Args: + paths (list[str]): Files or directories to walk. + + Returns: + tuple[list[Path], list[str]]: The files found and a message for each path + that was missing or of an unsupported type. + """ files: list[Path] = [] errors: list[str] = [] for raw_path in paths: @@ -92,6 +101,17 @@ def _paragraph_has_latex(lines: list[str]) -> bool: def unwrap_markdown_text(text: str, *, preserve_latex: bool = False) -> str: + """Join each Markdown paragraph onto a single line. + + Code fences, and optionally display-math blocks, are passed through unchanged. + + Args: + text (str): Markdown source to unwrap. + preserve_latex (bool): Leave display-math blocks unwrapped. + + Returns: + str: The unwrapped text, preserving the original line ending style. + """ if not text: return text @@ -232,6 +252,15 @@ def _split_notebook_source(text: str) -> list[str]: def process_markdown_file(path: Path, check: bool) -> bool: + """Unwrap the paragraphs of one Markdown file. + + Args: + path (Path): Markdown file to process. + check (bool): Report whether the file would change without writing. + + Returns: + bool: Whether the file changed, or would change under ``check``. + """ original = path.read_text(encoding="utf-8") updated = unwrap_markdown_text(original, preserve_latex=True) changed = updated != original @@ -241,6 +270,15 @@ def process_markdown_file(path: Path, check: bool) -> bool: def process_notebook(path: Path, check: bool) -> tuple[bool, int]: + """Unwrap the paragraphs of every Markdown cell in one notebook. + + Args: + path (Path): Notebook file to process. + check (bool): Report whether the notebook would change without writing. + + Returns: + tuple[bool, int]: Whether the notebook changed, and how many cells changed. + """ with path.open(encoding="utf-8") as handle: notebook = json.load(handle) @@ -280,23 +318,37 @@ def main() -> int: print("error: no .md or .ipynb files found", file=sys.stderr) return 2 - changed_files = 0 + changed_paths = [] changed_cells = 0 for path in targets: suffix = path.suffix.lower() if suffix == ".md": - changed = process_markdown_file(path, args.check) - changed_files += int(changed) + if process_markdown_file(path, args.check): + changed_paths.append(path) elif suffix == ".ipynb": changed, cell_count = process_notebook(path, args.check) - changed_files += int(changed) + if changed: + changed_paths.append(path) changed_cells += cell_count mode = "would update" if args.check else "updated" - print( - f"markdown unwrap {mode}: {changed_files} file(s), {changed_cells} markdown cell(s)", + summary = ( + f"markdown unwrap {mode}: {len(changed_paths)} file(s), " + f"{changed_cells} markdown cell(s)" ) - return 1 if args.check and changed_files else 0 + if changed_paths: + print() + print(" - " + summary + ":") + for path in sorted(changed_paths): + print(f" - {path}") + + if not changed_paths: + print(f"clean (0 of {len(targets)} files)") + elif args.check: + print(f"ERROR: {len(changed_paths)} would change ({len(changed_paths)} of {len(targets)} files)") + else: + print(f"{len(changed_paths)} changed ({len(changed_paths)} of {len(targets)} files)") + return 1 if args.check and changed_paths else 0 if __name__ == "__main__": diff --git a/test/README.md b/test/README.md index 0a243c785..73cc08d92 100644 --- a/test/README.md +++ b/test/README.md @@ -25,6 +25,52 @@ This document describes the available test targets in the Makefile for QMCSoftwa | `make delcoverage` | Reset coverage tracking | Instant | Start fresh coverage analysis | +## Test File Organization + +Unit tests live flat in `test/` (no subpackage subfolders). Every file is named: + +``` +test__.py +``` + +`` is a short code for the `qmcpy` subpackage under test, or a cross-cutting bucket: + +| area | scope | +|------|-------| +| `dd` | `qmcpy/discrete_distribution` | +| `ft` | `qmcpy/fast_transform` | +| `ig` | `qmcpy/integrand` | +| `kn` | `qmcpy/kernel` | +| `sc` | `qmcpy/stopping_criterion` | +| `tm` | `qmcpy/true_measure` | +| `ut` | `qmcpy/util` | +| `ee` | end-to-end / cross-cutting pipeline (`integrate()`, worked problems such as Keister and pi) | +| `sr` | `scripts/` tooling, packaging, and docs checks | + +This keeps related tests adjacent when the directory is sorted, and lets you run one area at a time: + +```bash +python -m pytest test/ -k test_tm_ # every true_measure test +make unittests PYTEST_EXTRA_ARGS="-k test_sc_" +``` + +When a test spans two areas (say a stopping criterion exercised against a particular kernel), file it under the component actually under test and name the other in `` — e.g. `test_sc_cubbayes_kernels.py`. Reserve `ee` for cases where neither side is the clear subject. Do not invent new area codes: only the prefixes in the table are accepted, and `make check_test_style STRICT=--strict` fails on anything else. + +Notebook tests are separate: they live in `test/booktests/` as `tb_*.py` and are generated from `demos/` (see `test/booktests/README.md`). + +### Conventions checked by `make check_test_style` + +1. **Area prefix** — the filename must start with a recognized `test__` prefix from the table above. +2. **Object class** — write a test file as one or more `unittest.TestCase` subclasses rather than bare `def test_*` pytest functions. A class groups related assertions under a name (so `pytest -k TestCubMCG` selects them and a failure report names the group), shares construction through `setUp` / `setUpClass` / `self.addCleanup`, and runs identically under `pytest`, `python -m unittest`, and the coverage and booktest runners without depending on pytest fixtures. Most of the suite already follows this; a few legacy files still use bare functions and new files should not. + +`make check_test_style` lists any violation and is informational (exit 0). It also runs as part of `make format`. To make it fail instead — for a pre-commit hook or CI gate — pass `--strict`: + +```bash +STRICT=--strict make check_test_style +``` + +`STRICT=--strict make check_test_style` also runs in CI (the `alltests` workflow), so both conventions are enforced on every pull request. + ## Detailed Descriptions ## Scope diff --git a/test/booktests/tb_makefile_dev_tools.py b/test/booktests/tb_makefile_dev_tools.py new file mode 100644 index 000000000..f04153a01 --- /dev/null +++ b/test/booktests/tb_makefile_dev_tools.py @@ -0,0 +1,12 @@ +import unittest +from testbook import testbook +from __init__ import TB_TIMEOUT, BaseNotebookTest + +class NotebookTests(BaseNotebookTest): + + @testbook('../../demos/makefile_dev_tools.ipynb', execute=True, timeout=TB_TIMEOUT) + def test_makefile_dev_tools_notebook(self, tb): + pass + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_check_links.py b/test/test_check_links.py deleted file mode 100644 index 6a2b3dc52..000000000 --- a/test/test_check_links.py +++ /dev/null @@ -1,178 +0,0 @@ -import ssl -import sys -import urllib.error -from unittest.mock import patch - -from scripts import check_links - - -def _http_error(url, code): - return urllib.error.HTTPError(url, code, "test response", {}, None) - - -def test_head_success_is_reachable(): - with patch.object(check_links.urllib.request, "urlopen", return_value=object()) as urlopen: - assert check_links._check_one("https://example.test", timeout=1) is None - - assert urlopen.call_count == 1 - assert urlopen.call_args.args[0].get_method() == "HEAD" - - -def test_get_success_after_head_failure_is_reachable(): - url = "https://example.test" - with patch.object( - check_links.urllib.request, - "urlopen", - side_effect=[_http_error(url, 405), object()], - ) as urlopen: - assert check_links._check_one(url, timeout=1) is None - - assert urlopen.call_count == 2 - assert urlopen.call_args_list[1].args[0].get_method() == "GET" - - -def test_not_found_and_gone_gets_are_broken(): - for code in (404, 410): - url = f"https://example.test/{code}" - with patch.object( - check_links.urllib.request, - "urlopen", - side_effect=[_http_error(url, code), _http_error(url, code)], - ): - assert check_links._check_one(url, timeout=1) == ( - "broken", - f"{url} -- HTTP {code}", - ) - - -def test_bot_block_and_rate_limit_are_warnings(): - for code in (403, 429): - url = f"https://example.test/{code}" - with patch.object( - check_links.urllib.request, - "urlopen", - side_effect=[_http_error(url, code), _http_error(url, code)], - ): - severity, message = check_links._check_one(url, timeout=1) - - assert severity == "warning" - assert f"HTTP {code}" in message - - -def test_tls_and_timeout_failures_are_warnings(): - failures = ( - ssl.SSLCertVerificationError("certificate verify failed"), - TimeoutError("timed out"), - ) - for failure in failures: - with patch.object( - check_links.urllib.request, - "urlopen", - side_effect=[failure, failure], - ): - severity, message = check_links._check_one( - "https://example.test", timeout=1 - ) - - assert severity == "warning" - assert str(failure) in message - - -def test_external_results_are_separated_and_duplicate_urls_checked_once(tmp_path): - (tmp_path / "page.html").write_text( - 'missing' - 'duplicate' - 'blocked', - encoding="utf-8", - ) - - def result_for(url, _timeout): - if url.endswith("/missing"): - return "broken", f"{url} -- HTTP 404" - return "warning", f"{url} -- HTTP 403" - - with patch.object(check_links, "_check_one", side_effect=result_for) as check_one: - broken, warnings = check_links.check_external(tmp_path, workers=1) - - assert check_one.call_count == 2 - assert broken == [ - "https://example.test/missing -- HTTP 404 (seen on page.html)" - ] - assert warnings == [ - "https://example.test/blocked -- HTTP 403 (seen on page.html)" - ] - - -def test_internal_links_strip_site_url_deployment_path(tmp_path): - target = tmp_path / "target" - target.mkdir() - (target / "index.html").write_text( - '

Target

', encoding="utf-8" - ) - (tmp_path / "index.html").write_text( - 'root-relative' - 'absolute', - encoding="utf-8", - ) - - assert ( - check_links.check_internal( - tmp_path, site_url="https://qmcsoftware.github.io/QMCSoftware/" - ) - == [] - ) - - -def test_external_check_skips_same_site_urls(tmp_path): - (tmp_path / "page.html").write_text( - 'same' - 'external', - encoding="utf-8", - ) - - with patch.object(check_links, "_check_one", return_value=None) as check_one: - broken, warnings = check_links.check_external( - tmp_path, - workers=1, - site_url="https://qmcsoftware.github.io/QMCSoftware/", - ) - - assert broken == [] - assert warnings == [] - assert check_one.call_count == 1 - assert check_one.call_args.args[0] == "https://example.test/target/" - - -def test_external_warnings_do_not_make_main_fail(tmp_path, monkeypatch, capsys): - monkeypatch.setattr(sys, "argv", ["check_links.py", str(tmp_path), "--external"]) - monkeypatch.setattr( - check_links, "check_internal", lambda _site_dir, site_url=None: [] - ) - monkeypatch.setattr( - check_links, - "check_external", - lambda _site_dir, site_url=None: ( - [], - ["https://example.test -- HTTP 403"], - ), - ) - - assert check_links.main() == 0 - assert "0 broken link(s), 1 warning(s)" in capsys.readouterr().out - - -def test_confirmed_external_breakage_makes_main_fail(tmp_path, monkeypatch): - monkeypatch.setattr(sys, "argv", ["check_links.py", str(tmp_path), "--external"]) - monkeypatch.setattr( - check_links, "check_internal", lambda _site_dir, site_url=None: [] - ) - monkeypatch.setattr( - check_links, - "check_external", - lambda _site_dir, site_url=None: ( - ["https://example.test -- HTTP 404"], - [], - ), - ) - - assert check_links.main() == 1 diff --git a/test/test_check_removed_urls.py b/test/test_check_removed_urls.py deleted file mode 100644 index 9374e34ad..000000000 --- a/test/test_check_removed_urls.py +++ /dev/null @@ -1,134 +0,0 @@ -import sys -import urllib.error -from unittest.mock import patch - -from scripts import check_removed_urls as cru - -SITE = "https://qmcsoftware.github.io/QMCSoftware/" - - -def _sitemap(*paths): - locs = "".join(f"{SITE}{path}" for path in paths) - return f'{locs}' - - -def _config(redirect_maps=None): - plugins = ["material/search", {"mkdocs-jupyter": {"execute": False}}] - if redirect_maps is not None: - plugins.append({"redirects": {"redirect_maps": redirect_maps}}) - return {"site_url": SITE, "plugins": plugins} - - -def _run(tmp_path, monkeypatch, sitemap_paths, redirect_maps=None, extra_argv=()): - """Run main() offline against a temp sitemap and a temp docs/ tree.""" - docs = tmp_path / "docs" - docs.mkdir(parents=True) - (docs / "README.md").write_text("home", encoding="utf-8") - (docs / "good_practices.md").write_text("page", encoding="utf-8") - sitemap = tmp_path / "sitemap.xml" - sitemap.write_text(_sitemap(*sitemap_paths), encoding="utf-8") - - monkeypatch.setattr(cru, "read_config", lambda *a, **k: _config(redirect_maps)) - monkeypatch.setattr(sys, "argv", [ - "check_removed_urls.py", "--sitemap", str(sitemap), "--docs-dir", str(docs), - *extra_argv, - ]) - return cru.main() - - -def test_url_path_and_source_round_trip(tmp_path): - for source, url_path in [("blogs/scipywrapper/index.md", "blogs/scipywrapper/"), - ("good_practices.md", "good_practices/"), - ("demos/quickstart.ipynb", "demos/quickstart/"), - ("index.md", ""), ("README.md", "")]: - assert cru.url_path_for_source(source) == url_path - - for source in ("README.md", "good_practices.md", "demos/quickstart.ipynb", - "api/index.md"): - path = tmp_path / source - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text("page", encoding="utf-8") - assert cru.source_exists(cru.url_path_for_source(source), tmp_path) - assert not cru.source_exists("blogs/scipywrapper/", tmp_path) - - -def test_redirect_maps_reads_the_plugin_and_tolerates_its_absence(): - entry = {"blogs/x/index.md": "https://qmcsoftware.org/blogs/x/"} - assert cru.redirect_maps(_config(entry)) == entry - assert cru.redirect_maps(_config()) == {} - assert cru.redirect_maps({}) == {} - - -def test_published_paths_separates_foreign_urls(): - sitemap = _sitemap("", "good_practices/").replace( - "", "https://example.test/other/") - - assert cru.published_paths(sitemap, SITE) == ( - ["", "good_practices/"], ["https://example.test/other/"]) - - -def test_http_status_falls_back_to_get_when_head_is_unsupported(): - url = "https://example.test" - error = urllib.error.HTTPError(url, 405, "test response", {}, None) - response = type("Response", (), {"status": 200, "__enter__": lambda s: s, - "__exit__": lambda s, *a: False})() - with patch.object(cru.urllib.request, "urlopen", - side_effect=[error, response]) as urlopen: - assert cru.http_status(url, timeout=1) == "200" - - assert urlopen.call_count == 2 - assert urlopen.call_args_list[1].args[0].get_method() == "GET" - - -def test_removed_page_without_redirect_is_flagged(tmp_path, monkeypatch, capsys): - code = _run(tmp_path, monkeypatch, ["", "good_practices/", "blogs/scipywrapper/"]) - out = capsys.readouterr().out - - assert code == 1 - assert "1 removed with no redirect" in out - assert f"[ORPHAN] {SITE}blogs/scipywrapper/" in out - assert "blogs/scipywrapper/index.md: " in out - - -def test_removed_page_covered_by_a_redirect_passes(tmp_path, monkeypatch, capsys): - code = _run( - tmp_path, monkeypatch, ["", "good_practices/", "blogs/scipywrapper/"], - redirect_maps={ - "blogs/scipywrapper/index.md": "https://qmcsoftware.org/blogs/scipywrapper/"}, - ) - out = capsys.readouterr().out - - assert code == 0 - assert "0 removed with no redirect" in out - assert "[redirect]" in out and "[ORPHAN]" not in out - - -def test_intact_site_passes(tmp_path, monkeypatch, capsys): - assert _run(tmp_path, monkeypatch, ["", "good_practices/"]) == 0 - assert "2 still have a page source" in capsys.readouterr().out - - -def test_verify_redirects_follows_the_target_status(tmp_path, monkeypatch, capsys): - redirects = {"blogs/x/index.md": "https://qmcsoftware.org/blogs/x/"} - for status, expected_code in [("200", 0), ("404", 1)]: - monkeypatch.setattr(cru, "http_status", lambda *a, **k: status) - code = _run(tmp_path / status, monkeypatch, ["", "blogs/x/"], - redirect_maps=redirects, extra_argv=("--verify-redirects",)) - out = capsys.readouterr().out - - assert code == expected_code - assert status in out - # The URL itself is covered, so a failure is the target, not an orphan. - assert "[ORPHAN]" not in out - - -def test_unreachable_sitemap_fails_unless_offline_is_allowed(tmp_path, monkeypatch, capsys): - monkeypatch.setattr(cru, "read_config", lambda *a, **k: _config()) - argv = ["check_removed_urls.py", "--sitemap", str(tmp_path / "absent.xml")] - - monkeypatch.setattr(sys, "argv", argv) - assert cru.main() == 1 - - monkeypatch.setattr(sys, "argv", argv + ["--allow-offline"]) - assert cru.main() == 0 - assert "skipping the check" in capsys.readouterr().out diff --git a/test/test_colab_notebooks.py b/test/test_colab_notebooks.py deleted file mode 100644 index 858c4c4b9..000000000 --- a/test/test_colab_notebooks.py +++ /dev/null @@ -1,314 +0,0 @@ -from __future__ import annotations - -import json -import os -import sys -from pathlib import Path - -import pytest - -from scripts import check_colab_notebooks as check -from scripts import harden_colab_notebook as harden -from scripts import smoke_test_colab_notebooks as smoke - - -def markdown_cell(source: str, cell_id: str = "markdown") -> dict: - return { - "cell_type": "markdown", - "id": cell_id, - "metadata": {}, - "source": source.splitlines(keepends=True), - } - - -def code_cell(source: str, cell_id: str = "code") -> dict: - return { - "cell_type": "code", - "execution_count": None, - "id": cell_id, - "metadata": {}, - "outputs": [], - "source": source.splitlines(keepends=True), - } - - -@pytest.fixture -def colab_repo(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): - demos_dir = tmp_path / "demos" - demos_dir.mkdir() - notebook_path = demos_dir / "example.ipynb" - notebook = { - "cells": [ - markdown_cell("# Example\n", "title"), - code_cell("import math\n", "imports"), - ], - "metadata": {}, - "nbformat": 4, - "nbformat_minor": 5, - } - notebook_path.write_text(json.dumps(notebook, indent=1) + "\n", encoding="utf-8") - - manifest_path = tmp_path / "manifest.json" - manifest = { - "repo": "QMCSoftware/QMCSoftware", - "git_ref": "develop", - "enabled": [], - "disabled": {}, - } - manifest_path.write_text(json.dumps(manifest, indent=1) + "\n", encoding="utf-8") - - monkeypatch.setattr(check, "REPO_ROOT", tmp_path) - monkeypatch.setattr(check, "DEMOS_DIR", demos_dir) - monkeypatch.setattr(harden, "REPO_ROOT", tmp_path) - monkeypatch.setattr(smoke, "REPO_ROOT", tmp_path) - return notebook_path, manifest_path - - -def test_badge_stripping_preserves_intro_and_drops_badge_only_cells(): - intro = markdown_cell( - "# ML Sensitivity Indices\n\n" - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)]" - "(https://colab.research.google.com/github/QMCSoftware/QMCSoftware/" - "blob/develop/demos/iris.ipynb)\n\n" - "This notebook demonstrates sensitivity indices.\n" - ) - badge_only = markdown_cell( - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)]" - "(https://colab.research.google.com/github/QMCSoftware/QMCSoftware/" - "blob/develop/demos/iris.ipynb)\n" - ) - - cleaned_intro = harden.badge_stripped_cell(intro) - assert cleaned_intro is not None - assert "# ML Sensitivity Indices" in check.cell_source_text(cleaned_intro) - assert "sensitivity indices" in check.cell_source_text(cleaned_intro) - assert "Open In Colab" not in check.cell_source_text(cleaned_intro) - assert harden.remove_any_badge_cells([badge_only, code_cell("pass\n")]) == [ - code_cell("pass\n") - ] - - -def test_is_any_badge_cell_rejects_spoofed_hostname(): - spoofed = markdown_cell( - "[click](https://evil.example/colab.research.google.com/assets/colab-badge.svg)\n" - ) - genuine = markdown_cell( - "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)]" - "(https://colab.research.google.com/github/QMCSoftware/QMCSoftware/" - "blob/develop/demos/iris.ipynb)\n" - ) - - assert not check.is_any_badge_cell(spoofed) - assert check.is_any_badge_cell(genuine) - - -def test_bootstrap_detection_uses_marker_and_real_install_command( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - misleading = code_cell( - '"""import google.colab\n# @title Execute this cell to install dependencies\n' - '!pip install qmcpy\n"""\n' - ) - comment_only = code_cell( - "# @title Execute this cell to install dependencies\n" - "# import google.colab\n" - "# !pip install qmcpy\n" - ) - assert not check.is_any_install_cell(misleading) - assert not check.is_bootstrap_cell(misleading) - assert check.is_any_install_cell(comment_only) - assert not check.is_bootstrap_cell(comment_only) - - monkeypatch.setattr(harden, "REPO_ROOT", tmp_path) - notebook_path = tmp_path / "demos" / "example.ipynb" - notebook_path.parent.mkdir() - source = "".join( - harden.bootstrap_cell_source( - notebook_path, - {"repo": "QMCSoftware/QMCSoftware"}, - [], - ) - ) - generated = code_cell(source) - assert check.is_bootstrap_cell(generated) - assert "except ImportError:" in source - assert "if IN_COLAB:" in source - assert "except:\n" not in source - compile(smoke.rewrite_shell_magics(source), "", "exec") - - -def test_extra_pip_packages_preserves_later_explicit_installs(): - cells = [ - code_cell("import qmcpy as qp\n"), - code_cell("import ipywidgets as widgets\n"), - code_cell( - "try:\n" - " import QuantLib as ql\n" - "except ModuleNotFoundError:\n" - " !pip install -q QuantLib\n" - ), - code_cell("!pip install -q seaborn\n"), - ] - - assert harden.extra_pip_packages(cells) == ["QuantLib", "ipywidgets", "seaborn"] - - -def test_needs_latex_setup_detects_tueplots(): - cells = [ - code_cell("import qmcpy as qp\n"), - code_cell( - "from tueplots import bundles\n" - "pyplot.rcParams.update(bundles.probnum2025())\n" - ), - ] - - assert harden.needs_latex_setup(cells) - - -def test_imported_modules_survives_magic_only_block_body(): - # A shell-magic line as the *only* statement in a block used to leave an - # empty `if:`/`try:` body, making ast.parse raise and silently hiding - # every import in the cell (not just the magic line itself). - source = ( - "import os\n" - "from util import helper\n" - "if True:\n" - " !echo hi\n" - ) - assert check.imported_modules(source) == {"os", "util"} - - -def test_local_module_matches_finds_ancestor_directory( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): - monkeypatch.setattr(check, "DEMOS_DIR", tmp_path) - (tmp_path / "util.py").write_text("", encoding="utf-8") - notebook_dir = tmp_path / "output" - notebook_dir.mkdir() - - matches = check.local_module_matches(notebook_dir, "util") - - assert matches == [tmp_path / "util.py"] - - -def test_extra_pip_packages_honors_colab_deps_marker(): - cells = [ - code_cell("import qmcpy as qp\n"), - code_cell( - "# colab-deps: plotly, some-package\n" - "import plotly\n" - ), - ] - - assert harden.extra_pip_packages(cells) == ["plotly", "some-package"] - - -def test_dump_notebook_preserves_existing_json_indent(tmp_path: Path): - notebook_path = tmp_path / "example.ipynb" - notebook = { - "cells": [code_cell("pass\n")], - "metadata": {}, - "nbformat": 4, - "nbformat_minor": 5, - } - original_source = json.dumps(notebook, indent=2) + "\n" - - harden.dump_notebook(notebook_path, notebook, original_source) - - assert notebook_path.read_text(encoding="utf-8") == original_source - - -def test_harden_check_smoke_round_trip_is_idempotent( - colab_repo, monkeypatch: pytest.MonkeyPatch -): - notebook_path, manifest_path = colab_repo - harden.harden_notebook(notebook_path, manifest_path) - - assert check.run_check(manifest_path, strict=True) == 0 - smoke_notebook, source_indices = smoke.build_smoke_notebook(notebook_path, 1) - assert len(smoke_notebook["cells"]) == len(source_indices) - - sentinel = object() - old_modules = { - name: sys.modules.get(name, sentinel) for name in ("google", "google.colab") - } - old_environment = { - name: os.environ.get(name, sentinel) - for name in ("QMC_COLAB_SMOKE", "QMC_COLAB_SMOKE_REPO_ROOT", "QMC_COLAB_SMOKE_NOTEBOOK_DIR") - } - namespace: dict = {} - try: - for cell in smoke_notebook["cells"]: - if cell["cell_type"] == "code": - exec(check.cell_source_text(cell), namespace) - finally: - for name, value in old_modules.items(): - if value is sentinel: - sys.modules.pop(name, None) - else: - sys.modules[name] = value - for name, value in old_environment.items(): - if value is sentinel: - os.environ.pop(name, None) - else: - os.environ[name] = value - - monkeypatch.setattr( - harden, - "dump_notebook", - lambda *_args, **_kwargs: pytest.fail("unchanged notebook was rewritten"), - ) - monkeypatch.setattr( - harden, - "dump_json", - lambda *_args, **_kwargs: pytest.fail("unchanged manifest was rewritten"), - ) - harden.harden_notebook(notebook_path, manifest_path) - - -def test_checker_rejects_wrong_badge(colab_repo): - notebook_path, manifest_path = colab_repo - harden.harden_notebook(notebook_path, manifest_path) - notebook = check.load_json(notebook_path) - badge = next(cell for cell in notebook["cells"] if check.is_any_badge_cell(cell)) - badge["source"] = [check.cell_source_text(badge).replace("develop", "wrong-ref")] - notebook_path.write_text(json.dumps(notebook, indent=1) + "\n", encoding="utf-8") - - assert check.run_check(manifest_path, strict=True) == 1 - - -def test_harden_failure_does_not_disable_notebook( - colab_repo, monkeypatch: pytest.MonkeyPatch -): - notebook_path, manifest_path = colab_repo - original_manifest = manifest_path.read_text(encoding="utf-8") - monkeypatch.setattr( - harden, - "harden_notebook", - lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("failure")), - ) - - successes, failures = harden.harden_batch([notebook_path], manifest_path) - - assert successes == [] - assert failures == [("demos/example.ipynb", "failure")] - assert manifest_path.read_text(encoding="utf-8") == original_manifest - - -def test_smoke_batch_continues_after_a_notebook_failure(monkeypatch: pytest.MonkeyPatch): - def fake_build(notebook_path: Path, cells_after_bootstrap: int): - return {"cells": []}, [] - - def fake_execute(notebook_path: Path, smoke_nb, source_indices, timeout): - if "broken" in notebook_path.as_posix(): - raise RuntimeError("boom") - - monkeypatch.setattr(smoke, "build_smoke_notebook", fake_build) - monkeypatch.setattr(smoke, "execute_smoke_notebook", fake_execute) - - passed, failed = smoke.smoke_test_batch( - ["demos/broken.ipynb", "demos/ok.ipynb"], cells_after_bootstrap=1, timeout=60 - ) - - assert passed == ["demos/ok.ipynb"] - assert failed == [("demos/broken.ipynb", "boom")] diff --git a/test/test_copulas.py b/test/test_copulas.py deleted file mode 100644 index 2f4bd06ba..000000000 --- a/test/test_copulas.py +++ /dev/null @@ -1,1362 +0,0 @@ -import warnings - -import numpy as np -import pytest -import scipy.stats as stats - -from qmcpy import ( - AbstractCopula, - ClaytonCopula, - DigitalNetB2, - FrankCopula, - GaussianCopula, - GumbelCopula, - StudentTCopula, -) - -from qmcpy.true_measure.copula import ( - AbstractCopula as ModuleAbstractCopula, - _apply_marginal_ppfs, - _build_marginal_range, - _clip_unit_interval, - _marginal_cdfs_and_logpdf, - _validate_correlation_matrix, - _validate_dimension, - _validate_marginals, -) - -from qmcpy.util import DimensionError, MethodImplementationError, ParameterError - - -class PPFOnlyMarginal: - def ppf(self, u): - return np.asarray(u, dtype=float) - - -class NonCallablePPFMarginal: - ppf = 1.0 - - -class UnitPDFMarginal: - def ppf(self, u): - return np.asarray(u, dtype=float) - - def cdf(self, x): - return np.asarray(x, dtype=float) - - def pdf(self, x): - return np.ones_like(np.asarray(x, dtype=float)) - - -class CDFOnlyMarginal(PPFOnlyMarginal): - def cdf(self, x): - return np.asarray(x, dtype=float) - - -class BadIntervalMarginal(PPFOnlyMarginal): - def interval(self, confidence): - raise ValueError("interval unavailable") - - -class BadRangeMarginal: - def ppf(self, u): - raise ValueError("ppf unavailable") - - -def _equicorrelation(d, rho): - corr = np.full((d, d), rho, dtype=float) - np.fill_diagonal(corr, 1.0) - return corr - - -def _make_copula(copula_cls, dimension=2, marginals=None, correlation=None, seed=7): - if marginals is None: - marginals = [stats.norm()] * dimension - if correlation is None: - correlation = np.eye(dimension) - - kwargs = {} - if copula_cls is StudentTCopula: - kwargs["df"] = 4 - if copula_cls is ClaytonCopula: - kwargs["theta"] = 2.0 - if copula_cls is FrankCopula: - kwargs["theta"] = 5.0 - if copula_cls is GumbelCopula: - kwargs["theta"] = 2.0 - - common = { - "sampler": DigitalNetB2(dimension, seed=seed), - "marginals": marginals, - **kwargs, - } - if copula_cls in [ClaytonCopula, FrankCopula, GumbelCopula]: - return copula_cls(**common) - return copula_cls(correlation=correlation, **common) - - -# Base AbstractCopula and helper tests - - -def test_abstract_copula_is_importable_from_public_module_path(): - assert ModuleAbstractCopula is AbstractCopula - - -def test_public_api_imports_and_normal_usage(): - for copula_cls in [ - GaussianCopula, - StudentTCopula, - ClaytonCopula, - FrankCopula, - GumbelCopula, - ]: - assert issubclass(copula_cls, AbstractCopula) - - tm = _make_copula(copula_cls) - x = tm(8) - x_gen = tm.gen_samples(8) - v = tm.gen_copula_samples(8) - - assert x.shape == (8, 2) - assert x_gen.shape == (8, 2) - assert v.shape == (8, 2) - assert np.all(np.isfinite(x)) - assert np.all(np.isfinite(x_gen)) - assert np.all((0 <= v) & (v <= 1)) - - -def test_abstract_copula_rejects_unimplemented_transform(): - tm = AbstractCopula( - DigitalNetB2(2, seed=101), - marginals=[stats.uniform(), stats.uniform()], - ) - - with pytest.raises(MethodImplementationError): - tm.copula_transform(np.full((3, 2), 0.5)) - - -def test_abstract_copula_rejects_invalid_sampler(): - with pytest.raises(ParameterError, match="sampler"): - AbstractCopula(object(), marginals=[stats.uniform()]) - - -def test_validate_marginals_error_branches(): - with pytest.raises(ParameterError, match="marginals"): - _validate_marginals(None) - - with pytest.raises(ParameterError, match="at least one"): - _validate_marginals([]) - - with pytest.raises(ParameterError, match="ppf"): - _validate_marginals([NonCallablePPFMarginal()]) - - -def test_validate_dimension_error_branches(): - with pytest.raises(DimensionError, match="integer dimension"): - _validate_dimension(object(), [stats.uniform()]) - - with pytest.raises(DimensionError, match="marginals"): - _validate_dimension(3, [stats.uniform(), stats.uniform()]) - - -def test_apply_marginal_ppfs_clips_endpoints_and_checks_dimension(): - transformed = _apply_marginal_ppfs( - np.array([[0.0, 1.0], [1.0, 0.0]]), - [stats.norm(), stats.norm()], - ) - - assert transformed.shape == (2, 2) - assert np.all(np.isfinite(transformed)) - - with pytest.raises(DimensionError, match="marginals"): - _apply_marginal_ppfs(np.full((2, 3), 0.5), [stats.uniform(), stats.uniform()]) - - -def test_marginal_range_falls_back_when_interval_or_ppf_fails(): - ranges = _build_marginal_range([BadIntervalMarginal(), BadRangeMarginal()]) - - assert ranges.shape == (2, 2) - assert np.all(np.isfinite(ranges[0])) - np.testing.assert_allclose(ranges[1], [-np.inf, np.inf]) - - -def test_marginal_cdfs_and_logpdf_pdf_branch_and_errors(): - x = np.array([[0.25, 0.75], [0.4, 0.6]]) - u, log_density = _marginal_cdfs_and_logpdf( - x, - [UnitPDFMarginal(), UnitPDFMarginal()], - ) - - np.testing.assert_allclose(u, x) - np.testing.assert_allclose(log_density, np.zeros(2)) - - with pytest.raises(ParameterError, match="cdf"): - _marginal_cdfs_and_logpdf(x, [PPFOnlyMarginal(), UnitPDFMarginal()]) - - with pytest.raises(ParameterError, match="pdf"): - _marginal_cdfs_and_logpdf(x, [CDFOnlyMarginal(), UnitPDFMarginal()]) - - -def test_validate_correlation_matrix_rejects_nonfinite_values(): - with pytest.raises(ValueError, match="finite"): - _validate_correlation_matrix([[1.0, np.nan], [np.nan, 1.0]], 2) - - -def test_clip_unit_interval_uses_machine_epsilon(): - clipped = _clip_unit_interval(np.array([0.0, 0.5, 1.0])) - eps = np.finfo(float).eps - - np.testing.assert_allclose(clipped, [eps, 0.5, 1.0 - eps]) - - -@pytest.mark.parametrize( - "copula_cls", - [GaussianCopula, StudentTCopula, ClaytonCopula, GumbelCopula, FrankCopula], -) -def test_copula_transform_outputs_dependent_uniforms_in_unit_cube(copula_cls): - tm = _make_copula(copula_cls, dimension=3) - u = np.array( - [ - [0.1, 0.3, 0.7], - [0.5, 0.5, 0.5], - [0.9, 0.8, 0.2], - ] - ) - - v = tm.copula_transform(u) - - assert v.shape == u.shape - assert np.all(np.isfinite(v)) - assert np.all((0.0 <= v) & (v <= 1.0)) - - -@pytest.mark.parametrize( - "copula_cls,dimension", - [ - (GaussianCopula, 3), - (StudentTCopula, 3), - (ClaytonCopula, 3), - (FrankCopula, 3), - (GumbelCopula, 3), - ], -) -def test_copula_sample_shapes_are_preserved(copula_cls, dimension): - tm = _make_copula(copula_cls, dimension=dimension, seed=9) - - one = tm(1) - many = tm(8) - batched_transform = tm._transform(np.full((2, 3, dimension), 0.5)) - - assert one.shape == (1, dimension) - assert many.shape == (8, dimension) - assert batched_transform.shape == (2, 3, dimension) - assert np.all(np.isfinite(one)) - assert np.all(np.isfinite(many)) - assert np.all(np.isfinite(batched_transform)) - - -# Elliptical copulas - - -def test_output_shape_with_nonnormal_marginals(): - tm = GaussianCopula( - sampler=DigitalNetB2(2, seed=7), - marginals=[stats.beta(a=2, b=5), stats.gamma(a=3, scale=2)], - correlation=[[1.0, 0.4], [0.4, 1.0]], - ) - - x = tm(16) - - assert x.shape == (16, 2) - - -def test_finite_output_for_normal_marginals(): - tm = GaussianCopula( - sampler=DigitalNetB2(2, seed=11), - marginals=[stats.norm(), stats.norm(loc=1.0, scale=2.0)], - correlation=[[1.0, -0.3], [-0.3, 1.0]], - ) - - x = tm(128) - - assert np.all(np.isfinite(x)) - - -def test_return_weights_shape_when_marginal_densities_available(): - tm = GaussianCopula( - sampler=DigitalNetB2(2, seed=12), - marginals=[stats.norm(), stats.gamma(a=2.0)], - correlation=[[1.0, 0.25], [0.25, 1.0]], - ) - - x, weights = tm(32, return_weights=True) - - assert x.shape == (32, 2) - assert weights.shape == (32,) - assert np.all(np.isfinite(weights)) - assert np.all(weights > 0.0) - - -def test_identity_correlation_matches_independent_marginal_transforms(): - marginals = [stats.norm(loc=-1.0, scale=2.0), stats.gamma(a=2.0, scale=3.0)] - tm = GaussianCopula( - sampler=DigitalNetB2(2, seed=13), - marginals=marginals, - correlation=np.eye(2), - ) - u = np.array([[0.2, 0.7], [0.4, 0.8], [0.9, 0.1]]) - - x = tm._transform(u) - expected = np.column_stack( - [marginals[j].ppf(u[:, j]) for j in range(len(marginals))] - ) - - np.testing.assert_allclose(x, expected, rtol=1e-12, atol=1e-12) - - -def test_positive_correlation_produces_positive_dependence(): - rho = 0.75 - tm = GaussianCopula( - sampler=DigitalNetB2(2, seed=17), - marginals=[stats.norm(), stats.norm()], - correlation=[[1.0, rho], [rho, 1.0]], - ) - - x = tm(4096) - empirical_corr = np.corrcoef(x.T)[0, 1] - - assert empirical_corr > 0.5 - assert abs(empirical_corr - rho) < 0.2 - - -@pytest.mark.parametrize("copula_cls", [GaussianCopula, StudentTCopula]) -@pytest.mark.parametrize("dimension", [1, 3, 5]) -def test_elliptical_copulas_support_general_dimensions(copula_cls, dimension): - correlation = _equicorrelation(dimension, 0.25) - tm = _make_copula( - copula_cls, - dimension=dimension, - marginals=[stats.norm()] * dimension, - correlation=correlation, - seed=19, - ) - - x = tm(16) - one = tm(1) - - assert x.shape == (16, dimension) - assert one.shape == (1, dimension) - assert np.all(np.isfinite(x)) - assert np.all(np.isfinite(one)) - - -@pytest.mark.parametrize("copula_cls", [GaussianCopula, StudentTCopula]) -def test_elliptical_copulas_handle_valid_near_singular_correlation(copula_cls): - dimension = 5 - tm = _make_copula( - copula_cls, - dimension=dimension, - marginals=[stats.norm()] * dimension, - correlation=_equicorrelation(dimension, 0.999), - seed=20, - ) - - x = tm(32) - - assert x.shape == (32, dimension) - assert np.all(np.isfinite(x)) - - -@pytest.mark.parametrize("copula_cls", [GaussianCopula, StudentTCopula]) -def test_elliptical_copulas_reject_singular_correlation(copula_cls): - with pytest.raises(ValueError, match="positive definite"): - _make_copula( - copula_cls, - dimension=3, - marginals=[stats.norm(), stats.norm(), stats.norm()], - correlation=np.ones((3, 3)), - seed=22, - ) - - -@pytest.mark.parametrize( - "copula_cls", - [GaussianCopula, StudentTCopula, ClaytonCopula, FrankCopula, GumbelCopula], -) -def test_distribution_dimension_matches_number_of_marginals(copula_cls): - tm = _make_copula( - copula_cls, - dimension=5, - marginals=[ - stats.norm(), - stats.beta(a=2, b=5), - stats.gamma(a=3), - stats.expon(), - stats.lognorm(s=0.5), - ], - correlation=np.eye(5), - ) - - x = tm(32) - - assert x.shape == (32, 5) - assert np.all(np.isfinite(x)) - - -@pytest.mark.parametrize("copula_cls", [GaussianCopula, StudentTCopula]) -def test_invalid_dimension_mismatches_raise(copula_cls): - with pytest.raises(DimensionError, match="marginals"): - _make_copula( - copula_cls, - dimension=2, - marginals=[stats.norm(), stats.norm(), stats.norm()], - correlation=np.eye(2), - ) - - with pytest.raises(ValueError, match="shape"): - _make_copula( - copula_cls, - dimension=2, - marginals=[stats.norm(), stats.norm()], - correlation=np.eye(3), - ) - - with pytest.raises(ValueError, match="square"): - _make_copula( - copula_cls, - dimension=2, - marginals=[stats.norm(), stats.norm()], - correlation=[[1.0, 0.2, 0.3], [0.2, 1.0, 0.4]], - ) - - -@pytest.mark.parametrize("copula_cls", [ClaytonCopula, FrankCopula, GumbelCopula]) -def test_archimedean_dimension_mismatch_raises_dimension_error(copula_cls): - with pytest.raises(DimensionError, match="marginals"): - _make_copula( - copula_cls, - dimension=2, - marginals=[stats.norm(), stats.norm(), stats.norm()], - ) - - -@pytest.mark.parametrize( - "copula_cls", - [GaussianCopula, StudentTCopula], -) -@pytest.mark.parametrize( - "correlation", - [ - [[1.0, 0.2], [0.3, 1.0]], - [[1.0, 0.2], [0.2, 0.9]], - [[1.0, 1.2], [1.2, 1.0]], - ], -) -def test_invalid_correlation_matrices_raise_value_error(copula_cls, correlation): - with pytest.raises(ValueError): - _make_copula( - copula_cls, - dimension=2, - marginals=[stats.norm(), stats.norm()], - correlation=correlation, - ) - - -def test_marginal_length_mismatch_raises_dimension_error(): - with pytest.raises(DimensionError, match="marginals"): - GaussianCopula( - sampler=DigitalNetB2(2, seed=21), - marginals=[stats.norm()], - correlation=np.eye(2), - ) - - -def test_marginal_without_ppf_raises_clear_error(): - class NoPPF: - pass - - with pytest.raises(ParameterError, match="ppf"): - GaussianCopula( - sampler=DigitalNetB2(1, seed=23), - marginals=[NoPPF()], - correlation=[[1.0]], - ) - - -@pytest.mark.parametrize( - "copula_cls", - [GaussianCopula, StudentTCopula, ClaytonCopula, FrankCopula, GumbelCopula], -) -def test_common_scipy_frozen_marginals_work(copula_cls): - tm = _make_copula( - copula_cls, - dimension=5, - marginals=[ - stats.norm(), - stats.beta(a=2, b=5), - stats.gamma(a=3), - stats.expon(), - stats.lognorm(s=0.5), - ], - correlation=np.eye(5), - seed=47, - ) - - x = tm(128) - - assert x.shape == (128, 5) - assert np.all(np.isfinite(x)) - - -@pytest.mark.parametrize( - "copula_cls", - [GaussianCopula, StudentTCopula, ClaytonCopula, FrankCopula, GumbelCopula], -) -def test_endpoint_uniforms_are_clipped_to_finite_outputs(copula_cls): - tm = _make_copula( - copula_cls, - dimension=5, - marginals=[ - stats.norm(), - stats.beta(a=2, b=5), - stats.gamma(a=3), - stats.expon(), - stats.lognorm(s=0.5), - ], - correlation=np.eye(5), - seed=53, - ) - u = np.array( - [ - [0.0, 1.0, 0.0, 1.0, 0.5], - [1.0, 0.0, 1.0, 0.0, 0.5], - ] - ) - - x = tm._transform(u) - - assert x.shape == (2, 5) - assert np.all(np.isfinite(x)) - - -def test_student_t_copula_output_shape_and_finite_values(): - tm = StudentTCopula( - sampler=DigitalNetB2(2, seed=29), - marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], - correlation=[[1.0, 0.5], [0.5, 1.0]], - df=4, - ) - - x = tm(128) - - assert x.shape == (128, 2) - assert np.all(np.isfinite(x)) - - -def test_student_t_copula_positive_correlation_produces_positive_dependence(): - tm = StudentTCopula( - sampler=DigitalNetB2(2, seed=31), - marginals=[stats.norm(), stats.norm()], - correlation=[[1.0, 0.7], [0.7, 1.0]], - df=5, - ) - - x = tm(4096) - empirical_corr = np.corrcoef(x.T)[0, 1] - - assert empirical_corr > 0.45 - - -def test_student_t_copula_has_stronger_joint_tail_than_gaussian_copula(): - rho = 0.7 - df = 4 - n = 2**12 - marginals = [stats.norm(), stats.norm()] - correlation = [[1.0, rho], [rho, 1.0]] - - gaussian = GaussianCopula( - sampler=DigitalNetB2(2, seed=101), - marginals=marginals, - correlation=correlation, - ) - student_t = StudentTCopula( - sampler=DigitalNetB2(2, seed=101), - marginals=marginals, - correlation=correlation, - df=df, - ) - - x_gaussian = gaussian(n) - x_student_t = student_t(n) - threshold = stats.norm.ppf(0.99) - - def joint_tail_rate(x): - tail_0 = x[:, 0] > threshold - return np.mean(x[tail_0, 1] > threshold) - - gaussian_tail = joint_tail_rate(x_gaussian) - student_t_tail = joint_tail_rate(x_student_t) - - assert student_t_tail > gaussian_tail + 0.08 - - -def test_student_t_copula_return_weights_shape_when_density_available(): - tm = StudentTCopula( - sampler=DigitalNetB2(2, seed=37), - marginals=[stats.norm(), stats.gamma(a=2.0)], - correlation=[[1.0, 0.3], [0.3, 1.0]], - df=6, - ) - - x, weights = tm(32, return_weights=True) - - assert x.shape == (32, 2) - assert weights.shape == (32,) - assert np.all(np.isfinite(weights)) - assert np.all(weights > 0.0) - - -@pytest.mark.parametrize("df", [1.0, 100.0]) -def test_student_t_copula_boundary_df_values_are_finite(df): - dimension = 3 - tm = StudentTCopula( - sampler=DigitalNetB2(dimension, seed=39), - marginals=[stats.norm()] * dimension, - correlation=_equicorrelation(dimension, 0.4), - df=df, - ) - - x = tm(128) - - assert x.shape == (128, dimension) - assert np.all(np.isfinite(x)) - - -def test_student_t_copula_large_df_is_close_to_gaussian_copula(): - rho = 0.6 - correlation = [[1.0, rho], [rho, 1.0]] - marginals = [stats.norm(), stats.norm()] - gaussian = GaussianCopula( - sampler=DigitalNetB2(2, seed=40), - marginals=marginals, - correlation=correlation, - ) - student_t = StudentTCopula( - sampler=DigitalNetB2(2, seed=40), - marginals=marginals, - correlation=correlation, - df=100, - ) - - x_gaussian = gaussian(4096) - x_student_t = student_t(4096) - corr_gaussian = np.corrcoef(x_gaussian.T)[0, 1] - corr_student_t = np.corrcoef(x_student_t.T)[0, 1] - - assert abs(corr_student_t - corr_gaussian) < 0.02 - - -@pytest.mark.parametrize("df", [0, -1, np.inf, "not-a-number"]) -def test_student_t_copula_invalid_df_raises_parameter_error(df): - with pytest.raises(ParameterError, match="df"): - StudentTCopula( - sampler=DigitalNetB2(2, seed=41), - marginals=[stats.norm(), stats.norm()], - correlation=np.eye(2), - df=df, - ) - - -def test_student_t_copula_marginal_without_ppf_raises_clear_error(): - class NoPPF: - pass - - with pytest.raises(ParameterError, match="ppf"): - StudentTCopula( - sampler=DigitalNetB2(1, seed=43), - marginals=[NoPPF()], - correlation=[[1.0]], - df=4, - ) - - -# Archimedean copulas - - -def test_clayton_copula_output_shape_and_finite_values(): - tm = ClaytonCopula( - sampler=DigitalNetB2(2, seed=57), - marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], - theta=2.0, - ) - - x = tm(128) - - assert x.shape == (128, 2) - assert np.all(np.isfinite(x)) - - -def test_clayton_copula_return_weights_shape_when_density_available(): - tm = ClaytonCopula( - sampler=DigitalNetB2(3, seed=59), - marginals=[stats.norm(), stats.gamma(a=2.0), stats.expon()], - theta=1.5, - ) - - x, weights = tm(32, return_weights=True) - - assert x.shape == (32, 3) - assert weights.shape == (32,) - assert np.all(np.isfinite(weights)) - assert np.all(weights > 0.0) - - -@pytest.mark.parametrize("theta", [0, -1, np.inf, "not-a-number"]) -def test_clayton_copula_invalid_theta_raises_parameter_error(theta): - with pytest.raises(ParameterError, match="theta"): - ClaytonCopula( - sampler=DigitalNetB2(2, seed=61), - marginals=[stats.norm(), stats.norm()], - theta=theta, - ) - - -@pytest.mark.parametrize("dimension", [2, 3, 5]) -def test_clayton_copula_supports_general_dimension(dimension): - tm = ClaytonCopula( - sampler=DigitalNetB2(dimension, seed=63), - marginals=[stats.norm()] * dimension, - theta=2.0, - ) - - x = tm(128) - - assert x.shape == (128, dimension) - assert np.all(np.isfinite(x)) - - -def test_clayton_copula_marginal_without_ppf_raises_clear_error(): - class NoPPF: - pass - - with pytest.raises(ParameterError, match="ppf"): - ClaytonCopula( - sampler=DigitalNetB2(2, seed=67), - marginals=[stats.norm(), NoPPF()], - theta=2.0, - ) - - -@pytest.mark.parametrize( - "marginals", - [ - [stats.norm(), stats.beta(a=2, b=5)], - [stats.gamma(a=3), stats.expon()], - [stats.lognorm(s=0.5), stats.norm()], - ], -) -def test_clayton_copula_common_scipy_frozen_marginals_work(marginals): - tm = ClaytonCopula( - sampler=DigitalNetB2(2, seed=69), - marginals=marginals, - theta=2.0, - ) - - x = tm(128) - - assert x.shape == (128, 2) - assert np.all(np.isfinite(x)) - - -def test_clayton_copula_endpoint_uniforms_are_clipped_to_finite_outputs(): - tm = ClaytonCopula( - sampler=DigitalNetB2(2, seed=70), - marginals=[stats.norm(), stats.lognorm(s=0.5)], - theta=2.0, - ) - u = np.array([[0.0, 1.0], [1.0, 0.0]]) - - x = tm._transform(u) - - assert x.shape == (2, 2) - assert np.all(np.isfinite(x)) - - -@pytest.mark.parametrize("dimension", [2, 3, 5]) -def test_clayton_copula_tiny_theta_is_near_independent(dimension): - marginals = [stats.uniform()] * dimension - tm = ClaytonCopula( - sampler=DigitalNetB2(dimension, seed=70), - marginals=marginals, - theta=1e-8, - ) - u = np.array( - [ - [0.2, 0.7, 0.4, 0.6, 0.8], - [0.4, 0.8, 0.9, 0.3, 0.2], - [0.9, 0.1, 0.3, 0.7, 0.5], - ] - )[:, :dimension] - - x = tm._transform(u) - - assert x.shape == (3, dimension) - assert np.all(np.isfinite(x)) - np.testing.assert_allclose(x, u, atol=5e-6) - - -@pytest.mark.parametrize("dimension", [2, 3, 5]) -@pytest.mark.parametrize("theta", [20.0, 50.0]) -def test_clayton_copula_large_theta_is_finite(theta, dimension): - tm = ClaytonCopula( - sampler=DigitalNetB2(dimension, seed=70), - marginals=[stats.norm()] * dimension, - theta=theta, - ) - - x = tm(128) - - assert x.shape == (128, dimension) - assert np.all(np.isfinite(x)) - - -def test_clayton_copula_positive_dependence_behavior(): - tm = ClaytonCopula( - sampler=DigitalNetB2(2, seed=71), - marginals=[stats.uniform(), stats.uniform()], - theta=2.0, - ) - - x = tm(4096) - empirical_corr = np.corrcoef(x.T)[0, 1] - - assert empirical_corr > 0.45 - - -def test_clayton_copula_has_stronger_lower_tail_than_gaussian_copula(): - theta = 2.0 - n = 2**12 - marginals = [stats.uniform(), stats.uniform()] - # Clayton Kendall tau is theta/(theta+2); convert to Gaussian rho. - rho = np.sin(np.pi * (theta / (theta + 2.0)) / 2.0) - - clayton = ClaytonCopula( - sampler=DigitalNetB2(2, seed=73), - marginals=marginals, - theta=theta, - ) - gaussian = GaussianCopula( - sampler=DigitalNetB2(2, seed=73), - marginals=marginals, - correlation=[[1.0, rho], [rho, 1.0]], - ) - - x_clayton = clayton(n) - x_gaussian = gaussian(n) - threshold = 0.05 - - def lower_tail_rate(x): - tail_0 = x[:, 0] < threshold - return np.mean(x[tail_0, 1] < threshold) - - clayton_tail = lower_tail_rate(x_clayton) - gaussian_tail = lower_tail_rate(x_gaussian) - - assert clayton_tail > gaussian_tail + 0.2 - - -def test_frank_copula_output_shape_for_two_dimensions(): - tm = FrankCopula( - sampler=DigitalNetB2(2, seed=75), - marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], - theta=5.0, - ) - - x = tm(128) - - assert x.shape == (128, 2) - assert np.all(np.isfinite(x)) - - -@pytest.mark.parametrize("dimension", [3, 5]) -def test_frank_copula_positive_theta_supports_higher_dimensions(dimension): - tm = FrankCopula( - sampler=DigitalNetB2(dimension, seed=76), - marginals=[stats.norm()] * dimension, - theta=5.0, - ) - - x = tm(128) - - assert x.shape == (128, dimension) - assert np.all(np.isfinite(x)) - - -def test_frank_copula_return_weights_shape_when_density_available(): - tm = FrankCopula( - sampler=DigitalNetB2(3, seed=77), - marginals=[stats.norm(), stats.gamma(a=2.0), stats.expon()], - theta=4.0, - ) - - x, weights = tm(32, return_weights=True) - - assert x.shape == (32, 3) - assert weights.shape == (32,) - assert np.all(np.isfinite(weights)) - assert np.all(weights > 0.0) - - -@pytest.mark.parametrize("theta", [0, np.inf, -np.inf, "not-a-number"]) -def test_frank_copula_invalid_theta_raises_parameter_error(theta): - with pytest.raises(ParameterError, match="theta"): - FrankCopula( - sampler=DigitalNetB2(2, seed=78), - marginals=[stats.norm(), stats.norm()], - theta=theta, - ) - - -def test_frank_copula_negative_theta_rejected_above_two_dimensions(): - with pytest.raises(ParameterError, match="d=2"): - FrankCopula( - sampler=DigitalNetB2(3, seed=79), - marginals=[stats.norm(), stats.norm(), stats.norm()], - theta=-2.0, - ) - - -def test_frank_copula_dimension_mismatch_raises_dimension_error(): - with pytest.raises(DimensionError, match="marginals"): - FrankCopula( - sampler=DigitalNetB2(2, seed=80), - marginals=[stats.norm(), stats.norm(), stats.norm()], - theta=5.0, - ) - - -def test_frank_copula_marginal_without_ppf_raises_clear_error(): - class NoPPF: - pass - - with pytest.raises(ParameterError, match="ppf"): - FrankCopula( - sampler=DigitalNetB2(2, seed=82), - marginals=[stats.norm(), NoPPF()], - theta=5.0, - ) - - -def test_frank_copula_positive_dependence_behavior(): - tm = FrankCopula( - sampler=DigitalNetB2(2, seed=84), - marginals=[stats.uniform(), stats.uniform()], - theta=6.0, - ) - - x = tm(4096) - empirical_corr = np.corrcoef(x.T)[0, 1] - - assert empirical_corr > 0.45 - - -@pytest.mark.parametrize( - "theta,dimension", - [ - (1e-8, 3), - (-1e-8, 2), - ], -) -def test_frank_copula_tiny_theta_is_close_to_independence(theta, dimension): - marginals = [stats.uniform()] * dimension - tm = FrankCopula( - sampler=DigitalNetB2(dimension, seed=86), - marginals=marginals, - theta=theta, - ) - u = np.array( - [ - [0.2, 0.7, 0.4, 0.6, 0.8], - [0.4, 0.8, 0.9, 0.3, 0.2], - [0.9, 0.1, 0.3, 0.7, 0.5], - ] - )[:, :dimension] - - x = tm._transform(u) - - assert x.shape == (3, dimension) - assert np.all(np.isfinite(x)) - np.testing.assert_allclose(x, u, atol=5e-6) - - -@pytest.mark.parametrize( - "theta,dimension", - [ - (50.0, 5), - (-50.0, 2), - ], -) -def test_frank_copula_large_theta_is_finite(theta, dimension): - tm = FrankCopula( - sampler=DigitalNetB2(dimension, seed=87), - marginals=[stats.norm()] * dimension, - theta=theta, - ) - - x = tm(128) - - assert x.shape == (128, dimension) - assert np.all(np.isfinite(x)) - - -def test_frank_copula_negative_theta_produces_negative_dependence_in_2d(): - tm = FrankCopula( - sampler=DigitalNetB2(2, seed=88), - marginals=[stats.uniform(), stats.uniform()], - theta=-6.0, - ) - - x = tm(4096) - empirical_corr = np.corrcoef(x.T)[0, 1] - - assert empirical_corr < -0.35 - - -def test_gumbel_copula_output_shape_and_finite_values(): - tm = GumbelCopula( - sampler=DigitalNetB2(2, seed=79), - marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], - theta=2.0, - ) - - x = tm(128) - - assert x.shape == (128, 2) - assert np.all(np.isfinite(x)) - - -def test_gumbel_copula_return_weights_shape_when_density_available(): - tm = GumbelCopula( - sampler=DigitalNetB2(3, seed=81), - marginals=[stats.norm(), stats.gamma(a=2.0), stats.expon()], - theta=1.5, - ) - - x, weights = tm(32, return_weights=True) - - assert x.shape == (32, 3) - assert weights.shape == (32,) - assert np.all(np.isfinite(weights)) - assert np.all(weights > 0.0) - - -@pytest.mark.parametrize("theta", [0, 0.5, -1, np.inf, "not-a-number"]) -def test_gumbel_copula_invalid_theta_raises_parameter_error(theta): - with pytest.raises(ParameterError, match="theta"): - GumbelCopula( - sampler=DigitalNetB2(2, seed=83), - marginals=[stats.norm(), stats.norm()], - theta=theta, - ) - - -def test_gumbel_copula_theta_one_is_independent_marginal_transform(): - marginals = [stats.norm(loc=-1.0, scale=2.0), stats.gamma(a=2.0, scale=3.0)] - tm = GumbelCopula( - sampler=DigitalNetB2(2, seed=85), - marginals=marginals, - theta=1.0, - ) - u = np.array([[0.2, 0.7], [0.4, 0.8], [0.9, 0.1]]) - - x = tm._transform(u) - expected = np.column_stack( - [marginals[j].ppf(u[:, j]) for j in range(len(marginals))] - ) - - np.testing.assert_allclose(x, expected, rtol=1e-12, atol=1e-12) - - -@pytest.mark.parametrize("dimension", [2, 3, 5]) -def test_gumbel_copula_theta_close_to_one_is_near_independent(dimension): - marginals = [stats.uniform()] * dimension - tm = GumbelCopula( - sampler=DigitalNetB2(dimension, seed=85), - marginals=marginals, - theta=1.000001, - ) - u = np.array( - [ - [0.2, 0.7, 0.4, 0.6, 0.8], - [0.4, 0.8, 0.9, 0.3, 0.2], - [0.9, 0.1, 0.3, 0.7, 0.5], - ] - )[:, :dimension] - - x = tm._transform(u) - - assert x.shape == (3, dimension) - assert np.all(np.isfinite(x)) - np.testing.assert_allclose(x, u, atol=5e-5) - - -@pytest.mark.parametrize("dimension", [2, 3, 5]) -@pytest.mark.parametrize("theta", [20.0, 50.0]) -def test_gumbel_copula_large_theta_is_finite(theta, dimension): - tm = GumbelCopula( - sampler=DigitalNetB2(dimension, seed=86), - marginals=[stats.norm()] * dimension, - theta=theta, - ) - - x = tm(128) - - assert x.shape == (128, dimension) - assert np.all(np.isfinite(x)) - - -@pytest.mark.parametrize("dimension", [2, 3, 5]) -def test_gumbel_copula_supports_general_dimension(dimension): - tm = GumbelCopula( - sampler=DigitalNetB2(dimension, seed=87), - marginals=[stats.norm()] * dimension, - theta=2.0, - ) - - x = tm(128) - - assert x.shape == (128, dimension) - assert np.all(np.isfinite(x)) - - -def test_gumbel_copula_marginal_without_ppf_raises_clear_error(): - class NoPPF: - pass - - with pytest.raises(ParameterError, match="ppf"): - GumbelCopula( - sampler=DigitalNetB2(2, seed=89), - marginals=[stats.norm(), NoPPF()], - theta=2.0, - ) - - -@pytest.mark.parametrize( - "marginals", - [ - [stats.norm(), stats.beta(a=2, b=5)], - [stats.gamma(a=3), stats.expon()], - [stats.lognorm(s=0.5), stats.norm()], - ], -) -def test_gumbel_copula_common_scipy_frozen_marginals_work(marginals): - tm = GumbelCopula( - sampler=DigitalNetB2(2, seed=91), - marginals=marginals, - theta=2.0, - ) - - x = tm(128) - - assert x.shape == (128, 2) - assert np.all(np.isfinite(x)) - - -def test_gumbel_copula_endpoint_uniforms_are_clipped_to_finite_outputs(): - tm = GumbelCopula( - sampler=DigitalNetB2(2, seed=93), - marginals=[stats.norm(), stats.lognorm(s=0.5)], - theta=2.0, - ) - u = np.array([[0.0, 1.0], [1.0, 0.0]]) - - x = tm._transform(u) - - assert x.shape == (2, 2) - assert np.all(np.isfinite(x)) - - -def test_gumbel_copula_positive_dependence_behavior(): - tm = GumbelCopula( - sampler=DigitalNetB2(2, seed=95), - marginals=[stats.uniform(), stats.uniform()], - theta=2.0, - ) - - x = tm(4096) - empirical_corr = np.corrcoef(x.T)[0, 1] - - assert empirical_corr > 0.45 - - -def test_gumbel_copula_has_stronger_upper_tail_than_gaussian_copula(): - theta = 2.0 - n = 2**12 - marginals = [stats.uniform(), stats.uniform()] - # Gumbel Kendall tau is 1 - 1/theta; convert to Gaussian rho. - rho = np.sin(np.pi * (1.0 - 1.0 / theta) / 2.0) - - gumbel = GumbelCopula( - sampler=DigitalNetB2(2, seed=97), - marginals=marginals, - theta=theta, - ) - gaussian = GaussianCopula( - sampler=DigitalNetB2(2, seed=97), - marginals=marginals, - correlation=[[1.0, rho], [rho, 1.0]], - ) - - x_gumbel = gumbel(n) - x_gaussian = gaussian(n) - threshold = 0.95 - - def upper_tail_rate(x): - tail_0 = x[:, 0] > threshold - return np.mean(x[tail_0, 1] > threshold) - - gumbel_tail = upper_tail_rate(x_gumbel) - gaussian_tail = upper_tail_rate(x_gaussian) - - assert gumbel_tail > gaussian_tail + 0.15 - - -# Weights, fallback behavior, spawn, and edge cases - - -@pytest.mark.parametrize( - "copula_cls", - [GaussianCopula, StudentTCopula, ClaytonCopula, GumbelCopula, FrankCopula], -) -def test_copula_weight_fallback_warns_once_when_density_methods_are_missing( - copula_cls, -): - tm = _make_copula( - copula_cls, - dimension=2, - marginals=[PPFOnlyMarginal(), PPFOnlyMarginal()], - ) - x = np.full((4, 2), 0.5) - expected_message = getattr( - tm, - "_missing_weight_warning_message", - f"{copula_cls.__name__} marginals must implement 'cdf' and " - "'pdf' or 'logpdf' to compute density weights. " - "Weights will be treated as 1.", - ) - - assert "_unit_weight_with_warning" not in copula_cls.__dict__ - assert ( - tm._unit_weight_with_warning.__func__ - is AbstractCopula._unit_weight_with_warning - ) - - with pytest.warns(UserWarning) as warning_info: - weights = tm._weight(x) - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - second_weights = tm._weight(x) - - np.testing.assert_allclose(weights, np.ones(4)) - np.testing.assert_allclose(second_weights, np.ones(4)) - assert str(warning_info[0].message) == expected_message - assert caught == [] - - -def test_student_t_weight_falls_back_when_multivariate_t_is_unavailable(): - tm = StudentTCopula( - DigitalNetB2(2, seed=115), - marginals=[stats.norm(), stats.norm()], - correlation=np.eye(2), - df=4, - ) - tm._mvt_scipy = None - - with pytest.warns(UserWarning, match="Weights will be treated as 1"): - weights = tm._weight(np.full((3, 2), 0.25)) - - np.testing.assert_allclose(weights, np.ones(3)) - - -def test_gaussian_weight_uses_pdf_branch_when_logpdf_is_unavailable(): - tm = GaussianCopula( - DigitalNetB2(2, seed=117), - marginals=[UnitPDFMarginal(), UnitPDFMarginal()], - correlation=[[1.0, 0.4], [0.4, 1.0]], - ) - - weights = tm._weight(np.array([[0.25, 0.5], [0.75, 0.5]])) - - assert weights.shape == (2,) - assert np.all(np.isfinite(weights)) - assert np.all(weights > 0.0) - - -def test_gumbel_theta_one_weight_is_independent_marginal_density(): - tm = GumbelCopula( - DigitalNetB2(2, seed=119), - marginals=[stats.gamma(a=2.0), stats.expon()], - theta=1.0, - ) - x = np.array([[1.0, 0.5], [2.0, 1.5]]) - expected = stats.gamma(a=2.0).pdf(x[:, 0]) * stats.expon().pdf(x[:, 1]) - - weights = tm._weight(x) - - np.testing.assert_allclose(weights, expected) - - -def test_gen_copula_samples_composed_transform_branch(): - inner = GaussianCopula( - DigitalNetB2(2, seed=121), - marginals=[stats.uniform(), stats.uniform()], - correlation=[[1.0, 0.3], [0.3, 1.0]], - ) - outer = ClaytonCopula(inner, marginals=[stats.uniform(), stats.uniform()], theta=1.5) - - v = outer.gen_copula_samples(n_min=4, n_max=8) - - assert v.shape == (4, 2) - assert np.all(np.isfinite(v)) - assert np.all((0.0 <= v) & (v <= 1.0)) - - -@pytest.mark.parametrize( - "copula_cls", - [GaussianCopula, StudentTCopula, ClaytonCopula, GumbelCopula, FrankCopula], -) -def test_copula_spawn_same_dimension_and_reject_different_dimension(copula_cls): - tm = _make_copula(copula_cls, dimension=2) - - spawned = tm.spawn(s=1, dimensions=[2]) - assert len(spawned) == 1 - assert isinstance(spawned[0], copula_cls) - assert spawned[0](4).shape == (4, 2) - - with pytest.raises(DimensionError): - tm._spawn(DigitalNetB2(3, seed=123), 3) - - -def test_frank_one_dimensional_weight_covers_zero_order_eulerian_term(): - tm = FrankCopula( - DigitalNetB2(1, seed=125), - marginals=[UnitPDFMarginal()], - theta=3.0, - ) - - weights = tm._weight(np.array([[0.25], [0.75]])) - - assert weights.shape == (2,) - assert np.all(np.isfinite(weights)) - assert np.all(weights > 0.0) - - -def test_frank_rejects_large_negative_theta_when_exponential_overflows(): - with np.errstate(over="ignore"): - with pytest.raises(ParameterError, match="too close to 0 or too large"): - FrankCopula( - DigitalNetB2(2, seed=127), - marginals=[stats.uniform(), stats.uniform()], - theta=-1000.0, - ) diff --git a/test/test_discrete_distribs.py b/test/test_dd_discrete_distribs.py similarity index 100% rename from test/test_discrete_distribs.py rename to test/test_dd_discrete_distribs.py diff --git a/test/test_dd_dummy_sampler.py b/test/test_dd_dummy_sampler.py new file mode 100644 index 000000000..50d5005b8 --- /dev/null +++ b/test/test_dd_dummy_sampler.py @@ -0,0 +1,107 @@ +import unittest + +import numpy as np + +from qmcpy import DummySampler +from qmcpy.util import ParameterError + + +PLACEHOLDER_ERROR = "construction placeholder" + + +class TestDummySampler(unittest.TestCase): + + def test_dummy_sampler_constructs_dimension_one(self): + sampler = DummySampler(1) + + self.assertEqual(sampler.d, 1) + self.assertEqual(sampler.replications, 1) + self.assertTrue(sampler.no_replications) + self.assertEqual(sampler.mimics, "StdUniform") + self.assertEqual(sampler.parameters, []) + + def test_dummy_sampler_constructs_larger_dimensions(self): + sampler = DummySampler(3, seed=7) + + self.assertEqual(sampler.d, 3) + self.assertEqual(sampler.replications, 1) + self.assertTrue(sampler.no_replications) + self.assertTrue(np.array_equal(sampler.dvec, np.arange(3))) + + def test_dummy_sampler_constructs_larger_dimension_with_replications(self): + sampler = DummySampler(4, replications=3, seed=7) + + self.assertEqual(sampler.d, 4) + self.assertEqual(sampler.replications, 3) + self.assertFalse(sampler.no_replications) + self.assertTrue(np.array_equal(sampler.dvec, np.arange(4))) + + def test_dummy_sampler_direct_sampling_raises_placeholder_error(self): + sampler = DummySampler(2) + + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler(8) + + def test_dummy_sampler_replicated_direct_sampling_raises_placeholder_error(self): + sampler = DummySampler(2, replications=3) + + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler(8) + + def test_dummy_sampler_supported_calling_conventions_raise_placeholder_error(self): + sampler = DummySampler(2) + + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler(n=4) + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler(n_min=2, n_max=6) + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler(n=2, n_min=6) + + def test_dummy_sampler_nonzero_n_min_raises_placeholder_error(self): + sampler = DummySampler(2) + + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler(n_min=5, n_max=9) + + def test_dummy_sampler_rejects_return_binary(self): + sampler = DummySampler(2) + + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler(4, return_binary=True) + + def test_dummy_sampler_internal_gen_samples_raises_placeholder_error(self): + sampler = DummySampler(2) + + with self.assertRaisesRegex(ParameterError, PLACEHOLDER_ERROR): + sampler._gen_samples(n_min=5, n_max=9, return_binary=False, warn=True) + + def test_dummy_sampler_spawn_preserves_relevant_fields(self): + sampler = DummySampler(2, replications=3, seed=11) + + spawned = sampler.spawn(s=2, dimensions=[1, 5]) + + self.assertEqual([spawn.d for spawn in spawned], [1, 5]) + self.assertEqual([spawn.replications for spawn in spawned], [3, 3]) + self.assertTrue(all(isinstance(spawn, DummySampler) for spawn in spawned)) + + def test_dummy_sampler_spawn_without_explicit_replications(self): + sampler = DummySampler(2, seed=11) + + spawned = sampler.spawn(s=1, dimensions=4)[0] + + self.assertEqual(spawned.d, 4) + self.assertEqual(spawned.replications, 1) + self.assertTrue(spawned.no_replications) + + def test_dummy_sampler_limits_are_enforced(self): + with self.assertRaisesRegex(ParameterError, "dimension greater than dimension limit"): + DummySampler(10_002) + + sampler = DummySampler(1) + with self.assertRaisesRegex(ParameterError, "n_limit"): + sampler(n_min=0, n_max=2**32 + 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_dd_mpmc_optional_imports.py b/test/test_dd_mpmc_optional_imports.py new file mode 100644 index 000000000..4628f2c93 --- /dev/null +++ b/test/test_dd_mpmc_optional_imports.py @@ -0,0 +1,114 @@ +import ast +import builtins +import unittest +from pathlib import Path + + +def _execute_optional_import(blocked_import): + repository_root = Path(__file__).resolve().parent.parent + init_path = repository_root / "qmcpy" / "__init__.py" + init_tree = ast.parse(init_path.read_text()) + optional_import = next( + node + for node in init_tree.body + if isinstance(node, ast.Try) + and any( + isinstance(statement, ast.ImportFrom) + and statement.module == "discrete_distribution.mpmc" + for statement in node.body + ) + ) + + import qmcpy + + real_import = builtins.__import__ + + def guarded_import(name, globals=None, locals=None, fromlist=(), level=0): + missing_module = blocked_import(name, fromlist, level) + if missing_module is not None: + raise ModuleNotFoundError( + "blocked optional dependency", + name=missing_module, + ) + return real_import(name, globals, locals, fromlist, level) + + test_builtins = vars(builtins).copy() + test_builtins["__import__"] = guarded_import + namespace = {"__builtins__": test_builtins, "__package__": "qmcpy"} + module = ast.Module(body=[optional_import], type_ignores=[]) + exec(compile(module, str(init_path), "exec"), namespace) + return namespace + + +class TestMPMCOptionalImports(unittest.TestCase): + + def test_mpmc_utils_remain_available_without_pyg(self): + try: + import torch # noqa: F401 + except ImportError: + self.skipTest("torch not available") + + def block_pyg_models(name, fromlist, level): + if level == 1 and name == "discrete_distribution.mpmc.models": + return "torch_geometric" + return None + + namespace = _execute_optional_import(block_pyg_models) + + import qmcpy + + self.assertIs(namespace["mpmc_utils"], qmcpy.mpmc_utils) + self.assertEqual( + namespace["mpmc_utils"].__name__, + "qmcpy.discrete_distribution.mpmc.utils", + ) + self.assertNotIn("utils", namespace) + + with self.assertRaisesRegex( + ModuleNotFoundError, "MPMC_net.*torch_geometric" + ) as cm: + namespace["MPMC_net"]() + self.assertEqual(cm.exception.name, "torch_geometric") + + def test_mpmc_placeholders_report_missing_torch(self): + def block_torch_utils(name, fromlist, level): + if ( + level == 1 + and name == "discrete_distribution.mpmc" + and "utils" in fromlist + ): + return "torch" + return None + + namespace = _execute_optional_import(block_torch_utils) + + with self.assertRaisesRegex(ModuleNotFoundError, "mpmc_utils.*torch") as cm: + namespace["mpmc_utils"].L2star + self.assertEqual(cm.exception.name, "torch") + + with self.assertRaisesRegex(ModuleNotFoundError, "MPMC_net.*torch") as cm: + namespace["MPMC_net"]() + self.assertEqual(cm.exception.name, "torch") + + def test_mpmc_placeholder_missing_torch_scatter(self): + try: + import torch # noqa: F401 + except ImportError: + self.skipTest("torch not available") + + def block_torch_scatter(name, fromlist, level): + if level == 1 and name == "discrete_distribution.mpmc.models": + return "torch_scatter" + return None + + namespace = _execute_optional_import(block_torch_scatter) + + with self.assertRaisesRegex( + ModuleNotFoundError, "MPMC_net.*torch_scatter" + ) as cm: + namespace["MPMC_net"]() + self.assertEqual(cm.exception.name, "torch_scatter") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_dummy_sampler.py b/test/test_dummy_sampler.py deleted file mode 100644 index 24bec84eb..000000000 --- a/test/test_dummy_sampler.py +++ /dev/null @@ -1,111 +0,0 @@ -import numpy as np -import pytest - -from qmcpy import DummySampler -from qmcpy.util import ParameterError - - -PLACEHOLDER_ERROR = "construction placeholder" - - -def test_dummy_sampler_constructs_dimension_one(): - sampler = DummySampler(1) - - assert sampler.d == 1 - assert sampler.replications == 1 - assert sampler.no_replications - assert sampler.mimics == "StdUniform" - assert sampler.parameters == [] - - -def test_dummy_sampler_constructs_larger_dimensions(): - sampler = DummySampler(3, seed=7) - - assert sampler.d == 3 - assert sampler.replications == 1 - assert sampler.no_replications - assert np.array_equal(sampler.dvec, np.arange(3)) - - -def test_dummy_sampler_constructs_larger_dimension_with_replications(): - sampler = DummySampler(4, replications=3, seed=7) - - assert sampler.d == 4 - assert sampler.replications == 3 - assert not sampler.no_replications - assert np.array_equal(sampler.dvec, np.arange(4)) - - -def test_dummy_sampler_direct_sampling_raises_placeholder_error(): - sampler = DummySampler(2) - - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler(8) - - -def test_dummy_sampler_replicated_direct_sampling_raises_placeholder_error(): - sampler = DummySampler(2, replications=3) - - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler(8) - - -def test_dummy_sampler_supported_calling_conventions_raise_placeholder_error(): - sampler = DummySampler(2) - - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler(n=4) - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler(n_min=2, n_max=6) - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler(n=2, n_min=6) - - -def test_dummy_sampler_nonzero_n_min_raises_placeholder_error(): - sampler = DummySampler(2) - - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler(n_min=5, n_max=9) - - -def test_dummy_sampler_rejects_return_binary(): - sampler = DummySampler(2) - - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler(4, return_binary=True) - - -def test_dummy_sampler_internal_gen_samples_raises_placeholder_error(): - sampler = DummySampler(2) - - with pytest.raises(ParameterError, match=PLACEHOLDER_ERROR): - sampler._gen_samples(n_min=5, n_max=9, return_binary=False, warn=True) - - -def test_dummy_sampler_spawn_preserves_relevant_fields(): - sampler = DummySampler(2, replications=3, seed=11) - - spawned = sampler.spawn(s=2, dimensions=[1, 5]) - - assert [spawn.d for spawn in spawned] == [1, 5] - assert [spawn.replications for spawn in spawned] == [3, 3] - assert all(isinstance(spawn, DummySampler) for spawn in spawned) - - -def test_dummy_sampler_spawn_without_explicit_replications(): - sampler = DummySampler(2, seed=11) - - spawned = sampler.spawn(s=1, dimensions=4)[0] - - assert spawned.d == 4 - assert spawned.replications == 1 - assert spawned.no_replications - - -def test_dummy_sampler_limits_are_enforced(): - with pytest.raises(ParameterError, match="dimension greater than dimension limit"): - DummySampler(10_002) - - sampler = DummySampler(1) - with pytest.raises(ParameterError, match="n_limit"): - sampler(n_min=0, n_max=2**32 + 1) diff --git a/test/test_integrate.py b/test/test_ee_integrate.py similarity index 100% rename from test/test_integrate.py rename to test/test_ee_integrate.py diff --git a/test/test_keister.py b/test/test_ee_keister.py similarity index 100% rename from test/test_keister.py rename to test/test_ee_keister.py diff --git a/test/test_pi_problem.py b/test/test_ee_pi_problem.py similarity index 100% rename from test/test_pi_problem.py rename to test/test_ee_pi_problem.py diff --git a/test/test_fast_transform_fallbacks.py b/test/test_fast_transform_fallbacks.py deleted file mode 100644 index 2dcbd2b7a..000000000 --- a/test/test_fast_transform_fallbacks.py +++ /dev/null @@ -1,46 +0,0 @@ -import numpy as np -import pytest - -from qmcpy import ( - fftbr, - fftbr_torch, - fwht, - fwht_torch, - ifftbr, - ifftbr_torch, - omega_fftbr, - omega_fftbr_torch, - omega_fwht, - omega_fwht_torch, -) - - -def test_non_torch_transforms_basic(): - rng = np.random.default_rng(11) - x = rng.random(8) + 1j * rng.random(8) - y = fftbr(x) - assert y.shape == x.shape - xr = ifftbr(y) - assert xr.shape == x.shape - - a = rng.random(8) - b = fwht(a) - assert b.shape == a.shape - - omega = omega_fftbr(3) - assert omega.shape[0] == 2**3 - omega2 = omega_fwht(3) - assert omega2.shape[0] == 2**3 - - -def test_torch_fallbacks_raise(): - with pytest.raises(Exception): - fftbr_torch() - with pytest.raises(Exception): - ifftbr_torch() - with pytest.raises(Exception): - fwht_torch() - with pytest.raises(Exception): - omega_fftbr_torch() - with pytest.raises(Exception): - omega_fwht_torch() diff --git a/test/test_financial_option_quick.py b/test/test_financial_option_quick.py deleted file mode 100644 index bd9f3022b..000000000 --- a/test/test_financial_option_quick.py +++ /dev/null @@ -1,94 +0,0 @@ -import numpy as np - -from qmcpy import FinancialOption -import qmcpy - - -class SmallSampler(qmcpy.AbstractDiscreteDistribution): - def __init__(self, d=3): - super().__init__( - dimension=d, replications=1, seed=123, d_limit=100, n_limit=1024 - ) - - def _gen_samples(self, n_min, n_max, return_binary=False, warn=True): - n = n_max - n_min - # return shape (replications, n, d) - arr = np.tile(np.linspace(0.1, 1.0, n)[:, None], (1, self.d)) - return arr.reshape(self.replications, n, self.d) - - -def test_financial_option_payoffs_and_exact(): - sampler = SmallSampler(d=3) - fo = FinancialOption( - sampler, - option="EUROPEAN", - call_put="CALL", - volatility=0.5, - start_price=30, - strike_price=25, - interest_rate=0.01, - t_final=1, - ) - gbm = np.array([[30.0, 28.0, 35.0]]) - c = fo.payoff_european_call(gbm) - p = fo.payoff_european_put(gbm) - assert c.shape == (1,) - assert p.shape == (1,) - - # Asian arithmetic trapezoidal - fo_asian = FinancialOption( - sampler, - option="ASIAN", - asian_mean="ARITHMETIC", - asian_mean_quadrature_rule="TRAPEZOIDAL", - ) - gbm2 = np.array([[30.0, 32.0, 34.0]]) - a_call = fo_asian.payoff_asian_arithmetic_trap_call(gbm2) - assert a_call.shape == (1,) - - # geometric right call - fo_geo = FinancialOption( - sampler, - option="ASIAN", - asian_mean="GEOMETRIC", - asian_mean_quadrature_rule="RIGHT", - ) - g_call = fo_geo.payoff_asian_geometric_right_call(np.array([[30.0, 30.0, 30.0]])) - assert g_call.shape == (1,) - - # barrier options: up and down behaviors - fo_barrier_up = FinancialOption( - sampler, option="BARRIER", barrier_in_out="IN", barrier_price=25, start_price=20 - ) - gbm_up = np.array([[20.0, 26.0, 27.0]]) - v = fo_barrier_up.payoff_barrier_in_up_call(gbm_up) - assert v.shape == (1,) - - fo_barrier_out = FinancialOption( - sampler, - option="BARRIER", - barrier_in_out="OUT", - barrier_price=40, - start_price=30, - ) - gbm_out = np.array([[30.0, 32.0, 33.0]]) - v2 = fo_barrier_out.payoff_barrier_out_up_call(gbm_out) - assert v2.shape == (1,) - - # lookback - fo_lb = FinancialOption(sampler, option="LOOKBACK") - lb = fo_lb.payoff_lookback_call(np.array([[10.0, 9.0, 12.0]])) - assert lb.shape == (1,) - - # digital - fo_dig = FinancialOption(sampler, option="DIGITAL", digital_payout=5) - dig = fo_dig.payoff_digital_call(np.array([[10.0, 11.0, 12.0]])) - assert dig.shape == (1,) - - # exact value for European should return a float - val = fo.get_exact_value() - assert np.isscalar(val) - - # exact value for Asian geometric right - val2 = fo_geo.get_exact_value() - assert np.isscalar(val2) diff --git a/test/test_flatten_qmcpy_imports.py b/test/test_flatten_qmcpy_imports.py deleted file mode 100644 index c7fc36fdb..000000000 --- a/test/test_flatten_qmcpy_imports.py +++ /dev/null @@ -1,330 +0,0 @@ -import json -from pathlib import Path - -from scripts.flatten_qmcpy_imports import ( - _load_qmcpy_public_names, - flatten_imports, - main, -) - - -def _nested_import(module, imported): - return f"from {'qmcpy.' + module} import {imported}" - - -def test_flatten_imports_basic(): - source = ( - _nested_import("integrand", "Keister") - + "\n" - + _nested_import("discrete_distribution.lattice", "Lattice as LD") - + "\nfrom qmcpy import DigitalNetB2\nimport qmcpy.util\n" - ).encode() - - updated, count = flatten_imports( - source, frozenset({"DigitalNetB2", "Keister", "Lattice"}) - ) - - assert count == 3 - assert updated == ( - b"from qmcpy import DigitalNetB2, Keister, Lattice as LD\n" - b"import qmcpy.util\n" - ) - - -def test_flatten_preserves_private(): - source = ( - _nested_import("_internal._helpers", "PublicHelper") - + "\n" - + _nested_import( - "true_measure.uniform_triangle", - "UniformTriangle, _UniformTriangleAdapter", - ) - + "\n" - + _nested_import( - "true_measure.copula", - "(\n AbstractCopula,\n _validate_dimension,\n)", - ) - + "\n" - + _nested_import("integrand", "Keister") - + "\n" - ).encode() - - updated, count = flatten_imports(source, frozenset({"Keister"})) - - assert count == 1 - assert updated == source.replace( - _nested_import("integrand", "Keister").encode(), - b"from qmcpy import Keister", - ) - - -def test_private_module_splits_groups(): - source = ( - b"from qmcpy import Zeta\n" - b"from qmcpy._internal._helpers import PublicHelper\n" - b"from qmcpy import Alpha\n" - ) - - updated, count = flatten_imports(source) - - assert count == 0 - assert updated == source - - -def test_flatten_preserves_util_imports(): - source = ( - b"from qmcpy.util import ParameterError\n" - b"from qmcpy.util.transforms import tf_exp\n" - ) - - updated, count = flatten_imports(source, frozenset({"ParameterError", "tf_exp"})) - - assert (updated, count) == (source, 0) - - -def test_flatten_keeps_nonpublic_names(): - source = b"from qmcpy.stopping_criterion.pf_gp_ci import PFGPCIData\n" - - updated, count = flatten_imports(source, frozenset({"PFGPCI"})) - - assert (updated, count) == (source, 0) - - -def test_flatten_no_public_api_noop(): - source = (_nested_import("integrand", "Keister") + "\n").encode() - - updated, count = flatten_imports(source) - - assert (updated, count) == (source, 0) - - -def test_flatten_preserve_str_literals(): - source = b'text = """\nfrom qmcpy.integrand import Keister\n"""\n' - - updated, count = flatten_imports(source, frozenset({"Keister"})) - - assert count == 0 - assert updated == source - - -def test_python_string_protection_applies_to_every_rewrite_stage(): - string_body = ( - b'text = """\n' - b"from qmcpy.integrand import Keister\n" - b"from qmcpy import Zeta,Beta\n" - b"from qmcpy import Alpha\n" - b"from qmcpy import *\n" - b"from qmcpy import *\n" - b'"""\n' - ) - source = string_body + b"from qmcpy.integrand import Keister\n" - - updated, count = flatten_imports(source, frozenset({"Keister"})) - - assert count == 1 - assert updated == string_body + b"from qmcpy import Keister\n" - - -def test_python_tokenize_failure_is_fail_closed(): - source = b'"""unterminated\nfrom qmcpy.integrand import Keister\n' - - assert flatten_imports(source, frozenset({"Keister"})) == (source, 0) - - -def test_flatten_skip_star_expansion(): - source = ( - b"from qmcpy import *\n\n" - b"def f(Lattice):\n" - b" return Lattice\n\n" - b"y = Keister(dimension=2)\n" - b"x = Lattice(dimension=2)\n" - ) - - updated, count = flatten_imports(source, frozenset({"Keister", "Lattice"})) - - assert count == 0 - assert updated == source - - -def test_notebook_star_dedup(): - notebook = { - "cells": [ - { - "cell_type": "code", - "source": [ - _nested_import("integrand", "*") + "\n", - _nested_import("true_measure", "*"), - ], - } - ] - } - source = json.dumps(notebook, indent=1).encode() - - updated, count = flatten_imports(source, frozenset({"Keister"})) - - assert count == 3 - assert json.loads(updated)["cells"][0]["source"] == ["from qmcpy import *"] - - -def test_named_imports_merge_sort(): - source = ( - b"from qmcpy import Zeta,Beta\n" - b"from qmcpy import Alpha\n" - b"\n" - b"from qmcpy import Gamma\n" - ) - - updated, count = flatten_imports(source) - - assert count == 1 - assert updated == ( - b"from qmcpy import Alpha, Beta, Zeta\n" - b"\n" - b"from qmcpy import Gamma\n" - ) - assert flatten_imports(updated) == (updated, 0) - - -def test_merge_paren_and_single_line(): - source = b"""from qmcpy import ( - KernelDigShiftInvar, - KernelDigShiftInvarAdaptiveAlpha, - KernelDigShiftInvarCombined, - KernelShiftInvar, - KernelShiftInvarCombined, -) -from qmcpy import tf_exp_eps, tf_exp_eps_inv -""" - - updated, count = flatten_imports(source) - - assert count == 1 - assert updated == b"""from qmcpy import ( - KernelDigShiftInvar, - KernelDigShiftInvarAdaptiveAlpha, - KernelDigShiftInvarCombined, - KernelShiftInvar, - KernelShiftInvarCombined, - tf_exp_eps, - tf_exp_eps_inv, -) -""" - assert flatten_imports(updated) == (updated, 0) - - -def test_merge_same_scope_only(): - source = ( - b"if enabled:\n" - b" from qmcpy import Zeta\n" - b" from qmcpy import Alpha as First\n" - b"else:\n" - b" from qmcpy import Beta\n" - b"from qmcpy import _Private\n" - b"from qmcpy import Gamma # keep this comment\n" - ) - - updated, count = flatten_imports(source) - - assert count == 1 - assert updated == ( - b"if enabled:\n" - b" from qmcpy import Alpha as First, Zeta\n" - b"else:\n" - b" from qmcpy import Beta\n" - b"from qmcpy import _Private\n" - b"from qmcpy import Gamma # keep this comment\n" - ) - - -def test_notebook_named_merge(): - notebook = { - "cells": [ - { - "cell_type": "code", - "source": [ - "from qmcpy import Zeta\n", - "from qmcpy import Alpha,Beta\n", - "print(Alpha)\n", - ], - } - ] - } - source = json.dumps(notebook, indent=1).encode() - - updated, count = flatten_imports(source) - - assert count == 1 - assert json.loads(updated)["cells"][0]["source"] == [ - "from qmcpy import Alpha, Beta, Zeta\n", - "print(Alpha)\n", - ] - assert flatten_imports(updated) == (updated, 0) - - -def test_notebook_flattens_nested_imports_only_in_code_cells(): - nested_import = _nested_import("integrand", "Keister") + "\n" - metadata_import = _nested_import("true_measure", "Gaussian") + "\n" - string_literal = f'text = "{nested_import.rstrip()}"\n' - multiline_string = ['text = """\n', nested_import, '"""\n'] - notebook = { - "metadata": {"source": [metadata_import]}, - "cells": [ - {"cell_type": "markdown", "source": [nested_import]}, - {"cell_type": "code", "source": [nested_import]}, - {"cell_type": "code", "source": [string_literal]}, - {"cell_type": "code", "source": multiline_string}, - ] - } - source = json.dumps(notebook, indent=1).encode() - - updated, count = flatten_imports(source, frozenset({"Keister"})) - - cells = json.loads(updated)["cells"] - assert count == 1 - assert json.loads(updated)["metadata"]["source"] == [metadata_import] - assert cells[0]["source"] == [nested_import] - assert cells[1]["source"] == ["from qmcpy import Keister\n"] - assert cells[2]["source"] == [string_literal] - assert cells[3]["source"] == multiline_string - assert flatten_imports(updated, frozenset({"Keister"})) == (updated, 0) - - -def test_markdown_import_examples_are_flattened(tmp_path): - path = tmp_path / "example.md" - path.write_bytes( - b'Example with unmatched prose delimiter: """\n\n' - b"```python\n" - b"from qmcpy.integrand import Keister\n" - b"```\n" - ) - - assert main([str(path)]) == 0 - assert b"from qmcpy import Keister" in path.read_bytes() - - -def test_check_mode_no_write(tmp_path): - path = tmp_path / "example.py" - original = (_nested_import("true_measure", "Gaussian") + "\n").encode() - path.write_bytes(original) - - assert main(["--check", str(path)]) == 1 - assert path.read_bytes() == original - - assert main([str(path)]) == 0 - assert path.read_bytes() == b"from qmcpy import Gaussian\n" - assert main(["--check", str(path)]) == 0 - - -def test_public_names_optional_free_stable(): - repository_root = Path(__file__).resolve().parent.parent - names = _load_qmcpy_public_names(repository_root) - - assert names is not None - assert "Gaussian" in names - assert "Keister" in names - # Optional dependencies are blocked in the probe context, so fallback - # exports are part of the deterministic name set. - assert "PFGPCI" in names - # Helpers that are deliberately not part of the top-level API. - assert "PFGPCIData" not in names - assert "TriangularDistribution" not in names \ No newline at end of file diff --git a/test/test_ft_fast_transform_fallbacks.py b/test/test_ft_fast_transform_fallbacks.py new file mode 100644 index 000000000..7bf9c3493 --- /dev/null +++ b/test/test_ft_fast_transform_fallbacks.py @@ -0,0 +1,68 @@ +import unittest + +import numpy as np + +try: + import torch +except ImportError: + torch = None + +from qmcpy import ( + fftbr, + fftbr_torch, + fwht, + fwht_torch, + ifftbr, + ifftbr_torch, + omega_fftbr, + omega_fftbr_torch, + omega_fwht, + omega_fwht_torch, +) + + +class TestFastTransformFallbacks(unittest.TestCase): + + def test_non_torch_transforms_basic(self): + rng = np.random.default_rng(11) + x = rng.random(8) + 1j * rng.random(8) + y = fftbr(x) + self.assertEqual(y.shape, x.shape) + xr = ifftbr(y) + self.assertEqual(xr.shape, x.shape) + + a = rng.random(8) + b = fwht(a) + self.assertEqual(b.shape, a.shape) + + omega = omega_fftbr(3) + self.assertEqual(omega.shape[0], 2**3) + omega2 = omega_fwht(3) + self.assertEqual(omega2.shape[0], 2**3) + + def test_torch_transforms_or_fallbacks(self): + if torch is None: + calls = ( + (fftbr_torch, np.zeros(8, dtype=complex)), + (ifftbr_torch, np.zeros(8, dtype=complex)), + (fwht_torch, np.zeros(8)), + (omega_fftbr_torch, 3), + (omega_fwht_torch, 3), + ) + for transform, argument in calls: + with self.subTest(transform=transform.__name__): + with self.assertRaisesRegex(ModuleNotFoundError, "requires torch"): + transform(argument) + return + + complex_x = torch.zeros(8, dtype=torch.complex64) + real_x = torch.zeros(8) + self.assertEqual(fftbr_torch(complex_x).shape, complex_x.shape) + self.assertEqual(ifftbr_torch(complex_x).shape, complex_x.shape) + self.assertEqual(fwht_torch(real_x).shape, real_x.shape) + self.assertEqual(omega_fftbr_torch(3).shape[0], 2**3) + self.assertEqual(omega_fwht_torch(3).shape[0], 2**3) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_ig_financial_option_quick.py b/test/test_ig_financial_option_quick.py new file mode 100644 index 000000000..b05b8708b --- /dev/null +++ b/test/test_ig_financial_option_quick.py @@ -0,0 +1,101 @@ +import unittest + +import numpy as np + +from qmcpy import AbstractDiscreteDistribution, FinancialOption + + +class SmallSampler(AbstractDiscreteDistribution): + def __init__(self, d=3): + super().__init__( + dimension=d, replications=1, seed=123, d_limit=100, n_limit=1024 + ) + + def _gen_samples(self, n_min, n_max, return_binary=False, warn=True): + n = n_max - n_min + # return shape (replications, n, d) + arr = np.tile(np.linspace(0.1, 1.0, n)[:, None], (1, self.d)) + return arr.reshape(self.replications, n, self.d) + + +class TestFinancialOptionPayoffs(unittest.TestCase): + + def test_financial_option_payoffs_and_exact(self): + sampler = SmallSampler(d=3) + fo = FinancialOption( + sampler, + option="EUROPEAN", + call_put="CALL", + volatility=0.5, + start_price=30, + strike_price=25, + interest_rate=0.01, + t_final=1, + ) + gbm = np.array([[30.0, 28.0, 35.0]]) + c = fo.payoff_european_call(gbm) + p = fo.payoff_european_put(gbm) + self.assertEqual(c.shape, (1,)) + self.assertEqual(p.shape, (1,)) + + # Asian arithmetic trapezoidal + fo_asian = FinancialOption( + sampler, + option="ASIAN", + asian_mean="ARITHMETIC", + asian_mean_quadrature_rule="TRAPEZOIDAL", + ) + gbm2 = np.array([[30.0, 32.0, 34.0]]) + a_call = fo_asian.payoff_asian_arithmetic_trap_call(gbm2) + self.assertEqual(a_call.shape, (1,)) + + # geometric right call + fo_geo = FinancialOption( + sampler, + option="ASIAN", + asian_mean="GEOMETRIC", + asian_mean_quadrature_rule="RIGHT", + ) + g_call = fo_geo.payoff_asian_geometric_right_call(np.array([[30.0, 30.0, 30.0]])) + self.assertEqual(g_call.shape, (1,)) + + # barrier options: up and down behaviors + fo_barrier_up = FinancialOption( + sampler, option="BARRIER", barrier_in_out="IN", barrier_price=25, start_price=20 + ) + gbm_up = np.array([[20.0, 26.0, 27.0]]) + v = fo_barrier_up.payoff_barrier_in_up_call(gbm_up) + self.assertEqual(v.shape, (1,)) + + fo_barrier_out = FinancialOption( + sampler, + option="BARRIER", + barrier_in_out="OUT", + barrier_price=40, + start_price=30, + ) + gbm_out = np.array([[30.0, 32.0, 33.0]]) + v2 = fo_barrier_out.payoff_barrier_out_up_call(gbm_out) + self.assertEqual(v2.shape, (1,)) + + # lookback + fo_lb = FinancialOption(sampler, option="LOOKBACK") + lb = fo_lb.payoff_lookback_call(np.array([[10.0, 9.0, 12.0]])) + self.assertEqual(lb.shape, (1,)) + + # digital + fo_dig = FinancialOption(sampler, option="DIGITAL", digital_payout=5) + dig = fo_dig.payoff_digital_call(np.array([[10.0, 11.0, 12.0]])) + self.assertEqual(dig.shape, (1,)) + + # exact value for European should return a float + val = fo.get_exact_value() + self.assertTrue(np.isscalar(val)) + + # exact value for Asian geometric right + val2 = fo_geo.get_exact_value() + self.assertTrue(np.isscalar(val2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_integrands.py b/test/test_ig_integrands.py similarity index 100% rename from test/test_integrands.py rename to test/test_ig_integrands.py diff --git a/test/test_option.py b/test/test_ig_option.py similarity index 100% rename from test/test_option.py rename to test/test_ig_option.py diff --git a/test/test_option_ml.py b/test/test_ig_option_ml.py similarity index 100% rename from test/test_option_ml.py rename to test/test_ig_option_ml.py diff --git a/test/test_install_mpmc_pyg.py b/test/test_install_mpmc_pyg.py deleted file mode 100644 index 41b1b4e0b..000000000 --- a/test/test_install_mpmc_pyg.py +++ /dev/null @@ -1,85 +0,0 @@ -"""Tests for the platform-specific MPMC dependency installer.""" - -import subprocess -from types import SimpleNamespace - -import pytest - -from qmcpy.util import install_mpmc_pyg - - -def _torch(version="2.12.1+cpu", cuda=None, hip=None): - return SimpleNamespace( - __version__=version, - version=SimpleNamespace(cuda=cuda, hip=hip), - ) - - -def test_torch_versions_include_baseline_fallback(): - """Wheel lookup tries an exact patch release, then its minor baseline.""" - assert install_mpmc_pyg.torch_versions("2.12.1+cpu") == ["2.12.1", "2.12.0"] - assert install_mpmc_pyg.torch_versions("2.12.0") == ["2.12.0"] - - with pytest.raises(RuntimeError, match="Unable to parse torch version"): - install_mpmc_pyg.torch_versions("development") - - -@pytest.mark.parametrize( - ("torch_module", "expected"), - [ - (_torch(), "cpu"), - (_torch(cuda="12.6"), "cu126"), - (_torch(cuda="13.0.1"), "cu130"), - ], -) -def test_accelerator_tag(torch_module, expected): - """PyTorch build metadata maps to the expected PyG wheel tag.""" - assert install_mpmc_pyg.accelerator_tag(torch_module) == expected - - -def test_accelerator_tag_rejects_rocm(): - """The installer directs unsupported ROCm users to upstream guidance.""" - with pytest.raises(RuntimeError, match="does not currently support ROCm"): - install_mpmc_pyg.accelerator_tag(_torch(hip="6.3")) - - -def test_main_retries_with_torch_minor_baseline(monkeypatch): - """A missing exact wheel page falls back to the minor baseline page.""" - calls = [] - - def fake_run(*args): - calls.append(args) - if args[-1].endswith("torch-2.12.1+cpu.html"): - raise subprocess.CalledProcessError(1, args) - - monkeypatch.setattr(install_mpmc_pyg, "run", fake_run) - - install_mpmc_pyg.main(_torch()) - - assert calls[0][-1] == "torch-geometric>=2.6.1" - assert calls[1][-1] == "https://data.pyg.org/whl/torch-2.12.1+cpu.html" - assert calls[2][-1] == "https://data.pyg.org/whl/torch-2.12.0+cpu.html" - assert "--only-binary" in calls[1] - - -def test_main_explains_that_torch_must_be_installed(monkeypatch): - """Running the helper before installing the extra gives a useful error.""" - def missing_torch(_name): - raise ModuleNotFoundError("No module named 'torch'", name="torch") - - monkeypatch.setattr(install_mpmc_pyg.importlib, "import_module", missing_torch) - - with pytest.raises(RuntimeError, match=r"install 'qmcpy\[mpmc\]'"): - install_mpmc_pyg.main() - - -def test_main_reports_missing_wheel(monkeypatch): - """Exhausting candidate wheel pages reports the build that failed.""" - def fail_pyg_lib(*args): - if "pyg_lib>=0.6.0" in args: - raise subprocess.CalledProcessError(1, args) - - monkeypatch.setattr(install_mpmc_pyg, "run", fail_pyg_lib) - - with pytest.raises(RuntimeError, match=r"torch 2\.12\.1\+cpu \(cpu\)"): - install_mpmc_pyg.main(_torch()) diff --git a/test/test_kernels.py b/test/test_kn_kernels.py similarity index 100% rename from test/test_kernels.py rename to test/test_kn_kernels.py diff --git a/test/test_mpmc_optional_imports.py b/test/test_mpmc_optional_imports.py deleted file mode 100644 index 33b71b841..000000000 --- a/test/test_mpmc_optional_imports.py +++ /dev/null @@ -1,100 +0,0 @@ -import ast -import builtins -from pathlib import Path - -import pytest - - -def _execute_optional_import(blocked_import): - repository_root = Path(__file__).resolve().parent.parent - init_path = repository_root / "qmcpy" / "__init__.py" - init_tree = ast.parse(init_path.read_text()) - optional_import = next( - node - for node in init_tree.body - if isinstance(node, ast.Try) - and any( - isinstance(statement, ast.ImportFrom) - and statement.module == "discrete_distribution.mpmc" - for statement in node.body - ) - ) - - import qmcpy - - real_import = builtins.__import__ - - def guarded_import(name, globals=None, locals=None, fromlist=(), level=0): - missing_module = blocked_import(name, fromlist, level) - if missing_module is not None: - raise ModuleNotFoundError( - "blocked optional dependency", - name=missing_module, - ) - return real_import(name, globals, locals, fromlist, level) - - test_builtins = vars(builtins).copy() - test_builtins["__import__"] = guarded_import - namespace = {"__builtins__": test_builtins, "__package__": "qmcpy"} - module = ast.Module(body=[optional_import], type_ignores=[]) - exec(compile(module, str(init_path), "exec"), namespace) - return namespace - - -def test_mpmc_utils_remain_available_without_pyg(): - pytest.importorskip("torch") - - def block_pyg_models(name, fromlist, level): - if level == 1 and name == "discrete_distribution.mpmc.models": - return "torch_geometric" - return None - - namespace = _execute_optional_import(block_pyg_models) - - import qmcpy - - assert namespace["mpmc_utils"] is qmcpy.mpmc_utils - assert namespace["mpmc_utils"].__name__ == ( - "qmcpy.discrete_distribution.mpmc.utils" - ) - assert "utils" not in namespace - - with pytest.raises(ModuleNotFoundError, match="MPMC_net.*torch_geometric") as error: - namespace["MPMC_net"]() - assert error.value.name == "torch_geometric" - - -def test_mpmc_placeholders_report_missing_torch(): - def block_torch_utils(name, fromlist, level): - if ( - level == 1 - and name == "discrete_distribution.mpmc" - and "utils" in fromlist - ): - return "torch" - return None - - namespace = _execute_optional_import(block_torch_utils) - - with pytest.raises(ModuleNotFoundError, match="mpmc_utils.*torch") as error: - namespace["mpmc_utils"].L2star - assert error.value.name == "torch" - - with pytest.raises(ModuleNotFoundError, match="MPMC_net.*torch") as error: - namespace["MPMC_net"]() - assert error.value.name == "torch" - - -def test_mpmc_placeholder_missing_torch_scatter(): - pytest.importorskip("torch") - - def block_torch_scatter(name, fromlist, level): - if level == 1 and name == "discrete_distribution.mpmc.models": - return "torch_scatter" - return None - - namespace = _execute_optional_import(block_torch_scatter) - - with pytest.raises(ModuleNotFoundError, match="MPMC_net.*torch_scatter") as error: - namespace["MPMC_net"]() - assert error.value.name == "torch_scatter" diff --git a/test/test_plot_and_stop.py b/test/test_plot_and_stop.py deleted file mode 100644 index dd564cd9d..000000000 --- a/test/test_plot_and_stop.py +++ /dev/null @@ -1,155 +0,0 @@ -import sys -import types -import numpy as np -import builtins -import pytest - -import qmcpy -from qmcpy import plot_proj -from qmcpy.util import stop_notebook - - -class FakeAxes: - def __init__(self): - self.removed = False - self.calls = [] - - def remove(self): - self.removed = True - - def set_xlim(self, *a, **k): - self.calls.append(("set_xlim", a)) - - def set_ylim(self, *a, **k): - self.calls.append(("set_ylim", a)) - - def set_xticks(self, *a, **k): - self.calls.append(("set_xticks", a)) - - def set_yticks(self, *a, **k): - self.calls.append(("set_yticks", a)) - - def set_aspect(self, *a, **k): - self.calls.append(("set_aspect", a)) - - def grid(self, *a, **k): - self.calls.append(("grid", a)) - - def tick_params(self, *a, **k): - self.calls.append(("tick_params", a)) - - def set_xlabel(self, *a, **k): - self.calls.append(("set_xlabel", a)) - - def set_ylabel(self, *a, **k): - self.calls.append(("set_ylabel", a)) - - def scatter(self, *a, **k): - self.calls.append(("scatter", a)) - - -class FakeFig: - def __init__(self): - self.tl = False - - def tight_layout(self, *a, **k): - self.tl = True - - -def make_fake_matplotlib(nrows, ncols): - plt = types.ModuleType("matplotlib.pyplot") - plt.style = types.SimpleNamespace() - plt.style.use = lambda *a, **k: None - plt.rcParams = { - "font.family": "sans-serif", - "axes.prop_cycle": types.SimpleNamespace( - by_key=lambda: {"color": ["k", "b", "r"]} - ), - } - - def subplots(nrows=1, ncols=1, figsize=None, squeeze=False): - fig = FakeFig() - ax = np.empty((nrows, ncols), dtype=object) - for i in range(nrows): - for j in range(ncols): - ax[i, j] = FakeAxes() - return fig, ax - - plt.subplots = subplots - plt.suptitle = lambda *a, **k: None - return plt - - -class DummySampler(qmcpy.AbstractDiscreteDistribution): - def __init__(self, d=2): - super().__init__(dimension=d, replications=1, seed=1, d_limit=10, n_limit=100) - - def _gen_samples(self, n_min, n_max, return_binary=False, warn=True): - n = n_max - n_min - return np.tile(np.arange(n)[:, None] / max(1, n - 1), (1, 1, self.d)).reshape( - self.replications, n, self.d - ) - - def __repr__(self): - return "DummySampler" - - -def test_plot_proj_with_fake_matplotlib_and_sampler(monkeypatch): - # Inject fake matplotlib.pyplot - fake_plt = make_fake_matplotlib(1, 1) - # Create a proper matplotlib package module with colors submodule - fake_matplotlib = types.ModuleType("matplotlib") - fake_matplotlib.pyplot = fake_plt - fake_matplotlib.colors = types.SimpleNamespace() - monkeypatch.setitem(sys.modules, "matplotlib.pyplot", fake_plt) - monkeypatch.setitem(sys.modules, "matplotlib", fake_matplotlib) - - sampler = DummySampler(d=3) - fig, ax = plot_proj( - sampler, - n=4, - d_horizontal=1, - d_vertical=2, - math_ind=True, - marker_size=1, - figfac=1, - ) - assert isinstance(fig, FakeFig) - assert isinstance(ax, np.ndarray) - # At least one axes should have scatter calls or be removed - found = False - for a in ax.flatten(): - if getattr(a, "removed", False) or any(c[0] == "scatter" for c in a.calls): - found = True - break - assert found - - -def test_plot_proj_with_callable_sampler(monkeypatch): - # sampler not instance of AbstractDiscreteDistribution -> uses t_i labels - fake_plt = make_fake_matplotlib(1, 1) - fake_matplotlib = types.ModuleType("matplotlib") - fake_matplotlib.pyplot = fake_plt - fake_matplotlib.colors = types.SimpleNamespace() - monkeypatch.setitem(sys.modules, "matplotlib.pyplot", fake_plt) - monkeypatch.setitem(sys.modules, "matplotlib", fake_matplotlib) - - def sampler_callable(n): - return np.zeros((n, 1)) - - fig, ax = plot_proj( - sampler_callable, n=3, d_horizontal=0, d_vertical=0, math_ind=False - ) - assert isinstance(fig, FakeFig) - - -def test_stop_notebook_yes_and_no(monkeypatch): - # When input is 'yes' nothing should happen - monkeypatch.setattr(builtins, "input", lambda prompt="": "yes") - # Should not raise - stop_notebook("prompt") - - # When input is not 'yes' should exit - monkeypatch.setattr(builtins, "input", lambda prompt="": "no") - with pytest.raises(SystemExit): - stop_notebook("prompt") diff --git a/test/test_accumulate_data.py b/test/test_sc_accumulate_data.py similarity index 100% rename from test/test_accumulate_data.py rename to test/test_sc_accumulate_data.py diff --git a/test/test_cubbayes_vec.py b/test/test_sc_cubbayes_vec.py similarity index 100% rename from test/test_cubbayes_vec.py rename to test/test_sc_cubbayes_vec.py diff --git a/test/test_stopping_criteria.py b/test/test_sc_stopping_criteria.py similarity index 100% rename from test/test_stopping_criteria.py rename to test/test_sc_stopping_criteria.py diff --git a/test/test_scipy_wrapper_custom.py b/test/test_scipy_wrapper_custom.py deleted file mode 100644 index dd713e934..000000000 --- a/test/test_scipy_wrapper_custom.py +++ /dev/null @@ -1,302 +0,0 @@ -import warnings - -import pytest -import numpy as np -import scipy.stats as stats - -from qmcpy import DigitalNetB2, SciPyWrapper, StudentT, ZeroInflatedExpUniform - -from qmcpy.true_measure.triangular import TriangularDistribution -from qmcpy.util import DimensionError, ParameterError - - -MISSING_PDF_WARNING = "no 'pdf' or 'logpdf'" - - -def _missing_pdf_warnings(caught): - return [ - warning - for warning in caught - if issubclass(warning.category, UserWarning) - and MISSING_PDF_WARNING in str(warning.message) - ] - - -def test_mvn_dependence_correlation_and_moment(): - """ - Check that passing a SciPy multivariate normal through SciPyWrapper - preserves correlation and the mixed moment E[X1 X2]. - """ - sampler = DigitalNetB2(2, seed=5) - rho_target = 0.7 - cov = [[1.0, rho_target], [rho_target, 1.0]] - mvn = stats.multivariate_normal(mean=[0.0, 0.0], cov=cov) - tm_mvn = SciPyWrapper(sampler, scipy_distribs=mvn) - - n = 4096 - x = tm_mvn(n) - - rho_hat = np.corrcoef(x.T)[0, 1] - est_moment = np.mean(x[:, 0] * x[:, 1]) - - assert np.isfinite(rho_hat) - assert np.isfinite(est_moment) - - assert abs(rho_hat - rho_target) < 0.05 - assert abs(est_moment - rho_target) < 0.05 - - -def test_triangular_custom_marginal_range_and_shape(): - """ - Make sure our custom triangular marginal behaves sensibly: - samples stay in the right interval and the empirical mean is close - to the analytic mean. - """ - tri = TriangularDistribution(c=0.3, loc=-1.0, scale=2.0) - tm = SciPyWrapper(DigitalNetB2(1, seed=11), scipy_distribs=tri) - - n = 4096 - x = tm(n).ravel() - - assert x.min() >= -1.1 - assert x.max() <= 1.1 - - a = -1.0 - b = 1.0 - m = -1.0 + 0.3 * 2.0 - true_mean = (a + b + m) / 3.0 - emp_mean = x.mean() - assert abs(emp_mean - true_mean) < 0.05 - - -def test_zero_inflated_zero_rate(): - """ - Check that the zero-inflated exponential distribution preserves the - specified probability mass at X = 0. - """ - p_zero = 0.4 - sampler = DigitalNetB2(1, seed=17) - tm = ZeroInflatedExpUniform(sampler, p_zero=p_zero, lam=1.5) - - n = 4096 - samples = tm(n) - x = samples.ravel() - zero_rate = np.mean(x == 0.0) - - assert samples.shape == (n, 1) - assert abs(zero_rate - p_zero) < 0.05 - - -def test_zero_inflated_replications_shape(): - tm = ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17, replications=2), - p_zero=0.4, - lam=1.5, - ) - - x = tm(8) - - assert x.shape == (2, 8, 1) - assert np.all(x >= 0.0) - - -@pytest.mark.parametrize("p_zero", [0.0, 1.0, -0.1, 1.1]) -def test_zero_inflated_rejects_invalid_p_zero(p_zero): - with pytest.raises(ParameterError, match="p_zero must be in"): - ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17), - p_zero=p_zero, - lam=1.5, - ) - - -@pytest.mark.parametrize("lam", [0.0, -1.0]) -def test_zero_inflated_rejects_nonpositive_lam(lam): - with pytest.raises(ParameterError, match="lam must be positive"): - ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17), - p_zero=0.4, - lam=lam, - ) - - -def test_zero_inflated_requires_one_dimensional_sampler(): - with pytest.raises( - DimensionError, - match="requires a one-dimensional sampler", - ): - ZeroInflatedExpUniform( - DigitalNetB2(2, seed=17), - p_zero=0.4, - lam=1.5, - ) - - -def test_zero_inflated_inverse_transform_exact_values(): - tm = ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17), - p_zero=0.4, - lam=2.0, - ) - u = np.array([[0.0], [0.2], [0.4], [0.7], [0.9]]) - - x = tm._transform(u) - - assert x.shape == (5, 1) - assert np.array_equal(x[:3], np.zeros((3, 1))) - assert np.all(x[3:] > 0.0) - - u_positive = u[3:, 0] - u_rescaled = (u_positive - 0.4) / 0.6 - expected = -np.log1p(-u_rescaled) / 2.0 - assert np.allclose(x[3:, 0], expected) - - -def test_zero_inflated_inverse_transform_all_zero_branch(): - tm = ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17), - p_zero=0.4, - lam=2.0, - ) - u = np.array([[0.0], [0.1], [0.4]]) - - x = tm._transform(u) - - assert x.shape == (3, 1) - assert np.array_equal(x, np.zeros((3, 1))) - - -def test_zero_inflated_inverse_transform_clips_one(): - tm = ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17), - p_zero=0.4, - lam=2.0, - ) - u = np.array([[1.0]]) - - x = tm._transform(u) - - assert x.shape == (1, 1) - assert np.isfinite(x).all() - assert x[0, 0] > 0.0 - - -def test_zero_inflated_construction_does_not_warn_about_missing_pdf(): - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - tm = ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17), - p_zero=0.4, - lam=1.5, - ) - - assert tm.d == 1 - assert _missing_pdf_warnings(caught) == [] - - -def test_zero_inflated_sampling_does_not_warn_about_missing_pdf(): - tm = ZeroInflatedExpUniform(DigitalNetB2(1, seed=17), p_zero=0.4, lam=1.5) - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - x = tm(8) - - assert x.shape == (8, 1) - assert _missing_pdf_warnings(caught) == [] - - -def test_zero_inflated_return_weights_warns_once_for_missing_pdf(): - tm = ZeroInflatedExpUniform(DigitalNetB2(1, seed=17), p_zero=0.4, lam=1.5) - - with pytest.warns(UserWarning, match=MISSING_PDF_WARNING): - x, jac = tm(8, return_weights=True) - - assert x.shape == (8, 1) - assert jac.shape == (8,) - assert np.allclose(jac, 1.0) - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - x_second, jac_second = tm(8, return_weights=True) - - assert x_second.shape == (8, 1) - assert np.allclose(jac_second, 1.0) - assert _missing_pdf_warnings(caught) == [] - - -def test_zero_inflated_y_split_warns_and_uses_one_dimensional_interface(): - with pytest.warns(DeprecationWarning, match="y_split"): - tm = ZeroInflatedExpUniform( - DigitalNetB2(1, seed=17), - p_zero=0.4, - lam=1.5, - y_split=0.5, - ) - - x = tm(4) - - assert x.shape == (4, 1) - assert np.all(x >= 0.0) - - -def test_zero_inflated_y_split_preserves_deprecated_two_dimensional_usage(): - with pytest.warns(DeprecationWarning, match="2D zero-inflated"): - tm = ZeroInflatedExpUniform( - DigitalNetB2(2, seed=17), - p_zero=0.4, - lam=1.5, - y_split=0.5, - ) - - x = tm(16) - - assert x.shape == (16, 2) - assert np.all(x[:, 0] >= 0.0) - assert np.all((0.0 <= x[:, 1]) & (x[:, 1] <= 1.0)) - assert np.all(x[x[:, 0] == 0.0, 1] <= 0.5) - assert np.all(x[x[:, 0] > 0.0, 1] >= 0.5) - - -def test_zero_inflated_y_split_preserves_replicated_two_dimensional_usage(): - with pytest.warns(DeprecationWarning, match="2D zero-inflated"): - tm = ZeroInflatedExpUniform( - DigitalNetB2(2, seed=17, replications=2), - p_zero=0.4, - lam=1.5, - y_split=0.5, - ) - - x = tm(16) - - assert x.shape == (2, 16, 2) - assert np.all(x[..., 0] >= 0.0) - assert np.all((0.0 <= x[..., 1]) & (x[..., 1] <= 1.0)) - assert np.all(x[..., 1][x[..., 0] == 0.0] <= 0.5) - assert np.all(x[..., 1][x[..., 0] > 0.0] >= 0.5) - - -def test_student_t_marginals_shape(): - tm = SciPyWrapper( - sampler=DigitalNetB2(2, seed=5), - scipy_distribs=stats.t(df=5), - ) - x = tm(8) - assert x.shape == (8, 2) - - -def test_multivariate_student_t_joint_corr_and_cov(): - if not hasattr(stats, "multivariate_t"): - pytest.skip("scipy.stats.multivariate_t not available in this SciPy version") - - df = 5.0 - rho = 0.8 - loc = np.array([0.0, 0.0]) - shape = np.array([[1.0, rho], [rho, 1.0]]) - - tm = StudentT(DigitalNetB2(2, seed=123), loc=loc, shape=shape, df=df) - - n = 4096 - x = tm(n) - emp_corr = np.corrcoef(x.T)[0, 1] - - assert abs(emp_corr - rho) < 0.05 diff --git a/test/test_sr_annotate_public_api_types.py b/test/test_sr_annotate_public_api_types.py new file mode 100644 index 000000000..e3f12344e --- /dev/null +++ b/test/test_sr_annotate_public_api_types.py @@ -0,0 +1,230 @@ +import shutil +import tempfile +import textwrap +import unittest +from contextlib import redirect_stdout +from io import StringIO +from pathlib import Path + +from scripts import annotate_public_api_types + + +class TestAnnotatePublicAPITypes(unittest.TestCase): + + def setUp(self): + self.tmp_path = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp_path, ignore_errors=True) + + def _write(self, source): + path = self.tmp_path / "sample.py" + path.write_text(textwrap.dedent(source).lstrip(), encoding="utf-8") + return path + + def test_annotates_public_method_inputs_and_output(self): + path = self._write( + ''' + import numpy as np + + class Model: + + def evaluate(self, x, scale=1.0): + """Evaluate the model. + + Args: + x (np.ndarray): Evaluation points. + scale (float): Output scale. + + Returns: + np.ndarray: Scaled values. + """ + return scale * x + ''' + ) + + result = annotate_public_api_types.update_file(path) + source = path.read_text(encoding="utf-8") + + self.assertTrue(result.changed) + self.assertEqual(len(result.updates), 3) + self.assertIn( + "def evaluate(self, x: np.ndarray, scale: float = 1.0) -> np.ndarray:", + source, + ) + + def test_annotates_constructor_and_adds_none_return(self): + path = self._write( + ''' + class Body: + + def __init__(self, mass): + """Initialize a body. + + Args: + mass (float): Body mass. + """ + self.mass = mass + ''' + ) + + annotate_public_api_types.update_file(path) + + self.assertIn( + "def __init__(self, mass: float) -> None:", + path.read_text(encoding="utf-8"), + ) + + def test_annotates_decorated_public_method(self): + path = self._write( + ''' + class Model: + + @staticmethod + def normalize(x): + """Normalize a value. + + Args: + x (float): Value to normalize. + + Returns: + float: Normalized value. + """ + return x + ''' + ) + + result = annotate_public_api_types.update_file(path) + + self.assertTrue(result.changed) + self.assertIn( + "def normalize(x: float) -> float:", + path.read_text(encoding="utf-8"), + ) + + def test_preserves_existing_annotation_and_reports_conflict(self): + path = self._write( + ''' + def scale(x: int) -> float: + """Scale a value. + + Args: + x (float): Value to scale. + + Returns: + float: Scaled value. + """ + return float(x) + ''' + ) + original = path.read_text(encoding="utf-8") + + result = annotate_public_api_types.update_file(path) + + self.assertFalse(result.changed) + self.assertEqual(len(result.conflicts), 1) + self.assertEqual(result.conflicts[0].slot, "x") + self.assertEqual(path.read_text(encoding="utf-8"), original) + + def test_skips_type_whose_name_is_not_available(self): + path = self._write( + ''' + def evaluate(x): + """Evaluate points. + + Args: + x (ArrayLike): Evaluation points. + """ + return x + ''' + ) + original = path.read_text(encoding="utf-8") + + result = annotate_public_api_types.update_file(path) + + self.assertFalse(result.changed) + self.assertEqual(len(result.skips), 1) + self.assertIn("not available", result.skips[0].reason) + self.assertEqual(path.read_text(encoding="utf-8"), original) + + def test_skips_types_that_contradict_literal_defaults(self): + path = self._write( + ''' + import numpy as np + + def evaluate(x=None, tolerance=0.5): + """Evaluate points. + + Args: + x (np.ndarray): Evaluation points. + tolerance (np.ndarray): Error tolerance. + """ + return x + ''' + ) + original = path.read_text(encoding="utf-8") + + result = annotate_public_api_types.update_file(path) + + self.assertFalse(result.changed) + self.assertEqual(len(result.skips), 2) + self.assertTrue(any("not optional" in skip.reason for skip in result.skips)) + self.assertTrue(any("conflicts" in skip.reason for skip in result.skips)) + self.assertEqual(path.read_text(encoding="utf-8"), original) + + def test_ignores_private_and_nested_functions(self): + path = self._write( + ''' + def _private(x): + """Private helper. + + Args: + x (int): Value. + """ + return x + + def public(): + """Return a nested callable.""" + + def nested(x): + """Nested helper. + + Args: + x (int): Value. + """ + return x + + return nested + ''' + ) + + result = annotate_public_api_types.update_file(path) + + self.assertFalse(result.changed) + self.assertEqual(result.updates, ()) + + def test_check_mode_reports_without_writing(self): + path = self._write( + ''' + def scale(x): + """Scale a value. + + Args: + x (float): Value to scale. + """ + return 2 * x + ''' + ) + original = path.read_text(encoding="utf-8") + output = StringIO() + + with redirect_stdout(output): + status = annotate_public_api_types.main( + ["--check", "--root", str(self.tmp_path), str(path)] + ) + + self.assertEqual(status, 1) + self.assertIn("1 file(s) would change", output.getvalue()) + self.assertEqual(path.read_text(encoding="utf-8"), original) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_sr_check_links.py b/test/test_sr_check_links.py new file mode 100644 index 000000000..49c148256 --- /dev/null +++ b/test/test_sr_check_links.py @@ -0,0 +1,198 @@ +import contextlib +import io +import shutil +import ssl +import sys +import tempfile +import unittest +import urllib.error +from pathlib import Path +from unittest.mock import patch + +from scripts import check_links + + +def _http_error(url, code): + return urllib.error.HTTPError(url, code, "test response", {}, None) + + +class TestCheckLinks(unittest.TestCase): + + def setUp(self): + self.tmp_path = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp_path, ignore_errors=True) + + def _patch(self, target, name, value): + """monkeypatch.setattr equivalent: set now, auto-restore at test end.""" + patcher = patch.object(target, name, value) + patcher.start() + self.addCleanup(patcher.stop) + + def test_head_success_is_reachable(self): + with patch.object(check_links.urllib.request, "urlopen", return_value=object()) as urlopen: + self.assertIsNone(check_links._check_one("https://example.test", timeout=1)) + + self.assertEqual(urlopen.call_count, 1) + self.assertEqual(urlopen.call_args.args[0].get_method(), "HEAD") + + def test_get_success_after_head_failure_is_reachable(self): + url = "https://example.test" + with patch.object( + check_links.urllib.request, + "urlopen", + side_effect=[_http_error(url, 405), object()], + ) as urlopen: + self.assertIsNone(check_links._check_one(url, timeout=1)) + + self.assertEqual(urlopen.call_count, 2) + self.assertEqual(urlopen.call_args_list[1].args[0].get_method(), "GET") + + def test_not_found_and_gone_gets_are_broken(self): + for code in (404, 410): + with self.subTest(code=code): + url = f"https://example.test/{code}" + with patch.object( + check_links.urllib.request, + "urlopen", + side_effect=[_http_error(url, code), _http_error(url, code)], + ): + self.assertEqual( + check_links._check_one(url, timeout=1), + ("broken", f"{url} -- HTTP {code}"), + ) + + def test_bot_block_and_rate_limit_are_warnings(self): + for code in (403, 429): + with self.subTest(code=code): + url = f"https://example.test/{code}" + with patch.object( + check_links.urllib.request, + "urlopen", + side_effect=[_http_error(url, code), _http_error(url, code)], + ): + severity, message = check_links._check_one(url, timeout=1) + + self.assertEqual(severity, "warning") + self.assertIn(f"HTTP {code}", message) + + def test_tls_and_timeout_failures_are_warnings(self): + failures = ( + ssl.SSLCertVerificationError("certificate verify failed"), + TimeoutError("timed out"), + ) + for failure in failures: + with self.subTest(failure=type(failure).__name__): + with patch.object( + check_links.urllib.request, + "urlopen", + side_effect=[failure, failure], + ): + severity, message = check_links._check_one( + "https://example.test", timeout=1 + ) + + self.assertEqual(severity, "warning") + self.assertIn(str(failure), message) + + def test_external_results_are_separated_and_duplicate_urls_checked_once(self): + (self.tmp_path / "page.html").write_text( + 'missing' + 'duplicate' + 'blocked', + encoding="utf-8", + ) + + def result_for(url, _timeout): + if url.endswith("/missing"): + return "broken", f"{url} -- HTTP 404" + return "warning", f"{url} -- HTTP 403" + + with patch.object(check_links, "_check_one", side_effect=result_for) as check_one: + broken, warnings = check_links.check_external(self.tmp_path, workers=1) + + self.assertEqual(check_one.call_count, 2) + self.assertEqual( + broken, + ["https://example.test/missing -- HTTP 404 (seen on page.html)"], + ) + self.assertEqual( + warnings, + ["https://example.test/blocked -- HTTP 403 (seen on page.html)"], + ) + + def test_internal_links_strip_site_url_deployment_path(self): + target = self.tmp_path / "target" + target.mkdir() + (target / "index.html").write_text( + '

Target

', encoding="utf-8" + ) + (self.tmp_path / "index.html").write_text( + 'root-relative' + 'absolute', + encoding="utf-8", + ) + + self.assertEqual( + check_links.check_internal( + self.tmp_path, site_url="https://qmcsoftware.github.io/QMCSoftware/" + ), + [], + ) + + def test_external_check_skips_same_site_urls(self): + (self.tmp_path / "page.html").write_text( + 'same' + 'external', + encoding="utf-8", + ) + + with patch.object(check_links, "_check_one", return_value=None) as check_one: + broken, warnings = check_links.check_external( + self.tmp_path, + workers=1, + site_url="https://qmcsoftware.github.io/QMCSoftware/", + ) + + self.assertEqual(broken, []) + self.assertEqual(warnings, []) + self.assertEqual(check_one.call_count, 1) + self.assertEqual(check_one.call_args.args[0], "https://example.test/target/") + + def test_external_warnings_do_not_make_main_fail(self): + self._patch(sys, "argv", ["check_links.py", str(self.tmp_path), "--external"]) + self._patch( + check_links, "check_internal", lambda _site_dir, site_url=None: [] + ) + self._patch( + check_links, + "check_external", + lambda _site_dir, site_url=None: ( + [], + ["https://example.test -- HTTP 403"], + ), + ) + + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + self.assertEqual(check_links.main(), 0) + self.assertIn("0 broken link(s), 1 warning(s)", buf.getvalue()) + + def test_confirmed_external_breakage_makes_main_fail(self): + self._patch(sys, "argv", ["check_links.py", str(self.tmp_path), "--external"]) + self._patch( + check_links, "check_internal", lambda _site_dir, site_url=None: [] + ) + self._patch( + check_links, + "check_external", + lambda _site_dir, site_url=None: ( + ["https://example.test -- HTTP 404"], + [], + ), + ) + + self.assertEqual(check_links.main(), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_sr_check_removed_urls.py b/test/test_sr_check_removed_urls.py new file mode 100644 index 000000000..57a30223d --- /dev/null +++ b/test/test_sr_check_removed_urls.py @@ -0,0 +1,165 @@ +import contextlib +import io +import shutil +import sys +import tempfile +import unittest +import urllib.error +from pathlib import Path +from unittest.mock import patch + +from scripts import check_removed_urls as cru + +SITE = "https://qmcsoftware.github.io/QMCSoftware/" + + +def _sitemap(*paths): + locs = "".join(f"{SITE}{path}" for path in paths) + return f'{locs}' + + +def _config(redirect_maps=None): + plugins = ["material/search", {"mkdocs-jupyter": {"execute": False}}] + if redirect_maps is not None: + plugins.append({"redirects": {"redirect_maps": redirect_maps}}) + return {"site_url": SITE, "plugins": plugins} + + +class TestCheckRemovedUrls(unittest.TestCase): + + def setUp(self): + self.tmp_path = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp_path, ignore_errors=True) + self._last_out = "" + + def _patch(self, target, name, value): + """monkeypatch.setattr equivalent: set now, auto-restore at test end.""" + patcher = patch.object(target, name, value) + patcher.start() + self.addCleanup(patcher.stop) + + def _run(self, sitemap_paths, redirect_maps=None, extra_argv=(), base=None): + """Run main() offline against a temp sitemap and a temp docs/ tree.""" + base = self.tmp_path if base is None else base + docs = base / "docs" + docs.mkdir(parents=True) + (docs / "README.md").write_text("home", encoding="utf-8") + (docs / "good_practices.md").write_text("page", encoding="utf-8") + sitemap = base / "sitemap.xml" + sitemap.write_text(_sitemap(*sitemap_paths), encoding="utf-8") + + self._patch(cru, "read_config", lambda *a, **k: _config(redirect_maps)) + self._patch(sys, "argv", [ + "check_removed_urls.py", "--sitemap", str(sitemap), "--docs-dir", str(docs), + *extra_argv, + ]) + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + code = cru.main() + self._last_out = buf.getvalue() + return code + + def test_url_path_and_source_round_trip(self): + for source, url_path in [("blogs/scipywrapper/index.md", "blogs/scipywrapper/"), + ("good_practices.md", "good_practices/"), + ("demos/quickstart.ipynb", "demos/quickstart/"), + ("index.md", ""), ("README.md", "")]: + self.assertEqual(cru.url_path_for_source(source), url_path) + + for source in ("README.md", "good_practices.md", "demos/quickstart.ipynb", + "api/index.md"): + path = self.tmp_path / source + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("page", encoding="utf-8") + self.assertTrue( + cru.source_exists(cru.url_path_for_source(source), self.tmp_path) + ) + self.assertFalse(cru.source_exists("blogs/scipywrapper/", self.tmp_path)) + + def test_redirect_maps_reads_the_plugin_and_tolerates_its_absence(self): + entry = {"blogs/x/index.md": "https://qmcsoftware.org/blogs/x/"} + self.assertEqual(cru.redirect_maps(_config(entry)), entry) + self.assertEqual(cru.redirect_maps(_config()), {}) + self.assertEqual(cru.redirect_maps({}), {}) + + def test_published_paths_separates_foreign_urls(self): + sitemap = _sitemap("", "good_practices/").replace( + "", "https://example.test/other/") + + self.assertEqual( + cru.published_paths(sitemap, SITE), + (["", "good_practices/"], ["https://example.test/other/"]), + ) + + def test_http_status_falls_back_to_get_when_head_is_unsupported(self): + url = "https://example.test" + error = urllib.error.HTTPError(url, 405, "test response", {}, None) + response = type("Response", (), {"status": 200, "__enter__": lambda s: s, + "__exit__": lambda s, *a: False})() + with patch.object(cru.urllib.request, "urlopen", + side_effect=[error, response]) as urlopen: + self.assertEqual(cru.http_status(url, timeout=1), "200") + + self.assertEqual(urlopen.call_count, 2) + self.assertEqual(urlopen.call_args_list[1].args[0].get_method(), "GET") + + def test_removed_page_without_redirect_is_flagged(self): + code = self._run(["", "good_practices/", "blogs/scipywrapper/"]) + out = self._last_out + + self.assertEqual(code, 1) + self.assertIn("1 removed with no redirect", out) + self.assertIn(f"[ORPHAN] {SITE}blogs/scipywrapper/", out) + self.assertIn("blogs/scipywrapper/index.md: ", out) + + def test_removed_page_covered_by_a_redirect_passes(self): + code = self._run( + ["", "good_practices/", "blogs/scipywrapper/"], + redirect_maps={ + "blogs/scipywrapper/index.md": "https://qmcsoftware.org/blogs/scipywrapper/"}, + ) + out = self._last_out + + self.assertEqual(code, 0) + self.assertIn("0 removed with no redirect", out) + self.assertIn("[redirect]", out) + self.assertNotIn("[ORPHAN]", out) + + def test_intact_site_passes(self): + self.assertEqual(self._run(["", "good_practices/"]), 0) + self.assertIn("2 still have a page source", self._last_out) + + def test_verify_redirects_follows_the_target_status(self): + redirects = {"blogs/x/index.md": "https://qmcsoftware.org/blogs/x/"} + for status, expected_code in [("200", 0), ("404", 1)]: + with self.subTest(status=status): + self._patch(cru, "http_status", lambda *a, **k: status) + code = self._run( + ["", "blogs/x/"], + redirect_maps=redirects, + extra_argv=("--verify-redirects",), + base=self.tmp_path / status, + ) + out = self._last_out + + self.assertEqual(code, expected_code) + self.assertIn(status, out) + # The URL itself is covered, so a failure is the target, not an orphan. + self.assertNotIn("[ORPHAN]", out) + + def test_unreachable_sitemap_fails_unless_offline_is_allowed(self): + self._patch(cru, "read_config", lambda *a, **k: _config()) + argv = ["check_removed_urls.py", "--sitemap", str(self.tmp_path / "absent.xml")] + + self._patch(sys, "argv", argv) + self.assertEqual(cru.main(), 1) + + self._patch(sys, "argv", argv + ["--allow-offline"]) + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + self.assertEqual(cru.main(), 0) + self.assertIn("skipping the check", buf.getvalue()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_sr_colab_notebooks.py b/test/test_sr_colab_notebooks.py new file mode 100644 index 000000000..f23bdd49e --- /dev/null +++ b/test/test_sr_colab_notebooks.py @@ -0,0 +1,345 @@ +from __future__ import annotations + +import json +import os +import shutil +import sys +import tempfile +import unittest +import unittest.mock as mock +from pathlib import Path + +from scripts import check_colab_notebooks as check +from scripts import harden_colab_notebook as harden +from scripts import smoke_test_colab_notebooks as smoke + + +def markdown_cell(source: str, cell_id: str = "markdown") -> dict: + """Build a minimal notebook Markdown cell. + + Args: + source (str): Cell source text. + cell_id (str): Notebook cell identifier. + + Returns: + dict: A Markdown cell in nbformat 4 shape. + """ + return { + "cell_type": "markdown", + "id": cell_id, + "metadata": {}, + "source": source.splitlines(keepends=True), + } + + +def code_cell(source: str, cell_id: str = "code") -> dict: + """Build a minimal notebook code cell. + + Args: + source (str): Cell source text. + cell_id (str): Notebook cell identifier. + + Returns: + dict: A code cell in nbformat 4 shape, with no outputs. + """ + return { + "cell_type": "code", + "execution_count": None, + "id": cell_id, + "metadata": {}, + "outputs": [], + "source": source.splitlines(keepends=True), + } + + +class TestColabNotebooks(unittest.TestCase): + + def _tmp_path(self) -> Path: + """Fresh temp directory, removed after the test (pytest ``tmp_path``).""" + path = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, path, ignore_errors=True) + return path + + def _setattr(self, target, name, value): + """Set ``target.name = value`` for the test only (pytest ``monkeypatch``).""" + patcher = mock.patch.object(target, name, value) + patcher.start() + self.addCleanup(patcher.stop) + + def _colab_repo(self): + """Build a throwaway repo layout and point the scripts at it. + + Returns ``(notebook_path, manifest_path)`` (pytest ``colab_repo``). + """ + tmp_path = self._tmp_path() + demos_dir = tmp_path / "demos" + demos_dir.mkdir() + notebook_path = demos_dir / "example.ipynb" + notebook = { + "cells": [ + markdown_cell("# Example\n", "title"), + code_cell("import math\n", "imports"), + ], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5, + } + notebook_path.write_text(json.dumps(notebook, indent=1) + "\n", encoding="utf-8") + + manifest_path = tmp_path / "manifest.json" + manifest = { + "repo": "QMCSoftware/QMCSoftware", + "git_ref": "develop", + "enabled": [], + "disabled": {}, + } + manifest_path.write_text(json.dumps(manifest, indent=1) + "\n", encoding="utf-8") + + self._setattr(check, "REPO_ROOT", tmp_path) + self._setattr(check, "DEMOS_DIR", demos_dir) + self._setattr(harden, "REPO_ROOT", tmp_path) + self._setattr(smoke, "REPO_ROOT", tmp_path) + return notebook_path, manifest_path + + def test_badge_stripping_preserves_intro_and_drops_badge_only_cells(self): + intro = markdown_cell( + "# ML Sensitivity Indices\n\n" + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)]" + "(https://colab.research.google.com/github/QMCSoftware/QMCSoftware/" + "blob/develop/demos/iris.ipynb)\n\n" + "This notebook demonstrates sensitivity indices.\n" + ) + badge_only = markdown_cell( + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)]" + "(https://colab.research.google.com/github/QMCSoftware/QMCSoftware/" + "blob/develop/demos/iris.ipynb)\n" + ) + + cleaned_intro = harden.badge_stripped_cell(intro) + self.assertIsNotNone(cleaned_intro) + self.assertIn("# ML Sensitivity Indices", check.cell_source_text(cleaned_intro)) + self.assertIn("sensitivity indices", check.cell_source_text(cleaned_intro)) + self.assertNotIn("Open In Colab", check.cell_source_text(cleaned_intro)) + self.assertEqual( + harden.remove_any_badge_cells([badge_only, code_cell("pass\n")]), + [code_cell("pass\n")], + ) + + def test_is_any_badge_cell_rejects_spoofed_hostname(self): + spoofed = markdown_cell( + "[click](https://evil.example/colab.research.google.com/assets/colab-badge.svg)\n" + ) + genuine = markdown_cell( + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)]" + "(https://colab.research.google.com/github/QMCSoftware/QMCSoftware/" + "blob/develop/demos/iris.ipynb)\n" + ) + + self.assertFalse(check.is_any_badge_cell(spoofed)) + self.assertTrue(check.is_any_badge_cell(genuine)) + + def test_bootstrap_detection_uses_marker_and_real_install_command(self): + misleading = code_cell( + '"""import google.colab\n# @title Execute this cell to install dependencies\n' + '!pip install qmcpy\n"""\n' + ) + comment_only = code_cell( + "# @title Execute this cell to install dependencies\n" + "# import google.colab\n" + "# !pip install qmcpy\n" + ) + self.assertFalse(check.is_any_install_cell(misleading)) + self.assertFalse(check.is_bootstrap_cell(misleading)) + self.assertTrue(check.is_any_install_cell(comment_only)) + self.assertFalse(check.is_bootstrap_cell(comment_only)) + + tmp_path = self._tmp_path() + self._setattr(harden, "REPO_ROOT", tmp_path) + notebook_path = tmp_path / "demos" / "example.ipynb" + notebook_path.parent.mkdir() + source = "".join( + harden.bootstrap_cell_source( + notebook_path, + {"repo": "QMCSoftware/QMCSoftware"}, + [], + ) + ) + generated = code_cell(source) + self.assertTrue(check.is_bootstrap_cell(generated)) + self.assertIn("except ImportError:", source) + self.assertIn("if IN_COLAB:", source) + self.assertNotIn("except:\n", source) + compile(smoke.rewrite_shell_magics(source), "", "exec") + + def test_extra_pip_packages_preserves_later_explicit_installs(self): + cells = [ + code_cell("import qmcpy as qp\n"), + code_cell("import ipywidgets as widgets\n"), + code_cell( + "try:\n" + " import QuantLib as ql\n" + "except ModuleNotFoundError:\n" + " !pip install -q QuantLib\n" + ), + code_cell("!pip install -q seaborn\n"), + ] + + self.assertEqual( + harden.extra_pip_packages(cells), ["QuantLib", "ipywidgets", "seaborn"] + ) + + def test_needs_latex_setup_detects_tueplots(self): + cells = [ + code_cell("import qmcpy as qp\n"), + code_cell( + "from tueplots import bundles\n" + "pyplot.rcParams.update(bundles.probnum2025())\n" + ), + ] + + self.assertTrue(harden.needs_latex_setup(cells)) + + def test_imported_modules_survives_magic_only_block_body(self): + # A shell-magic line as the *only* statement in a block used to leave an + # empty `if:`/`try:` body, making ast.parse raise and silently hiding + # every import in the cell (not just the magic line itself). + source = ( + "import os\n" + "from util import helper\n" + "if True:\n" + " !echo hi\n" + ) + self.assertEqual(check.imported_modules(source), {"os", "util"}) + + def test_local_module_matches_finds_ancestor_directory(self): + tmp_path = self._tmp_path() + self._setattr(check, "DEMOS_DIR", tmp_path) + (tmp_path / "util.py").write_text("", encoding="utf-8") + notebook_dir = tmp_path / "output" + notebook_dir.mkdir() + + matches = check.local_module_matches(notebook_dir, "util") + + self.assertEqual(matches, [tmp_path / "util.py"]) + + def test_extra_pip_packages_honors_colab_deps_marker(self): + cells = [ + code_cell("import qmcpy as qp\n"), + code_cell( + "# colab-deps: plotly, some-package\n" + "import plotly\n" + ), + ] + + self.assertEqual( + harden.extra_pip_packages(cells), ["plotly", "some-package"] + ) + + def test_dump_notebook_preserves_existing_json_indent(self): + tmp_path = self._tmp_path() + notebook_path = tmp_path / "example.ipynb" + notebook = { + "cells": [code_cell("pass\n")], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 5, + } + original_source = json.dumps(notebook, indent=2) + "\n" + + harden.dump_notebook(notebook_path, notebook, original_source) + + self.assertEqual( + notebook_path.read_text(encoding="utf-8"), original_source + ) + + def test_harden_check_smoke_round_trip_is_idempotent(self): + notebook_path, manifest_path = self._colab_repo() + harden.harden_notebook(notebook_path, manifest_path) + + self.assertEqual(check.run_check(manifest_path, strict=True), 0) + smoke_notebook, source_indices = smoke.build_smoke_notebook(notebook_path, 1) + self.assertEqual(len(smoke_notebook["cells"]), len(source_indices)) + + sentinel = object() + old_modules = { + name: sys.modules.get(name, sentinel) for name in ("google", "google.colab") + } + old_environment = { + name: os.environ.get(name, sentinel) + for name in ("QMC_COLAB_SMOKE", "QMC_COLAB_SMOKE_REPO_ROOT", "QMC_COLAB_SMOKE_NOTEBOOK_DIR") + } + namespace: dict = {} + try: + for cell in smoke_notebook["cells"]: + if cell["cell_type"] == "code": + exec(check.cell_source_text(cell), namespace) + finally: + for name, value in old_modules.items(): + if value is sentinel: + sys.modules.pop(name, None) + else: + sys.modules[name] = value + for name, value in old_environment.items(): + if value is sentinel: + os.environ.pop(name, None) + else: + os.environ[name] = value + + self._setattr( + harden, + "dump_notebook", + lambda *_args, **_kwargs: self.fail("unchanged notebook was rewritten"), + ) + self._setattr( + harden, + "dump_json", + lambda *_args, **_kwargs: self.fail("unchanged manifest was rewritten"), + ) + harden.harden_notebook(notebook_path, manifest_path) + + def test_checker_rejects_wrong_badge(self): + notebook_path, manifest_path = self._colab_repo() + harden.harden_notebook(notebook_path, manifest_path) + notebook = check.load_json(notebook_path) + badge = next(cell for cell in notebook["cells"] if check.is_any_badge_cell(cell)) + badge["source"] = [check.cell_source_text(badge).replace("develop", "wrong-ref")] + notebook_path.write_text(json.dumps(notebook, indent=1) + "\n", encoding="utf-8") + + self.assertEqual(check.run_check(manifest_path, strict=True), 1) + + def test_harden_failure_does_not_disable_notebook(self): + notebook_path, manifest_path = self._colab_repo() + original_manifest = manifest_path.read_text(encoding="utf-8") + self._setattr( + harden, + "harden_notebook", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("failure")), + ) + + successes, failures = harden.harden_batch([notebook_path], manifest_path) + + self.assertEqual(successes, []) + self.assertEqual(failures, [("demos/example.ipynb", "failure")]) + self.assertEqual(manifest_path.read_text(encoding="utf-8"), original_manifest) + + def test_smoke_batch_continues_after_a_notebook_failure(self): + def fake_build(notebook_path: Path, cells_after_bootstrap: int): + return {"cells": []}, [] + + def fake_execute(notebook_path: Path, smoke_nb, source_indices, timeout): + if "broken" in notebook_path.as_posix(): + raise RuntimeError("boom") + + self._setattr(smoke, "build_smoke_notebook", fake_build) + self._setattr(smoke, "execute_smoke_notebook", fake_execute) + + passed, failed = smoke.smoke_test_batch( + ["demos/broken.ipynb", "demos/ok.ipynb"], cells_after_bootstrap=1, timeout=60 + ) + + self.assertEqual(passed, ["demos/ok.ipynb"]) + self.assertEqual(failed, [("demos/broken.ipynb", "boom")]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_sr_convert_asserts.py b/test/test_sr_convert_asserts.py new file mode 100644 index 000000000..4260453a7 --- /dev/null +++ b/test/test_sr_convert_asserts.py @@ -0,0 +1,138 @@ +import shutil +import tempfile +import textwrap +import unittest +from contextlib import redirect_stderr, redirect_stdout +from io import StringIO +from pathlib import Path + +from scripts import convert_asserts + + +class TestConvertAsserts(unittest.TestCase): + + def setUp(self): + self.tmp_path = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp_path, ignore_errors=True) + + def _write(self, source): + path = self.tmp_path / "sample.py" + path.write_text(textwrap.dedent(source).lstrip(), encoding="utf-8") + return path + + def test_converts_assert_message_and_preserves_inline_comment(self): + source = textwrap.dedent( + ''' + def positive(x): + assert x > 0, f"expected positive x, got {x}" # public input + return x + ''' + ).lstrip() + + result = convert_asserts.transform_source(source) + + self.assertEqual(result.converted_lines, (2,)) + self.assertEqual(result.skipped_lines, ()) + self.assertIn("if not (x > 0): # public input", result.source) + self.assertIn( + 'raise AssertionError(f"expected positive x, got {x}")', + result.source, + ) + + namespace = {} + exec(result.source, namespace) + self.assertEqual(namespace["positive"](2), 2) + with self.assertRaisesRegex(AssertionError, "expected positive x, got -1"): + namespace["positive"](-1) + + def test_preserves_multiline_condition_and_message(self): + source = textwrap.dedent( + ''' + def bounded(x): + assert ( + 0 <= x <= 1 + ), ( + f"x outside [0, 1]: {x}" + ) + ''' + ).lstrip() + + result = convert_asserts.transform_source(source) + + self.assertIn("if not (\n 0 <= x <= 1\n ):", result.source) + self.assertIn( + 'raise AssertionError(\n f"x outside [0, 1]: {x}"\n )', + result.source, + ) + compile(result.source, "sample.py", "exec") + + def test_supports_an_explicit_exception_already_in_scope(self): + source = "def positive(x):\n assert x > 0, 'positive required'\n" + + result = convert_asserts.transform_source(source, exception="ValueError") + + namespace = {} + exec(result.source, namespace) + with self.assertRaisesRegex(ValueError, "positive required"): + namespace["positive"](0) + + def test_preserves_a_tuple_as_one_exception_argument(self): + source = "def f():\n assert False, ('left', 'right')\n" + + result = convert_asserts.transform_source(source) + + namespace = {} + exec(result.source, namespace) + with self.assertRaises(AssertionError) as context: + namespace["f"]() + self.assertEqual(context.exception.args, (("left", "right"),)) + + def test_check_mode_reports_without_writing(self): + path = self._write( + ''' + def positive(x): + assert x > 0 + ''' + ) + original = path.read_text(encoding="utf-8") + output = StringIO() + + with redirect_stdout(output): + status = convert_asserts.main(["--check", str(path)]) + + self.assertEqual(status, 1) + self.assertIn("1 file(s) would change", output.getvalue()) + self.assertEqual(path.read_text(encoding="utf-8"), original) + + def test_skips_assert_mixed_with_other_one_line_statements(self): + source = "def f(x):\n assert x; return x\n" + + result = convert_asserts.transform_source(source) + + self.assertEqual(result.source, source) + self.assertEqual(result.converted_lines, ()) + self.assertEqual(result.skipped_lines, (2,)) + + def test_skips_assert_in_a_one_line_compound_suite(self): + source = "def f(x):\n if x: assert x > 0\n" + + result = convert_asserts.transform_source(source) + + self.assertEqual(result.source, source) + self.assertEqual(result.converted_lines, ()) + self.assertEqual(result.skipped_lines, (2,)) + + def test_rejects_an_exception_expression(self): + error = StringIO() + + with redirect_stderr(error): + status = convert_asserts.main( + ["--exception", "ValueError()", "unused.py"] + ) + + self.assertEqual(status, 2) + self.assertIn("exception must be a name already in scope", error.getvalue()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_sr_docstring_arg_types.py b/test/test_sr_docstring_arg_types.py new file mode 100644 index 000000000..ddc1a7dbe --- /dev/null +++ b/test/test_sr_docstring_arg_types.py @@ -0,0 +1,253 @@ +import shutil +import tempfile +import textwrap +import unittest +from contextlib import redirect_stdout +from io import StringIO +from pathlib import Path + +from scripts import add_docstring_arg_types + + +class TestAddDocstringArgTypes(unittest.TestCase): + + def setUp(self): + self.tmp_path = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp_path, ignore_errors=True) + + def _write(self, source): + path = self.tmp_path / "sample.py" + path.write_text(textwrap.dedent(source).lstrip(), encoding="utf-8") + return path + + def test_adds_annotation_types_to_existing_google_args(self): + path = self._write( + ''' + def calculate_velocity( + distance: float, + time: float, + acceleration: float = 0.0, + ) -> float: + """Calculate velocity. + + Args: + distance: Distance traveled. + time: Elapsed time. + acceleration: Constant acceleration. + + Returns: + float: Velocity. + """ + return distance / time + acceleration * time + ''' + ) + + result = add_docstring_arg_types.update_file(path) + source = path.read_text(encoding="utf-8") + + self.assertTrue(result.changed) + self.assertEqual( + [(update.argument, update.annotation) for update in result.updates], + [ + ("distance", "float"), + ("time", "float"), + ("acceleration", "float"), + ], + ) + self.assertIn("distance (float): Distance traveled.", source) + self.assertIn("time (float): Elapsed time.", source) + self.assertIn("acceleration (float): Constant acceleration.", source) + + def test_check_mode_reports_without_writing(self): + path = self._write( + ''' + def scale(x: int) -> int: + """Scale x. + + Args: + x: Value to scale. + + Returns: + Scaled value. + """ + return 2 * x + ''' + ) + original = path.read_text(encoding="utf-8") + + output = StringIO() + with redirect_stdout(output): + status = add_docstring_arg_types.main( + ["--check", "--include-outputs", str(path)] + ) + + self.assertEqual(status, 1) + self.assertIn("1 file(s) would change", output.getvalue()) + self.assertIn("1 output type update(s)", output.getvalue()) + self.assertEqual(path.read_text(encoding="utf-8"), original) + + def test_preserves_existing_type_unless_overwrite_is_requested(self): + path = self._write( + ''' + def scale(x: float): + """Scale x. + + Args: + x (int): Value to scale. + """ + return 2 * x + ''' + ) + + result = add_docstring_arg_types.update_file(path) + self.assertFalse(result.changed) + self.assertIn("x (int):", path.read_text(encoding="utf-8")) + + result = add_docstring_arg_types.update_file(path, overwrite_existing=True) + self.assertTrue(result.changed) + self.assertIn("x (float):", path.read_text(encoding="utf-8")) + + def test_updates_public_constructor_without_documenting_self(self): + path = self._write( + ''' + class Body: + + def __init__(self, mass: float): + """Initialize a body. + + Args: + mass: Body mass. + """ + self.mass = mass + ''' + ) + + result = add_docstring_arg_types.update_file(path) + + self.assertEqual(len(result.updates), 1) + self.assertEqual(result.updates[0].argument, "mass") + self.assertIn("mass (float): Body mass.", path.read_text(encoding="utf-8")) + + def test_normalizes_multiline_annotations_and_colon_spacing(self): + path = self._write( + ''' + def first( + values: list[ + float + ], + ): + """Return the first value. + + Args: + values : Values to inspect. + """ + return values[0] + ''' + ) + + result = add_docstring_arg_types.update_file(path) + source = path.read_text(encoding="utf-8") + + self.assertTrue(result.changed) + self.assertIn("values (list[float]): Values to inspect.", source) + + def test_does_not_infer_a_type_for_an_unannotated_argument(self): + path = self._write( + ''' + def scale(x): + """Scale a value. + + Args: + x: Value to scale. + """ + return 2 * x + ''' + ) + original = path.read_text(encoding="utf-8") + + result = add_docstring_arg_types.update_file(path) + + self.assertFalse(result.changed) + self.assertEqual(result.updates, []) + self.assertEqual(path.read_text(encoding="utf-8"), original) + + def test_adds_return_type_when_output_sync_is_requested(self): + path = self._write( + ''' + def norm(x: float) -> float: + """Compute a norm. + + Args: + x: Input value. + + Returns: + Computed norm. + """ + return abs(x) + ''' + ) + + result = add_docstring_arg_types.update_file(path, include_outputs=True) + source = path.read_text(encoding="utf-8") + + self.assertTrue(result.changed) + self.assertIn("x (float): Input value.", source) + self.assertIn("float: Computed norm.", source) + self.assertEqual( + [update.section for update in result.updates], + ["Args", "Returns"], + ) + + def test_replaces_existing_output_type_only_when_requested(self): + path = self._write( + ''' + def norm(x) -> float: + """Compute a norm. + + Returns: + int: Computed norm. + """ + return abs(x) + ''' + ) + + result = add_docstring_arg_types.update_file(path, include_outputs=True) + self.assertFalse(result.changed) + self.assertIn("int: Computed norm.", path.read_text(encoding="utf-8")) + + result = add_docstring_arg_types.update_file( + path, + include_outputs=True, + overwrite_existing=True, + ) + self.assertTrue(result.changed) + self.assertIn("float: Computed norm.", path.read_text(encoding="utf-8")) + + def test_extracts_item_type_for_yields_section(self): + path = self._write( + ''' + from typing import Iterator + + def indices(n: int) -> Iterator[int]: + """Yield indices. + + Args: + n: Number of indices. + + Yields: + Next index. + """ + yield from range(n) + ''' + ) + + result = add_docstring_arg_types.update_file(path, include_outputs=True) + source = path.read_text(encoding="utf-8") + + self.assertTrue(result.changed) + self.assertIn("n (int): Number of indices.", source) + self.assertIn("int: Next index.", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_sr_flatten_qmcpy_imports.py b/test/test_sr_flatten_qmcpy_imports.py new file mode 100644 index 000000000..5a295121a --- /dev/null +++ b/test/test_sr_flatten_qmcpy_imports.py @@ -0,0 +1,351 @@ +import json +import shutil +import tempfile +import unittest +from pathlib import Path + +from scripts.flatten_qmcpy_imports import ( + _load_qmcpy_public_names, + flatten_imports, + main, +) + + +def _nested_import(module, imported): + return f"from {'qmcpy.' + module} import {imported}" + + +class TestFlattenQmcpyImports(unittest.TestCase): + + def setUp(self): + self.tmp_path = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp_path, ignore_errors=True) + + def test_flatten_imports_basic(self): + source = ( + _nested_import("integrand", "Keister") + + "\n" + + _nested_import("discrete_distribution.lattice", "Lattice as LD") + + "\nfrom qmcpy import DigitalNetB2\nimport qmcpy.util\n" + ).encode() + + updated, count = flatten_imports( + source, frozenset({"DigitalNetB2", "Keister", "Lattice"}) + ) + + self.assertEqual(count, 3) + self.assertEqual( + updated, + ( + b"from qmcpy import DigitalNetB2, Keister, Lattice as LD\n" + b"import qmcpy.util\n" + ), + ) + + def test_flatten_preserves_private(self): + source = ( + _nested_import("_internal._helpers", "PublicHelper") + + "\n" + + _nested_import( + "true_measure.uniform_triangle", + "UniformTriangle, _UniformTriangleAdapter", + ) + + "\n" + + _nested_import( + "true_measure.copula", + "(\n AbstractCopula,\n _validate_dimension,\n)", + ) + + "\n" + + _nested_import("integrand", "Keister") + + "\n" + ).encode() + + updated, count = flatten_imports(source, frozenset({"Keister"})) + + self.assertEqual(count, 1) + self.assertEqual( + updated, + source.replace( + _nested_import("integrand", "Keister").encode(), + b"from qmcpy import Keister", + ), + ) + + def test_private_module_splits_groups(self): + source = ( + b"from qmcpy import Zeta\n" + b"from qmcpy._internal._helpers import PublicHelper\n" + b"from qmcpy import Alpha\n" + ) + + updated, count = flatten_imports(source) + + self.assertEqual(count, 0) + self.assertEqual(updated, source) + + def test_flatten_preserves_util_imports(self): + source = ( + b"from qmcpy.util import ParameterError\n" + b"from qmcpy.util.transforms import tf_exp\n" + ) + + updated, count = flatten_imports(source, frozenset({"ParameterError", "tf_exp"})) + + self.assertEqual((updated, count), (source, 0)) + + def test_flatten_keeps_nonpublic_names(self): + source = b"from qmcpy.stopping_criterion.pf_gp_ci import PFGPCIData\n" + + updated, count = flatten_imports(source, frozenset({"PFGPCI"})) + + self.assertEqual((updated, count), (source, 0)) + + def test_flatten_no_public_api_noop(self): + source = (_nested_import("integrand", "Keister") + "\n").encode() + + updated, count = flatten_imports(source) + + self.assertEqual((updated, count), (source, 0)) + + def test_flatten_preserve_str_literals(self): + source = b'text = """\nfrom qmcpy.integrand import Keister\n"""\n' + + updated, count = flatten_imports(source, frozenset({"Keister"})) + + self.assertEqual(count, 0) + self.assertEqual(updated, source) + + def test_python_string_protection_applies_to_every_rewrite_stage(self): + string_body = ( + b'text = """\n' + b"from qmcpy.integrand import Keister\n" + b"from qmcpy import Zeta,Beta\n" + b"from qmcpy import Alpha\n" + b"from qmcpy import *\n" + b"from qmcpy import *\n" + b'"""\n' + ) + source = string_body + b"from qmcpy.integrand import Keister\n" + + updated, count = flatten_imports(source, frozenset({"Keister"})) + + self.assertEqual(count, 1) + self.assertEqual(updated, string_body + b"from qmcpy import Keister\n") + + def test_python_tokenize_failure_is_fail_closed(self): + source = b'"""unterminated\nfrom qmcpy.integrand import Keister\n' + + self.assertEqual( + flatten_imports(source, frozenset({"Keister"})), (source, 0) + ) + + def test_flatten_skip_star_expansion(self): + source = ( + b"from qmcpy import *\n\n" + b"def f(Lattice):\n" + b" return Lattice\n\n" + b"y = Keister(dimension=2)\n" + b"x = Lattice(dimension=2)\n" + ) + + updated, count = flatten_imports(source, frozenset({"Keister", "Lattice"})) + + self.assertEqual(count, 0) + self.assertEqual(updated, source) + + def test_notebook_star_dedup(self): + notebook = { + "cells": [ + { + "cell_type": "code", + "source": [ + _nested_import("integrand", "*") + "\n", + _nested_import("true_measure", "*"), + ], + } + ] + } + source = json.dumps(notebook, indent=1).encode() + + updated, count = flatten_imports(source, frozenset({"Keister"})) + + self.assertEqual(count, 3) + self.assertEqual( + json.loads(updated)["cells"][0]["source"], ["from qmcpy import *"] + ) + + def test_named_imports_merge_sort(self): + source = ( + b"from qmcpy import Zeta,Beta\n" + b"from qmcpy import Alpha\n" + b"\n" + b"from qmcpy import Gamma\n" + ) + + updated, count = flatten_imports(source) + + self.assertEqual(count, 1) + self.assertEqual( + updated, + ( + b"from qmcpy import Alpha, Beta, Zeta\n" + b"\n" + b"from qmcpy import Gamma\n" + ), + ) + self.assertEqual(flatten_imports(updated), (updated, 0)) + + def test_merge_paren_and_single_line(self): + source = b"""from qmcpy import ( + KernelDigShiftInvar, + KernelDigShiftInvarAdaptiveAlpha, + KernelDigShiftInvarCombined, + KernelShiftInvar, + KernelShiftInvarCombined, +) +from qmcpy import tf_exp_eps, tf_exp_eps_inv +""" + + updated, count = flatten_imports(source) + + self.assertEqual(count, 1) + self.assertEqual( + updated, + b"""from qmcpy import ( + KernelDigShiftInvar, + KernelDigShiftInvarAdaptiveAlpha, + KernelDigShiftInvarCombined, + KernelShiftInvar, + KernelShiftInvarCombined, + tf_exp_eps, + tf_exp_eps_inv, +) +""", + ) + self.assertEqual(flatten_imports(updated), (updated, 0)) + + def test_merge_same_scope_only(self): + source = ( + b"if enabled:\n" + b" from qmcpy import Zeta\n" + b" from qmcpy import Alpha as First\n" + b"else:\n" + b" from qmcpy import Beta\n" + b"from qmcpy import _Private\n" + b"from qmcpy import Gamma # keep this comment\n" + ) + + updated, count = flatten_imports(source) + + self.assertEqual(count, 1) + self.assertEqual( + updated, + ( + b"if enabled:\n" + b" from qmcpy import Alpha as First, Zeta\n" + b"else:\n" + b" from qmcpy import Beta\n" + b"from qmcpy import _Private\n" + b"from qmcpy import Gamma # keep this comment\n" + ), + ) + + def test_notebook_named_merge(self): + notebook = { + "cells": [ + { + "cell_type": "code", + "source": [ + "from qmcpy import Zeta\n", + "from qmcpy import Alpha,Beta\n", + "print(Alpha)\n", + ], + } + ] + } + source = json.dumps(notebook, indent=1).encode() + + updated, count = flatten_imports(source) + + self.assertEqual(count, 1) + self.assertEqual( + json.loads(updated)["cells"][0]["source"], + [ + "from qmcpy import Alpha, Beta, Zeta\n", + "print(Alpha)\n", + ], + ) + self.assertEqual(flatten_imports(updated), (updated, 0)) + + def test_notebook_flattens_nested_imports_only_in_code_cells(self): + nested_import = _nested_import("integrand", "Keister") + "\n" + metadata_import = _nested_import("true_measure", "Gaussian") + "\n" + string_literal = f'text = "{nested_import.rstrip()}"\n' + multiline_string = ['text = """\n', nested_import, '"""\n'] + notebook = { + "metadata": {"source": [metadata_import]}, + "cells": [ + {"cell_type": "markdown", "source": [nested_import]}, + {"cell_type": "code", "source": [nested_import]}, + {"cell_type": "code", "source": [string_literal]}, + {"cell_type": "code", "source": multiline_string}, + ] + } + source = json.dumps(notebook, indent=1).encode() + + updated, count = flatten_imports(source, frozenset({"Keister"})) + + cells = json.loads(updated)["cells"] + self.assertEqual(count, 1) + self.assertEqual( + json.loads(updated)["metadata"]["source"], [metadata_import] + ) + self.assertEqual(cells[0]["source"], [nested_import]) + self.assertEqual(cells[1]["source"], ["from qmcpy import Keister\n"]) + self.assertEqual(cells[2]["source"], [string_literal]) + self.assertEqual(cells[3]["source"], multiline_string) + self.assertEqual( + flatten_imports(updated, frozenset({"Keister"})), (updated, 0) + ) + + def test_markdown_import_examples_are_flattened(self): + path = self.tmp_path / "example.md" + path.write_bytes( + b'Example with unmatched prose delimiter: """\n\n' + b"```python\n" + b"from qmcpy.integrand import Keister\n" + b"```\n" + ) + + self.assertEqual(main([str(path)]), 0) + self.assertIn(b"from qmcpy import Keister", path.read_bytes()) + + def test_check_mode_no_write(self): + path = self.tmp_path / "example.py" + original = (_nested_import("true_measure", "Gaussian") + "\n").encode() + path.write_bytes(original) + + self.assertEqual(main(["--check", str(path)]), 1) + self.assertEqual(path.read_bytes(), original) + + self.assertEqual(main([str(path)]), 0) + self.assertEqual(path.read_bytes(), b"from qmcpy import Gaussian\n") + self.assertEqual(main(["--check", str(path)]), 0) + + def test_public_names_optional_free_stable(self): + repository_root = Path(__file__).resolve().parent.parent + names = _load_qmcpy_public_names(repository_root) + + self.assertIsNotNone(names) + self.assertIn("Gaussian", names) + self.assertIn("Keister", names) + # Optional dependencies are blocked in the probe context, so fallback + # exports are part of the deterministic name set. + self.assertIn("PFGPCI", names) + # Helpers that are deliberately not part of the top-level API. + self.assertNotIn("PFGPCIData", names) + self.assertNotIn("TriangularDistribution", names) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_sr_unwrap_markdown.py b/test/test_sr_unwrap_markdown.py new file mode 100644 index 000000000..4115d9ad2 --- /dev/null +++ b/test/test_sr_unwrap_markdown.py @@ -0,0 +1,99 @@ +import unittest + +from scripts.unwrap_markdown import unwrap_markdown_text + + +class TestUnwrapMarkdown(unittest.TestCase): + + def test_unwraps_list_item_continuations(self): + cases = [ + ( + "- unordered first\n unordered second\n", + "- unordered first unordered second\n", + ), + ( + "- [ ] task first\n task second\n", + "- [ ] task first task second\n", + ), + ( + "10. ordered first\n ordered second\n", + "10. ordered first ordered second\n", + ), + ] + for source, expected in cases: + with self.subTest(source=source): + updated = unwrap_markdown_text(source) + + self.assertEqual(updated, expected) + self.assertEqual(unwrap_markdown_text(updated), updated) + + def test_unwraps_adjacent_and_nested_list_items_separately(self): + source = ( + "- parent first\n" + " parent second\n" + " - child first\n" + " child second\n" + "- sibling first\n" + " sibling second\n" + ) + + self.assertEqual( + unwrap_markdown_text(source), + ( + "- parent first parent second\n" + " - child first child second\n" + "- sibling first sibling second\n" + ), + ) + + def test_preserves_list_item_blocks_and_explicit_hard_breaks(self): + source = ( + "- first paragraph\n" + " continuation\n" + "\n" + " second paragraph\n" + " continuation\n" + "\n" + "- item before code\n" + " indented code\n" + "\n" + "- explicit hard break \n" + " remains separate\n" + ) + + self.assertEqual( + unwrap_markdown_text(source), + ( + "- first paragraph continuation\n" + "\n" + " second paragraph continuation\n" + "\n" + "- item before code\n" + " indented code\n" + "\n" + "- explicit hard break \n" + " remains separate\n" + ), + ) + + def test_unwraps_ordinary_paragraphs(self): + self.assertEqual( + unwrap_markdown_text("first line\nsecond line\n"), + "first line second line\n", + ) + + def test_preserves_horizontal_rules(self): + for rule in ["- - -", "* * *", "_ _ _"]: + with self.subTest(rule=rule): + source = f"{rule}\nfollowing paragraph\n" + + self.assertEqual(unwrap_markdown_text(source), source) + + def test_preserves_indented_code_that_looks_like_a_list(self): + source = " - code first\n code second\n" + + self.assertEqual(unwrap_markdown_text(source), source) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_tm_copulas.py b/test/test_tm_copulas.py new file mode 100644 index 000000000..ed29969fb --- /dev/null +++ b/test/test_tm_copulas.py @@ -0,0 +1,1271 @@ +import unittest +import warnings + +import numpy as np +import scipy.stats as stats + +from qmcpy import ( + AbstractCopula, + ClaytonCopula, + DigitalNetB2, + FrankCopula, + GaussianCopula, + GumbelCopula, + StudentTCopula, +) + +from qmcpy.true_measure.copula import ( + AbstractCopula as ModuleAbstractCopula, + _apply_marginal_ppfs, + _build_marginal_range, + _clip_unit_interval, + _marginal_cdfs_and_logpdf, + _validate_correlation_matrix, + _validate_dimension, + _validate_marginals, +) + +from qmcpy.util import DimensionError, MethodImplementationError, ParameterError + + +class PPFOnlyMarginal: + def ppf(self, u): + return np.asarray(u, dtype=float) + + +class NonCallablePPFMarginal: + ppf = 1.0 + + +class UnitPDFMarginal: + def ppf(self, u): + return np.asarray(u, dtype=float) + + def cdf(self, x): + return np.asarray(x, dtype=float) + + def pdf(self, x): + return np.ones_like(np.asarray(x, dtype=float)) + + +class CDFOnlyMarginal(PPFOnlyMarginal): + def cdf(self, x): + return np.asarray(x, dtype=float) + + +class BadIntervalMarginal(PPFOnlyMarginal): + def interval(self, confidence): + raise ValueError("interval unavailable") + + +class BadRangeMarginal: + def ppf(self, u): + raise ValueError("ppf unavailable") + + +def _equicorrelation(d, rho): + corr = np.full((d, d), rho, dtype=float) + np.fill_diagonal(corr, 1.0) + return corr + + +def _make_copula(copula_cls, dimension=2, marginals=None, correlation=None, seed=7): + if marginals is None: + marginals = [stats.norm()] * dimension + if correlation is None: + correlation = np.eye(dimension) + + kwargs = {} + if copula_cls is StudentTCopula: + kwargs["df"] = 4 + if copula_cls is ClaytonCopula: + kwargs["theta"] = 2.0 + if copula_cls is FrankCopula: + kwargs["theta"] = 5.0 + if copula_cls is GumbelCopula: + kwargs["theta"] = 2.0 + + common = { + "sampler": DigitalNetB2(dimension, seed=seed), + "marginals": marginals, + **kwargs, + } + if copula_cls in [ClaytonCopula, FrankCopula, GumbelCopula]: + return copula_cls(**common) + return copula_cls(correlation=correlation, **common) + + +class TestAbstractCopulaAndHelpers(unittest.TestCase): + + def test_abstract_copula_is_importable_from_public_module_path(self): + self.assertIs(ModuleAbstractCopula, AbstractCopula) + + def test_public_api_imports_and_normal_usage(self): + for copula_cls in [ + GaussianCopula, + StudentTCopula, + ClaytonCopula, + FrankCopula, + GumbelCopula, + ]: + with self.subTest(copula_cls=copula_cls.__name__): + self.assertTrue(issubclass(copula_cls, AbstractCopula)) + + tm = _make_copula(copula_cls) + x = tm(8) + x_gen = tm.gen_samples(8) + v = tm.gen_copula_samples(8) + + self.assertEqual(x.shape, (8, 2)) + self.assertEqual(x_gen.shape, (8, 2)) + self.assertEqual(v.shape, (8, 2)) + self.assertTrue(np.all(np.isfinite(x))) + self.assertTrue(np.all(np.isfinite(x_gen))) + self.assertTrue(np.all((0 <= v) & (v <= 1))) + + def test_abstract_copula_rejects_unimplemented_transform(self): + tm = AbstractCopula( + DigitalNetB2(2, seed=101), + marginals=[stats.uniform(), stats.uniform()], + ) + + with self.assertRaises(MethodImplementationError): + tm.copula_transform(np.full((3, 2), 0.5)) + + def test_abstract_copula_rejects_invalid_sampler(self): + with self.assertRaisesRegex(ParameterError, "sampler"): + AbstractCopula(object(), marginals=[stats.uniform()]) + + def test_validate_marginals_error_branches(self): + with self.assertRaisesRegex(ParameterError, "marginals"): + _validate_marginals(None) + + with self.assertRaisesRegex(ParameterError, "at least one"): + _validate_marginals([]) + + with self.assertRaisesRegex(ParameterError, "ppf"): + _validate_marginals([NonCallablePPFMarginal()]) + + def test_validate_dimension_error_branches(self): + with self.assertRaisesRegex(DimensionError, "integer dimension"): + _validate_dimension(object(), [stats.uniform()]) + + with self.assertRaisesRegex(DimensionError, "marginals"): + _validate_dimension(3, [stats.uniform(), stats.uniform()]) + + def test_apply_marginal_ppfs_clips_endpoints_and_checks_dimension(self): + transformed = _apply_marginal_ppfs( + np.array([[0.0, 1.0], [1.0, 0.0]]), + [stats.norm(), stats.norm()], + ) + + self.assertEqual(transformed.shape, (2, 2)) + self.assertTrue(np.all(np.isfinite(transformed))) + + with self.assertRaisesRegex(DimensionError, "marginals"): + _apply_marginal_ppfs(np.full((2, 3), 0.5), [stats.uniform(), stats.uniform()]) + + def test_marginal_range_falls_back_when_interval_or_ppf_fails(self): + ranges = _build_marginal_range([BadIntervalMarginal(), BadRangeMarginal()]) + + self.assertEqual(ranges.shape, (2, 2)) + self.assertTrue(np.all(np.isfinite(ranges[0]))) + np.testing.assert_allclose(ranges[1], [-np.inf, np.inf]) + + def test_marginal_cdfs_and_logpdf_pdf_branch_and_errors(self): + x = np.array([[0.25, 0.75], [0.4, 0.6]]) + u, log_density = _marginal_cdfs_and_logpdf( + x, + [UnitPDFMarginal(), UnitPDFMarginal()], + ) + + np.testing.assert_allclose(u, x) + np.testing.assert_allclose(log_density, np.zeros(2)) + + with self.assertRaisesRegex(ParameterError, "cdf"): + _marginal_cdfs_and_logpdf(x, [PPFOnlyMarginal(), UnitPDFMarginal()]) + + with self.assertRaisesRegex(ParameterError, "pdf"): + _marginal_cdfs_and_logpdf(x, [CDFOnlyMarginal(), UnitPDFMarginal()]) + + def test_validate_correlation_matrix_rejects_nonfinite_values(self): + with self.assertRaisesRegex(ValueError, "finite"): + _validate_correlation_matrix([[1.0, np.nan], [np.nan, 1.0]], 2) + + def test_clip_unit_interval_uses_machine_epsilon(self): + clipped = _clip_unit_interval(np.array([0.0, 0.5, 1.0])) + eps = np.finfo(float).eps + + np.testing.assert_allclose(clipped, [eps, 0.5, 1.0 - eps]) + + def test_copula_transform_outputs_dependent_uniforms_in_unit_cube(self): + for copula_cls in [GaussianCopula, StudentTCopula, ClaytonCopula, GumbelCopula, FrankCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + tm = _make_copula(copula_cls, dimension=3) + u = np.array( + [ + [0.1, 0.3, 0.7], + [0.5, 0.5, 0.5], + [0.9, 0.8, 0.2], + ] + ) + + v = tm.copula_transform(u) + + self.assertEqual(v.shape, u.shape) + self.assertTrue(np.all(np.isfinite(v))) + self.assertTrue(np.all((0.0 <= v) & (v <= 1.0))) + + def test_copula_sample_shapes_are_preserved(self): + for copula_cls, dimension in [ + (GaussianCopula, 3), + (StudentTCopula, 3), + (ClaytonCopula, 3), + (FrankCopula, 3), + (GumbelCopula, 3), + ]: + with self.subTest(copula_cls=copula_cls.__name__, dimension=dimension): + tm = _make_copula(copula_cls, dimension=dimension, seed=9) + + one = tm(1) + many = tm(8) + batched_transform = tm._transform(np.full((2, 3, dimension), 0.5)) + + self.assertEqual(one.shape, (1, dimension)) + self.assertEqual(many.shape, (8, dimension)) + self.assertEqual(batched_transform.shape, (2, 3, dimension)) + self.assertTrue(np.all(np.isfinite(one))) + self.assertTrue(np.all(np.isfinite(many))) + self.assertTrue(np.all(np.isfinite(batched_transform))) + + +class TestEllipticalCopulas(unittest.TestCase): + + def test_output_shape_with_nonnormal_marginals(self): + tm = GaussianCopula( + sampler=DigitalNetB2(2, seed=7), + marginals=[stats.beta(a=2, b=5), stats.gamma(a=3, scale=2)], + correlation=[[1.0, 0.4], [0.4, 1.0]], + ) + + x = tm(16) + + self.assertEqual(x.shape, (16, 2)) + + def test_finite_output_for_normal_marginals(self): + tm = GaussianCopula( + sampler=DigitalNetB2(2, seed=11), + marginals=[stats.norm(), stats.norm(loc=1.0, scale=2.0)], + correlation=[[1.0, -0.3], [-0.3, 1.0]], + ) + + x = tm(128) + + self.assertTrue(np.all(np.isfinite(x))) + + def test_return_weights_shape_when_marginal_densities_available(self): + tm = GaussianCopula( + sampler=DigitalNetB2(2, seed=12), + marginals=[stats.norm(), stats.gamma(a=2.0)], + correlation=[[1.0, 0.25], [0.25, 1.0]], + ) + + x, weights = tm(32, return_weights=True) + + self.assertEqual(x.shape, (32, 2)) + self.assertEqual(weights.shape, (32,)) + self.assertTrue(np.all(np.isfinite(weights))) + self.assertTrue(np.all(weights > 0.0)) + + def test_identity_correlation_matches_independent_marginal_transforms(self): + marginals = [stats.norm(loc=-1.0, scale=2.0), stats.gamma(a=2.0, scale=3.0)] + tm = GaussianCopula( + sampler=DigitalNetB2(2, seed=13), + marginals=marginals, + correlation=np.eye(2), + ) + u = np.array([[0.2, 0.7], [0.4, 0.8], [0.9, 0.1]]) + + x = tm._transform(u) + expected = np.column_stack( + [marginals[j].ppf(u[:, j]) for j in range(len(marginals))] + ) + + np.testing.assert_allclose(x, expected, rtol=1e-12, atol=1e-12) + + def test_positive_correlation_produces_positive_dependence(self): + rho = 0.75 + tm = GaussianCopula( + sampler=DigitalNetB2(2, seed=17), + marginals=[stats.norm(), stats.norm()], + correlation=[[1.0, rho], [rho, 1.0]], + ) + + x = tm(4096) + empirical_corr = np.corrcoef(x.T)[0, 1] + + self.assertGreater(empirical_corr, 0.5) + self.assertLess(abs(empirical_corr - rho), 0.2) + + def test_elliptical_copulas_support_general_dimensions(self): + for copula_cls in [GaussianCopula, StudentTCopula]: + for dimension in [1, 3, 5]: + with self.subTest(copula_cls=copula_cls.__name__, dimension=dimension): + correlation = _equicorrelation(dimension, 0.25) + tm = _make_copula( + copula_cls, + dimension=dimension, + marginals=[stats.norm()] * dimension, + correlation=correlation, + seed=19, + ) + + x = tm(16) + one = tm(1) + + self.assertEqual(x.shape, (16, dimension)) + self.assertEqual(one.shape, (1, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + self.assertTrue(np.all(np.isfinite(one))) + + def test_elliptical_copulas_handle_valid_near_singular_correlation(self): + for copula_cls in [GaussianCopula, StudentTCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + dimension = 5 + tm = _make_copula( + copula_cls, + dimension=dimension, + marginals=[stats.norm()] * dimension, + correlation=_equicorrelation(dimension, 0.999), + seed=20, + ) + + x = tm(32) + + self.assertEqual(x.shape, (32, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_elliptical_copulas_reject_singular_correlation(self): + for copula_cls in [GaussianCopula, StudentTCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + with self.assertRaisesRegex(ValueError, "positive definite"): + _make_copula( + copula_cls, + dimension=3, + marginals=[stats.norm(), stats.norm(), stats.norm()], + correlation=np.ones((3, 3)), + seed=22, + ) + + def test_distribution_dimension_matches_number_of_marginals(self): + for copula_cls in [GaussianCopula, StudentTCopula, ClaytonCopula, FrankCopula, GumbelCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + tm = _make_copula( + copula_cls, + dimension=5, + marginals=[ + stats.norm(), + stats.beta(a=2, b=5), + stats.gamma(a=3), + stats.expon(), + stats.lognorm(s=0.5), + ], + correlation=np.eye(5), + ) + + x = tm(32) + + self.assertEqual(x.shape, (32, 5)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_invalid_dimension_mismatches_raise(self): + for copula_cls in [GaussianCopula, StudentTCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + with self.assertRaisesRegex(DimensionError, "marginals"): + _make_copula( + copula_cls, + dimension=2, + marginals=[stats.norm(), stats.norm(), stats.norm()], + correlation=np.eye(2), + ) + + with self.assertRaisesRegex(ValueError, "shape"): + _make_copula( + copula_cls, + dimension=2, + marginals=[stats.norm(), stats.norm()], + correlation=np.eye(3), + ) + + with self.assertRaisesRegex(ValueError, "square"): + _make_copula( + copula_cls, + dimension=2, + marginals=[stats.norm(), stats.norm()], + correlation=[[1.0, 0.2, 0.3], [0.2, 1.0, 0.4]], + ) + + def test_archimedean_dimension_mismatch_raises_dimension_error(self): + for copula_cls in [ClaytonCopula, FrankCopula, GumbelCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + with self.assertRaisesRegex(DimensionError, "marginals"): + _make_copula( + copula_cls, + dimension=2, + marginals=[stats.norm(), stats.norm(), stats.norm()], + ) + + def test_invalid_correlation_matrices_raise_value_error(self): + correlations = [ + [[1.0, 0.2], [0.3, 1.0]], + [[1.0, 0.2], [0.2, 0.9]], + [[1.0, 1.2], [1.2, 1.0]], + ] + for copula_cls in [GaussianCopula, StudentTCopula]: + for correlation in correlations: + with self.subTest(copula_cls=copula_cls.__name__, correlation=correlation): + with self.assertRaises(ValueError): + _make_copula( + copula_cls, + dimension=2, + marginals=[stats.norm(), stats.norm()], + correlation=correlation, + ) + + def test_marginal_length_mismatch_raises_dimension_error(self): + with self.assertRaisesRegex(DimensionError, "marginals"): + GaussianCopula( + sampler=DigitalNetB2(2, seed=21), + marginals=[stats.norm()], + correlation=np.eye(2), + ) + + def test_marginal_without_ppf_raises_clear_error(self): + class NoPPF: + pass + + with self.assertRaisesRegex(ParameterError, "ppf"): + GaussianCopula( + sampler=DigitalNetB2(1, seed=23), + marginals=[NoPPF()], + correlation=[[1.0]], + ) + + def test_common_scipy_frozen_marginals_work(self): + for copula_cls in [GaussianCopula, StudentTCopula, ClaytonCopula, FrankCopula, GumbelCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + tm = _make_copula( + copula_cls, + dimension=5, + marginals=[ + stats.norm(), + stats.beta(a=2, b=5), + stats.gamma(a=3), + stats.expon(), + stats.lognorm(s=0.5), + ], + correlation=np.eye(5), + seed=47, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, 5)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_endpoint_uniforms_are_clipped_to_finite_outputs(self): + for copula_cls in [GaussianCopula, StudentTCopula, ClaytonCopula, FrankCopula, GumbelCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + tm = _make_copula( + copula_cls, + dimension=5, + marginals=[ + stats.norm(), + stats.beta(a=2, b=5), + stats.gamma(a=3), + stats.expon(), + stats.lognorm(s=0.5), + ], + correlation=np.eye(5), + seed=53, + ) + u = np.array( + [ + [0.0, 1.0, 0.0, 1.0, 0.5], + [1.0, 0.0, 1.0, 0.0, 0.5], + ] + ) + + x = tm._transform(u) + + self.assertEqual(x.shape, (2, 5)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_student_t_copula_output_shape_and_finite_values(self): + tm = StudentTCopula( + sampler=DigitalNetB2(2, seed=29), + marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], + correlation=[[1.0, 0.5], [0.5, 1.0]], + df=4, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_student_t_copula_positive_correlation_produces_positive_dependence(self): + tm = StudentTCopula( + sampler=DigitalNetB2(2, seed=31), + marginals=[stats.norm(), stats.norm()], + correlation=[[1.0, 0.7], [0.7, 1.0]], + df=5, + ) + + x = tm(4096) + empirical_corr = np.corrcoef(x.T)[0, 1] + + self.assertGreater(empirical_corr, 0.45) + + def test_student_t_copula_has_stronger_joint_tail_than_gaussian_copula(self): + rho = 0.7 + df = 4 + n = 2**12 + marginals = [stats.norm(), stats.norm()] + correlation = [[1.0, rho], [rho, 1.0]] + + gaussian = GaussianCopula( + sampler=DigitalNetB2(2, seed=101), + marginals=marginals, + correlation=correlation, + ) + student_t = StudentTCopula( + sampler=DigitalNetB2(2, seed=101), + marginals=marginals, + correlation=correlation, + df=df, + ) + + x_gaussian = gaussian(n) + x_student_t = student_t(n) + threshold = stats.norm.ppf(0.99) + + def joint_tail_rate(x): + tail_0 = x[:, 0] > threshold + return np.mean(x[tail_0, 1] > threshold) + + gaussian_tail = joint_tail_rate(x_gaussian) + student_t_tail = joint_tail_rate(x_student_t) + + self.assertGreater(student_t_tail, gaussian_tail + 0.08) + + def test_student_t_copula_return_weights_shape_when_density_available(self): + tm = StudentTCopula( + sampler=DigitalNetB2(2, seed=37), + marginals=[stats.norm(), stats.gamma(a=2.0)], + correlation=[[1.0, 0.3], [0.3, 1.0]], + df=6, + ) + + x, weights = tm(32, return_weights=True) + + self.assertEqual(x.shape, (32, 2)) + self.assertEqual(weights.shape, (32,)) + self.assertTrue(np.all(np.isfinite(weights))) + self.assertTrue(np.all(weights > 0.0)) + + def test_student_t_copula_boundary_df_values_are_finite(self): + for df in [1.0, 100.0]: + with self.subTest(df=df): + dimension = 3 + tm = StudentTCopula( + sampler=DigitalNetB2(dimension, seed=39), + marginals=[stats.norm()] * dimension, + correlation=_equicorrelation(dimension, 0.4), + df=df, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_student_t_copula_large_df_is_close_to_gaussian_copula(self): + rho = 0.6 + correlation = [[1.0, rho], [rho, 1.0]] + marginals = [stats.norm(), stats.norm()] + gaussian = GaussianCopula( + sampler=DigitalNetB2(2, seed=40), + marginals=marginals, + correlation=correlation, + ) + student_t = StudentTCopula( + sampler=DigitalNetB2(2, seed=40), + marginals=marginals, + correlation=correlation, + df=100, + ) + + x_gaussian = gaussian(4096) + x_student_t = student_t(4096) + corr_gaussian = np.corrcoef(x_gaussian.T)[0, 1] + corr_student_t = np.corrcoef(x_student_t.T)[0, 1] + + self.assertLess(abs(corr_student_t - corr_gaussian), 0.02) + + def test_student_t_copula_invalid_df_raises_parameter_error(self): + for df in [0, -1, np.inf, "not-a-number"]: + with self.subTest(df=df): + with self.assertRaisesRegex(ParameterError, "df"): + StudentTCopula( + sampler=DigitalNetB2(2, seed=41), + marginals=[stats.norm(), stats.norm()], + correlation=np.eye(2), + df=df, + ) + + def test_student_t_copula_marginal_without_ppf_raises_clear_error(self): + class NoPPF: + pass + + with self.assertRaisesRegex(ParameterError, "ppf"): + StudentTCopula( + sampler=DigitalNetB2(1, seed=43), + marginals=[NoPPF()], + correlation=[[1.0]], + df=4, + ) + + +class TestArchimedeanCopulas(unittest.TestCase): + + def test_clayton_copula_output_shape_and_finite_values(self): + tm = ClaytonCopula( + sampler=DigitalNetB2(2, seed=57), + marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], + theta=2.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_clayton_copula_return_weights_shape_when_density_available(self): + tm = ClaytonCopula( + sampler=DigitalNetB2(3, seed=59), + marginals=[stats.norm(), stats.gamma(a=2.0), stats.expon()], + theta=1.5, + ) + + x, weights = tm(32, return_weights=True) + + self.assertEqual(x.shape, (32, 3)) + self.assertEqual(weights.shape, (32,)) + self.assertTrue(np.all(np.isfinite(weights))) + self.assertTrue(np.all(weights > 0.0)) + + def test_clayton_copula_invalid_theta_raises_parameter_error(self): + for theta in [0, -1, np.inf, "not-a-number"]: + with self.subTest(theta=theta): + with self.assertRaisesRegex(ParameterError, "theta"): + ClaytonCopula( + sampler=DigitalNetB2(2, seed=61), + marginals=[stats.norm(), stats.norm()], + theta=theta, + ) + + def test_clayton_copula_supports_general_dimension(self): + for dimension in [2, 3, 5]: + with self.subTest(dimension=dimension): + tm = ClaytonCopula( + sampler=DigitalNetB2(dimension, seed=63), + marginals=[stats.norm()] * dimension, + theta=2.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_clayton_copula_marginal_without_ppf_raises_clear_error(self): + class NoPPF: + pass + + with self.assertRaisesRegex(ParameterError, "ppf"): + ClaytonCopula( + sampler=DigitalNetB2(2, seed=67), + marginals=[stats.norm(), NoPPF()], + theta=2.0, + ) + + def test_clayton_copula_common_scipy_frozen_marginals_work(self): + for marginals in [ + [stats.norm(), stats.beta(a=2, b=5)], + [stats.gamma(a=3), stats.expon()], + [stats.lognorm(s=0.5), stats.norm()], + ]: + with self.subTest(marginals=[type(m.dist).__name__ for m in marginals]): + tm = ClaytonCopula( + sampler=DigitalNetB2(2, seed=69), + marginals=marginals, + theta=2.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_clayton_copula_endpoint_uniforms_are_clipped_to_finite_outputs(self): + tm = ClaytonCopula( + sampler=DigitalNetB2(2, seed=70), + marginals=[stats.norm(), stats.lognorm(s=0.5)], + theta=2.0, + ) + u = np.array([[0.0, 1.0], [1.0, 0.0]]) + + x = tm._transform(u) + + self.assertEqual(x.shape, (2, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_clayton_copula_tiny_theta_is_near_independent(self): + for dimension in [2, 3, 5]: + with self.subTest(dimension=dimension): + marginals = [stats.uniform()] * dimension + tm = ClaytonCopula( + sampler=DigitalNetB2(dimension, seed=70), + marginals=marginals, + theta=1e-8, + ) + u = np.array( + [ + [0.2, 0.7, 0.4, 0.6, 0.8], + [0.4, 0.8, 0.9, 0.3, 0.2], + [0.9, 0.1, 0.3, 0.7, 0.5], + ] + )[:, :dimension] + + x = tm._transform(u) + + self.assertEqual(x.shape, (3, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + np.testing.assert_allclose(x, u, atol=5e-6) + + def test_clayton_copula_large_theta_is_finite(self): + for dimension in [2, 3, 5]: + for theta in [20.0, 50.0]: + with self.subTest(dimension=dimension, theta=theta): + tm = ClaytonCopula( + sampler=DigitalNetB2(dimension, seed=70), + marginals=[stats.norm()] * dimension, + theta=theta, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_clayton_copula_positive_dependence_behavior(self): + tm = ClaytonCopula( + sampler=DigitalNetB2(2, seed=71), + marginals=[stats.uniform(), stats.uniform()], + theta=2.0, + ) + + x = tm(4096) + empirical_corr = np.corrcoef(x.T)[0, 1] + + self.assertGreater(empirical_corr, 0.45) + + def test_clayton_copula_has_stronger_lower_tail_than_gaussian_copula(self): + theta = 2.0 + n = 2**12 + marginals = [stats.uniform(), stats.uniform()] + # Clayton Kendall tau is theta/(theta+2); convert to Gaussian rho. + rho = np.sin(np.pi * (theta / (theta + 2.0)) / 2.0) + + clayton = ClaytonCopula( + sampler=DigitalNetB2(2, seed=73), + marginals=marginals, + theta=theta, + ) + gaussian = GaussianCopula( + sampler=DigitalNetB2(2, seed=73), + marginals=marginals, + correlation=[[1.0, rho], [rho, 1.0]], + ) + + x_clayton = clayton(n) + x_gaussian = gaussian(n) + threshold = 0.05 + + def lower_tail_rate(x): + tail_0 = x[:, 0] < threshold + return np.mean(x[tail_0, 1] < threshold) + + clayton_tail = lower_tail_rate(x_clayton) + gaussian_tail = lower_tail_rate(x_gaussian) + + self.assertGreater(clayton_tail, gaussian_tail + 0.2) + + def test_frank_copula_output_shape_for_two_dimensions(self): + tm = FrankCopula( + sampler=DigitalNetB2(2, seed=75), + marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], + theta=5.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_frank_copula_positive_theta_supports_higher_dimensions(self): + for dimension in [3, 5]: + with self.subTest(dimension=dimension): + tm = FrankCopula( + sampler=DigitalNetB2(dimension, seed=76), + marginals=[stats.norm()] * dimension, + theta=5.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_frank_copula_return_weights_shape_when_density_available(self): + tm = FrankCopula( + sampler=DigitalNetB2(3, seed=77), + marginals=[stats.norm(), stats.gamma(a=2.0), stats.expon()], + theta=4.0, + ) + + x, weights = tm(32, return_weights=True) + + self.assertEqual(x.shape, (32, 3)) + self.assertEqual(weights.shape, (32,)) + self.assertTrue(np.all(np.isfinite(weights))) + self.assertTrue(np.all(weights > 0.0)) + + def test_frank_copula_invalid_theta_raises_parameter_error(self): + for theta in [0, np.inf, -np.inf, "not-a-number"]: + with self.subTest(theta=theta): + with self.assertRaisesRegex(ParameterError, "theta"): + FrankCopula( + sampler=DigitalNetB2(2, seed=78), + marginals=[stats.norm(), stats.norm()], + theta=theta, + ) + + def test_frank_copula_negative_theta_rejected_above_two_dimensions(self): + with self.assertRaisesRegex(ParameterError, "d=2"): + FrankCopula( + sampler=DigitalNetB2(3, seed=79), + marginals=[stats.norm(), stats.norm(), stats.norm()], + theta=-2.0, + ) + + def test_frank_copula_dimension_mismatch_raises_dimension_error(self): + with self.assertRaisesRegex(DimensionError, "marginals"): + FrankCopula( + sampler=DigitalNetB2(2, seed=80), + marginals=[stats.norm(), stats.norm(), stats.norm()], + theta=5.0, + ) + + def test_frank_copula_marginal_without_ppf_raises_clear_error(self): + class NoPPF: + pass + + with self.assertRaisesRegex(ParameterError, "ppf"): + FrankCopula( + sampler=DigitalNetB2(2, seed=82), + marginals=[stats.norm(), NoPPF()], + theta=5.0, + ) + + def test_frank_copula_positive_dependence_behavior(self): + tm = FrankCopula( + sampler=DigitalNetB2(2, seed=84), + marginals=[stats.uniform(), stats.uniform()], + theta=6.0, + ) + + x = tm(4096) + empirical_corr = np.corrcoef(x.T)[0, 1] + + self.assertGreater(empirical_corr, 0.45) + + def test_frank_copula_tiny_theta_is_close_to_independence(self): + for theta, dimension in [(1e-8, 3), (-1e-8, 2)]: + with self.subTest(theta=theta, dimension=dimension): + marginals = [stats.uniform()] * dimension + tm = FrankCopula( + sampler=DigitalNetB2(dimension, seed=86), + marginals=marginals, + theta=theta, + ) + u = np.array( + [ + [0.2, 0.7, 0.4, 0.6, 0.8], + [0.4, 0.8, 0.9, 0.3, 0.2], + [0.9, 0.1, 0.3, 0.7, 0.5], + ] + )[:, :dimension] + + x = tm._transform(u) + + self.assertEqual(x.shape, (3, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + np.testing.assert_allclose(x, u, atol=5e-6) + + def test_frank_copula_large_theta_is_finite(self): + for theta, dimension in [(50.0, 5), (-50.0, 2)]: + with self.subTest(theta=theta, dimension=dimension): + tm = FrankCopula( + sampler=DigitalNetB2(dimension, seed=87), + marginals=[stats.norm()] * dimension, + theta=theta, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_frank_copula_negative_theta_produces_negative_dependence_in_2d(self): + tm = FrankCopula( + sampler=DigitalNetB2(2, seed=88), + marginals=[stats.uniform(), stats.uniform()], + theta=-6.0, + ) + + x = tm(4096) + empirical_corr = np.corrcoef(x.T)[0, 1] + + self.assertLess(empirical_corr, -0.35) + + def test_gumbel_copula_output_shape_and_finite_values(self): + tm = GumbelCopula( + sampler=DigitalNetB2(2, seed=79), + marginals=[stats.norm(), stats.gamma(a=3.0, scale=2.0)], + theta=2.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_gumbel_copula_return_weights_shape_when_density_available(self): + tm = GumbelCopula( + sampler=DigitalNetB2(3, seed=81), + marginals=[stats.norm(), stats.gamma(a=2.0), stats.expon()], + theta=1.5, + ) + + x, weights = tm(32, return_weights=True) + + self.assertEqual(x.shape, (32, 3)) + self.assertEqual(weights.shape, (32,)) + self.assertTrue(np.all(np.isfinite(weights))) + self.assertTrue(np.all(weights > 0.0)) + + def test_gumbel_copula_invalid_theta_raises_parameter_error(self): + for theta in [0, 0.5, -1, np.inf, "not-a-number"]: + with self.subTest(theta=theta): + with self.assertRaisesRegex(ParameterError, "theta"): + GumbelCopula( + sampler=DigitalNetB2(2, seed=83), + marginals=[stats.norm(), stats.norm()], + theta=theta, + ) + + def test_gumbel_copula_theta_one_is_independent_marginal_transform(self): + marginals = [stats.norm(loc=-1.0, scale=2.0), stats.gamma(a=2.0, scale=3.0)] + tm = GumbelCopula( + sampler=DigitalNetB2(2, seed=85), + marginals=marginals, + theta=1.0, + ) + u = np.array([[0.2, 0.7], [0.4, 0.8], [0.9, 0.1]]) + + x = tm._transform(u) + expected = np.column_stack( + [marginals[j].ppf(u[:, j]) for j in range(len(marginals))] + ) + + np.testing.assert_allclose(x, expected, rtol=1e-12, atol=1e-12) + + def test_gumbel_copula_theta_close_to_one_is_near_independent(self): + for dimension in [2, 3, 5]: + with self.subTest(dimension=dimension): + marginals = [stats.uniform()] * dimension + tm = GumbelCopula( + sampler=DigitalNetB2(dimension, seed=85), + marginals=marginals, + theta=1.000001, + ) + u = np.array( + [ + [0.2, 0.7, 0.4, 0.6, 0.8], + [0.4, 0.8, 0.9, 0.3, 0.2], + [0.9, 0.1, 0.3, 0.7, 0.5], + ] + )[:, :dimension] + + x = tm._transform(u) + + self.assertEqual(x.shape, (3, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + np.testing.assert_allclose(x, u, atol=5e-5) + + def test_gumbel_copula_large_theta_is_finite(self): + for dimension in [2, 3, 5]: + for theta in [20.0, 50.0]: + with self.subTest(dimension=dimension, theta=theta): + tm = GumbelCopula( + sampler=DigitalNetB2(dimension, seed=86), + marginals=[stats.norm()] * dimension, + theta=theta, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_gumbel_copula_supports_general_dimension(self): + for dimension in [2, 3, 5]: + with self.subTest(dimension=dimension): + tm = GumbelCopula( + sampler=DigitalNetB2(dimension, seed=87), + marginals=[stats.norm()] * dimension, + theta=2.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, dimension)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_gumbel_copula_marginal_without_ppf_raises_clear_error(self): + class NoPPF: + pass + + with self.assertRaisesRegex(ParameterError, "ppf"): + GumbelCopula( + sampler=DigitalNetB2(2, seed=89), + marginals=[stats.norm(), NoPPF()], + theta=2.0, + ) + + def test_gumbel_copula_common_scipy_frozen_marginals_work(self): + for marginals in [ + [stats.norm(), stats.beta(a=2, b=5)], + [stats.gamma(a=3), stats.expon()], + [stats.lognorm(s=0.5), stats.norm()], + ]: + with self.subTest(marginals=[type(m.dist).__name__ for m in marginals]): + tm = GumbelCopula( + sampler=DigitalNetB2(2, seed=91), + marginals=marginals, + theta=2.0, + ) + + x = tm(128) + + self.assertEqual(x.shape, (128, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_gumbel_copula_endpoint_uniforms_are_clipped_to_finite_outputs(self): + tm = GumbelCopula( + sampler=DigitalNetB2(2, seed=93), + marginals=[stats.norm(), stats.lognorm(s=0.5)], + theta=2.0, + ) + u = np.array([[0.0, 1.0], [1.0, 0.0]]) + + x = tm._transform(u) + + self.assertEqual(x.shape, (2, 2)) + self.assertTrue(np.all(np.isfinite(x))) + + def test_gumbel_copula_positive_dependence_behavior(self): + tm = GumbelCopula( + sampler=DigitalNetB2(2, seed=95), + marginals=[stats.uniform(), stats.uniform()], + theta=2.0, + ) + + x = tm(4096) + empirical_corr = np.corrcoef(x.T)[0, 1] + + self.assertGreater(empirical_corr, 0.45) + + def test_gumbel_copula_has_stronger_upper_tail_than_gaussian_copula(self): + theta = 2.0 + n = 2**12 + marginals = [stats.uniform(), stats.uniform()] + # Gumbel Kendall tau is 1 - 1/theta; convert to Gaussian rho. + rho = np.sin(np.pi * (1.0 - 1.0 / theta) / 2.0) + + gumbel = GumbelCopula( + sampler=DigitalNetB2(2, seed=97), + marginals=marginals, + theta=theta, + ) + gaussian = GaussianCopula( + sampler=DigitalNetB2(2, seed=97), + marginals=marginals, + correlation=[[1.0, rho], [rho, 1.0]], + ) + + x_gumbel = gumbel(n) + x_gaussian = gaussian(n) + threshold = 0.95 + + def upper_tail_rate(x): + tail_0 = x[:, 0] > threshold + return np.mean(x[tail_0, 1] > threshold) + + gumbel_tail = upper_tail_rate(x_gumbel) + gaussian_tail = upper_tail_rate(x_gaussian) + + self.assertGreater(gumbel_tail, gaussian_tail + 0.15) + + +class TestCopulaWeightsFallbackAndSpawn(unittest.TestCase): + + def test_copula_weight_fallback_warns_once_when_density_methods_are_missing(self): + for copula_cls in [GaussianCopula, StudentTCopula, ClaytonCopula, GumbelCopula, FrankCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + tm = _make_copula( + copula_cls, + dimension=2, + marginals=[PPFOnlyMarginal(), PPFOnlyMarginal()], + ) + x = np.full((4, 2), 0.5) + expected_message = getattr( + tm, + "_missing_weight_warning_message", + f"{copula_cls.__name__} marginals must implement 'cdf' and " + "'pdf' or 'logpdf' to compute density weights. " + "Weights will be treated as 1.", + ) + + self.assertNotIn("_unit_weight_with_warning", copula_cls.__dict__) + self.assertIs( + tm._unit_weight_with_warning.__func__, + AbstractCopula._unit_weight_with_warning, + ) + + with self.assertWarns(UserWarning) as wcm: + weights = tm._weight(x) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + second_weights = tm._weight(x) + + np.testing.assert_allclose(weights, np.ones(4)) + np.testing.assert_allclose(second_weights, np.ones(4)) + self.assertEqual(str(wcm.warning), expected_message) + self.assertEqual(caught, []) + + def test_student_t_weight_falls_back_when_multivariate_t_is_unavailable(self): + tm = StudentTCopula( + DigitalNetB2(2, seed=115), + marginals=[stats.norm(), stats.norm()], + correlation=np.eye(2), + df=4, + ) + tm._mvt_scipy = None + + with self.assertWarnsRegex(UserWarning, "Weights will be treated as 1"): + weights = tm._weight(np.full((3, 2), 0.25)) + + np.testing.assert_allclose(weights, np.ones(3)) + + def test_gaussian_weight_uses_pdf_branch_when_logpdf_is_unavailable(self): + tm = GaussianCopula( + DigitalNetB2(2, seed=117), + marginals=[UnitPDFMarginal(), UnitPDFMarginal()], + correlation=[[1.0, 0.4], [0.4, 1.0]], + ) + + weights = tm._weight(np.array([[0.25, 0.5], [0.75, 0.5]])) + + self.assertEqual(weights.shape, (2,)) + self.assertTrue(np.all(np.isfinite(weights))) + self.assertTrue(np.all(weights > 0.0)) + + def test_gumbel_theta_one_weight_is_independent_marginal_density(self): + tm = GumbelCopula( + DigitalNetB2(2, seed=119), + marginals=[stats.gamma(a=2.0), stats.expon()], + theta=1.0, + ) + x = np.array([[1.0, 0.5], [2.0, 1.5]]) + expected = stats.gamma(a=2.0).pdf(x[:, 0]) * stats.expon().pdf(x[:, 1]) + + weights = tm._weight(x) + + np.testing.assert_allclose(weights, expected) + + def test_gen_copula_samples_composed_transform_branch(self): + inner = GaussianCopula( + DigitalNetB2(2, seed=121), + marginals=[stats.uniform(), stats.uniform()], + correlation=[[1.0, 0.3], [0.3, 1.0]], + ) + outer = ClaytonCopula(inner, marginals=[stats.uniform(), stats.uniform()], theta=1.5) + + v = outer.gen_copula_samples(n_min=4, n_max=8) + + self.assertEqual(v.shape, (4, 2)) + self.assertTrue(np.all(np.isfinite(v))) + self.assertTrue(np.all((0.0 <= v) & (v <= 1.0))) + + def test_copula_spawn_same_dimension_and_reject_different_dimension(self): + for copula_cls in [GaussianCopula, StudentTCopula, ClaytonCopula, GumbelCopula, FrankCopula]: + with self.subTest(copula_cls=copula_cls.__name__): + tm = _make_copula(copula_cls, dimension=2) + + spawned = tm.spawn(s=1, dimensions=[2]) + self.assertEqual(len(spawned), 1) + self.assertIsInstance(spawned[0], copula_cls) + self.assertEqual(spawned[0](4).shape, (4, 2)) + + with self.assertRaises(DimensionError): + tm._spawn(DigitalNetB2(3, seed=123), 3) + + def test_frank_one_dimensional_weight_covers_zero_order_eulerian_term(self): + tm = FrankCopula( + DigitalNetB2(1, seed=125), + marginals=[UnitPDFMarginal()], + theta=3.0, + ) + + weights = tm._weight(np.array([[0.25], [0.75]])) + + self.assertEqual(weights.shape, (2,)) + self.assertTrue(np.all(np.isfinite(weights))) + self.assertTrue(np.all(weights > 0.0)) + + def test_frank_rejects_large_negative_theta_when_exponential_overflows(self): + with np.errstate(over="ignore"): + with self.assertRaisesRegex(ParameterError, "too close to 0 or too large"): + FrankCopula( + DigitalNetB2(2, seed=127), + marginals=[stats.uniform(), stats.uniform()], + theta=-1000.0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_product_measure.py b/test/test_tm_product_measure.py similarity index 100% rename from test/test_product_measure.py rename to test/test_tm_product_measure.py diff --git a/test/test_tm_scipy_wrapper_custom.py b/test/test_tm_scipy_wrapper_custom.py new file mode 100644 index 000000000..105b984e5 --- /dev/null +++ b/test/test_tm_scipy_wrapper_custom.py @@ -0,0 +1,292 @@ +import unittest +import warnings + +import numpy as np +import scipy.stats as stats + +from qmcpy import DigitalNetB2, SciPyWrapper, StudentT, ZeroInflatedExpUniform + +from qmcpy.true_measure.triangular import TriangularDistribution +from qmcpy.util import DimensionError, ParameterError + + +MISSING_PDF_WARNING = "no 'pdf' or 'logpdf'" + + +def _missing_pdf_warnings(caught): + return [ + warning + for warning in caught + if issubclass(warning.category, UserWarning) + and MISSING_PDF_WARNING in str(warning.message) + ] + + +class TestSciPyWrapperCustom(unittest.TestCase): + + def test_mvn_dependence_correlation_and_moment(self): + """ + Check that passing a SciPy multivariate normal through SciPyWrapper + preserves correlation and the mixed moment E[X1 X2]. + """ + sampler = DigitalNetB2(2, seed=5) + rho_target = 0.7 + cov = [[1.0, rho_target], [rho_target, 1.0]] + mvn = stats.multivariate_normal(mean=[0.0, 0.0], cov=cov) + tm_mvn = SciPyWrapper(sampler, scipy_distribs=mvn) + + n = 4096 + x = tm_mvn(n) + + rho_hat = np.corrcoef(x.T)[0, 1] + est_moment = np.mean(x[:, 0] * x[:, 1]) + + self.assertTrue(np.isfinite(rho_hat)) + self.assertTrue(np.isfinite(est_moment)) + + self.assertLess(abs(rho_hat - rho_target), 0.05) + self.assertLess(abs(est_moment - rho_target), 0.05) + + def test_triangular_custom_marginal_range_and_shape(self): + """ + Make sure our custom triangular marginal behaves sensibly: + samples stay in the right interval and the empirical mean is close + to the analytic mean. + """ + tri = TriangularDistribution(c=0.3, loc=-1.0, scale=2.0) + tm = SciPyWrapper(DigitalNetB2(1, seed=11), scipy_distribs=tri) + + n = 4096 + x = tm(n).ravel() + + self.assertGreaterEqual(x.min(), -1.1) + self.assertLessEqual(x.max(), 1.1) + + a = -1.0 + b = 1.0 + m = -1.0 + 0.3 * 2.0 + true_mean = (a + b + m) / 3.0 + emp_mean = x.mean() + self.assertLess(abs(emp_mean - true_mean), 0.05) + + def test_zero_inflated_zero_rate(self): + """ + Check that the zero-inflated exponential distribution preserves the + specified probability mass at X = 0. + """ + p_zero = 0.4 + sampler = DigitalNetB2(1, seed=17) + tm = ZeroInflatedExpUniform(sampler, p_zero=p_zero, lam=1.5) + + n = 4096 + samples = tm(n) + x = samples.ravel() + zero_rate = np.mean(x == 0.0) + + self.assertEqual(samples.shape, (n, 1)) + self.assertLess(abs(zero_rate - p_zero), 0.05) + + def test_zero_inflated_replications_shape(self): + tm = ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17, replications=2), + p_zero=0.4, + lam=1.5, + ) + + x = tm(8) + + self.assertEqual(x.shape, (2, 8, 1)) + self.assertTrue(np.all(x >= 0.0)) + + def test_zero_inflated_rejects_invalid_p_zero(self): + for p_zero in [0.0, 1.0, -0.1, 1.1]: + with self.subTest(p_zero=p_zero): + with self.assertRaisesRegex(ParameterError, "p_zero must be in"): + ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17), + p_zero=p_zero, + lam=1.5, + ) + + def test_zero_inflated_rejects_nonpositive_lam(self): + for lam in [0.0, -1.0]: + with self.subTest(lam=lam): + with self.assertRaisesRegex(ParameterError, "lam must be positive"): + ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17), + p_zero=0.4, + lam=lam, + ) + + def test_zero_inflated_requires_one_dimensional_sampler(self): + with self.assertRaisesRegex( + DimensionError, "requires a one-dimensional sampler" + ): + ZeroInflatedExpUniform( + DigitalNetB2(2, seed=17), + p_zero=0.4, + lam=1.5, + ) + + def test_zero_inflated_inverse_transform_exact_values(self): + tm = ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17), + p_zero=0.4, + lam=2.0, + ) + u = np.array([[0.0], [0.2], [0.4], [0.7], [0.9]]) + + x = tm._transform(u) + + self.assertEqual(x.shape, (5, 1)) + self.assertTrue(np.array_equal(x[:3], np.zeros((3, 1)))) + self.assertTrue(np.all(x[3:] > 0.0)) + + u_positive = u[3:, 0] + u_rescaled = (u_positive - 0.4) / 0.6 + expected = -np.log1p(-u_rescaled) / 2.0 + self.assertTrue(np.allclose(x[3:, 0], expected)) + + def test_zero_inflated_inverse_transform_all_zero_branch(self): + tm = ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17), + p_zero=0.4, + lam=2.0, + ) + u = np.array([[0.0], [0.1], [0.4]]) + + x = tm._transform(u) + + self.assertEqual(x.shape, (3, 1)) + self.assertTrue(np.array_equal(x, np.zeros((3, 1)))) + + def test_zero_inflated_inverse_transform_clips_one(self): + tm = ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17), + p_zero=0.4, + lam=2.0, + ) + u = np.array([[1.0]]) + + x = tm._transform(u) + + self.assertEqual(x.shape, (1, 1)) + self.assertTrue(np.isfinite(x).all()) + self.assertGreater(x[0, 0], 0.0) + + def test_zero_inflated_construction_does_not_warn_about_missing_pdf(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + tm = ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17), + p_zero=0.4, + lam=1.5, + ) + + self.assertEqual(tm.d, 1) + self.assertEqual(_missing_pdf_warnings(caught), []) + + def test_zero_inflated_sampling_does_not_warn_about_missing_pdf(self): + tm = ZeroInflatedExpUniform(DigitalNetB2(1, seed=17), p_zero=0.4, lam=1.5) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + x = tm(8) + + self.assertEqual(x.shape, (8, 1)) + self.assertEqual(_missing_pdf_warnings(caught), []) + + def test_zero_inflated_return_weights_warns_once_for_missing_pdf(self): + tm = ZeroInflatedExpUniform(DigitalNetB2(1, seed=17), p_zero=0.4, lam=1.5) + + with self.assertWarnsRegex(UserWarning, MISSING_PDF_WARNING): + x, jac = tm(8, return_weights=True) + + self.assertEqual(x.shape, (8, 1)) + self.assertEqual(jac.shape, (8,)) + self.assertTrue(np.allclose(jac, 1.0)) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + x_second, jac_second = tm(8, return_weights=True) + + self.assertEqual(x_second.shape, (8, 1)) + self.assertTrue(np.allclose(jac_second, 1.0)) + self.assertEqual(_missing_pdf_warnings(caught), []) + + def test_zero_inflated_y_split_warns_and_uses_one_dimensional_interface(self): + with self.assertWarnsRegex(DeprecationWarning, "y_split"): + tm = ZeroInflatedExpUniform( + DigitalNetB2(1, seed=17), + p_zero=0.4, + lam=1.5, + y_split=0.5, + ) + + x = tm(4) + + self.assertEqual(x.shape, (4, 1)) + self.assertTrue(np.all(x >= 0.0)) + + def test_zero_inflated_y_split_preserves_deprecated_two_dimensional_usage(self): + with self.assertWarnsRegex(DeprecationWarning, "2D zero-inflated"): + tm = ZeroInflatedExpUniform( + DigitalNetB2(2, seed=17), + p_zero=0.4, + lam=1.5, + y_split=0.5, + ) + + x = tm(16) + + self.assertEqual(x.shape, (16, 2)) + self.assertTrue(np.all(x[:, 0] >= 0.0)) + self.assertTrue(np.all((0.0 <= x[:, 1]) & (x[:, 1] <= 1.0))) + self.assertTrue(np.all(x[x[:, 0] == 0.0, 1] <= 0.5)) + self.assertTrue(np.all(x[x[:, 0] > 0.0, 1] >= 0.5)) + + def test_zero_inflated_y_split_preserves_replicated_two_dimensional_usage(self): + with self.assertWarnsRegex(DeprecationWarning, "2D zero-inflated"): + tm = ZeroInflatedExpUniform( + DigitalNetB2(2, seed=17, replications=2), + p_zero=0.4, + lam=1.5, + y_split=0.5, + ) + + x = tm(16) + + self.assertEqual(x.shape, (2, 16, 2)) + self.assertTrue(np.all(x[..., 0] >= 0.0)) + self.assertTrue(np.all((0.0 <= x[..., 1]) & (x[..., 1] <= 1.0))) + self.assertTrue(np.all(x[..., 1][x[..., 0] == 0.0] <= 0.5)) + self.assertTrue(np.all(x[..., 1][x[..., 0] > 0.0] >= 0.5)) + + def test_student_t_marginals_shape(self): + tm = SciPyWrapper( + sampler=DigitalNetB2(2, seed=5), + scipy_distribs=stats.t(df=5), + ) + x = tm(8) + self.assertEqual(x.shape, (8, 2)) + + def test_multivariate_student_t_joint_corr_and_cov(self): + if not hasattr(stats, "multivariate_t"): + self.skipTest("scipy.stats.multivariate_t not available in this SciPy version") + + df = 5.0 + rho = 0.8 + loc = np.array([0.0, 0.0]) + shape = np.array([[1.0, rho], [rho, 1.0]]) + + tm = StudentT(DigitalNetB2(2, seed=123), loc=loc, shape=shape, df=df) + + n = 4096 + x = tm(n) + emp_corr = np.corrcoef(x.T)[0, 1] + + self.assertLess(abs(emp_corr - rho), 0.05) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_true_measures.py b/test/test_tm_true_measures.py similarity index 100% rename from test/test_true_measures.py rename to test/test_tm_true_measures.py diff --git a/test/test_unwrap_markdown.py b/test/test_unwrap_markdown.py deleted file mode 100644 index e227b1807..000000000 --- a/test/test_unwrap_markdown.py +++ /dev/null @@ -1,89 +0,0 @@ -import pytest - -from scripts.unwrap_markdown import unwrap_markdown_text - - -@pytest.mark.parametrize( - ("source", "expected"), - [ - ( - "- unordered first\n unordered second\n", - "- unordered first unordered second\n", - ), - ( - "- [ ] task first\n task second\n", - "- [ ] task first task second\n", - ), - ( - "10. ordered first\n ordered second\n", - "10. ordered first ordered second\n", - ), - ], -) -def test_unwraps_list_item_continuations(source, expected): - updated = unwrap_markdown_text(source) - - assert updated == expected - assert unwrap_markdown_text(updated) == updated - - -def test_unwraps_adjacent_and_nested_list_items_separately(): - source = ( - "- parent first\n" - " parent second\n" - " - child first\n" - " child second\n" - "- sibling first\n" - " sibling second\n" - ) - - assert unwrap_markdown_text(source) == ( - "- parent first parent second\n" - " - child first child second\n" - "- sibling first sibling second\n" - ) - - -def test_preserves_list_item_blocks_and_explicit_hard_breaks(): - source = ( - "- first paragraph\n" - " continuation\n" - "\n" - " second paragraph\n" - " continuation\n" - "\n" - "- item before code\n" - " indented code\n" - "\n" - "- explicit hard break \n" - " remains separate\n" - ) - - assert unwrap_markdown_text(source) == ( - "- first paragraph continuation\n" - "\n" - " second paragraph continuation\n" - "\n" - "- item before code\n" - " indented code\n" - "\n" - "- explicit hard break \n" - " remains separate\n" - ) - - -def test_unwraps_ordinary_paragraphs(): - assert unwrap_markdown_text("first line\nsecond line\n") == "first line second line\n" - - -@pytest.mark.parametrize("rule", ["- - -", "* * *", "_ _ _"]) -def test_preserves_horizontal_rules(rule): - source = f"{rule}\nfollowing paragraph\n" - - assert unwrap_markdown_text(source) == source - - -def test_preserves_indented_code_that_looks_like_a_list(): - source = " - code first\n code second\n" - - assert unwrap_markdown_text(source) == source diff --git a/test/test_ut_install_mpmc_pyg.py b/test/test_ut_install_mpmc_pyg.py new file mode 100644 index 000000000..25e810ec1 --- /dev/null +++ b/test/test_ut_install_mpmc_pyg.py @@ -0,0 +1,92 @@ +"""Tests for the platform-specific MPMC dependency installer.""" + +import subprocess +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from qmcpy.util import install_mpmc_pyg + + +def _torch(version="2.12.1+cpu", cuda=None, hip=None): + return SimpleNamespace( + __version__=version, + version=SimpleNamespace(cuda=cuda, hip=hip), + ) + + +class TestInstallMPMCPyG(unittest.TestCase): + + def test_torch_versions_include_baseline_fallback(self): + """Wheel lookup tries an exact patch release, then its minor baseline.""" + self.assertEqual( + install_mpmc_pyg.torch_versions("2.12.1+cpu"), ["2.12.1", "2.12.0"] + ) + self.assertEqual(install_mpmc_pyg.torch_versions("2.12.0"), ["2.12.0"]) + + with self.assertRaisesRegex(RuntimeError, "Unable to parse torch version"): + install_mpmc_pyg.torch_versions("development") + + def test_accelerator_tag(self): + """PyTorch build metadata maps to the expected PyG wheel tag.""" + cases = [ + (_torch(), "cpu"), + (_torch(cuda="12.6"), "cu126"), + (_torch(cuda="13.0.1"), "cu130"), + ] + for torch_module, expected in cases: + with self.subTest(expected=expected): + self.assertEqual( + install_mpmc_pyg.accelerator_tag(torch_module), expected + ) + + def test_accelerator_tag_rejects_rocm(self): + """The installer directs unsupported ROCm users to upstream guidance.""" + with self.assertRaisesRegex(RuntimeError, "does not currently support ROCm"): + install_mpmc_pyg.accelerator_tag(_torch(hip="6.3")) + + def test_main_retries_with_torch_minor_baseline(self): + """A missing exact wheel page falls back to the minor baseline page.""" + calls = [] + + def fake_run(*args): + calls.append(args) + if args[-1].endswith("torch-2.12.1+cpu.html"): + raise subprocess.CalledProcessError(1, args) + + with patch.object(install_mpmc_pyg, "run", fake_run): + install_mpmc_pyg.main(_torch()) + + self.assertEqual(calls[0][-1], "torch-geometric>=2.6.1") + self.assertEqual( + calls[1][-1], "https://data.pyg.org/whl/torch-2.12.1+cpu.html" + ) + self.assertEqual( + calls[2][-1], "https://data.pyg.org/whl/torch-2.12.0+cpu.html" + ) + self.assertIn("--only-binary", calls[1]) + + def test_main_explains_that_torch_must_be_installed(self): + """Running the helper before installing the extra gives a useful error.""" + def missing_torch(_name): + raise ModuleNotFoundError("No module named 'torch'", name="torch") + + with patch.object( + install_mpmc_pyg.importlib, "import_module", missing_torch + ): + with self.assertRaisesRegex(RuntimeError, r"install 'qmcpy\[mpmc\]'"): + install_mpmc_pyg.main() + + def test_main_reports_missing_wheel(self): + """Exhausting candidate wheel pages reports the build that failed.""" + def fail_pyg_lib(*args): + if "pyg_lib>=0.6.0" in args: + raise subprocess.CalledProcessError(1, args) + + with patch.object(install_mpmc_pyg, "run", fail_pyg_lib): + with self.assertRaisesRegex(RuntimeError, r"torch 2\.12\.1\+cpu \(cpu\)"): + install_mpmc_pyg.main(_torch()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_ut_plot_and_stop.py b/test/test_ut_plot_and_stop.py new file mode 100644 index 000000000..4ce81545b --- /dev/null +++ b/test/test_ut_plot_and_stop.py @@ -0,0 +1,166 @@ +import builtins +import sys +import types +import unittest +from unittest.mock import patch + +import numpy as np + +from qmcpy import AbstractDiscreteDistribution, plot_proj +from qmcpy.util import stop_notebook + + +class FakeAxes: + def __init__(self): + self.removed = False + self.calls = [] + + def remove(self): + self.removed = True + + def set_xlim(self, *a, **k): + self.calls.append(("set_xlim", a)) + + def set_ylim(self, *a, **k): + self.calls.append(("set_ylim", a)) + + def set_xticks(self, *a, **k): + self.calls.append(("set_xticks", a)) + + def set_yticks(self, *a, **k): + self.calls.append(("set_yticks", a)) + + def set_aspect(self, *a, **k): + self.calls.append(("set_aspect", a)) + + def grid(self, *a, **k): + self.calls.append(("grid", a)) + + def tick_params(self, *a, **k): + self.calls.append(("tick_params", a)) + + def set_xlabel(self, *a, **k): + self.calls.append(("set_xlabel", a)) + + def set_ylabel(self, *a, **k): + self.calls.append(("set_ylabel", a)) + + def scatter(self, *a, **k): + self.calls.append(("scatter", a)) + + +class FakeFig: + def __init__(self): + self.tl = False + + def tight_layout(self, *a, **k): + self.tl = True + + +def make_fake_matplotlib(nrows, ncols): + plt = types.ModuleType("matplotlib.pyplot") + plt.style = types.SimpleNamespace() + plt.style.use = lambda *a, **k: None + plt.rcParams = { + "font.family": "sans-serif", + "axes.prop_cycle": types.SimpleNamespace( + by_key=lambda: {"color": ["k", "b", "r"]} + ), + } + + def subplots(nrows=1, ncols=1, figsize=None, squeeze=False): + fig = FakeFig() + ax = np.empty((nrows, ncols), dtype=object) + for i in range(nrows): + for j in range(ncols): + ax[i, j] = FakeAxes() + return fig, ax + + plt.subplots = subplots + plt.suptitle = lambda *a, **k: None + return plt + + +class DummySampler(AbstractDiscreteDistribution): + def __init__(self, d=2): + super().__init__(dimension=d, replications=1, seed=1, d_limit=10, n_limit=100) + + def _gen_samples(self, n_min, n_max, return_binary=False, warn=True): + n = n_max - n_min + return np.tile(np.arange(n)[:, None] / max(1, n - 1), (1, 1, self.d)).reshape( + self.replications, n, self.d + ) + + def __repr__(self): + return "DummySampler" + + +class TestPlotProjAndStopNotebook(unittest.TestCase): + + def test_plot_proj_with_fake_matplotlib_and_sampler(self): + # Inject fake matplotlib.pyplot + fake_plt = make_fake_matplotlib(1, 1) + # Create a proper matplotlib package module with colors submodule + fake_matplotlib = types.ModuleType("matplotlib") + fake_matplotlib.pyplot = fake_plt + fake_matplotlib.colors = types.SimpleNamespace() + + with patch.dict( + sys.modules, + {"matplotlib.pyplot": fake_plt, "matplotlib": fake_matplotlib}, + ): + sampler = DummySampler(d=3) + fig, ax = plot_proj( + sampler, + n=4, + d_horizontal=1, + d_vertical=2, + math_ind=True, + marker_size=1, + figfac=1, + ) + + self.assertIsInstance(fig, FakeFig) + self.assertIsInstance(ax, np.ndarray) + # At least one axes should have scatter calls or be removed + found = False + for a in ax.flatten(): + if getattr(a, "removed", False) or any(c[0] == "scatter" for c in a.calls): + found = True + break + self.assertTrue(found) + + def test_plot_proj_with_callable_sampler(self): + # sampler not instance of AbstractDiscreteDistribution -> uses t_i labels + fake_plt = make_fake_matplotlib(1, 1) + fake_matplotlib = types.ModuleType("matplotlib") + fake_matplotlib.pyplot = fake_plt + fake_matplotlib.colors = types.SimpleNamespace() + + with patch.dict( + sys.modules, + {"matplotlib.pyplot": fake_plt, "matplotlib": fake_matplotlib}, + ): + def sampler_callable(n): + return np.zeros((n, 1)) + + fig, ax = plot_proj( + sampler_callable, n=3, d_horizontal=0, d_vertical=0, math_ind=False + ) + + self.assertIsInstance(fig, FakeFig) + + def test_stop_notebook_yes_and_no(self): + # When input is 'yes' nothing should happen + with patch.object(builtins, "input", lambda prompt="": "yes"): + # Should not raise + stop_notebook("prompt") + + # When input is not 'yes' should exit + with patch.object(builtins, "input", lambda prompt="": "no"): + with self.assertRaises(SystemExit): + stop_notebook("prompt") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_util.py b/test/test_ut_util.py similarity index 100% rename from test/test_util.py rename to test/test_ut_util.py