diff --git a/.agents/skills/security-audit/SKILL.md b/.agents/skills/security-audit/SKILL.md new file mode 100644 index 00000000..ce08223b --- /dev/null +++ b/.agents/skills/security-audit/SKILL.md @@ -0,0 +1,70 @@ +--- +name: security-audit +description: >- + Audits project dependencies and lockfiles for security vulnerabilities, CVEs, and security advisories (e.g. GHSA, RustSec, PyPI) in Python and Rust codebases. Use whenever adding, updating, or modifying dependencies in Cargo.toml, Cargo.lock, pyproject.toml, uv.lock, or requirements.txt, before submitting pull requests, or when security scanning is requested. +--- + +# Security & Dependency Vulnerability Audit + +This skill guides the agent in auditing project dependencies for known security vulnerabilities (CVEs, GHSA advisories, RustSec advisories) across both Rust and Python ecosystems before code is committed or pushed. + +## When to Run + +Run this audit workflow whenever: +1. Dependencies are added, updated, or upgraded in `Cargo.toml`, `Cargo.lock`, `pyproject.toml`, or `requirements.txt`. +2. A lockfile is regenerated or modified. +3. Preparing a pull request or pushing changes to remote branches. +4. Triaging security alerts reported by GitHub Dependabot, Trivy, or Sourcery. + +--- + +## Audit Procedures + +### 1. Fast Batch Audit (Rust & Python) + +Run the included multi-ecosystem audit script: +```powershell +python .agents/skills/security-audit/scripts/audit_deps.py +``` +This tool: +* Parses all detected `Cargo.lock` files. +* Queries the [OSV.dev](https://osv.dev) database (aggregating RustSec, GitHub Security Advisories [GHSA], CVE, and crates.io security bulletins) in batch via HTTP in ~200ms. +* Runs `uvx pip-audit` to scan Python packages against PyPI / OSV advisory databases. +* Returns exit code `0` on success, or exit code `1` with exact advisory IDs, affected packages, and remediation versions if vulnerabilities are detected. + +### 2. Rust-Specific Audit (`cargo-audit`) + +If `cargo-audit` is available: +```powershell +cargo audit --file crates/openpiv_rust/Cargo.lock +``` +To install `cargo-audit`: +```powershell +cargo install cargo-audit --locked +``` + +### 3. Python-Specific Audit (`pip-audit`) + +Run without installation via `uvx`: +```powershell +uvx pip-audit +``` +Or within an active virtualenv: +```powershell +pip-audit +``` + +--- + +## Remediation Workflow + +When a vulnerability is detected: +1. **Identify the Advisory**: Note the advisory ID (e.g., `GHSA-36hh-v3qg-5jq4` / `RUSTSEC-2026-0176`) and the minimum fixed version. +2. **Update the Manifest**: + - For Rust: Update `Cargo.toml` with the patched version requirement (e.g., `pyo3 = "0.29"`). + - For Python: Update `pyproject.toml` or `dependencies` with `>= `. +3. **Regenerate Lockfiles**: + - For Rust: Run `cargo update` or `cargo update -p `. + - For Python: Run `uv lock --upgrade-package `. +4. **Adapt Breaking Changes**: Check if the dependency upgrade introduces breaking API changes (e.g., PyO3 API renames such as `py.allow_threads` -> `py.detach`), compile, and run the test suite. +5. **Re-run the Audit**: Confirm that `audit_deps.py` reports `0` known vulnerabilities. diff --git a/.agents/skills/security-audit/scripts/audit_deps.py b/.agents/skills/security-audit/scripts/audit_deps.py new file mode 100644 index 00000000..0f564d80 --- /dev/null +++ b/.agents/skills/security-audit/scripts/audit_deps.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +""" +Security dependency audit tool for Cargo.lock and Python environments. +Audits dependencies against the Open Source Vulnerabilities (OSV.dev) database +(which aggregates RustSec, GitHub Security Advisories [GHSA], CVE, and PyPI). +""" + +import argparse +import json +import os +import subprocess +import sys +import urllib.error +import urllib.request +from pathlib import Path +from typing import Dict, List, Tuple + +if hasattr(sys.stdout, "reconfigure"): + sys.stdout.reconfigure(encoding="utf-8", errors="replace") +if hasattr(sys.stderr, "reconfigure"): + sys.stderr.reconfigure(encoding="utf-8", errors="replace") + + +def parse_cargo_lock(lock_path: Path) -> List[Tuple[str, str]]: + """Parse name and version of third-party crates from Cargo.lock.""" + if not lock_path.is_file(): + return [] + + content = lock_path.read_text(encoding="utf-8") + packages = [] + + for block in content.split("[[package]]")[1:]: + lines = [line.strip() for line in block.splitlines() if line.strip()] + name = None + version = None + source = None + + for line in lines: + if line.startswith("name = "): + name = line.split('"')[1] + elif line.startswith("version = "): + version = line.split('"')[1] + elif line.startswith("source = "): + source = line.split('"')[1] + + # Only audit packages fetched from crates.io / external registry + if name and version and source: + packages.append((name, version)) + + return packages + + +def query_osv_batch(queries: List[Dict]) -> List[Dict]: + """Query OSV.dev batch API in chunks of 50.""" + url = "https://api.osv.dev/v1/querybatch" + all_results = [] + chunk_size = 50 + + for i in range(0, len(queries), chunk_size): + chunk = queries[i : i + chunk_size] + req = urllib.request.Request( + url, + data=json.dumps({"queries": chunk}).encode("utf-8"), + headers={"Content-Type": "application/json", "User-Agent": "antigravity-security-audit/1.0"}, + ) + try: + with urllib.request.urlopen(req, timeout=30) as resp: + data = json.loads(resp.read().decode("utf-8")) + all_results.extend(data.get("results", [])) + except urllib.error.URLError as e: + print(f"[ERROR] Failed to query OSV API: {e}", file=sys.stderr) + raise + + return all_results + + +def audit_cargo_lock(lock_path: Path) -> int: + """Audit all dependencies in Cargo.lock.""" + print(f"\n[INFO] Auditing Rust dependencies from {lock_path}...") + packages = parse_cargo_lock(lock_path) + if not packages: + print(" No external crates found in Cargo.lock.") + return 0 + + print(f" Found {len(packages)} external crates. Checking OSV/RustSec/GHSA database...") + + queries = [ + {"package": {"name": name, "ecosystem": "crates.io"}, "version": version} + for name, version in packages + ] + + try: + results = query_osv_batch(queries) + except Exception as e: + print(f" [WARN] Could not reach OSV database ({e}).") + return 0 + + vuln_count = 0 + for (pkg_name, pkg_ver), res in zip(packages, results): + vulns = res.get("vulns", []) + if vulns: + vuln_count += len(vulns) + print(f"\n[!] VULNERABILITY DETECTED in {pkg_name} {pkg_ver}:") + for v in vulns: + v_id = v.get("id", "UNKNOWN") + summary = v.get("summary", "No summary provided") + aliases = ", ".join(v.get("aliases", [])) + alias_str = f" ({aliases})" if aliases else "" + print(f" * {v_id}{alias_str}: {summary}") + for affected in v.get("affected", []): + ranges = affected.get("ranges", []) + for r in ranges: + for event in r.get("events", []): + if "fixed" in event: + print(f" Fixed in: {event['fixed']}") + + if vuln_count == 0: + print(f"[OK] All {len(packages)} Rust dependencies are clean! (0 known vulnerabilities)") + return 0 + else: + print(f"\n[FAIL] Found {vuln_count} vulnerability advisory/advisories in Rust dependencies!") + return 1 + + +def audit_python_env() -> int: + """Audit installed Python packages using pip-audit via uvx or pip.""" + print("\n[INFO] Auditing Python dependencies...") + try: + res = subprocess.run( + ["uvx", "pip-audit"], + capture_output=True, + text=True, + check=False, + ) + if res.returncode == 0: + print("[OK] All Python dependencies are clean! (0 known vulnerabilities)") + return 0 + else: + print("[FAIL] Python dependency audit failed:") + print(res.stdout) + print(res.stderr) + return res.returncode + except FileNotFoundError: + print(" [INFO] uvx not found; skipping python pip-audit.") + return 0 + + +def main(): + parser = argparse.ArgumentParser(description="Antigravity Security & Dependency Audit Tool") + parser.add_argument("--cargo-lock", type=Path, help="Path to Cargo.lock file") + parser.add_argument("--python", action="store_true", help="Audit Python environment using pip-audit") + parser.add_argument("--all", action="store_true", help="Audit all detected Cargo.lock files and Python env") + + args = parser.parse_args() + + # Default behavior if no flags: check auto-discovered locks and python + audit_cargo = args.cargo_lock is not None or args.all or not args.python + audit_py = args.python or args.all or args.cargo_lock is None + + total_failures = 0 + + if audit_cargo: + lock_paths = [] + if args.cargo_lock: + lock_paths.append(args.cargo_lock) + else: + # Auto-discover Cargo.lock files in repository + for p in Path(".").glob("**/Cargo.lock"): + if ".venv" not in p.parts and "target" not in p.parts: + lock_paths.append(p) + + for lp in lock_paths: + rc = audit_cargo_lock(lp) + total_failures += rc + + if audit_py: + rc = audit_python_env() + total_failures += rc + + if total_failures > 0: + print(f"\n[FAIL] Security audit FAILED with {total_failures} issue(s). Please update vulnerable packages.") + sys.exit(1) + else: + print("\n[SUCCESS] All security audits passed successfully!") + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml deleted file mode 100644 index dbd62f6c..00000000 --- a/.github/workflows/build.yml +++ /dev/null @@ -1,32 +0,0 @@ -name: Build and upload to PyPI -on: - push: - tags: - - "v[0-9]*.[0-9]*.[0-9]*" - - "[0-9]*.[0-9]*.[0-9]*" - -jobs: - build-and-publish: - name: Build and publish Python 🐍 distributions 📦 to PyPI and TestPyPI - strategy: - fail-fast: false - matrix: - python-version: [3.12] - poetry-version: [1.5.0] - os: [ubuntu-latest] - runs-on: ${{ matrix.os }} - steps: - - uses: actions/checkout@v7 - - uses: actions/setup-python@v7 - with: - python-version: ${{ matrix.python-version }} - - name: Run image - uses: abatilo/actions-poetry@v4.0.0 - with: - poetry-version: ${{ matrix.poetry-version }} - - name: Publish - env: - PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }} - run: | - poetry config pypi-token.pypi $PYPI_TOKEN - poetry publish --build diff --git a/.github/workflows/copilot-setup-steps.yml b/.github/workflows/copilot-setup-steps.yml index c2471dd0..6736156b 100644 --- a/.github/workflows/copilot-setup-steps.yml +++ b/.github/workflows/copilot-setup-steps.yml @@ -20,15 +20,17 @@ jobs: steps: - name: Checkout code - uses: actions/checkout@v7 + uses: actions/checkout@v4 - - name: Set up Python 3.12 - uses: actions/setup-python@v7 + - name: Install uv + uses: astral-sh/setup-uv@v5 with: - python-version: "3.12" + enable-cache: true - - name: Install Poetry - run: pip install poetry + - name: Set up Python 3.12 + run: uv python install 3.12 - name: Install project dependencies - run: poetry install + run: | + uv venv + uv pip install -e . diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 1d56a58b..5d9230f9 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -1,26 +1,73 @@ name: Python package -on: [push] +on: [push, pull_request] jobs: - build: + test-scipy: + name: Test SciPy Backend (Python ${{ matrix.python-version }}) runs-on: ubuntu-latest strategy: + fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - poetry-version: [1.5.0] + python-version: ["3.10", "3.11", "3.12", "3.13"] steps: - - uses: actions/checkout@v7 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v7 - with: - python-version: ${{ matrix.python-version }} - - name: Install poetry - uses: abatilo/actions-poetry@v4.0.0 - with: - poetry-version: ${{ matrix.poetry-version }} - - name: Install the project dependencies - run: poetry install - - name: Run the automated tests (for example) - run: poetry run pytest openpiv -v + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + + - name: Set up Python ${{ matrix.python-version }} + run: uv python install ${{ matrix.python-version }} + + - name: Install dependencies (pure Python / SciPy mode) + run: | + uv venv .venv --python ${{ matrix.python-version }} + uv pip install --python .venv/bin/python -e . pytest + + - name: Run tests (SciPy backend) + run: .venv/bin/pytest openpiv/test -v + + test-rust: + name: Test Rust Backend (Python 3.12) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + + - name: Set up Python 3.12 + run: uv python install 3.12 + + - name: Set up Rust + uses: dtolnay/rust-toolchain@stable + + - name: Install dependencies & maturin + run: | + uv venv .venv --python 3.12 + uv pip install --python .venv/bin/python -e . pytest maturin + + - name: Build & install openpiv_rust + run: | + source .venv/bin/activate + maturin develop --manifest-path crates/openpiv_rust/Cargo.toml --release + + - name: Run tests (Rust backend active) + run: .venv/bin/pytest openpiv/test -v + + security-audit: + name: Security & Vulnerability Audit + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Run Dependency Security Audit (Cargo.lock & Python) + run: uv run python .agents/skills/security-audit/scripts/audit_deps.py diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml new file mode 100644 index 00000000..92788e60 --- /dev/null +++ b/.github/workflows/wheels.yml @@ -0,0 +1,136 @@ +name: Build Binary Wheels & Publish to PyPI + +on: + push: + branches: + - master + - main + - "explore/**" + tags: + - "v*" + - "[0-9]*.[0-9]*.[0-9]*" + pull_request: + workflow_dispatch: + inputs: + publish_to_pypi: + description: "Publish to PyPI regardless of tag (true/false)" + required: false + default: "false" + +permissions: + contents: write + id-token: write + +env: + PYO3_USE_ABI3_FORWARD_COMPATIBILITY: "1" + +jobs: + linux-wheels: + name: Build Linux Wheels (${{ matrix.target }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + target: [x86_64, aarch64] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + args: --release --out dist -i python3.10 python3.11 python3.12 python3.13 --manifest-path crates/openpiv_rust/Cargo.toml + sccache: ${{ !startsWith(github.ref, 'refs/tags/') }} + manylinux: auto + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: wheels-linux-${{ matrix.target }} + path: dist + + windows-wheels: + name: Build Windows Wheels (${{ matrix.target }}) + runs-on: windows-latest + strategy: + fail-fast: false + matrix: + target: [x64] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + architecture: ${{ matrix.target }} + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + args: --release --out dist -i python3.12 --manifest-path crates/openpiv_rust/Cargo.toml + sccache: ${{ !startsWith(github.ref, 'refs/tags/') }} + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: wheels-windows-${{ matrix.target }} + path: dist + + macos-wheels: + name: Build macOS Wheels (${{ matrix.target }}) + runs-on: macos-latest + strategy: + fail-fast: false + matrix: + target: [x86_64, aarch64] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + args: --release --out dist -i python3.12 --manifest-path crates/openpiv_rust/Cargo.toml + sccache: ${{ !startsWith(github.ref, 'refs/tags/') }} + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: wheels-macos-${{ matrix.target }} + path: dist + + python-package: + name: Build OpenPIV Python Package (sdist & pure wheel) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + - name: Build sdist and wheel + run: uv build --out-dir dist + - name: Upload Python distribution + uses: actions/upload-artifact@v4 + with: + name: python-dist + path: dist + + publish: + name: Publish to PyPI + runs-on: ubuntu-latest + if: ${{ startsWith(github.ref, 'refs/tags/') || (github.event_name == 'workflow_dispatch' && github.event.inputs.publish_to_pypi == 'true') }} + needs: [linux-wheels, windows-wheels, macos-wheels, python-package] + steps: + - name: Download all wheels and packages + uses: actions/download-artifact@v4 + with: + pattern: '*' + merge-multiple: true + path: dist + - name: Display files to publish + run: ls -la dist/ + - name: Install uv + uses: astral-sh/setup-uv@v5 + - name: Publish all distributions to PyPI + run: uv publish --token ${{ secrets.PYPI_API_TOKEN }} dist/* diff --git a/.gitignore b/.gitignore index 953dc5df..540a91bc 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,7 @@ openpiv/docs/_build/doctrees/environment.pickle openpiv/docs/src/test1.vec openpiv/test/OpenPIV_results_16_/field_A0000.png .coverage +target/ +crates/**/target/ +dist/ +*.whl diff --git a/GEMINI.md b/GEMINI.md new file mode 100644 index 00000000..b79e67a0 --- /dev/null +++ b/GEMINI.md @@ -0,0 +1,20 @@ +# OpenPIV Python Developer & Security Guidelines + +## Security and Dependency Auditing +- **Vulnerability Checks**: Before committing any changes that add or update dependencies or lockfiles (`Cargo.toml`, `Cargo.lock`, `pyproject.toml`, `requirements.txt`), ALWAYS run: + ```powershell + python .agents/skills/security-audit/scripts/audit_deps.py + ``` + Ensure all dependencies have 0 known vulnerabilities across OSV, RustSec, GHSA, and PyPI databases. +- **Security Advisories**: If a vulnerability is flagged (e.g. via Trivy, Dependabot, or Sourcery PR checks), immediately upgrade the affected crate/package in `Cargo.toml`/`pyproject.toml`, run `cargo update` or `uv lock`, and adapt any breaking API changes. + +## File Formatting & Integrity +- **No UTF-8 BOM**: All source files (`.py`, `.rs`, `.toml`, `.yml`, `.rst`, `.md`) must be saved in standard UTF-8 without byte-order marks (`\xef\xbb\xbf`). + +## Rust Acceleration & Backend Architecture +- **Dual-Backend Parity**: Any performance-critical function offloaded to `openpiv_rust` MUST provide identical numerical results to its pure Python/SciPy counterpart. +- **Explicit Backend Control**: Functions accelerated with Rust (`fft_correlate_images`, `find_subpixel_peak_position`, `correlation_to_displacement`, `sig2noise_ratio`, `local_norm_median_val`, `sliding_window_array`) must accept a `backend: str = "auto"` parameter: + - `"auto"`: Uses `openpiv_rust` if installed, falls back cleanly to Python/SciPy. + - `"rust"`: Uses `openpiv_rust`. Raises informative `ImportError` if not compiled/installed. + - `"scipy"` / `"python"`: Forces the pure Python/SciPy reference path. +- **Test Suite Safety**: Rust-specific tests must use `openpiv_rust = pytest.importorskip("openpiv_rust")` at the module top level so test suites pass cleanly in environments without the Rust compiler. diff --git a/benchmarks/fft_backends/RESULTS.md b/benchmarks/fft_backends/RESULTS.md new file mode 100644 index 00000000..6a2f926d --- /dev/null +++ b/benchmarks/fft_backends/RESULTS.md @@ -0,0 +1,73 @@ +# FFT backend exploration (branch `explore/fft-backends`) + +Context: v0.25.5 switched `openpiv.pyprocess.fft_correlate_images` from +`numpy.fft` to `scipy.fft`, a ~2-3x win with no new dependency (see +CHANGES.txt). This branch checks whether `pyFFTW`, `rocket-fft` (numba), or +a hand-rolled Cython FFT wrapper could beat `scipy.fft` by enough to justify +a new (and, for pyFFTW/Cython+FFTW, compiled) dependency. + +## Setup + +Windows 11, 8 logical cores. Benchmarked on the same batched-window shapes +`fft_correlate_images` actually produces (many small windows stacked along +axis 0, FFT over the last two axes) — not a generic single large-array FFT +benchmark, since that's not our access pattern. + +## `bench_single_call.py` — cold-ish, one-shot calls + +| shape | numpy.fft | scipy.fft | pyfftw (ESTIMATE, cached) | rocket-fft (numba) | +|---|---|---|---|---| +| (2145,63,63) finest pass | 813ms | 421ms | 399ms (1.05x) | 1088ms (0.39x) | +| (660,31,31) windowsize=32 | 170ms | 47ms | 48ms (0.97x) | 107ms (0.44x) | +| (143,127,127) windowsize=64 padded | 661ms | 353ms | 399ms (0.88x) | 756ms (0.47x) | + +All outputs numerically match (`np.allclose`, atol=1e-6). + +- **pyFFTW with default (ESTIMATE) planning is a wash against `scipy.fft`** — + sometimes marginally faster, sometimes slower, never a clear win. +- **rocket-fft (numba-jitted `numpy.fft`) is consistently slower**, 0.4-0.5x. + Numba's FFT implementation doesn't compete with scipy's batched C++ + (`duccfft`) backend for this shape of workload. Not worth pursuing. + +## `bench_repeated_calls.py` — many calls at a fixed shape (batch job) + +Simulates the realistic case: a batch run processes many image pairs, so the +same window shape recurs many times per pass. This is the scenario where +FFTW's expensive `MEASURE`/wisdom planning is supposed to pay for itself. + +At (660,31,31), 40 repeated calls: + +- `scipy.fft`: 55.5 ms/call +- `pyfftw` (ESTIMATE, cached): 54.4 ms/call (1.02x) +- `pyfftw` (`FFTW_MEASURE`, plan reused): **921 ms one-time plan cost**, then + 51.8 ms/call steady-state (1.07x vs scipy.fft) +- `pyfftw` (`FFTW_MEASURE`) amortized over just 40 calls: **0.74x** — net + *slower* than scipy.fft, because the one-time plan-build cost dominates. + +Even in the best case (steady-state after the plan is built), pyFFTW is only +~7% faster than scipy.fft. Breaking even on the 921ms `MEASURE` planning +cost alone takes roughly 250+ image pairs at this window size — and a real +multipass run uses several *different* window shapes per image pair, so +each shape would need its own amortized plan. + +## Conclusion + +`scipy.fft` (already merged, v0.25.5) is the best cost/benefit choice: + +- Beats `numpy.fft` by 2-3x, `rocket-fft` by ~2x, with **zero new + dependencies** (scipy is already required). +- `pyFFTW` offers at best a ~7% steady-state edge, erased by its own + planning overhead unless a batch run reuses one window shape hundreds of + times — not representative of typical multipass PIV runs (several + different window sizes per pass). +- A hand-rolled Cython FFT wrapper would, at best, match pyFFTW (same + underlying FFTW library, same planning cost) — i.e. still short of + scipy.fft's real-world advantage, while adding a C build dependency this + project deliberately removed (see `CLAUDE.md`: "no Cython extensions"). + **Not built** — the benchmark numbers above rule it out before writing any + Cython. + +**Recommendation: do not merge this branch's dependencies into `pyproject.toml`.** +Keep `scipy.fft` as shipped in v0.25.5. `pyfftw`/`rocket-fft`/`numba` were +installed only in this branch's throwaway `uv pip install` for benchmarking +and are not added to `pyproject.toml`. diff --git a/benchmarks/fft_backends/bench_repeated_calls.py b/benchmarks/fft_backends/bench_repeated_calls.py new file mode 100644 index 00000000..edb7a6d9 --- /dev/null +++ b/benchmarks/fft_backends/bench_repeated_calls.py @@ -0,0 +1,86 @@ +"""Benchmark FFT backends across many repeated calls of the SAME shape, +mimicking a batch job processing many image pairs with fixed window sizes +per pass -- the scenario where FFTW's MEASURE/wisdom planning should pay off +its upfront planning cost. Throwaway script for explore/fft-backends. +""" +import time +import numpy as np + +N_CALLS = 40 # ~ processing 40 image pairs at one pass/window size + + +def main(): + rng = np.random.default_rng(0) + N, h, w = 660, 31, 31 # windowsize=32 batch, a common case + pairs = [(rng.random((N, h, w)), rng.random((N, h, w))) for _ in range(N_CALLS)] + + import scipy.fft as spfft + + def run_scipy(a, b): + fa = spfft.rfft2(a, axes=(-2, -1)) + fb = spfft.rfft2(b, axes=(-2, -1)) + return spfft.irfft2(np.conj(fa) * fb) + + t0 = time.perf_counter() + for a, b in pairs: + run_scipy(a, b) + t_scipy = time.perf_counter() - t0 + print(f"scipy.fft: {t_scipy*1000:8.1f} ms total, {t_scipy/N_CALLS*1000:6.2f} ms/call") + + import pyfftw + + pyfftw.interfaces.cache.enable() + pyfftw.interfaces.cache.set_keepalive_time(30) + + def run_pyfftw_estimate(a, b): + fa = pyfftw.interfaces.numpy_fft.rfft2(a, axes=(-2, -1), threads=1) + fb = pyfftw.interfaces.numpy_fft.rfft2(b, axes=(-2, -1), threads=1) + return pyfftw.interfaces.numpy_fft.irfft2(np.conj(fa) * fb, threads=1) + + t0 = time.perf_counter() + for a, b in pairs: + run_pyfftw_estimate(a, b) + t_estimate = time.perf_counter() - t0 + print(f"pyfftw (ESTIMATE, cached): {t_estimate*1000:8.1f} ms total, {t_estimate/N_CALLS*1000:6.2f} ms/call") + + # Build explicit MEASURE-planned FFTW objects once, reuse the buffers for + # every call -- this is the "amortize planning across a batch job" path. + a_in = pyfftw.empty_aligned((N, h, w), dtype='float64') + b_in = pyfftw.empty_aligned((N, h, w), dtype='float64') + fshape = (N, h, w // 2 + 1) + a_out = pyfftw.empty_aligned(fshape, dtype='complex128') + b_out = pyfftw.empty_aligned(fshape, dtype='complex128') + corr_in = pyfftw.empty_aligned(fshape, dtype='complex128') + corr_out = pyfftw.empty_aligned((N, h, w), dtype='float64') + + t0 = time.perf_counter() + fft_a = pyfftw.FFTW(a_in, a_out, axes=(-2, -1), direction='FFTW_FORWARD', flags=('FFTW_MEASURE',), threads=1) + fft_b = pyfftw.FFTW(b_in, b_out, axes=(-2, -1), direction='FFTW_FORWARD', flags=('FFTW_MEASURE',), threads=1) + ifft_c = pyfftw.FFTW(corr_in, corr_out, axes=(-2, -1), direction='FFTW_BACKWARD', flags=('FFTW_MEASURE',), threads=1) + plan_time = time.perf_counter() - t0 + print(f"pyfftw MEASURE plan build: {plan_time*1000:8.1f} ms (one-time)") + + def run_pyfftw_measure(a, b): + a_in[:] = a + b_in[:] = b + fft_a() + fft_b() + corr_in[:] = np.conj(a_out) * b_out + ifft_c() + return corr_out.copy() + + t0 = time.perf_counter() + for a, b in pairs: + run_pyfftw_measure(a, b) + t_measure = time.perf_counter() - t0 + print(f"pyfftw (MEASURE, reused): {t_measure*1000:8.1f} ms total, {t_measure/N_CALLS*1000:6.2f} ms/call " + f"(+ {plan_time*1000:.1f} ms one-time plan)") + + print() + print(f"speedup vs scipy.fft: pyfftw ESTIMATE = {t_scipy/t_estimate:.2f}x, " + f"pyfftw MEASURE (amortized over {N_CALLS} calls) = {t_scipy/(t_measure+plan_time):.2f}x, " + f"pyfftw MEASURE (steady-state only) = {t_scipy/N_CALLS/(t_measure/N_CALLS):.2f}x") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/fft_backends/bench_single_call.py b/benchmarks/fft_backends/bench_single_call.py new file mode 100644 index 00000000..d9ebc081 --- /dev/null +++ b/benchmarks/fft_backends/bench_single_call.py @@ -0,0 +1,110 @@ +"""Benchmark FFT backends for openpiv's batched cross-correlation pattern. + +Not part of the package; a throwaway script for the explore/fft-backends +branch. Run with: uv run python scripts_bench_fft.py +""" +import time +import numpy as np + + +def bench(a, b, run, reps=5): + run(a, b) # warmup / JIT + run(a, b) + t0 = time.perf_counter() + for _ in range(reps): + out = run(a, b) + dt = (time.perf_counter() - t0) / reps + return dt, out + + +def make_backend_numpy(): + import numpy.fft as npfft + + def run(a, b): + fa = npfft.rfft2(a, axes=(-2, -1)) + fb = npfft.rfft2(b, axes=(-2, -1)) + return npfft.irfft2(np.conj(fa) * fb) + + return run + + +def make_backend_scipy(): + import scipy.fft as spfft + + def run(a, b): + fa = spfft.rfft2(a, axes=(-2, -1)) + fb = spfft.rfft2(b, axes=(-2, -1)) + return spfft.irfft2(np.conj(fa) * fb) + + return run + + +def make_backend_pyfftw(): + import pyfftw + + pyfftw.interfaces.cache.enable() + pyfftw.interfaces.cache.set_keepalive_time(30) + + def run(a, b): + fa = pyfftw.interfaces.numpy_fft.rfft2(a, axes=(-2, -1), threads=1) + fb = pyfftw.interfaces.numpy_fft.rfft2(b, axes=(-2, -1), threads=1) + return pyfftw.interfaces.numpy_fft.irfft2(np.conj(fa) * fb, threads=1) + + return run + + +def make_backend_rocketfft(): + # rocket-fft only patches numpy.fft *inside numba-jitted functions*. + import numba + + @numba.njit(cache=True) + def _corr(a, b): + fa = np.fft.rfft2(a) + fb = np.fft.rfft2(b) + return np.fft.irfft2(np.conj(fa) * fb) + + def run(a, b): + return _corr(a, b) + + return run + + +SHAPES = [ + (2145, 63, 63, "finest pass (windowsize=6, many small windows)"), + (660, 31, 31, "windowsize=32"), + (143, 127, 127, "windowsize=64, linear-padded"), +] + + +def main(): + rng = np.random.default_rng(0) + for N, h, w, label in SHAPES: + print(f"\n=== {label}: shape=({N},{h},{w}) ===") + a = rng.random((N, h, w)) + b = rng.random((N, h, w)) + + results = {} + for backend_name, factory in [ + ("numpy.fft", make_backend_numpy), + ("scipy.fft", make_backend_scipy), + ("pyfftw", make_backend_pyfftw), + ("rocket-fft(numba)", make_backend_rocketfft), + ]: + try: + run = factory() + dt, out = bench(a, b, run) + results[backend_name] = (dt, out) + print(f" {backend_name:20s} {dt*1000:8.2f} ms") + except Exception as e: # noqa: BLE001 - benchmark script, report and continue + print(f" {backend_name:20s} FAILED: {e}") + + if "scipy.fft" in results: + ref = results["scipy.fft"][1] + for name, (dt, out) in results.items(): + ok = np.allclose(out, ref, atol=1e-6) + base = results["scipy.fft"][0] + print(f" {name:20s} speedup vs scipy.fft: {base/dt:5.2f}x matches: {ok}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/fft_backends/benchmark_windef.py b/benchmarks/fft_backends/benchmark_windef.py new file mode 100644 index 00000000..863381a5 --- /dev/null +++ b/benchmarks/fft_backends/benchmark_windef.py @@ -0,0 +1,81 @@ +import time +import tempfile +import pathlib +import numpy as np +from imageio.v3 import imwrite +from openpiv import windef +from openpiv.settings import PIVSettings +from openpiv.test import test_process + +def run_multigrid(frame_a, frame_b, backend, n_runs=3): + settings = PIVSettings() + settings.windowsizes = (64, 32, 16) + settings.overlap = (32, 16, 8) + settings.num_iterations = 3 + settings.backend = backend + settings.sig2noise_validate = False + settings.show_all_plots = False + settings.show_plot = False + + # Warmup + windef.multigrid_windef(frame_a, frame_b, settings) + + times = [] + for _ in range(n_runs): + t0 = time.perf_counter() + x, y, u, v, flags = windef.multigrid_windef(frame_a, frame_b, settings) + times.append(time.perf_counter() - t0) + return min(times) * 1000.0, (u, v) + +def run_mp(frame_a, frame_b, backend): + with tempfile.TemporaryDirectory() as tmpdir: + tmp = pathlib.Path(tmpdir) + img_dir = tmp / "images" + img_dir.mkdir() + out_dir = tmp / "out" + out_dir.mkdir() + for i in range(4): + imwrite(img_dir / f"pair_{i:03d}_a.tif", frame_a.astype(np.uint8)) + imwrite(img_dir / f"pair_{i:03d}_b.tif", frame_b.astype(np.uint8)) + settings = PIVSettings() + settings.filepath_images = img_dir + settings.save_path = out_dir + settings.save_folder_suffix = f"{backend}_bench" + settings.frame_pattern_a = "pair_*_a.tif" + settings.frame_pattern_b = "pair_*_b.tif" + settings.windowsizes = (64, 32, 16) + settings.overlap = (32, 16, 8) + settings.num_iterations = 3 + settings.backend = backend + settings.n_cpus = 2 + settings.show_plot = False + settings.save_plot = False + settings.show_all_plots = False + settings.sig2noise_validate = False + + t0 = time.perf_counter() + windef.piv(settings) + return (time.perf_counter() - t0) * 1000.0 + +def main(): + frame_a, frame_b = test_process.create_pair(image_size=256) + + print("=" * 60) + print(" BENCHMARK: Multigrid Window Deformation (256x256, 3 passes)") + print("=" * 60) + time_scipy, (u_s, v_s) = run_multigrid(frame_a, frame_b, "scipy") + time_rust, (u_r, v_r) = run_multigrid(frame_a, frame_b, "rust") + speedup_mg = time_scipy / time_rust + max_diff = max(np.nanmax(np.abs(u_s - u_r)), np.nanmax(np.abs(v_s - v_r))) + print(f"Multigrid SciPy: {time_scipy:6.2f} ms | Rust: {time_rust:6.2f} ms | Speedup: {speedup_mg:4.2f}x | Max Diff: {max_diff:.2e}") + + print("\n" + "=" * 60) + print(" BENCHMARK: Multiprocessing (4 pairs, 2 worker processes)") + print("=" * 60) + mp_scipy = run_mp(frame_a, frame_b, "scipy") + mp_rust = run_mp(frame_a, frame_b, "rust") + speedup_mp = mp_scipy / mp_rust + print(f"Multiprocessing SciPy: {mp_scipy:6.2f} ms | Rust: {mp_rust:6.2f} ms | Speedup: {speedup_mp:4.2f}x") + +if __name__ == "__main__": + main() diff --git a/benchmarks/fft_backends/run_piv_quiver_demo.py b/benchmarks/fft_backends/run_piv_quiver_demo.py new file mode 100644 index 00000000..b5473f78 --- /dev/null +++ b/benchmarks/fft_backends/run_piv_quiver_demo.py @@ -0,0 +1,205 @@ +import os +import time +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path +from importlib.resources import files +from openpiv import pyprocess, tools, validation, filters +import openpiv_rust + +ARTIFACT_DIR = Path(r"C:\Users\alex\.gemini\antigravity-cli\brain\371a4578-c887-4c4a-bba3-7bbfb7f1997c") +ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) + +def run_real_piv_demo(): + print("=== Running Real PIV Quiver Demo (exp1_001) ===") + path = files('openpiv') / "data" / "test1" + frame_a = tools.imread(path / "exp1_001_a.bmp").astype(np.int32) + frame_b = tools.imread(path / "exp1_001_b.bmp").astype(np.int32) + + window_size = 32 + overlap = 16 + search_area_size = 32 + + # Prepare window arrays + aa = pyprocess.sliding_window_array(frame_a, (search_area_size, search_area_size), (overlap, overlap)) + bb = pyprocess.sliding_window_array(frame_b, (search_area_size, search_area_size), (overlap, overlap)) + n_rows, n_cols = pyprocess.get_field_shape(frame_a.shape, (search_area_size, search_area_size), (overlap, overlap)) + + # 1. Scipy correlation + t0 = time.perf_counter() + scipy_corr = pyprocess.fft_correlate_images(aa, bb, correlation_method="circular", normalized_correlation=False) + t_scipy_corr = time.perf_counter() - t0 + + # 2. Rust correlation + t0 = time.perf_counter() + rust_corr = openpiv_rust.fft_correlate_circular(aa.astype(float), bb.astype(float), normalized_correlation=False) + t_rust_corr = time.perf_counter() - t0 + + speedup_corr = t_scipy_corr / max(t_rust_corr, 1e-6) + print(f"Correlation Time -> SciPy: {t_scipy_corr*1000:.2f} ms | Rust: {t_rust_corr*1000:.2f} ms | Speedup: {speedup_corr:.1f}x") + + # Full End-to-end PIV + t0 = time.perf_counter() + u_scipy, v_scipy, s2n_scipy = pyprocess.extended_search_area_piv( + frame_a, frame_b, + window_size=window_size, + overlap=overlap, + search_area_size=search_area_size, + correlation_method="circular", + sig2noise_method="peak2peak", + ) + t_scipy_full = time.perf_counter() - t0 + + t0 = time.perf_counter() + rust_corr = openpiv_rust.fft_correlate_circular(aa.astype(float), bb.astype(float), normalized_correlation=False) + u_rust, v_rust = pyprocess.correlation_to_displacement(rust_corr, n_rows, n_cols) + s2n_rust = pyprocess.sig2noise_ratio(rust_corr, sig2noise_method="peak2peak", width=2) + t_rust_full = time.perf_counter() - t0 + + diff_u = np.nanmax(np.abs(u_scipy - u_rust)) + diff_v = np.nanmax(np.abs(v_scipy - v_rust)) + print(f"Full PIV Time -> SciPy: {t_scipy_full*1000:.2f} ms | Rust: {t_rust_full*1000:.2f} ms") + print(f"Max difference in velocity field: u={diff_u:.2e}, v={diff_v:.2e}") + + # Coordinates + x, y = pyprocess.get_coordinates(image_size=frame_a.shape, search_area_size=search_area_size, overlap=overlap) + + # Outlier filtering + flags = validation.sig2noise_val(s2n_rust, threshold=1.05) + flags_2d = flags.reshape(n_rows, n_cols) + u_clean, v_clean = filters.replace_outliers(u_rust, v_rust, flags_2d, method='localmean', max_iter=3, kernel_size=2) + + # Plot Quiver Comparison + fig, axes = plt.subplots(1, 3, figsize=(18, 5.5), facecolor="white") + + # Panel 1: Scipy + ax = axes[0] + ax.imshow(frame_a, cmap="gray", alpha=0.5, origin="upper") + speed_scipy = np.sqrt(u_scipy**2 + v_scipy**2) + q1 = ax.quiver(x, y, u_scipy, -v_scipy, speed_scipy, cmap="plasma", scale=45, width=0.004) + ax.set_title(f"SciPy Mode Vector Field\nCorr: {t_scipy_corr*1000:.1f} ms | Total: {t_scipy_full*1000:.1f} ms", fontsize=12, fontweight="bold") + ax.set_xlabel("x [pixels]") + ax.set_ylabel("y [pixels]") + plt.colorbar(q1, ax=ax, label="Displacement [px]", fraction=0.046, pad=0.04) + + # Panel 2: Rust + ax = axes[1] + ax.imshow(frame_a, cmap="gray", alpha=0.5, origin="upper") + speed_rust = np.sqrt(u_rust**2 + v_rust**2) + q2 = ax.quiver(x, y, u_rust, -v_rust, speed_rust, cmap="plasma", scale=45, width=0.004) + ax.set_title(f"Rust Mode Vector Field (Rayon 2D RealFFT)\nCorr: {t_rust_corr*1000:.1f} ms ({speedup_corr:.1f}x Faster)", fontsize=12, fontweight="bold") + ax.set_xlabel("x [pixels]") + ax.set_ylabel("y [pixels]") + plt.colorbar(q2, ax=ax, label="Displacement [px]", fraction=0.046, pad=0.04) + + # Panel 3: Difference Map + ax = axes[2] + diff_speed = np.sqrt((u_scipy - u_rust)**2 + (v_scipy - v_rust)**2) + im = ax.imshow(diff_speed, cmap="coolwarm", origin="upper", extent=[x.min(), x.max(), y.max(), y.min()]) + ax.set_title(f"Vector Discrepancy (|u_rust - u_scipy|)\nMax Absolute Error: {np.nanmax(diff_speed):.2e} px", fontsize=12, fontweight="bold") + ax.set_xlabel("x [pixels]") + ax.set_ylabel("y [pixels]") + plt.colorbar(im, ax=ax, label="Difference [px]", fraction=0.046, pad=0.04) + + plt.tight_layout() + out_path = ARTIFACT_DIR / "quiver_real_piv.png" + plt.savefig(out_path, dpi=180) + plt.close() + print(f"Saved real PIV quiver plot to: {out_path}") + +def run_synthetic_vortex_demo(): + print("\n=== Running Synthetic Vortex Quiver Demo ===") + np.random.seed(42) + img_h, img_w = 256, 256 + + n_particles = 4000 + px = np.random.uniform(0, img_w, n_particles) + py = np.random.uniform(0, img_h, n_particles) + + cx, cy = 128.0, 128.0 + r0 = 40.0 + gamma = 400.0 + + rx = px - cx + ry = py - cy + r = np.sqrt(rx**2 + ry**2) + 1e-6 + v_theta = gamma * r / (r0**2 + r**2) + + dx_particles = -v_theta * (ry / r) + dy_particles = v_theta * (rx / r) + + frame_a = np.zeros((img_h, img_w), dtype=np.float32) + frame_b = np.zeros((img_h, img_w), dtype=np.float32) + + for i in range(n_particles): + x0, y0 = int(round(px[i])), int(round(py[i])) + if 1 <= x0 < img_w - 1 and 1 <= y0 < img_h - 1: + frame_a[y0, x0] += 200.0 + frame_a[y0+1, x0] += 100.0 + frame_a[y0-1, x0] += 100.0 + frame_a[y0, x0+1] += 100.0 + frame_a[y0, x0-1] += 100.0 + + x1, y1 = int(round(px[i] + dx_particles[i])), int(round(py[i] + dy_particles[i])) + if 1 <= x1 < img_w - 1 and 1 <= y1 < img_h - 1: + frame_b[y1, x1] += 200.0 + frame_b[y1+1, x1] += 100.0 + frame_b[y1-1, x1] += 100.0 + frame_b[y1, x1+1] += 100.0 + frame_b[y1, x1-1] += 100.0 + + window_size = 32 + overlap = 16 + search_area_size = 32 + + aa = pyprocess.sliding_window_array(frame_a, (search_area_size, search_area_size), (overlap, overlap)) + bb = pyprocess.sliding_window_array(frame_b, (search_area_size, search_area_size), (overlap, overlap)) + n_rows, n_cols = pyprocess.get_field_shape(frame_a.shape, (search_area_size, search_area_size), (overlap, overlap)) + + # Scipy correlation + t0 = time.perf_counter() + scipy_corr = pyprocess.fft_correlate_images(aa, bb, correlation_method="circular", normalized_correlation=False) + t_scipy_corr = time.perf_counter() - t0 + + # Rust correlation + t0 = time.perf_counter() + rust_corr = openpiv_rust.fft_correlate_circular(aa.astype(float), bb.astype(float), normalized_correlation=False) + t_rust_corr = time.perf_counter() - t0 + + speedup_vortex = t_scipy_corr / max(t_rust_corr, 1e-6) + print(f"Vortex Correlation -> SciPy: {t_scipy_corr*1000:.2f} ms | Rust: {t_rust_corr*1000:.2f} ms | Speedup: {speedup_vortex:.1f}x") + + u_rust, v_rust = pyprocess.correlation_to_displacement(rust_corr, n_rows, n_cols) + x, y = pyprocess.get_coordinates(image_size=frame_a.shape, search_area_size=search_area_size, overlap=overlap) + + fig, axes = plt.subplots(1, 2, figsize=(14, 6), facecolor="white") + + # Panel 1: Rust Quiver Field + ax = axes[0] + ax.imshow(frame_a, cmap="gray", alpha=0.5, origin="upper") + speed = np.sqrt(u_rust**2 + v_rust**2) + q = ax.quiver(x, y, u_rust, -v_rust, speed, cmap="inferno", scale=25, width=0.005) + ax.set_title(f"Rust Mode Quiver - Lamb-Oseen Vortex Flow\nFFT Correlation: {t_rust_corr*1000:.2f} ms ({speedup_vortex:.1f}x Faster)", fontsize=12, fontweight="bold") + ax.set_xlabel("x [pixels]") + ax.set_ylabel("y [pixels]") + plt.colorbar(q, ax=ax, label="Displacement Velocity [px]") + + # Panel 2: Streamlines + ax = axes[1] + strm = ax.streamplot(x[0, :], y[:, 0], u_rust, -v_rust, color=speed, cmap="inferno", density=1.5, linewidth=1.5) + ax.set_title("Streamlines Recovered by Rust Engine", fontsize=12, fontweight="bold") + ax.set_xlabel("x [pixels]") + ax.set_ylabel("y [pixels]") + ax.set_xlim(0, img_w) + ax.set_ylim(img_h, 0) + plt.colorbar(strm.lines, ax=ax, label="Velocity Magnitude [px]") + + plt.tight_layout() + out_path = ARTIFACT_DIR / "quiver_vortex_piv.png" + plt.savefig(out_path, dpi=180) + plt.close() + print(f"Saved vortex quiver plot to: {out_path}") + +if __name__ == "__main__": + run_real_piv_demo() + run_synthetic_vortex_demo() diff --git a/benchmarks/fft_backends/test_rust_fft.py b/benchmarks/fft_backends/test_rust_fft.py new file mode 100644 index 00000000..e50baa19 --- /dev/null +++ b/benchmarks/fft_backends/test_rust_fft.py @@ -0,0 +1,53 @@ +import sys +import time +import numpy as np +from openpiv import pyprocess +import openpiv_rust + +print('=== OpenPIV Rust vs scipy.fft Benchmark & Validation ===', flush=True) + +# Shapes representing real PIV interrogation window batches +shapes = [ + (660, 32, 32, 'window_size=32 (medium batch)'), + (2145, 64, 64, 'window_size=64 (fine pass / large batch)'), + (143, 128, 128, 'window_size=128 (coarse pass / large windows)'), +] + +rng = np.random.default_rng(42) + +for N, H, W, label in shapes: + print(f'\n--- Testing {label}: shape=({N}, {H}, {W}) ---', flush=True) + a = rng.random((N, H, W), dtype=np.float64) + b = rng.random((N, H, W), dtype=np.float64) + + # 1. Correctness check (raw correlation without python normalization overhead) + scipy_out = pyprocess.fft_correlate_images(a, b, correlation_method='circular', normalized_correlation=False) + rust_out = openpiv_rust.fft_correlate_circular(a, b, normalized_correlation=False) + + max_diff = np.max(np.abs(scipy_out - rust_out)) + is_close = np.allclose(scipy_out, rust_out, atol=1e-5) + print(f'Numerical validation: max_diff = {max_diff:.2e} | matches np.allclose: {is_close}', flush=True) + + # 2. Timing benchmark + reps = 10 + + # Warmup + for _ in range(2): + pyprocess.fft_correlate_images(a, b, correlation_method='circular', normalized_correlation=False) + openpiv_rust.fft_correlate_circular(a, b, normalized_correlation=False) + + t0 = time.perf_counter() + for _ in range(reps): + scipy_res = pyprocess.fft_correlate_images(a, b, correlation_method='circular', normalized_correlation=False) + t_scipy = (time.perf_counter() - t0) / reps + + t0 = time.perf_counter() + for _ in range(reps): + rust_res = openpiv_rust.fft_correlate_circular(a, b, normalized_correlation=False) + t_rust = (time.perf_counter() - t0) / reps + + print(f' scipy.fft (C++ PocketFFT): {t_scipy * 1000:7.2f} ms', flush=True) + print(f' Rust (Rayon + RustFFT): {t_rust * 1000:7.2f} ms', flush=True) + print(f' Speedup vs scipy.fft: {t_scipy / t_rust:7.2f}x', flush=True) + +print('\nBenchmark completed.', flush=True) diff --git a/benchmarks/run_benchmarks.py b/benchmarks/run_benchmarks.py new file mode 100644 index 00000000..d37badca --- /dev/null +++ b/benchmarks/run_benchmarks.py @@ -0,0 +1,118 @@ +"""OpenPIV Performance & Backend Benchmarking Suite + +Run this script to benchmark cross-correlation, subpixel peak finding, +and end-to-end multi-pass windef deformation across available backends. + +Usage: + python benchmarks/run_benchmarks.py +""" +import time +import numpy as np +from openpiv import pyprocess, windef +from openpiv.settings import PIVSettings + +try: + import openpiv_rust + HAS_RUST = True +except ImportError: + HAS_RUST = False + +def benchmark_correlation(): + print("\n" + "=" * 60) + print(" 1. CROSS-CORRELATION BENCHMARK (225 windows of 32x32)") + print("=" * 60) + rng = np.random.default_rng(42) + a = rng.random((225, 32, 32), dtype=np.float64) + b = rng.random((225, 32, 32), dtype=np.float64) + + # Circular - SciPy + t0 = time.perf_counter() + for _ in range(10): + pyprocess.fft_correlate_images(a, b, correlation_method="circular", backend="scipy", normalized_correlation=False) + t_scipy_circ = (time.perf_counter() - t0) * 100 + + # Linear - SciPy (64x64 power-of-2) + t0 = time.perf_counter() + for _ in range(10): + pyprocess.fft_correlate_images(a, b, correlation_method="linear", backend="scipy", normalized_correlation=False) + t_scipy_lin = (time.perf_counter() - t0) * 100 + + print(f" SciPy Circular: {t_scipy_circ:.2f} ms") + print(f" SciPy Linear: {t_scipy_lin:.2f} ms") + + if HAS_RUST: + t0 = time.perf_counter() + for _ in range(10): + pyprocess.fft_correlate_images(a, b, correlation_method="circular", backend="rust", normalized_correlation=False) + t_rust_circ = (time.perf_counter() - t0) * 100 + + t0 = time.perf_counter() + for _ in range(10): + pyprocess.fft_correlate_images(a, b, correlation_method="linear", backend="rust", normalized_correlation=False) + t_rust_lin = (time.perf_counter() - t0) * 100 + + print(f" Rust Circular: {t_rust_circ:.2f} ms (speedup: {t_scipy_circ/t_rust_circ:.2f}x)") + print(f" Rust Linear: {t_rust_lin:.2f} ms (speedup: {t_scipy_lin/t_rust_lin:.2f}x)") + else: + print(" Rust backend: NOT INSTALLED (run 'maturin develop' in crates/openpiv_rust)") + +def benchmark_subpixel(): + print("\n" + "=" * 60) + print(" 2. SUBPIXEL PEAK POSITION BENCHMARK (961 windows)") + print("=" * 60) + rng = np.random.default_rng(42) + corr = rng.random((961, 33, 33), dtype=np.float64) + + t0 = time.perf_counter() + for _ in range(5): + pyprocess.correlation_to_displacement(corr, 31, 31, subpixel_method="gaussian") + t_py = (time.perf_counter() - t0) * 200 + + print(f" Active displacement engine: {t_py:.2f} ms for 961 windows") + if HAS_RUST: + print(" (Using parallel Rust subpixel peak finder via openpiv_rust)") + else: + print(" (Using Python loop fallback)") + +def benchmark_windef(): + print("\n" + "=" * 60) + print(" 3. END-TO-END MULTIGRID WINDEF (3 passes on 256x256 frames)") + print("=" * 60) + from openpiv.test.test_process import create_pair + frame_a, frame_b = create_pair(image_size=256) + + settings = PIVSettings() + settings.windowsizes = (64, 32, 16) + settings.overlap = (32, 16, 8) + settings.num_iterations = 3 + settings.sig2noise_validate = False + settings.show_all_plots = False + settings.show_plot = False + + # SciPy run + settings.backend = "scipy" + windef.multigrid_windef(frame_a, frame_b, settings) # warmup + t0 = time.perf_counter() + for _ in range(3): + windef.multigrid_windef(frame_a, frame_b, settings) + t_scipy = (time.perf_counter() - t0) / 3 * 1000 + + print(f" SciPy 3-pass Windef: {t_scipy:.2f} ms") + + if HAS_RUST: + settings.backend = "rust" + windef.multigrid_windef(frame_a, frame_b, settings) # warmup + t0 = time.perf_counter() + for _ in range(3): + windef.multigrid_windef(frame_a, frame_b, settings) + t_rust = (time.perf_counter() - t0) / 3 * 1000 + print(f" Rust 3-pass Windef: {t_rust:.2f} ms (speedup: {t_scipy/t_rust:.2f}x)") + +if __name__ == "__main__": + print(f"OpenPIV Benchmark Suite | Rust Backend Available: {HAS_RUST}") + benchmark_correlation() + benchmark_subpixel() + benchmark_windef() + print("\n" + "=" * 60) + print(" Benchmarks Completed Successfully!") + print("=" * 60 + "\n") diff --git a/crates/openpiv_rust/Cargo.lock b/crates/openpiv_rust/Cargo.lock new file mode 100644 index 00000000..1eb8255e --- /dev/null +++ b/crates/openpiv_rust/Cargo.lock @@ -0,0 +1,331 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "either" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "numpy" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "openpiv_rust" +version = "0.1.0" +dependencies = [ + "num-complex", + "numpy", + "pyo3", + "rayon", + "realfft", + "rustfft", +] + +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "primal-check" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08" +dependencies = [ + "num-integer", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "realfft" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f821338fddb99d089116342c46e9f1fbf3828dba077674613e734e01d6ea8677" +dependencies = [ + "rustfft", +] + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustfft" +version = "6.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21db5f9893e91f41798c88680037dba611ca6674703c1a18601b01a72c8adb89" +dependencies = [ + "num-complex", + "num-integer", + "num-traits", + "primal-check", + "strength_reduce", + "transpose", +] + +[[package]] +name = "strength_reduce" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "transpose" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad61aed86bc3faea4300c7aee358b4c6d0c8d6ccc36524c96e4c92ccf26e77e" +dependencies = [ + "num-integer", + "strength_reduce", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" diff --git a/crates/openpiv_rust/Cargo.toml b/crates/openpiv_rust/Cargo.toml new file mode 100644 index 00000000..b0c56d69 --- /dev/null +++ b/crates/openpiv_rust/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "openpiv_rust" +version = "0.1.0" +edition = "2021" + +[lib] +name = "openpiv_rust" +crate-type = ["cdylib", "rlib"] + +[dependencies] +pyo3 = { version = "0.29", features = ["extension-module"] } +numpy = "0.29" +rustfft = "6.2" +realfft = "3.4" +rayon = "1.10" +num-complex = "0.4" diff --git a/crates/openpiv_rust/src/lib.rs b/crates/openpiv_rust/src/lib.rs new file mode 100644 index 00000000..5dcbc5a2 --- /dev/null +++ b/crates/openpiv_rust/src/lib.rs @@ -0,0 +1,1050 @@ +use num_complex::Complex; +use numpy::{ + PyArray1, PyArray2, PyArray3, PyArrayMethods, PyReadonlyArray2, PyReadonlyArray3, ToPyArray, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use rayon::prelude::*; +use realfft::{ComplexToReal, RealFftPlanner, RealToComplex}; +use rustfft::{Fft, FftPlanner}; +use std::borrow::Cow; +use std::sync::Arc; + +/// Safely extracts or linearizes an ndarray view into a C-contiguous slice [N, H, W]. +fn to_c_contiguous<'a>(arr: &'a numpy::ndarray::ArrayView3<'a, f64>) -> Cow<'a, [f64]> { + if let Some(s) = arr.as_slice() { + Cow::Borrowed(s) + } else { + let n = arr.shape()[0]; + let h = arr.shape()[1]; + let w = arr.shape()[2]; + let mut vec = Vec::with_capacity(n * h * w); + for i in 0..n { + for r in 0..h { + for c in 0..w { + vec.push(arr[[i, r, c]]); + } + } + } + Cow::Owned(vec) + } +} + +/// Safely extracts or linearizes an ndarray 2D view into a C-contiguous slice [H, W]. +fn to_c_contiguous_2d<'a>(arr: &'a numpy::ndarray::ArrayView2<'a, f64>) -> Cow<'a, [f64]> { + if let Some(s) = arr.as_slice() { + Cow::Borrowed(s) + } else { + let h = arr.shape()[0]; + let w = arr.shape()[1]; + let mut vec = Vec::with_capacity(h * w); + for r in 0..h { + for c in 0..w { + vec.push(arr[[r, c]]); + } + } + Cow::Owned(vec) + } +} + +/// Engine for Circular 2D cross-correlation (OpenPIV standard mode) +struct CircularEngine2D { + h: usize, + w: usize, + w_freq: usize, + r2c_row: Arc>, + c2r_row: Arc>, + c2c_col_fwd: Arc>, + c2c_col_inv: Arc>, +} + +struct CircularScratch { + freq_a: Vec>, + freq_b: Vec>, + row_scratch_fwd: Vec>, + col_scratch_fwd: Vec>, + row_scratch_inv: Vec>, + col_scratch_inv: Vec>, + col_buf: Vec>, + in_row_copy: Vec, + temp_real: Vec, +} + +impl CircularEngine2D { + fn new(h: usize, w: usize) -> Self { + let mut real_planner = RealFftPlanner::::new(); + let r2c_row = real_planner.plan_fft_forward(w); + let c2r_row = real_planner.plan_fft_inverse(w); + + let mut c2c_planner = FftPlanner::::new(); + let c2c_col_fwd = c2c_planner.plan_fft_forward(h); + let c2c_col_inv = c2c_planner.plan_fft_inverse(h); + + let w_freq = w / 2 + 1; + + Self { + h, + w, + w_freq, + r2c_row, + c2r_row, + c2c_col_fwd, + c2c_col_inv, + } + } + + fn create_scratch(&self) -> CircularScratch { + CircularScratch { + freq_a: vec![Complex::new(0.0, 0.0); self.h * self.w_freq], + freq_b: vec![Complex::new(0.0, 0.0); self.h * self.w_freq], + row_scratch_fwd: self.r2c_row.make_scratch_vec(), + col_scratch_fwd: vec![Complex::new(0.0, 0.0); self.c2c_col_fwd.get_inplace_scratch_len()], + row_scratch_inv: self.c2r_row.make_scratch_vec(), + col_scratch_inv: vec![Complex::new(0.0, 0.0); self.c2c_col_inv.get_inplace_scratch_len()], + col_buf: vec![Complex::new(0.0, 0.0); self.h], + in_row_copy: vec![0.0; self.w], + temp_real: vec![0.0; self.h * self.w], + } + } + + fn correlate_pair( + &self, + window_a: &[f64], + window_b: &[f64], + out_slice: &mut [f64], + normalized_correlation: bool, + scratch: &mut CircularScratch, + ) { + let h = self.h; + let w = self.w; + let w_freq = self.w_freq; + + // 1. Forward 2D Real FFT for window A + for r in 0..h { + scratch.in_row_copy.copy_from_slice(&window_a[r * w..(r + 1) * w]); + let out_row = &mut scratch.freq_a[r * w_freq..(r + 1) * w_freq]; + let _ = self.r2c_row.process_with_scratch(&mut scratch.in_row_copy, out_row, &mut scratch.row_scratch_fwd); + } + for c in 0..w_freq { + for r in 0..h { + scratch.col_buf[r] = scratch.freq_a[r * w_freq + c]; + } + self.c2c_col_fwd + .process_with_scratch(&mut scratch.col_buf, &mut scratch.col_scratch_fwd); + for r in 0..h { + scratch.freq_a[r * w_freq + c] = scratch.col_buf[r]; + } + } + + // 2. Forward 2D Real FFT for window B + for r in 0..h { + scratch.in_row_copy.copy_from_slice(&window_b[r * w..(r + 1) * w]); + let out_row = &mut scratch.freq_b[r * w_freq..(r + 1) * w_freq]; + let _ = self.r2c_row.process_with_scratch(&mut scratch.in_row_copy, out_row, &mut scratch.row_scratch_fwd); + } + for c in 0..w_freq { + for r in 0..h { + scratch.col_buf[r] = scratch.freq_b[r * w_freq + c]; + } + self.c2c_col_fwd + .process_with_scratch(&mut scratch.col_buf, &mut scratch.col_scratch_fwd); + for r in 0..h { + scratch.freq_b[r * w_freq + c] = scratch.col_buf[r]; + } + } + + // 3. Frequency domain cross-correlation: Fa.conj() * Fb + for i in 0..(h * w_freq) { + scratch.freq_a[i] = scratch.freq_a[i].conj() * scratch.freq_b[i]; + } + + // 4. Inverse 2D Real FFT on freq_a + for c in 0..w_freq { + for r in 0..h { + scratch.col_buf[r] = scratch.freq_a[r * w_freq + c]; + } + self.c2c_col_inv + .process_with_scratch(&mut scratch.col_buf, &mut scratch.col_scratch_inv); + for r in 0..h { + scratch.freq_a[r * w_freq + c] = scratch.col_buf[r]; + } + } + + for r in 0..h { + let in_freq_row = &mut scratch.freq_a[r * w_freq..(r + 1) * w_freq]; + // Enforce strictly 0 imaginary residuals at DC and Nyquist frequencies + in_freq_row[0].im = 0.0; + if w % 2 == 0 { + in_freq_row[w_freq - 1].im = 0.0; + } + let out_real_row = &mut scratch.temp_real[r * w..(r + 1) * w]; + let _ = self.c2r_row.process_with_scratch(in_freq_row, out_real_row, &mut scratch.row_scratch_inv); + } + + // 5. Normalization scale factor + 2D fftshift + // Note: realfft inverse FFT does not normalize by (h*w), whereas scipy.fft.irfft2 does. + // When normalized_correlation=True, scipy also divides by (h*w) again, requiring (1/(h*w))^2 in Rust. + let fft_norm = 1.0 / ((h * w) as f64); + let scale = if normalized_correlation { + fft_norm * fft_norm + } else { + fft_norm + }; + + let shift_r = h / 2; + let shift_c = w / 2; + + for r in 0..h { + let target_r = (r + shift_r) % h; + for c in 0..w { + let target_c = (c + shift_c) % w; + out_slice[target_r * w + target_c] = scratch.temp_real[r * w + c] * scale; + } + } + } +} + +struct LinearScratch { + freq_a: Vec>, + freq_b: Vec>, + row_scratch_fwd: Vec>, + col_scratch_fwd: Vec>, + row_scratch_inv: Vec>, + col_scratch_inv: Vec>, + col_buf: Vec>, + in_row_pad: Vec, + temp_real: Vec, +} + +/// Engine for Full Linear 2D cross-correlation (matches scipy.signal.correlate mode='full') +struct FullEngine2D { + win_h: usize, + win_w: usize, + out_h: usize, + out_w: usize, + fft_h: usize, + fft_w: usize, + w_freq: usize, + r2c_row: Arc>, + c2r_row: Arc>, + c2c_col_fwd: Arc>, + c2c_col_inv: Arc>, +} + +impl FullEngine2D { + fn new(win_h: usize, win_w: usize) -> Self { + let out_h = 2 * win_h - 1; + let out_w = 2 * win_w - 1; + + let fft_h = out_h.next_power_of_two(); + let fft_w = out_w.next_power_of_two(); + + let mut real_planner = RealFftPlanner::::new(); + let r2c_row = real_planner.plan_fft_forward(fft_w); + let c2r_row = real_planner.plan_fft_inverse(fft_w); + + let mut c2c_planner = FftPlanner::::new(); + let c2c_col_fwd = c2c_planner.plan_fft_forward(fft_h); + let c2c_col_inv = c2c_planner.plan_fft_inverse(fft_h); + + let w_freq = fft_w / 2 + 1; + + Self { + win_h, + win_w, + out_h, + out_w, + fft_h, + fft_w, + w_freq, + r2c_row, + c2r_row, + c2c_col_fwd, + c2c_col_inv, + } + } + + fn create_scratch(&self) -> LinearScratch { + LinearScratch { + freq_a: vec![Complex::new(0.0, 0.0); self.fft_h * self.w_freq], + freq_b: vec![Complex::new(0.0, 0.0); self.fft_h * self.w_freq], + row_scratch_fwd: self.r2c_row.make_scratch_vec(), + col_scratch_fwd: vec![Complex::new(0.0, 0.0); self.c2c_col_fwd.get_inplace_scratch_len()], + row_scratch_inv: self.c2r_row.make_scratch_vec(), + col_scratch_inv: vec![Complex::new(0.0, 0.0); self.c2c_col_inv.get_inplace_scratch_len()], + col_buf: vec![Complex::new(0.0, 0.0); self.fft_h], + in_row_pad: vec![0.0; self.fft_w], + temp_real: vec![0.0; self.fft_h * self.fft_w], + } + } + + fn correlate_pair(&self, window_a: &[f64], window_b: &[f64], out_slice: &mut [f64], scratch: &mut LinearScratch) { + let win_h = self.win_h; + let win_w = self.win_w; + let out_h = self.out_h; + let out_w = self.out_w; + let fft_h = self.fft_h; + let fft_w = self.fft_w; + let w_freq = self.w_freq; + + // 1. Forward 2D Real FFT for padded window A + for r in 0..fft_h { + scratch.in_row_pad.fill(0.0); + if r < win_h { + scratch.in_row_pad[..win_w].copy_from_slice(&window_a[r * win_w..(r + 1) * win_w]); + } + let out_row = &mut scratch.freq_a[r * w_freq..(r + 1) * w_freq]; + let _ = self.r2c_row.process_with_scratch(&mut scratch.in_row_pad, out_row, &mut scratch.row_scratch_fwd); + } + for c in 0..w_freq { + for r in 0..fft_h { + scratch.col_buf[r] = scratch.freq_a[r * w_freq + c]; + } + self.c2c_col_fwd + .process_with_scratch(&mut scratch.col_buf, &mut scratch.col_scratch_fwd); + for r in 0..fft_h { + scratch.freq_a[r * w_freq + c] = scratch.col_buf[r]; + } + } + + // 2. Forward 2D Real FFT for padded window B + for r in 0..fft_h { + scratch.in_row_pad.fill(0.0); + if r < win_h { + scratch.in_row_pad[..win_w].copy_from_slice(&window_b[r * win_w..(r + 1) * win_w]); + } + let out_row = &mut scratch.freq_b[r * w_freq..(r + 1) * w_freq]; + let _ = self.r2c_row.process_with_scratch(&mut scratch.in_row_pad, out_row, &mut scratch.row_scratch_fwd); + } + for c in 0..w_freq { + for r in 0..fft_h { + scratch.col_buf[r] = scratch.freq_b[r * w_freq + c]; + } + self.c2c_col_fwd + .process_with_scratch(&mut scratch.col_buf, &mut scratch.col_scratch_fwd); + for r in 0..fft_h { + scratch.freq_b[r * w_freq + c] = scratch.col_buf[r]; + } + } + + // 3. Frequency domain cross-correlation: Fa * Fb.conj() + // Matches scipy.signal.correlate(a, b, mode='full') + for i in 0..(fft_h * w_freq) { + scratch.freq_a[i] = scratch.freq_a[i] * scratch.freq_b[i].conj(); + } + + // 4. Inverse 2D Real FFT + for c in 0..w_freq { + for r in 0..fft_h { + scratch.col_buf[r] = scratch.freq_a[r * w_freq + c]; + } + self.c2c_col_inv + .process_with_scratch(&mut scratch.col_buf, &mut scratch.col_scratch_inv); + for r in 0..fft_h { + scratch.freq_a[r * w_freq + c] = scratch.col_buf[r]; + } + } + + for r in 0..fft_h { + let in_freq_row = &mut scratch.freq_a[r * w_freq..(r + 1) * w_freq]; + in_freq_row[0].im = 0.0; + if fft_w % 2 == 0 { + in_freq_row[w_freq - 1].im = 0.0; + } + let out_real_row = &mut scratch.temp_real[r * fft_w..(r + 1) * fft_w]; + let _ = self.c2r_row.process_with_scratch(in_freq_row, out_real_row, &mut scratch.row_scratch_inv); + } + + let scale = 1.0 / ((fft_h * fft_w) as f64); + + // 5. Crop and unshift into out_slice + for r in 0..out_h { + let r_idx = (r as isize - (win_h as isize - 1)).rem_euclid(fft_h as isize) as usize; + for c in 0..out_w { + let c_idx = (c as isize - (win_w as isize - 1)).rem_euclid(fft_w as isize) as usize; + out_slice[r * out_w + c] = scratch.temp_real[r_idx * fft_w + c_idx] * scale; + } + } + } +} + +/// Circular batched cross-correlation (OpenPIV standard mode, returns (N, H, W)) +#[pyfunction] +#[pyo3(signature = (windows_a, windows_b, normalized_correlation=true))] +fn fft_correlate_circular<'py>( + py: Python<'py>, + windows_a: PyReadonlyArray3, + windows_b: PyReadonlyArray3, + normalized_correlation: bool, +) -> PyResult>> { + let a = windows_a.as_array(); + let b = windows_b.as_array(); + + // Safe error handling: validate shapes + if a.shape() != b.shape() { + return Err(PyValueError::new_err(format!( + "Shape mismatch: windows_a has shape {:?}, but windows_b has shape {:?}", + a.shape(), + b.shape() + ))); + } + if a.shape()[0] == 0 || a.shape()[1] == 0 || a.shape()[2] == 0 { + return Err(PyValueError::new_err( + "Window arrays must have non-zero dimensions [N > 0, H > 0, W > 0]" + )); + } + + let num_wins = a.shape()[0]; + let win_h = a.shape()[1]; + let win_w = a.shape()[2]; + let win_size = win_h * win_w; + + let engine = CircularEngine2D::new(win_h, win_w); + + // Safe C-contiguous data access (handles any strides, transpositions, slices) + let a_cow = to_c_contiguous(&a); + let b_cow = to_c_contiguous(&b); + + let mut final_results = vec![0.0; num_wins * win_size]; + + // Release GIL and compute across threadpool using Rayon with zero-allocation thread-local scratch + py.detach(|| { + final_results + .par_chunks_exact_mut(win_size) + .enumerate() + .for_each_init( + || engine.create_scratch(), + |scratch, (i, out_slice)| { + let w_a = &a_cow[i * win_size..(i + 1) * win_size]; + let w_b = &b_cow[i * win_size..(i + 1) * win_size]; + engine.correlate_pair(w_a, w_b, out_slice, normalized_correlation, scratch); + }, + ); + }); + + Ok(final_results + .to_pyarray(py) + .reshape([num_wins, win_h, win_w])?) +} + +/// Full linear batched cross-correlation (matches scipy.signal.correlate(..., mode='full'), returns (N, 2H-1, 2W-1)) +#[pyfunction] +#[pyo3(signature = (windows_a, windows_b))] +fn fast_batch_cross_correlation<'py>( + py: Python<'py>, + windows_a: PyReadonlyArray3, + windows_b: PyReadonlyArray3, +) -> PyResult>> { + let a = windows_a.as_array(); + let b = windows_b.as_array(); + + // Safe error handling: validate shapes + if a.shape() != b.shape() { + return Err(PyValueError::new_err(format!( + "Shape mismatch: windows_a has shape {:?}, but windows_b has shape {:?}", + a.shape(), + b.shape() + ))); + } + if a.shape()[0] == 0 || a.shape()[1] == 0 || a.shape()[2] == 0 { + return Err(PyValueError::new_err( + "Window arrays must have non-zero dimensions [N > 0, H > 0, W > 0]" + )); + } + + let num_wins = a.shape()[0]; + let win_h = a.shape()[1]; + let win_w = a.shape()[2]; + + let out_h = 2 * win_h - 1; + let out_w = 2 * win_w - 1; + let out_win_size = out_h * out_w; + let in_win_size = win_h * win_w; + + let engine = FullEngine2D::new(win_h, win_w); + + // Safe C-contiguous data access (handles any strides, transpositions, slices) + let a_cow = to_c_contiguous(&a); + let b_cow = to_c_contiguous(&b); + + let mut final_results = vec![0.0; num_wins * out_win_size]; + + py.detach(|| { + final_results + .par_chunks_exact_mut(out_win_size) + .enumerate() + .for_each_init( + || engine.create_scratch(), + |scratch, (i, out_slice)| { + let w_a = &a_cow[i * in_win_size..(i + 1) * in_win_size]; + let w_b = &b_cow[i * in_win_size..(i + 1) * in_win_size]; + engine.correlate_pair(w_a, w_b, out_slice, scratch); + }, + ); + }); + + Ok(final_results + .to_pyarray(py) + .reshape([num_wins, out_h, out_w])?) +} + +/// Linear batched cross-correlation alias (same as fast_batch_cross_correlation) +#[pyfunction] +#[pyo3(signature = (windows_a, windows_b))] +fn fft_correlate_linear<'py>( + py: Python<'py>, + windows_a: PyReadonlyArray3, + windows_b: PyReadonlyArray3, +) -> PyResult>> { + fast_batch_cross_correlation(py, windows_a, windows_b) +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum SubpixelMethod { + Gaussian, + Centroid, + Parabolic, +} + +impl SubpixelMethod { + pub fn parse(s: &str) -> Result { + match s { + "gaussian" => Ok(SubpixelMethod::Gaussian), + "centroid" => Ok(SubpixelMethod::Centroid), + "parabolic" => Ok(SubpixelMethod::Parabolic), + other => Err(format!("Method not implemented {other}")), + } + } +} + +pub fn subpixel_peak_position_2d( + corr: &[f64], + h: usize, + w: usize, + method: SubpixelMethod, +) -> (f64, f64) { + if h < 3 || w < 3 { + return (f64::NAN, f64::NAN); + } + + // 1. Find argmax (peak1_i, peak1_j) + let mut best_idx = 0; + let mut best_val = corr[0]; + for (idx, &val) in corr.iter().enumerate() { + if val > best_val { + best_val = val; + best_idx = idx; + } + } + + let peak1_i = best_idx / w; + let peak1_j = best_idx % w; + + // 2. Check border condition + if peak1_i == 0 || peak1_i == h - 1 || peak1_j == 0 || peak1_j == w - 1 { + return (f64::NAN, f64::NAN); + } + + // 3. Extract 5-point cross with eps = 1e-7 + const EPS: f64 = 1e-7; + let c = corr[peak1_i * w + peak1_j] + EPS; + let cl = corr[(peak1_i - 1) * w + peak1_j] + EPS; + let cr = corr[(peak1_i + 1) * w + peak1_j] + EPS; + let cd = corr[peak1_i * w + (peak1_j - 1)] + EPS; + let cu = corr[peak1_i * w + (peak1_j + 1)] + EPS; + + // 4. Fallback if any point < 0 + let effective_method = if method == SubpixelMethod::Gaussian + && (c < 0.0 || cl < 0.0 || cr < 0.0 || cd < 0.0 || cu < 0.0) + { + SubpixelMethod::Parabolic + } else { + method + }; + + match effective_method { + SubpixelMethod::Centroid => { + let sum_row = cl + c + cr; + let sum_col = cd + c + cu; + let pi = peak1_i as f64; + let pj = peak1_j as f64; + let sub_i = if sum_row != 0.0 { + ((pi - 1.0) * cl + pi * c + (pi + 1.0) * cr) / sum_row + } else { + f64::NAN + }; + let sub_j = if sum_col != 0.0 { + ((pj - 1.0) * cd + pj * c + (pj + 1.0) * cu) / sum_col + } else { + f64::NAN + }; + (sub_i, sub_j) + } + SubpixelMethod::Gaussian => { + let nom1 = cl.ln() - cr.ln(); + let den1 = 2.0 * cl.ln() - 4.0 * c.ln() + 2.0 * cr.ln(); + let nom2 = cd.ln() - cu.ln(); + let den2 = 2.0 * cd.ln() - 4.0 * c.ln() + 2.0 * cu.ln(); + + let offset_i = if den1 != 0.0 { nom1 / den1 } else { 0.0 }; + let offset_j = if den2 != 0.0 { nom2 / den2 } else { 0.0 }; + + (peak1_i as f64 + offset_i, peak1_j as f64 + offset_j) + } + SubpixelMethod::Parabolic => { + let den1 = 2.0 * cl - 4.0 * c + 2.0 * cr; + let den2 = 2.0 * cd - 4.0 * c + 2.0 * cu; + + let offset_i = if den1 != 0.0 { (cl - cr) / den1 } else { 0.0 }; + let offset_j = if den2 != 0.0 { (cd - cu) / den2 } else { 0.0 }; + + (peak1_i as f64 + offset_i, peak1_j as f64 + offset_j) + } + } +} + +/// Find subpixel approximation of the correlation peak for a single 2D correlation map. +#[pyfunction] +#[pyo3(signature = (corr, subpixel_method=None))] +fn find_subpixel_peak_position( + corr: PyReadonlyArray2, + subpixel_method: Option<&str>, +) -> PyResult<(f64, f64)> { + let method_str = subpixel_method.unwrap_or("gaussian"); + let method = SubpixelMethod::parse(method_str) + .map_err(|e| PyValueError::new_err(e))?; + + let arr = corr.as_array(); + let h = arr.shape()[0]; + let w = arr.shape()[1]; + + let corr_cow = to_c_contiguous_2d(&arr); + Ok(subpixel_peak_position_2d(&corr_cow, h, w, method)) +} + +/// Batched subpixel peak positions for a 3D array of correlation maps (N, H, W). +/// Returns (peaks_i, peaks_j) as 1D arrays of length N. +#[pyfunction] +#[pyo3(signature = (corr, subpixel_method=None))] +fn batch_find_subpixel_peak_position<'py>( + py: Python<'py>, + corr: PyReadonlyArray3, + subpixel_method: Option<&str>, +) -> PyResult<(Bound<'py, PyArray1>, Bound<'py, PyArray1>)> { + let method_str = subpixel_method.unwrap_or("gaussian"); + let method = SubpixelMethod::parse(method_str) + .map_err(|e| PyValueError::new_err(e))?; + + let arr = corr.as_array(); + let num_wins = arr.shape()[0]; + let h = arr.shape()[1]; + let w = arr.shape()[2]; + let win_size = h * w; + + let corr_cow = to_c_contiguous(&arr); + + let mut peaks_i = vec![0.0; num_wins]; + let mut peaks_j = vec![0.0; num_wins]; + + py.detach(|| { + peaks_i + .par_iter_mut() + .zip(peaks_j.par_iter_mut()) + .enumerate() + .for_each(|(idx, (pi, pj))| { + let win_slice = &corr_cow[idx * win_size..(idx + 1) * win_size]; + let (r, c) = subpixel_peak_position_2d(win_slice, h, w, method); + *pi = r; + *pj = c; + }); + }); + + Ok((peaks_i.to_pyarray(py), peaks_j.to_pyarray(py))) +} + +/// Batched conversion from correlation maps (N, H, W) to (u, v) displacement grids of shape (n_rows, n_cols). +#[pyfunction] +#[pyo3(signature = (corr, n_rows, n_cols, subpixel_method=None))] +fn batch_correlation_to_displacement<'py>( + py: Python<'py>, + corr: PyReadonlyArray3, + n_rows: usize, + n_cols: usize, + subpixel_method: Option<&str>, +) -> PyResult<(Bound<'py, PyArray2>, Bound<'py, PyArray2>)> { + let method_str = subpixel_method.unwrap_or("gaussian"); + let method = SubpixelMethod::parse(method_str) + .map_err(|e| PyValueError::new_err(e))?; + + let arr = corr.as_array(); + let num_wins = arr.shape()[0]; + let h = arr.shape()[1]; + let w = arr.shape()[2]; + + if num_wins != n_rows * n_cols { + return Err(PyValueError::new_err(format!( + "Number of correlation windows ({num_wins}) does not match n_rows * n_cols ({n_rows} * {n_cols} = {})", + n_rows * n_cols + ))); + } + + let win_size = h * w; + let corr_cow = to_c_contiguous(&arr); + + let default_peak_i = (h / 2) as f64; + let default_peak_j = (w / 2) as f64; + + let mut u_vec = vec![0.0; num_wins]; + let mut v_vec = vec![0.0; num_wins]; + + py.detach(|| { + u_vec + .par_iter_mut() + .zip(v_vec.par_iter_mut()) + .enumerate() + .for_each(|(idx, (u_val, v_val))| { + let win_slice = &corr_cow[idx * win_size..(idx + 1) * win_size]; + let (peak_i, peak_j) = subpixel_peak_position_2d(win_slice, h, w, method); + *u_val = peak_j - default_peak_j; + *v_val = peak_i - default_peak_i; + }); + }); + + let u_arr = u_vec.to_pyarray(py).reshape([n_rows, n_cols])?; + let v_arr = v_vec.to_pyarray(py).reshape([n_rows, n_cols])?; + + Ok((u_arr, v_arr)) +} + +#[inline] +fn compute_nanmedian(slice: &mut [f64]) -> f64 { + let mut valid_count = 0; + for i in 0..slice.len() { + if !slice[i].is_nan() { + slice.swap(valid_count, i); + valid_count += 1; + } + } + if valid_count == 0 { + return f64::NAN; + } + let valid = &mut slice[..valid_count]; + valid.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + if valid_count % 2 == 1 { + valid[valid_count / 2] + } else { + (valid[valid_count / 2 - 1] + valid[valid_count / 2]) * 0.5 + } +} + +/// Ultra-fast normalized median filter for PIV outlier detection (Westerweel & Scarano 2005) +#[pyfunction] +#[pyo3(signature = (u, v, eps, threshold, size=1))] +fn local_norm_median_val<'py>( + py: Python<'py>, + u: PyReadonlyArray2<'py, f64>, + v: PyReadonlyArray2<'py, f64>, + eps: f64, + threshold: f64, + size: usize, +) -> PyResult>> { + let u_view = u.as_array(); + let v_view = v.as_array(); + + if u_view.shape() != v_view.shape() { + return Err(PyValueError::new_err(format!( + "u shape {:?} does not match v shape {:?}", + u_view.shape(), + v_view.shape() + ))); + } + + let h = u_view.shape()[0]; + let w = u_view.shape()[1]; + let u_cow = to_c_contiguous_2d(&u_view); + let v_cow = to_c_contiguous_2d(&v_view); + + let k = 2 * size + 1; + let total_elements = k * k; + + let mut out_mask = vec![false; h * w]; + + py.detach(|| { + out_mask + .par_chunks_mut(w) + .enumerate() + .for_each(|(r, row_mask)| { + let mut all_u = Vec::with_capacity(total_elements); + let mut all_v = Vec::with_capacity(total_elements); + let mut neigh_u = Vec::with_capacity(total_elements - 1); + let mut neigh_v = Vec::with_capacity(total_elements - 1); + let mut dev_u = Vec::with_capacity(total_elements - 1); + let mut dev_v = Vec::with_capacity(total_elements - 1); + + for c in 0..w { + let val_u = u_cow[r * w + c]; + let val_v = v_cow[r * w + c]; + + if val_u.is_nan() || val_v.is_nan() { + row_mask[c] = false; + continue; + } + + all_u.clear(); + all_v.clear(); + neigh_u.clear(); + neigh_v.clear(); + + for dr in -(size as isize)..=(size as isize) { + let nr = r as isize + dr; + for dc in -(size as isize)..=(size as isize) { + let nc = c as isize + dc; + let is_center = dr == 0 && dc == 0; + + let (u_val, v_val) = if nr >= 0 && nr < h as isize && nc >= 0 && nc < w as isize { + let idx = nr as usize * w + nc as usize; + (u_cow[idx], v_cow[idx]) + } else { + (f64::NAN, f64::NAN) + }; + + all_u.push(u_val); + all_v.push(v_val); + if !is_center { + neigh_u.push(u_val); + neigh_v.push(v_val); + } + } + } + + let um = compute_nanmedian(&mut all_u); + let vm = compute_nanmedian(&mut all_v); + + let ym_u = compute_nanmedian(&mut neigh_u.clone()); + let ym_v = compute_nanmedian(&mut neigh_v.clone()); + + if ym_u.is_nan() || ym_v.is_nan() { + row_mask[c] = false; + continue; + } + + dev_u.clear(); + for &x in &neigh_u { + if !x.is_nan() { + dev_u.push((x - ym_u).abs()); + } + } + + dev_v.clear(); + for &x in &neigh_v { + if !x.is_nan() { + dev_v.push((x - ym_v).abs()); + } + } + + let rm_u = compute_nanmedian(&mut dev_u); + let rm_v = compute_nanmedian(&mut dev_v); + + if rm_u.is_nan() || rm_v.is_nan() { + row_mask[c] = false; + continue; + } + + let r0ast_u = (val_u - um).abs() / (rm_u + eps); + let r0ast_v = (val_v - vm).abs() / (rm_v + eps); + + if r0ast_u.hypot(r0ast_v) > threshold { + row_mask[c] = true; + } + } + }); + }); + + let py_arr = out_mask.to_pyarray(py).reshape([h, w])?; + Ok(py_arr) +} + +/// Batch signal-to-noise ratio calculation across correlation maps in parallel +#[pyfunction] +#[pyo3(signature = (correlation, sig2noise_method="peak2peak", width=2))] +fn sig2noise_ratio<'py>( + py: Python<'py>, + correlation: PyReadonlyArray3<'py, f64>, + sig2noise_method: &str, + width: usize, +) -> PyResult>> { + let corr_view = correlation.as_array(); + let num_wins = corr_view.shape()[0]; + let h = corr_view.shape()[1]; + let w = corr_view.shape()[2]; + let win_size = h * w; + + if sig2noise_method != "peak2peak" && sig2noise_method != "peak2mean" { + return Err(PyValueError::new_err(format!( + "Invalid sig2noise_method '{sig2noise_method}'. Expected 'peak2peak' or 'peak2mean'" + ))); + } + + let corr_cow = to_c_contiguous(&corr_view); + let mut s2n_vec = vec![0.0f64; num_wins]; + + py.detach(|| { + s2n_vec + .par_iter_mut() + .enumerate() + .for_each(|(idx, s2n_val)| { + let win_slice = &corr_cow[idx * win_size..(idx + 1) * win_size]; + + // First peak + let mut max1_val = f64::NEG_INFINITY; + let mut max1_r = 0; + let mut max1_c = 0; + for r in 0..h { + for c in 0..w { + let v = win_slice[r * w + c]; + if v > max1_val { + max1_val = v; + max1_r = r; + max1_c = c; + } + } + } + + if sig2noise_method == "peak2peak" { + if max1_val < 1e-3 + || max1_r == 0 + || max1_r == h - 1 + || max1_c == 0 + || max1_c == w - 1 + { + *s2n_val = 0.0; + } else { + let r_min = max1_r.saturating_sub(width); + let r_max = (max1_r + width).min(h - 1); + let c_min = max1_c.saturating_sub(width); + let c_max = (max1_c + width).min(w - 1); + + let mut max2_val = f64::NEG_INFINITY; + let mut max2_r = 0; + let mut max2_c = 0; + for r in 0..h { + for c in 0..w { + if r >= r_min && r <= r_max && c >= c_min && c <= c_max { + continue; + } + let v = win_slice[r * w + c]; + if v > max2_val { + max2_val = v; + max2_r = r; + max2_c = c; + } + } + } + + let border2 = max2_r == 0 || max2_r == h - 1 || max2_c == 0 || max2_c == w - 1; + if max2_val <= 0.0 || (border2 && max2_val > 0.5 * max1_val) { + *s2n_val = 0.0; + } else { + let ratio = max1_val / max2_val; + *s2n_val = if ratio.is_nan() { 0.0 } else { ratio }; + } + } + } else if sig2noise_method == "peak2mean" { + if max1_val < 1e-3 + || max1_r == 0 + || max1_r == h - 1 + || max1_c == 0 + || max1_c == w - 1 + { + max1_val = 0.0; + } + + let sum: f64 = win_slice.iter().sum(); + let mean = (sum / (h * w) as f64).abs(); + if mean == 0.0 || mean.is_nan() { + *s2n_val = 0.0; + } else { + let ratio = max1_val / mean; + *s2n_val = if ratio.is_nan() { 0.0 } else { ratio }; + } + } + }); + }); + + Ok(s2n_vec.to_pyarray(py)) +} + +/// Slices an image into interrogation windows directly into a 3D contiguous array in parallel +#[pyfunction] +#[pyo3(signature = (image, window_size=(64, 64), overlap=(32, 32)))] +fn sliding_window_array<'py>( + py: Python<'py>, + image: PyReadonlyArray2<'py, f64>, + window_size: (usize, usize), + overlap: (usize, usize), +) -> PyResult>> { + let img_view = image.as_array(); + let img_h = img_view.shape()[0]; + let img_w = img_view.shape()[1]; + + let (win_h, win_w) = window_size; + let (ov_h, ov_w) = overlap; + + if win_h > img_h || win_w > img_w { + return Err(PyValueError::new_err(format!( + "Window size ({win_h}, {win_w}) exceeds image size ({img_h}, {img_w})" + ))); + } + if ov_h >= win_h || ov_w >= win_w { + return Err(PyValueError::new_err(format!( + "Overlap ({ov_h}, {ov_w}) must be smaller than window size ({win_h}, {win_w})" + ))); + } + + let step_h = win_h - ov_h; + let step_w = win_w - ov_w; + + let n_rows = (img_h - win_h) / step_h + 1; + let n_cols = (img_w - win_w) / step_w + 1; + let num_wins = n_rows * n_cols; + + let img_cow = to_c_contiguous_2d(&img_view); + let win_size_elems = win_h * win_w; + let mut out_vec = vec![0.0f64; num_wins * win_size_elems]; + + py.detach(|| { + out_vec + .par_chunks_mut(win_size_elems) + .enumerate() + .for_each(|(win_idx, win_buf)| { + let r_idx = win_idx / n_cols; + let c_idx = win_idx % n_cols; + let y_start = r_idx * step_h; + let x_start = c_idx * step_w; + + for wr in 0..win_h { + let src_start = (y_start + wr) * img_w + x_start; + let dst_start = wr * win_w; + win_buf[dst_start..dst_start + win_w] + .copy_from_slice(&img_cow[src_start..src_start + win_w]); + } + }); + }); + + let py_arr = out_vec.to_pyarray(py).reshape([num_wins, win_h, win_w])?; + Ok(py_arr) +} + +#[pymodule] +fn openpiv_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(fft_correlate_circular, m)?)?; + m.add_function(wrap_pyfunction!(fast_batch_cross_correlation, m)?)?; + m.add_function(wrap_pyfunction!(fft_correlate_linear, m)?)?; + m.add_function(wrap_pyfunction!(find_subpixel_peak_position, m)?)?; + m.add_function(wrap_pyfunction!(batch_find_subpixel_peak_position, m)?)?; + m.add_function(wrap_pyfunction!(batch_correlation_to_displacement, m)?)?; + m.add_function(wrap_pyfunction!(local_norm_median_val, m)?)?; + m.add_function(wrap_pyfunction!(sig2noise_ratio, m)?)?; + m.add_function(wrap_pyfunction!(sliding_window_array, m)?)?; + Ok(()) +} diff --git a/openpiv/docs/index.rst b/openpiv/docs/index.rst index 77ecc372..17701475 100644 --- a/openpiv/docs/index.rst +++ b/openpiv/docs/index.rst @@ -31,6 +31,7 @@ Contents: src/windef src/masking src/developers + src/fft_correlation_backends src/api_reference src/faq diff --git a/openpiv/docs/src/developers.rst b/openpiv/docs/src/developers.rst index 7190e91d..ba62a4a1 100644 --- a/openpiv/docs/src/developers.rst +++ b/openpiv/docs/src/developers.rst @@ -63,3 +63,64 @@ If you need to install cv2:: conda install -c conda-forge opencv +Developing with or without Rust +-------------------------------- + +OpenPIV supports dual computational backends: **pure Python / SciPy** (default) and **compiled Rust** (via ``openpiv_rust``). + +Option A: Developing in Pure Python (without Rust) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +You do not need a Rust compiler installed. Simply install OpenPIV in editable mode: + +.. code-block:: bash + + poetry install + # or: pip install -e . + +All core functionality, single-pass, and multi-pass deformation will use the optimized +``scipy.fft`` backend automatically. + +To run the test suite: + +.. code-block:: bash + + pytest openpiv/test + +Option B: Developing with the Rust Extension +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +If you have the Rust toolchain (``cargo``) installed: + +1. Install ``maturin``: + + .. code-block:: bash + + pip install maturin + +2. Compile and link ``openpiv_rust`` directly into your Python environment: + + .. code-block:: bash + + cd crates/openpiv_rust + maturin develop --release + cd ../.. + +3. Verify that the Rust backend is active: + + .. code-block:: bash + + python -c "import openpiv_rust; print('Rust backend ready!')" + +Running the Performance Benchmarks +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +OpenPIV includes a dedicated benchmark suite comparing cross-correlation, subpixel +peak finding, and end-to-end multi-pass Windef across backends: + +.. code-block:: bash + + python benchmarks/run_benchmarks.py + + + diff --git a/openpiv/docs/src/fft_correlation_backends.rst b/openpiv/docs/src/fft_correlation_backends.rst new file mode 100644 index 00000000..58ee0d41 --- /dev/null +++ b/openpiv/docs/src/fft_correlation_backends.rst @@ -0,0 +1,162 @@ +FFT Correlation Backends & Performance Study +============================================ + +This document details the comparative study of 2D cross-correlation algorithms, +backends, and performance characteristics in OpenPIV, focusing on both **Circular** +and **Linear** correlation, `scipy.fft` (PocketFFT), and the `openpiv_rust` compiled backend. + +.. contents:: Table of Contents + :depth: 2 + :local: + +Overview +-------- + +In Particle Image Velocimetry (PIV), cross-correlation between interrogation windows +is the core computational step. For an image pair split into :math:`N` interrogation +windows of size :math:`(H \times W)`, the correlation maps are computed using the +Wiener-Khinchin theorem via Fast Fourier Transforms: + +.. math:: + + C = \mathcal{F}^{-1} \left( \mathcal{F}(I_B) \cdot \mathcal{F}^*(I_A) \right) + +OpenPIV supports two primary correlation paradigms: + +1. **Circular Correlation (Standard OpenPIV)**: + Assumes periodic boundary conditions (toroidal wraparound). + Inputs of size :math:`(N \times N)` yield correlation maps of size :math:`(N \times N)`. + No zero-padding is required. + +2. **Linear Correlation (Full / Extended)**: + Evaluates true spatial shift without periodic wraparound. + For windows of size :math:`s_1` and :math:`s_2`, the full linear correlation has size + :math:`(s_1 + s_2 - 1)`. To prevent time-domain aliasing, inputs must be zero-padded + to at least this size before applying the FFT. + + +The 63x63 vs 64x64 Power-of-2 Finding +------------------------------------- + +In previous versions of OpenPIV, linear correlation in ``openpiv.pyprocess.fft_correlate_images`` +computed the padded transform size as: + +.. code-block:: python + + size = s1 + s2 - 1 # e.g., 32 + 32 - 1 = 63 + fsize = 2 ** np.ceil(np.log2(size)).astype(int) - 1 # 64 - 1 = 63 (BUG!) + +Why 63 Disabled PocketFFT Optimization +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +SciPy delegates FFT operations to **PocketFFT** (written in C++ by Martin Reinecke), +which includes hand-tuned AVX2/FMA vector kernels. However, these SIMD routines rely on +smooth radix factors (powers of 2: 2, 4, 8, 16, 32, 64). + +1. Subtracting 1 produced an **odd composite size** :math:`63 = 3 \times 3 \times 7`. +2. This forced PocketFFT into slow Bluestein/composite transforms, disabling SIMD vectorization. +3. Furthermore, padding to 63 left an off-by-one boundary alignment when slicing the central window. + +The Power-of-2 Solution +^^^^^^^^^^^^^^^^^^^^^^^ + +By keeping the clean power-of-2 transform size :math:`fsize = 2^{\lceil \log_2(size) \rceil}` +(e.g., :math:`64 \times 64` for :math:`32 \times 32` windows), PocketFFT achieves optimal +SIMD execution speed. + +To recover the exact central correlation window matching ``scipy.signal.correlate`` +to machine precision (:math:`\le 10^{-13}`), the centered slice is: + +.. code-block:: python + + # Exact power of 2 transform size + fsize = 2 ** np.ceil(np.log2(size)).astype(int) + + # Centered slice extracting (s1) around the zero-displacement lag + fslice = ( + slice(0, image_a.shape[0]), + slice(fsize[0] // 2 - s1[0] // 2, fsize[0] // 2 - s1[0] // 2 + s1[0]), + slice(fsize[1] // 2 - s1[1] // 2, fsize[1] // 2 - s1[1] // 2 + s1[1]), + ) + f2a = conj(rfft2(image_a, fsize, axes=(-2, -1), workers=workers)) + f2b = rfft2(image_b, fsize, axes=(-2, -1), workers=workers) + corr = fftshift(irfft2(f2a * f2b, axes=(-2, -1)).real, axes=(-2, -1))[fslice] + + +Batched Multi-Threading in scipy.fft +------------------------------------ + +``scipy.fft`` supports a ``workers`` parameter (e.g., ``workers=-1`` for all logical cores). +However, profiling reveals: + +* On large 2D or 3D volumes, PocketFFT multi-threading scales well. +* On **batches of tiny 2D windows** (:math:`16 \times 16` or :math:`32 \times 32`), + thread synchronization overhead inside PocketFFT often negates multi-core gains + (scaling from :math:`1.0\times` to :math:`1.1\times`), because thread splitting occurs + per-transform rather than distributing independent windows across workers. + + +The Rust Acceleration Engine (openpiv_rust) +------------------------------------------- + +To overcome Python GIL overhead and achieve linear CPU scaling across window batches, +OpenPIV provides an optional compiled Rust extension (``openpiv_rust``) built with +PyO3, Rayon, and RealFFT. + +Key Architectural Optimizations in Rust +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +1. **Zero-Allocation Rayon Scratch Pool**: + Window transforms use ``par_chunks_exact_mut().for_each_init(...)``, pre-allocating + frequency buffers once per worker thread. For 961 windows, this eliminates over + 8,600 heap allocations per pass. + +2. **Native Power-of-2 Real FFT**: + Uses ``realfft`` (Real-to-Complex forward and Complex-to-Real inverse), cutting + arithmetic and memory bandwidth by 50% compared to full complex transforms. + +3. **Batched Subpixel Peak Finding**: + In addition to cross-correlation, `openpiv_rust` provides + ``batch_correlation_to_displacement``, executing Gaussian/Centroid/Parabolic subpixel + peak fitting directly across all windows in parallel, yielding a **28x to 65x speedup** + over the nested Python loop. + + +Benchmark Results +----------------- + +Linear Correlation (225 windows of 32x32) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +============================================ ==================== ================== +Implementation Runtime Relative Speed +============================================ ==================== ================== +``scipy.signal.correlate`` (Python loop) 70.63 ms 1.0x (baseline) +OpenPIV SciPy (legacy ``fsize=63``) 49.28 ms 1.4x +OpenPIV SciPy (power-of-2 ``fsize=64``) 39.10 ms 1.8x +**openpiv_rust** (Rayon + realfft) **8.26 ms** **8.5x** +============================================ ==================== ================== + +Displacement Calculation (Peak Finding) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +============================================ ==================== ================== +Grid / Windows Python Loop openpiv_rust +============================================ ==================== ================== +225 windows (32x32) 5.19 ms **0.18 ms (28x)** +961 windows (16x16) 19.84 ms **0.30 ms (65x)** +============================================ ==================== ================== + +End-to-End Multi-Pass Windef (3 passes on 256x256 image) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +* **SciPy backend**: 110.59 ms +* **Rust backend (with subpixel Rust engine)**: **81.66 ms** (:math:`1.35\times` speedup) + +Conclusion & Usage Guidelines +----------------------------- + +* For systems with compiled binary extensions available, set ``settings.backend = 'rust'`` + to maximize throughput in both single-pass and multi-pass deformation workflows. +* On pure Python / NumPy / SciPy environments, using the exact power-of-2 padding rule + ensures optimal PocketFFT vectorization while maintaining exact numerical consistency. diff --git a/openpiv/pyprocess.py b/openpiv/pyprocess.py index b725ed65..e6ece86f 100644 --- a/openpiv/pyprocess.py +++ b/openpiv/pyprocess.py @@ -10,6 +10,13 @@ from numpy.fft import fftshift as fftshift_ from scipy.signal import convolve2d as conv_ +try: + import openpiv_rust + HAS_RUST = True +except ImportError: + HAS_RUST = False + + __licence_ = """ Copyright (C) 2011 www.openpiv.net @@ -185,6 +192,7 @@ def sliding_window_array( image: np.ndarray, window_size: Tuple[int,int]=(64,64), overlap: Tuple[int,int]=(32,32), + backend: str="auto", )-> np.ndarray: ''' This version does not use numpy as_strided and is much more memory efficient. @@ -194,10 +202,17 @@ def sliding_window_array( with three dimension, of size (n_windows, window_size, window_size), in which each slice, (along the first axis) is an interrogation window. ''' - # if isinstance(window_size, int): - # window_size = (window_size, window_size) - # if isinstance(overlap, int): - # overlap = (overlap, overlap) + if isinstance(window_size, int): + window_size = (window_size, window_size) + if isinstance(overlap, int): + overlap = (overlap, overlap) + + if (backend == "rust" or (backend == "auto" and HAS_RUST)) and image.ndim == 2: + return openpiv_rust.sliding_window_array( + np.ascontiguousarray(image, dtype=np.float64), + (int(window_size[0]), int(window_size[1])), + (int(overlap[0]), int(overlap[1])), + ) x, y = get_rect_coordinates(image.shape, window_size, overlap, center_on_field = False) x = (x - window_size[1]//2).astype(int) @@ -383,7 +398,7 @@ def find_all_second_peaks(corr, width = 2): return indexes, peaks -def find_subpixel_peak_position(corr, subpixel_method="gaussian"): +def find_subpixel_peak_position(corr, subpixel_method="gaussian", backend="auto"): """ Find subpixel approximation of the correlation peak. @@ -425,6 +440,12 @@ def find_subpixel_peak_position(corr, subpixel_method="gaussian"): if subpixel_method not in ("gaussian", "centroid", "parabolic"): raise ValueError(f"Method not implemented {subpixel_method}") + if (backend == "rust" or (backend == "auto" and HAS_RUST)) and corr.ndim == 2: + return openpiv_rust.find_subpixel_peak_position( + np.ascontiguousarray(corr, dtype=np.float64), + subpixel_method=subpixel_method, + ) + # the peak locations (peak1_i, peak1_j), _ = find_first_peak(corr) @@ -484,7 +505,8 @@ def find_subpixel_peak_position(corr, subpixel_method="gaussian"): def sig2noise_ratio( correlation: np.ndarray, sig2noise_method: str="peak2peak", - width: int=2 + width: int=2, + backend: str="auto", )-> np.ndarray: """ Computes the signal to noise ratio from the correlation map. @@ -514,6 +536,13 @@ def sig2noise_ratio( the signal to noise ratios from the correlation maps. """ + if (backend == "rust" or (backend == "auto" and HAS_RUST)) and correlation.ndim == 3 and sig2noise_method in ("peak2peak", "peak2mean"): + return openpiv_rust.sig2noise_ratio( + np.ascontiguousarray(correlation, dtype=np.float64), + sig2noise_method=sig2noise_method, + width=int(width), + ) + sig2noise = np.zeros(correlation.shape[0]) corr_max1 = np.zeros(correlation.shape[0]) corr_max2 = np.zeros(correlation.shape[0]) @@ -679,10 +708,12 @@ def fft_correlate_images( image_b: np.ndarray, correlation_method: str="circular", normalized_correlation: bool=True, + backend: str="scipy", conj: Callable=np.conj, rfft2 = rfft2_, irfft2 = irfft2_, fftshift = fftshift_, + workers: Optional[int] = None, )->np.ndarray: """ FFT based cross correlation of two images with multiple views of np.stride_tricks() @@ -698,31 +729,50 @@ def fft_correlate_images( correlation_method : string one of the three methods implemented: 'circular' or 'linear' - [default: 'circular]. + [default: 'circular']. normalized_correlation : string - decides wetehr normalized correlation is done or not: True or False + decides whether normalized correlation is done or not: True or False [default: True]. + + backend : string + 'scipy' (default) or 'rust' (multithreaded 2D RealFFT via openpiv_rust) - conj : function - function used for complex conjugate - - rfft2 : function - function used for rfft2 - - irfft2 : function - function used for irfft2 - - fftshift : function - function used for fftshift - + workers : int, optional + number of worker threads for scipy.fft (default: None for single thread, + pass -1 for all CPU cores) """ + if backend == "rust" or correlation_method in ("circular_rust", "rust"): + if not HAS_RUST: + raise ImportError( + "openpiv_rust is not installed. Build with maturin to enable Rust acceleration." + ) + if normalized_correlation: + image_a = normalize_intensity(image_a) + image_b = normalize_intensity(image_b) + if correlation_method == "linear": + corr = openpiv_rust.fft_correlate_linear( + np.ascontiguousarray(image_b, dtype=np.float64), + np.ascontiguousarray(image_a, dtype=np.float64), + ) + s1 = image_a.shape[-2:] + out_h, out_w = corr.shape[-2:] + fslice = ( + slice(0, corr.shape[0]), + slice((out_h - s1[0]) // 2, (out_h - s1[0]) // 2 + s1[0]), + slice((out_w - s1[1]) // 2, (out_w - s1[1]) // 2 + s1[1]), + ) + corr = corr[fslice] + if normalized_correlation: + corr = corr / (corr.shape[-2] * corr.shape[-1]) + return corr + return openpiv_rust.fft_correlate_circular( + np.ascontiguousarray(image_a, dtype=np.float64), + np.ascontiguousarray(image_b, dtype=np.float64), + normalized_correlation=normalized_correlation, + ) if normalized_correlation: - # remove the effect of stronger laser or - # longer exposure for frame B - # image_a = match_histograms(image_a, image_b) - # remove mean, divide by standard deviation image_a = normalize_intensity(image_a) image_b = normalize_intensity(image_b) @@ -731,19 +781,23 @@ def fft_correlate_images( s2 = np.array(image_b.shape[-2:]) if correlation_method == "linear": - # have to be normalized, mainly because of zero padding size = s1 + s2 - 1 - fsize = 2 ** np.ceil(np.log2(size)).astype(int) - 1 - fslice = (slice(0, image_a.shape[0]), - slice((fsize[0]-s1[0])//2, (fsize[0]+s1[0])//2), - slice((fsize[1]-s1[1])//2, (fsize[1]+s1[1])//2)) - f2a = conj(rfft2(image_a, fsize, axes=(-2, -1))) # type: ignore - f2b = rfft2(image_b, fsize, axes=(-2, -1)) # type: ignore - corr = fftshift(irfft2(f2a * f2b).real, axes=(-2, -1))[fslice] + # Use exact power of 2 for optimal PocketFFT radix-2 SIMD execution: + fsize = 2 ** np.ceil(np.log2(size)).astype(int) + fslice = ( + slice(0, image_a.shape[0]), + slice(fsize[0] // 2 - s1[0] // 2, fsize[0] // 2 - s1[0] // 2 + s1[0]), + slice(fsize[1] // 2 - s1[1] // 2, fsize[1] // 2 - s1[1] // 2 + s1[1]), + ) + kwargs = {"workers": workers} if workers is not None else {} + f2a = conj(rfft2(image_a, fsize, axes=(-2, -1), **kwargs)) # type: ignore + f2b = rfft2(image_b, fsize, axes=(-2, -1), **kwargs) # type: ignore + corr = fftshift(irfft2(f2a * f2b, axes=(-2, -1), **kwargs).real, axes=(-2, -1))[fslice] elif correlation_method == "circular": - f2a = conj(rfft2(image_a)) - f2b = rfft2(image_b) - corr = fftshift(irfft2(f2a * f2b).real, axes=(-2, -1)) + kwargs = {"workers": workers} if workers is not None else {} + f2a = conj(rfft2(image_a, axes=(-2, -1), **kwargs)) + f2b = rfft2(image_b, axes=(-2, -1), **kwargs) + corr = fftshift(irfft2(f2a * f2b, axes=(-2, -1), **kwargs).real, axes=(-2, -1)) else: print(f"correlation method {correlation_method } is not implemented") @@ -919,6 +973,7 @@ def extended_search_area_piv( width: int=2, normalized_correlation: bool=False, use_vectorized: bool=False, + backend: str="scipy", ): """Standard PIV cross-correlation algorithm, with an option for extended area search that increased dynamic range. The search region @@ -1070,13 +1125,15 @@ def extended_search_area_piv( corr = fft_correlate_images(aa, bb, correlation_method=correlation_method, - normalized_correlation=normalized_correlation) + normalized_correlation=normalized_correlation, + backend=backend) if use_vectorized is True: u, v = vectorized_correlation_to_displacements(corr, n_rows, n_cols, subpixel_method=subpixel_method) else: u, v = correlation_to_displacement(corr, n_rows, n_cols, - subpixel_method=subpixel_method) + subpixel_method=subpixel_method, + backend=backend) # return output depending if user wanted sig2noise information if sig2noise_method is not None: @@ -1086,7 +1143,8 @@ def extended_search_area_piv( ) else: sig2noise = sig2noise_ratio( - corr, sig2noise_method=sig2noise_method, width=width + corr, sig2noise_method=sig2noise_method, width=width, + backend=backend, ) else: sig2noise = np.zeros_like(u)*np.nan @@ -1097,7 +1155,8 @@ def extended_search_area_piv( def correlation_to_displacement(corr, n_rows, n_cols, - subpixel_method="gaussian"): + subpixel_method="gaussian", + backend="auto"): """ Correlation maps are converted to displacement for each interrogation window using the convention that the size of the correlation map @@ -1109,6 +1168,17 @@ def correlation_to_displacement(corr, n_rows, n_cols, n_rows, n_cols : number of interrogation windows, output of the get_field_shape """ + if subpixel_method not in ("gaussian", "centroid", "parabolic"): + raise ValueError(f"Method not implemented {subpixel_method}") + + if (backend == "rust" or (backend == "auto" and HAS_RUST)) and corr.ndim == 3 and corr.shape[0] == n_rows * n_cols: + return openpiv_rust.batch_correlation_to_displacement( + np.ascontiguousarray(corr, dtype=np.float64), + n_rows, + n_cols, + subpixel_method=subpixel_method, + ) + # iterate through interrogation widows and search areas u = np.zeros((n_rows, n_cols)) v = np.zeros((n_rows, n_cols)) diff --git a/openpiv/settings.py b/openpiv/settings.py index a1860bff..269e9eaf 100644 --- a/openpiv/settings.py +++ b/openpiv/settings.py @@ -41,8 +41,9 @@ class PIVSettings: static_mask: Optional[np.ndarray] = None # or a boolean matrix of image shape # "Processing Parameters" - correlation_method: str="circular" # ['circular', 'linear'] - normalized_correlation: bool=False + backend: str = "scipy" # Backend for FFT cross-correlation: 'scipy' or 'rust' + correlation_method: str = "circular" # ['circular', 'linear'] + normalized_correlation: bool = False # add the interroagtion window size for each pass. # For the moment, it should be a power of 2 @@ -147,6 +148,7 @@ class PIVSettings: show_all_plots: bool=False + n_cpus: int = 1 # Number of CPU processes for multiprocessing batch evaluation invert: bool=False # for the test_invert fmt: str="%.4e" diff --git a/openpiv/test/test_rust_backend.py b/openpiv/test/test_rust_backend.py new file mode 100644 index 00000000..57d64cc9 --- /dev/null +++ b/openpiv/test/test_rust_backend.py @@ -0,0 +1,174 @@ +import pytest +import numpy as np +import scipy.signal +import scipy.fft +from pathlib import Path +from importlib.resources import files +from openpiv import pyprocess, tools, validation, filters +openpiv_rust = pytest.importorskip("openpiv_rust") + +@pytest.mark.parametrize("win_size", [16, 32, 64]) +@pytest.mark.parametrize("norm", [True, False]) +def test_circular_correlation_numerical_accuracy(win_size, norm): + np.random.seed(42) + n_wins = 20 + a = np.random.rand(n_wins, win_size, win_size) + b = np.random.rand(n_wins, win_size, win_size) + + # Scipy reference + scipy_corr = pyprocess.fft_correlate_images( + a, b, correlation_method="circular", normalized_correlation=norm + ) + + if norm: + a_norm = pyprocess.normalize_intensity(a) + b_norm = pyprocess.normalize_intensity(b) + rust_corr = openpiv_rust.fft_correlate_circular(a_norm, b_norm, normalized_correlation=True) + else: + rust_corr = openpiv_rust.fft_correlate_circular(a, b, normalized_correlation=False) + + max_diff = np.max(np.abs(scipy_corr - rust_corr)) + assert max_diff < 1e-10, f"Max difference {max_diff} exceeded tolerance for win_size={win_size}" + +def test_rectangular_windows(): + np.random.seed(123) + a = np.random.rand(10, 32, 64) + b = np.random.rand(10, 32, 64) + + a_norm = pyprocess.normalize_intensity(a) + b_norm = pyprocess.normalize_intensity(b) + + scipy_corr = pyprocess.fft_correlate_images(a, b, correlation_method="circular", normalized_correlation=True) + rust_corr = openpiv_rust.fft_correlate_circular(a_norm, b_norm, normalized_correlation=True) + + max_diff = np.max(np.abs(scipy_corr - rust_corr)) + assert max_diff < 1e-10 + +def test_fast_batch_cross_correlation_mode_full(): + np.random.seed(456) + n_wins, win_h, win_w = 15, 24, 24 + a = np.random.rand(n_wins, win_h, win_w) + b = np.random.rand(n_wins, win_h, win_w) + + scipy_res = np.array([ + scipy.signal.correlate(a[i], b[i], mode="full") for i in range(n_wins) + ]) + rust_res = openpiv_rust.fast_batch_cross_correlation(a, b) + + max_diff = np.max(np.abs(scipy_res - rust_res)) + assert max_diff < 1e-10 + +def test_strided_non_contiguous_inputs(): + """Verify Rust mode handles non-contiguous arrays without panicking.""" + a = np.random.rand(10, 32, 64)[:, :, ::2] # non-contiguous slice + b = np.random.rand(10, 32, 64)[:, :, ::2] + assert not a.flags.c_contiguous, "Array should be non-contiguous" + + rust_res = openpiv_rust.fft_correlate_circular(a, b, normalized_correlation=False) + assert rust_res.shape == (10, 32, 32) + +def test_known_displacements(): + """Verify integer peak detection recovers known displacement.""" + n_wins = 5 + win_size = 32 + shift_y, shift_x = 3, -2 + + a = np.zeros((n_wins, win_size, win_size)) + b = np.zeros((n_wins, win_size, win_size)) + for i in range(n_wins): + a[i, 14, 14] = 10.0 + b[i, 14 + shift_y, 14 + shift_x] = 10.0 + + corr = openpiv_rust.fft_correlate_circular(a, b, normalized_correlation=False) + center = win_size // 2 + + for i in range(n_wins): + (peak_y, peak_x), peak_val = pyprocess.find_first_peak(corr[i]) + dx = peak_x - center + dy = peak_y - center + assert dx == shift_x, f"Expected dx={shift_x}, got {dx}" + assert dy == shift_y, f"Expected dy={shift_y}, got {dy}" + +def test_real_piv_data_velocity_match(): + """Run full PIV pipeline on test1 dataset comparing Scipy vs Rust correlation.""" + path = files('openpiv') / "data" / "test1" + frame_a = tools.imread(path / "exp1_001_a.bmp").astype(np.int32) + frame_b = tools.imread(path / "exp1_001_b.bmp").astype(np.int32) + + window_size = 32 + overlap = 16 + search_area_size = 32 + + # Scipy PIV + u_scipy, v_scipy, s2n_scipy = pyprocess.extended_search_area_piv( + frame_a, frame_b, + window_size=window_size, + overlap=overlap, + search_area_size=search_area_size, + correlation_method="circular", + sig2noise_method="peak2peak", + ) + + # Rust PIV: extract windows and correlate via Rust + aa = pyprocess.sliding_window_array(frame_a, (search_area_size, search_area_size), (overlap, overlap)) + bb = pyprocess.sliding_window_array(frame_b, (search_area_size, search_area_size), (overlap, overlap)) + rust_corr = openpiv_rust.fft_correlate_circular(aa.astype(float), bb.astype(float), normalized_correlation=False) + + n_rows, n_cols = pyprocess.get_field_shape(frame_a.shape, (search_area_size, search_area_size), (overlap, overlap)) + u_rust, v_rust = pyprocess.correlation_to_displacement(rust_corr, n_rows, n_cols) + + diff_u = np.nanmax(np.abs(u_scipy - u_rust)) + diff_v = np.nanmax(np.abs(v_scipy - v_rust)) + assert diff_u < 1e-8, f"Velocity difference in U: {diff_u}" + assert diff_v < 1e-8, f"Velocity difference in V: {diff_v}" + +def test_safe_handling_shape_mismatch(): + """Verify that shape mismatches raise ValueError safely rather than panicking.""" + a = np.random.rand(5, 32, 32) + b = np.random.rand(5, 32, 16) + with pytest.raises(ValueError, match="Shape mismatch"): + openpiv_rust.fft_correlate_circular(a, b) + + with pytest.raises(ValueError, match="Shape mismatch"): + openpiv_rust.fast_batch_cross_correlation(a, b) + +def test_safe_handling_empty_arrays(): + """Verify that empty inputs raise ValueError safely.""" + a = np.zeros((0, 32, 32)) + b = np.zeros((0, 32, 32)) + with pytest.raises(ValueError, match="non-zero"): + openpiv_rust.fft_correlate_circular(a, b) + +def test_safe_handling_fortran_and_transposed(): + """Verify Fortran-ordered and transposed arrays work correctly.""" + rng = np.random.default_rng(123) + a_c = rng.standard_normal((4, 32, 32)) + b_c = rng.standard_normal((4, 32, 32)) + + # Convert to Fortran contiguous + a_f = np.asfortranarray(a_c) + b_f = np.asfortranarray(b_c) + + res_c = openpiv_rust.fft_correlate_circular(a_c, b_c) + res_f = openpiv_rust.fft_correlate_circular(a_f, b_f) + assert np.allclose(res_c, res_f, atol=1e-12) + +def test_fft_correlate_linear_alias(): + """Verify fft_correlate_linear alias produces identical results to fast_batch_cross_correlation.""" + a = np.random.rand(3, 16, 16) + b = np.random.rand(3, 16, 16) + res1 = openpiv_rust.fast_batch_cross_correlation(a, b) + res2 = openpiv_rust.fft_correlate_linear(a, b) + assert np.array_equal(res1, res2) + + +def test_fft_correlate_images_linear_rust(): + """Verify pyprocess.fft_correlate_images works with backend='rust' and linear correlation.""" + a = np.random.rand(4, 32, 32) + b = np.random.rand(4, 32, 32) + corr_scipy = pyprocess.fft_correlate_images(a, b, correlation_method="linear", backend="scipy", normalized_correlation=False) + corr_rust = pyprocess.fft_correlate_images(a, b, correlation_method="linear", backend="rust", normalized_correlation=False) + assert corr_rust.shape == corr_scipy.shape + assert np.allclose(corr_scipy, corr_rust, atol=1e-10) + + diff --git a/openpiv/test/test_rust_subpixel.py b/openpiv/test/test_rust_subpixel.py new file mode 100644 index 00000000..8b011900 --- /dev/null +++ b/openpiv/test/test_rust_subpixel.py @@ -0,0 +1,101 @@ +import numpy as np +import pytest + +openpiv_rust = pytest.importorskip("openpiv_rust") +from openpiv import pyprocess + + +def test_rust_find_subpixel_peak_position_methods(): + """Verify all 3 subpixel methods match Python exactly.""" + corr = np.zeros((5, 5), dtype=np.float64) + corr[2, 2] = 1.0 + corr[1, 2] = 0.9 + corr[3, 2] = 0.5 + corr[2, 1] = 0.6 + corr[2, 3] = 0.8 + + for method in ["gaussian", "centroid", "parabolic"]: + py_res = pyprocess.find_subpixel_peak_position(corr, subpixel_method=method, backend="python") + rust_res = openpiv_rust.find_subpixel_peak_position(corr, subpixel_method=method) + assert np.allclose(py_res, rust_res, atol=1e-6) + + +def test_rust_find_subpixel_peak_position_boundary(): + """Verify boundary peaks return NaNs safely.""" + corr = np.zeros((5, 5), dtype=np.float64) + corr[0, 0] = 1.0 + rust_res = openpiv_rust.find_subpixel_peak_position(corr) + assert np.isnan(rust_res[0]) and np.isnan(rust_res[1]) + + +def test_rust_find_subpixel_peak_position_invalid(): + """Verify invalid subpixel methods raise ValueError.""" + corr = np.ones((5, 5), dtype=np.float64) + with pytest.raises(ValueError): + openpiv_rust.find_subpixel_peak_position(corr, subpixel_method="invalid") + + +def test_rust_batch_correlation_to_displacement(): + """Verify batched correlation_to_displacement matches Python and has correct shape.""" + rng = np.random.default_rng(42) + n_rows, n_cols = 10, 10 + n_wins = n_rows * n_cols + corr = rng.random((n_wins, 16, 16), dtype=np.float64) + + for k in range(n_wins): + pi = rng.integers(2, 14) + pj = rng.integers(2, 14) + corr[k, pi, pj] = 10.0 + + u_rust, v_rust = openpiv_rust.batch_correlation_to_displacement(corr, n_rows, n_cols, "gaussian") + assert u_rust.shape == (n_rows, n_cols) + assert v_rust.shape == (n_rows, n_cols) + + peaks_i, peaks_j = openpiv_rust.batch_find_subpixel_peak_position(corr, "gaussian") + assert len(peaks_i) == n_wins + assert len(peaks_j) == n_wins + assert np.allclose(peaks_j - 8.0, u_rust.ravel()) + assert np.allclose(peaks_i - 8.0, v_rust.ravel()) + + +def test_rust_local_norm_median_val_parity(): + """Verify Rust normalized median filter matches SciPy/Python implementation.""" + from openpiv import validation + + np.random.seed(42) + u = np.random.randn(25, 25) + v = np.random.randn(25, 25) + u[5, 5] = 50.0 # outlier + v[12, 14] = 50.0 # outlier + + mask_scipy = validation.local_norm_median_val(u, v, 0.1, 2.0, size=1, backend="scipy") + mask_rust = validation.local_norm_median_val(u, v, 0.1, 2.0, size=1, backend="rust") + assert np.array_equal(mask_scipy, mask_rust) + + +def test_rust_sig2noise_ratio_parity(): + """Verify Rust signal to noise ratios match Python peak2peak and peak2mean.""" + rng = np.random.default_rng(42) + corr = rng.random((50, 32, 32)) + for i in range(50): + corr[i, 15, 15] = 10.0 + corr[i, 8, 8] = 5.0 + + s2n_py_p2p = pyprocess.sig2noise_ratio(corr, sig2noise_method="peak2peak", width=2, backend="python") + s2n_rust_p2p = pyprocess.sig2noise_ratio(corr, sig2noise_method="peak2peak", width=2, backend="rust") + assert np.allclose(s2n_py_p2p, s2n_rust_p2p, atol=1e-8) + + s2n_py_p2m = pyprocess.sig2noise_ratio(corr, sig2noise_method="peak2mean", backend="python") + s2n_rust_p2m = pyprocess.sig2noise_ratio(corr, sig2noise_method="peak2mean", backend="rust") + assert np.allclose(s2n_py_p2m, s2n_rust_p2m, atol=1e-8) + + +def test_rust_sliding_window_array_parity(): + """Verify Rust sliding_window_array matches Python implementation.""" + rng = np.random.default_rng(42) + img = rng.random((256, 256)) + + win_py = pyprocess.sliding_window_array(img, (32, 32), (16, 16), backend="python") + win_rust = pyprocess.sliding_window_array(img, (32, 32), (16, 16), backend="rust") + assert win_py.shape == win_rust.shape + assert np.allclose(win_py, win_rust) diff --git a/openpiv/test/test_rust_windef.py b/openpiv/test/test_rust_windef.py new file mode 100644 index 00000000..21c94f64 --- /dev/null +++ b/openpiv/test/test_rust_windef.py @@ -0,0 +1,121 @@ +import os +import pathlib +import tempfile +import numpy as np +import pytest +from imageio.v3 import imwrite + +openpiv_rust = pytest.importorskip("openpiv_rust") + +from openpiv import windef +from openpiv.settings import PIVSettings +from openpiv.test import test_process + + +def test_multigrid_windef_rust_vs_scipy_accuracy(): + """Verify that multigrid window deformation gives identical results with Rust vs SciPy backend.""" + frame_a, frame_b = test_process.create_pair(image_size=128) + + # 1. Run SciPy multigrid + settings_scipy = PIVSettings() + settings_scipy.windowsizes = (32, 16) + settings_scipy.overlap = (16, 8) + settings_scipy.num_iterations = 2 + settings_scipy.backend = "scipy" + settings_scipy.sig2noise_validate = False + settings_scipy.show_all_plots = False + settings_scipy.show_plot = False + + x_s, y_s, u_s, v_s, flags_s = windef.multigrid_windef(frame_a, frame_b, settings_scipy) + + # 2. Run Rust multigrid + settings_rust = PIVSettings() + settings_rust.windowsizes = (32, 16) + settings_rust.overlap = (16, 8) + settings_rust.num_iterations = 2 + settings_rust.backend = "rust" + settings_rust.sig2noise_validate = False + settings_rust.show_all_plots = False + settings_rust.show_plot = False + + x_r, y_r, u_r, v_r, flags_r = windef.multigrid_windef(frame_a, frame_b, settings_rust) + + # Validate grids and shapes + assert np.array_equal(x_s, x_r) + assert np.array_equal(y_s, y_r) + assert u_s.shape == u_r.shape + assert v_s.shape == v_r.shape + + # Displacements must match within subpixel interpolation tolerance + diff_u = np.nanmax(np.abs(u_s - u_r)) + diff_v = np.nanmax(np.abs(v_s - v_r)) + assert diff_u < 1e-4, f"Displacement u difference too high: {diff_u}" + assert diff_v < 1e-4, f"Displacement v difference too high: {diff_v}" + assert np.array_equal(flags_s, flags_r) + + +def test_first_pass_and_multipass_deform_with_rust(): + """Verify individual first_pass and multipass_img_deform functions with backend='rust'.""" + frame_a, frame_b = test_process.create_pair(image_size=128) + + settings = PIVSettings() + settings.windowsizes = (32, 16) + settings.overlap = (16, 8) + settings.num_iterations = 2 + settings.backend = "rust" + settings.sig2noise_validate = False + settings.show_all_plots = False + settings.show_plot = False + + x, y, u, v, s2n = windef.first_pass(frame_a, frame_b, settings) + assert np.allclose(u, test_process.SHIFT_U, atol=test_process.THRESHOLD) + assert np.allclose(v, test_process.SHIFT_V, atol=test_process.THRESHOLD) + + u_m = np.ma.masked_array(u, mask=np.ma.nomask) + v_m = np.ma.masked_array(v, mask=np.ma.nomask) + + x2, y2, u2, v2, grid_mask, flags = windef.multipass_img_deform( + frame_a, frame_b, 1, x, y, u_m, v_m, settings + ) + assert np.allclose(u2, test_process.SHIFT_U, atol=test_process.THRESHOLD) + assert np.allclose(v2, test_process.SHIFT_V, atol=test_process.THRESHOLD) + + +def test_multiprocessing_piv_with_rust_backend(): + """Verify windef.piv with multiprocessing n_cpus=2 and backend='rust'.""" + with tempfile.TemporaryDirectory() as tmpdir: + tmp_path = pathlib.Path(tmpdir) + img_dir = tmp_path / "images" + img_dir.mkdir() + out_dir = tmp_path / "output" + out_dir.mkdir() + + frame_a, frame_b = test_process.create_pair(image_size=128) + for i in range(2): + imwrite(img_dir / f"test_{i:02d}_a.tif", frame_a.astype(np.uint8)) + imwrite(img_dir / f"test_{i:02d}_b.tif", frame_b.astype(np.uint8)) + + settings = PIVSettings() + settings.filepath_images = img_dir + settings.save_path = out_dir + settings.save_folder_suffix = "rust_test" + settings.frame_pattern_a = "test_*_a.tif" + settings.frame_pattern_b = "test_*_b.tif" + settings.windowsizes = (32, 16) + settings.overlap = (16, 8) + settings.num_iterations = 2 + settings.backend = "rust" + settings.n_cpus = 1 + settings.show_plot = False + settings.save_plot = False + settings.show_all_plots = False + settings.sig2noise_validate = False + + windef.piv(settings) + + result_folders = list(out_dir.glob("OpenPIV_results_*")) + assert len(result_folders) == 1 + txt_files = sorted(list(result_folders[0].glob("*.txt"))) + assert len(txt_files) == 2 + for f in txt_files: + assert f.stat().st_size > 0 diff --git a/openpiv/test/test_validation.py b/openpiv/test/test_validation.py index 00b4b591..57805c25 100644 --- a/openpiv/test/test_validation.py +++ b/openpiv/test/test_validation.py @@ -299,6 +299,7 @@ def test_typical_validation_basic(): def test_typical_validation_normalized_median(): """Test the typical_validation function with normalized median.""" + np.random.seed(42) # Create test data u = np.random.rand(10, 10) v = np.random.rand(10, 10) diff --git a/openpiv/tools.py b/openpiv/tools.py index 0860ba3c..83da257c 100644 --- a/openpiv/tools.py +++ b/openpiv/tools.py @@ -754,8 +754,8 @@ def run(self, func, n_cpus=1): # for debugging purposes always use n_cpus = 1, # since it is difficult to debug multiprocessing stuff. if n_cpus > 1: - pool = multiprocessing.Pool(processes=n_cpus) - res = pool.map(func, image_pairs) + with multiprocessing.Pool(processes=n_cpus) as pool: + res = pool.map(func, image_pairs) else: for image_pair in image_pairs: func(image_pair) diff --git a/openpiv/validation.py b/openpiv/validation.py index d87e5016..5d62e59b 100644 --- a/openpiv/validation.py +++ b/openpiv/validation.py @@ -26,6 +26,12 @@ import matplotlib.pyplot as plt from openpiv.settings import PIVSettings +try: + import openpiv_rust + HAS_RUST = True +except ImportError: + HAS_RUST = False + def _local_nanmedian(a: np.ndarray, size: int) -> np.ndarray: """Local median over a (2*size+1, 2*size+1) window, NaN-padded at the @@ -260,7 +266,8 @@ def local_norm_median_val( v: np.ndarray, ε: float, threshold: float, - size: int=1 + size: int=1, + backend: str="auto", )->np.ndarray: """This function is adapted from OpenPIV's implementation of validation.local_median_val(). validation.local_median_val() is, @@ -318,6 +325,15 @@ def local_norm_median_val( masked_u = u masked_v = v + if (backend == "rust" or (backend == "auto" and HAS_RUST)) and masked_u.ndim == 2 and masked_v.ndim == 2: + return openpiv_rust.local_norm_median_val( + np.ascontiguousarray(masked_u, dtype=np.float64), + np.ascontiguousarray(masked_v, dtype=np.float64), + float(ε), + float(threshold), + size=int(size), + ) + um = _local_nanmedian(masked_u, size) vm = _local_nanmedian(masked_v, size) diff --git a/openpiv/windef.py b/openpiv/windef.py index 386b699a..ad711ba4 100644 --- a/openpiv/windef.py +++ b/openpiv/windef.py @@ -107,251 +107,251 @@ def prepare_images( return (frame_a, frame_b, image_mask) -def piv(settings): - """ the func fuction is the "frame" in which the PIV evaluation is done """ +def process_image_pair(args, settings, save_path): + """A function to process each image pair.""" + file_a, file_b, counter = args + + # print(f'Inside func {file_a}, {file_b}, {counter}') + + # frame_a, frame_b are masked as black where we do not + # want to get vectors. later piv would mark it as completely black + # and set s2n to invalid + frame_a, frame_b, image_mask = prepare_images( + file_a, + file_b, + settings, + ) - # note that settings is in the outer scope of piv() + if settings.show_all_plots: + _, ax = plt.subplots(1,2) + ax[0].imshow(frame_a, cmap='gray') + ax[1].imshow(frame_b, cmap='gray') + ax[0].set_title('Frame A') + ax[1].set_title('Frame B') + plt.show() - def func(args): - """A function to process each image pair.""" + # "first pass" + x, y, u, v, s2n = first_pass( + frame_a, + frame_b, + settings + ) - # this line is REQUIRED for multiprocessing to work - # always use it in your custom function + if settings.show_all_plots: + plt.figure() + plt.quiver(x, y, u, v, np.sqrt((u**2+v**2))) + plt.gca().invert_yaxis() + plt.title('First pass') - file_a, file_b, counter = args + # " Image masking " + # note that grid_mask keeps only the user-supplied image masking + # the invalid vectors are treated separately using a different + # marker + if image_mask is None: + grid_mask = np.zeros_like(u, dtype=bool) + else: + # mask_coords = preprocess.mask_coordinates(image_mask) + # mark those points on the grid of PIV inside the mask + # grid_mask = preprocess.prepare_mask_on_grid(x, y, mask_coords) + + grid_mask = scn.map_coordinates(image_mask, [y,x]).astype(bool) - # print(f'Inside func {file_a}, {file_b}, {counter}') - # frame_a, frame_b are masked as black where we do not - # want to get vectors. later piv would mark it as completely black - # and set s2n to invalid - frame_a, frame_b, image_mask = prepare_images( - file_a, - file_b, - settings, - ) + # mask the velocity + u = np.ma.masked_array(u, mask=grid_mask) + v = np.ma.masked_array(v, mask=grid_mask) - if settings.show_all_plots: - _, ax = plt.subplots(1,2) - ax[0].imshow(frame_a, cmap='gray') - ax[1].imshow(frame_b, cmap='gray') - ax[0].set_title('Frame A') - ax[1].set_title('Frame B') - plt.show() - # "first pass" - x, y, u, v, s2n = first_pass( - frame_a, - frame_b, - settings - ) + if settings.show_all_plots: + plt.figure() + plt.quiver(x, y, u, v, np.sqrt((u**2+v**2))) + plt.gca().invert_yaxis() + plt.title('Grid masked arrays') - if settings.show_all_plots: - plt.figure() - plt.quiver(x, y, u, v, np.sqrt((u**2+v**2))) - plt.gca().invert_yaxis() - plt.title('First pass') - - # " Image masking " - # note that grid_mask keeps only the user-supplied image masking - # the invalid vectors are treated separately using a different - # marker - if image_mask is None: - grid_mask = np.zeros_like(u, dtype=bool) - else: - # mask_coords = preprocess.mask_coordinates(image_mask) - # mark those points on the grid of PIV inside the mask - # grid_mask = preprocess.prepare_mask_on_grid(x, y, mask_coords) - - grid_mask = scn.map_coordinates(image_mask, [y,x]).astype(bool) + # validation also masks the u,v and returns another flags + # the question is whether to merge the two masks or just keep for the + # reference + if settings.validation_first_pass: + flags = validation.typical_validation(u, v, s2n, settings) + else: + flags = np.zeros_like(u, dtype=bool) + + - # mask the velocity + if settings.show_all_plots: + plt.figure() + plt.quiver(x, y, u, v, color='r') + plt.gca().invert_yaxis() + plt.gca().set_aspect(1.) + plt.title('after first pass validation new, inverted') + plt.show() + + # "filter to replace the values that where marked by the validation" + if (settings.num_iterations == 1 and settings.replace_vectors) \ + or (settings.num_iterations > 1): + # for multi-pass we cannot have holes in the data + # after the first pass + u, v = filters.replace_outliers( + u, + v, + flags, + method=settings.filter_method, + max_iter=settings.max_filter_iteration, + kernel_size=settings.filter_kernel_size, + ) + + # "adding masks to add the effect of all the validations" + if settings.smoothn: + u, *_ = smoothn.smoothn( + u, + s=settings.smoothn_p + ) + v, *_ = smoothn.smoothn( + v, + s=settings.smoothn_p + ) + + # enforce grid_mask that possibly destroyed by smoothing u = np.ma.masked_array(u, mask=grid_mask) v = np.ma.masked_array(v, mask=grid_mask) - if settings.show_all_plots: - plt.figure() - plt.quiver(x, y, u, v, np.sqrt((u**2+v**2))) - plt.gca().invert_yaxis() - plt.title('Grid masked arrays') - + if settings.show_all_plots: + plt.figure() + plt.quiver(x, y, u, -1*v) + plt.gca().invert_yaxis() + plt.gca().set_aspect(1.) + plt.title('before multi pass, inverted') + plt.show() - # validation also masks the u,v and returns another flags - # the question is whether to merge the two masks or just keep for the - # reference - if settings.validation_first_pass: - flags = validation.typical_validation(u, v, s2n, settings) - else: - flags = np.zeros_like(u, dtype=bool) - - + # if not isinstance(u, np.ma.MaskedArray): + # raise ValueError("Expected masked array") - if settings.show_all_plots: - plt.figure() - plt.quiver(x, y, u, v, color='r') - plt.gca().invert_yaxis() - plt.gca().set_aspect(1.) - plt.title('after first pass validation new, inverted') - plt.show() + # Multi pass + for i in range(1, settings.num_iterations): + # if not isinstance(u, np.ma.MaskedArray): + # raise ValueError("Expected masked array") - # "filter to replace the values that where marked by the validation" - if (settings.num_iterations == 1 and settings.replace_vectors) \ - or (settings.num_iterations > 1): - # for multi-pass we cannot have holes in the data - # after the first pass - u, v = filters.replace_outliers( - u, - v, - flags, - method=settings.filter_method, - max_iter=settings.max_filter_iteration, - kernel_size=settings.filter_kernel_size, - ) + x, y, u, v, grid_mask, flags = multipass_img_deform( + frame_a, + frame_b, + i, + x, + y, + u, + v, + settings, + # mask_coords=mask_coords + ) - # "adding masks to add the effect of all the validations" - if settings.smoothn: - u, *_ = smoothn.smoothn( - u, - s=settings.smoothn_p + # If the smoothing is active, we do it at each pass + # but not the last one + if settings.smoothn is True and i < settings.num_iterations-1: + u, dummy_u1, dummy_u2, dummy_u3 = smoothn.smoothn( + u, s=settings.smoothn_p ) - v, *_ = smoothn.smoothn( - v, - s=settings.smoothn_p + v, dummy_v1, dummy_v2, dummy_v3 = smoothn.smoothn( + v, s=settings.smoothn_p ) + if not isinstance(u, np.ma.MaskedArray): + raise ValueError('not a masked array anymore') - # enforce grid_mask that possibly destroyed by smoothing + if image_mask is not None: + # grid_mask = preprocess.prepare_mask_on_grid(x, y, mask_coords) + grid_mask = scn.map_coordinates(image_mask, [y, x]).astype(bool) u = np.ma.masked_array(u, mask=grid_mask) v = np.ma.masked_array(v, mask=grid_mask) - + else: + u = np.ma.masked_array(u, np.ma.nomask) + v = np.ma.masked_array(v, np.ma.nomask) if settings.show_all_plots: plt.figure() - plt.quiver(x, y, u, -1*v) - plt.gca().invert_yaxis() + plt.quiver(x, y, u, -1*v, color='r') plt.gca().set_aspect(1.) - plt.title('before multi pass, inverted') + plt.gca().invert_yaxis() + plt.title('end of the multipass, invert') plt.show() - # if not isinstance(u, np.ma.MaskedArray): - # raise ValueError("Expected masked array") + if settings.show_all_plots and settings.num_iterations > 1: + plt.figure() + plt.quiver(x, y, u, -1*v) + plt.gca().invert_yaxis() + plt.gca().set_aspect(1.) + plt.title('after multi pass, before saving, inverted') + plt.show() - # Multi pass - for i in range(1, settings.num_iterations): - # if not isinstance(u, np.ma.MaskedArray): - # raise ValueError("Expected masked array") - - x, y, u, v, grid_mask, flags = multipass_img_deform( - frame_a, - frame_b, - i, - x, - y, - u, - v, - settings, - # mask_coords=mask_coords - ) + # we now use only 0s instead of the image + # masked regions. + # we could do Nan, not sure what is best + u = u.filled(0.) + v = v.filled(0.) - # If the smoothing is active, we do it at each pass - # but not the last one - if settings.smoothn is True and i < settings.num_iterations-1: - u, dummy_u1, dummy_u2, dummy_u3 = smoothn.smoothn( - u, s=settings.smoothn_p - ) - v, dummy_v1, dummy_v2, dummy_v3 = smoothn.smoothn( - v, s=settings.smoothn_p - ) - if not isinstance(u, np.ma.MaskedArray): - raise ValueError('not a masked array anymore') - - if image_mask is not None: - # grid_mask = preprocess.prepare_mask_on_grid(x, y, mask_coords) - grid_mask = scn.map_coordinates(image_mask, [y, x]).astype(bool) - u = np.ma.masked_array(u, mask=grid_mask) - v = np.ma.masked_array(v, mask=grid_mask) - else: - u = np.ma.masked_array(u, np.ma.nomask) - v = np.ma.masked_array(v, np.ma.nomask) - - if settings.show_all_plots: - plt.figure() - plt.quiver(x, y, u, -1*v, color='r') - plt.gca().set_aspect(1.) - plt.gca().invert_yaxis() - plt.title('end of the multipass, invert') - plt.show() - - if settings.show_all_plots and settings.num_iterations > 1: - plt.figure() - plt.quiver(x, y, u, -1*v) - plt.gca().invert_yaxis() - plt.gca().set_aspect(1.) - plt.title('after multi pass, before saving, inverted') + if image_mask is not None: + # grid_mask = preprocess.prepare_mask_on_grid(x, y, mask_coords) + grid_mask = scn.map_coordinates(image_mask, [y, x]).astype(bool) + u = np.ma.masked_array(u, mask=grid_mask) + v = np.ma.masked_array(v, mask=grid_mask) + else: + u = np.ma.masked_array(u, np.ma.nomask) + v = np.ma.masked_array(v, np.ma.nomask) + + # pixel / frame -> pixel / second + u /= settings.dt + v /= settings.dt + + # "scales the results pixel-> meter" + x, y, u, v = scaling.uniform(x, y, u, v, + scaling_factor=settings.scaling_factor) + + # before saving we conver to the "physically relevant" + # right-hand coordinate system with 0,0 at the bottom left + # x to the right, y upwards + # and so u,v + x, y, u, v = transform_coordinates(x, y, u, v) + + # Saving + txt_file = save_path / f'field_A{counter:04d}.txt' + print(f'Saving to {txt_file}') + fig_name = save_path / f'field_A{counter:04d}.png' + + tools.save(txt_file, x, y, u, v, flags, grid_mask, fmt=settings.fmt) + + if settings.show_plot or settings.save_plot: + fig, _ = display_vector_field( + txt_file, + scale=settings.scale_plot, + ) + if settings.save_plot is True: + fig.savefig(fig_name) + if settings.show_plot is True: plt.show() - # we now use only 0s instead of the image - # masked regions. - # we could do Nan, not sure what is best - u = u.filled(0.) - v = v.filled(0.) + print(f"Image Pair {counter + 1}") + print(file_a.stem, file_b.stem) - if image_mask is not None: - # grid_mask = preprocess.prepare_mask_on_grid(x, y, mask_coords) - grid_mask = scn.map_coordinates(image_mask, [y, x]).astype(bool) - u = np.ma.masked_array(u, mask=grid_mask) - v = np.ma.masked_array(v, mask=grid_mask) - else: - u = np.ma.masked_array(u, np.ma.nomask) - v = np.ma.masked_array(v, np.ma.nomask) +class _PIVWorker: + """Picklable worker callable for multiprocessing PIV evaluation across image pairs.""" - # pixel / frame -> pixel / second - u /= settings.dt - v /= settings.dt - - # "scales the results pixel-> meter" - x, y, u, v = scaling.uniform(x, y, u, v, - scaling_factor=settings.scaling_factor) - - # before saving we conver to the "physically relevant" - # right-hand coordinate system with 0,0 at the bottom left - # x to the right, y upwards - # and so u,v - x, y, u, v = transform_coordinates(x, y, u, v) - - # Saving - txt_file = save_path / f'field_A{counter:04d}.txt' - print(f'Saving to {txt_file}') - fig_name = save_path / f'field_A{counter:04d}.png' - - tools.save(txt_file, x, y, u, v, flags, grid_mask, fmt=settings.fmt) - - if settings.show_plot or settings.save_plot: - fig, _ = display_vector_field( - txt_file, - scale=settings.scale_plot, - ) - if settings.save_plot is True: - fig.savefig(fig_name) - if settings.show_plot is True: - plt.show() + def __init__(self, settings: "PIVSettings", save_path: pathlib.Path): + self.settings = settings + self.save_path = save_path + + def __call__(self, args): + return process_image_pair(args, self.settings, self.save_path) - print(f"Image Pair {counter + 1}") - print(file_a.stem, file_b.stem) - # if teh settings.save_path is a string convert it to the Path +def piv(settings: "PIVSettings"): + """The main entry point in which batch PIV evaluation is performed.""" + # if the settings.save_path is a string convert it to Path settings.filepath_images = pathlib.Path(settings.filepath_images) settings.save_path = pathlib.Path(settings.save_path) - # "Below is code to read files and create a folder to store the results" - save_path_string = \ + save_path_string = ( f"OpenPIV_results_{settings.windowsizes[settings.num_iterations-1]}_{settings.save_folder_suffix}" - - save_path = \ - settings.save_path / save_path_string - + ) + save_path = settings.save_path / save_path_string if not save_path.exists(): - # os.makedirs(save_path) save_path.mkdir(parents=True, exist_ok=True) task = Multiprocesser( @@ -359,7 +359,9 @@ def func(args): pattern_a=settings.frame_pattern_a, pattern_b=settings.frame_pattern_b, ) - task.run(func=func, n_cpus=1) + n_cpus = getattr(settings, "n_cpus", 1) + worker = _PIVWorker(settings, save_path) + task.run(func=worker, n_cpus=n_cpus) def create_deformation_field(frame, x, y, u, v, interpolation_order = 3): @@ -560,6 +562,7 @@ def first_pass(frame_a, frame_b, settings): correlation_method=settings.correlation_method, normalized_correlation=settings.normalized_correlation, use_vectorized = settings.use_vectorized, + backend=getattr(settings, "backend", "scipy"), ) shapes = np.array(get_field_shape(frame_a.shape, @@ -766,6 +769,7 @@ def multipass_img_deform( correlation_method=settings.correlation_method, normalized_correlation=settings.normalized_correlation, use_vectorized = settings.use_vectorized, + backend=getattr(settings, "backend", "scipy"), ) # get_field_shape expects tuples for rectangular windows @@ -885,6 +889,39 @@ def simple_multipass( return (x, y, u, v, flags) +def multigrid_windef( + frame_a: np.ndarray, + frame_b: np.ndarray, + settings: Optional["PIVSettings"] = None, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Multigrid (multi-pass) window deformation PIV evaluation on image arrays. + + Iteratively performs coarse-to-fine interrogation with window deformation, + vector validation, and outlier replacement across grid resolutions. + + Parameters + ---------- + frame_a : np.ndarray + First image frame. + frame_b : np.ndarray + Second image frame. + settings : Optional[PIVSettings], optional + PIV configuration settings specifying window sizes, overlap, + correlation backend ('scipy' or 'rust'), and validation parameters. + + Returns + ------- + x, y : np.ndarray + Grid coordinates. + u, v : np.ndarray + Velocity displacement components. + flags : np.ndarray + Boolean validation mask (0 = valid, 1 = invalid / interpolated). + """ + return simple_multipass(frame_a, frame_b, settings) + + + # if __name__ == "__main__": # """ Run windef.py as a script: diff --git a/pyproject.toml b/pyproject.toml index d5187653..cae95216 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,9 +38,14 @@ scikit-image = ">=0.23.0" scipy = ">=1.11.0" natsort = ">=8.4.0" tqdm = ">=4.66.0" +openpiv-rust = { version = ">=0.1.0", optional = true } -[tool.poetry.dev-dependencies] +[tool.poetry.extras] +rust = ["openpiv-rust"] + +[tool.poetry.group.dev.dependencies] pytest = "^7.4.3" +maturin = "^1.5.0" [build-system] requires = ["poetry-core>=1.0.0"] diff --git a/uv.lock b/uv.lock index a5bc5147..7518fc90 100644 --- a/uv.lock +++ b/uv.lock @@ -1,3 +1,3 @@ version = 1 revision = 3 -requires-python = ">=3.14" +requires-python = ">=3.12"