diff --git a/.github/workflows/gpu-ci.yml b/.github/workflows/gpu-ci.yml
new file mode 100644
index 0000000..1a9ee4e
--- /dev/null
+++ b/.github/workflows/gpu-ci.yml
@@ -0,0 +1,526 @@
+name: GPU CI (RunPod on-demand)
+
+on:
+ workflow_dispatch:
+ inputs:
+ ref:
+ description: "Git ref to test (branch, tag, or SHA). Empty = current ref."
+ required: false
+ default: ""
+ run_benchmarks:
+ description: "Run backend benchmark suite"
+ required: false
+ type: boolean
+ default: true
+ benchmark_sizes:
+ description: "Comma-separated grid sizes (e.g. 128,192,512)"
+ required: false
+ default: "128,192"
+ benchmark_steps:
+ description: "Propagation steps per benchmark run"
+ required: false
+ default: "25"
+ benchmark_repeats:
+ description: "Benchmark repeats"
+ required: false
+ default: "3"
+ benchmark_warmup:
+ description: "Benchmark warmup iterations"
+ required: false
+ default: "1"
+
+concurrency:
+ group: gpu-ci-${{ inputs.ref != '' && inputs.ref || github.ref }}
+ cancel-in-progress: true
+
+permissions:
+ contents: read
+
+jobs:
+ reject-unauthorized:
+ if: ${{ github.actor != 'savowe' }}
+ runs-on: ubuntu-latest
+ steps:
+ - name: Block unauthorized trigger
+ run: |
+ echo "::error::This workflow may only be started by @savowe."
+ exit 1
+
+ cupy-tests-runpod:
+ name: CuPy backend tests (RunPod)
+ if: ${{ github.actor == 'savowe' }}
+ runs-on: ubuntu-latest
+ timeout-minutes: 90
+ environment: gpu-ci
+
+ env:
+ RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+ RUNPOD_GPU_TYPE_ID: ${{ vars.RUNPOD_GPU_TYPE_ID != '' && vars.RUNPOD_GPU_TYPE_ID || 'NVIDIA GeForce RTX 4090' }}
+ RUNPOD_IMAGE_NAME: ${{ vars.RUNPOD_IMAGE_NAME != '' && vars.RUNPOD_IMAGE_NAME || 'runpod/pytorch:1.0.3-cu1290-torch290-ubuntu2204' }}
+ RUNPOD_CLOUD_TYPE: ${{ vars.RUNPOD_CLOUD_TYPE != '' && vars.RUNPOD_CLOUD_TYPE || 'ALL' }}
+ RUNPOD_VOLUME_MOUNT_PATH: ${{ vars.RUNPOD_VOLUME_MOUNT_PATH != '' && vars.RUNPOD_VOLUME_MOUNT_PATH || '/workspace' }}
+ RUNPOD_MIN_VCPU: ${{ vars.RUNPOD_MIN_VCPU != '' && vars.RUNPOD_MIN_VCPU || '4' }}
+ RUNPOD_MIN_MEMORY_GB: ${{ vars.RUNPOD_MIN_MEMORY_GB != '' && vars.RUNPOD_MIN_MEMORY_GB || '15' }}
+ RUNPOD_VOLUME_GB: ${{ vars.RUNPOD_VOLUME_GB != '' && vars.RUNPOD_VOLUME_GB || '30' }}
+ RUNPOD_CONTAINER_DISK_GB: ${{ vars.RUNPOD_CONTAINER_DISK_GB != '' && vars.RUNPOD_CONTAINER_DISK_GB || '30' }}
+ RUN_BENCHMARKS: ${{ inputs.run_benchmarks && 'true' || 'false' }}
+ BENCHMARK_SIZES: ${{ inputs.benchmark_sizes != '' && inputs.benchmark_sizes || '128,192' }}
+ BENCHMARK_STEPS: ${{ inputs.benchmark_steps != '' && inputs.benchmark_steps || '25' }}
+ BENCHMARK_REPEATS: ${{ inputs.benchmark_repeats != '' && inputs.benchmark_repeats || '3' }}
+ BENCHMARK_WARMUP: ${{ inputs.benchmark_warmup != '' && inputs.benchmark_warmup || '1' }}
+
+ steps:
+ - name: Checkout repository
+ uses: actions/checkout@v4
+ with:
+ ref: ${{ inputs.ref != '' && inputs.ref || github.ref }}
+
+ - name: Resolve checkout SHA
+ id: resolve-sha
+ run: |
+ set -euo pipefail
+ echo "sha=$(git rev-parse HEAD)" >> "$GITHUB_OUTPUT"
+
+ - name: Validate required secrets
+ run: |
+ set -euo pipefail
+ test -n "${RUNPOD_API_KEY}" || { echo "Missing secret RUNPOD_API_KEY (set it in environment 'gpu-ci')."; exit 1; }
+ test -n "${{ secrets.RUNPOD_SSH_PRIVATE_KEY }}" || { echo "Missing secret RUNPOD_SSH_PRIVATE_KEY (set it in environment 'gpu-ci')."; exit 1; }
+
+ - name: Validate RunPod image reference
+ run: |
+ set -euo pipefail
+ case "${RUNPOD_IMAGE_NAME}" in
+ *:*) ;;
+ *)
+ echo "RUNPOD_IMAGE_NAME must include an explicit tag (e.g. runpod/pytorch:1.0.3-cu1290-torch290-ubuntu2204)."
+ exit 1
+ ;;
+ esac
+
+ - name: Validate benchmark parameters
+ run: |
+ set -euo pipefail
+ [[ "${BENCHMARK_SIZES}" =~ ^[0-9]+(,[0-9]+)*$ ]] || { echo "benchmark_sizes must be comma-separated positive integers"; exit 1; }
+ [[ "${BENCHMARK_STEPS}" =~ ^[0-9]+$ ]] || { echo "benchmark_steps must be a positive integer"; exit 1; }
+ [[ "${BENCHMARK_REPEATS}" =~ ^[0-9]+$ ]] || { echo "benchmark_repeats must be a positive integer"; exit 1; }
+ [[ "${BENCHMARK_WARMUP}" =~ ^[0-9]+$ ]] || { echo "benchmark_warmup must be a non-negative integer"; exit 1; }
+
+ [ "${BENCHMARK_STEPS}" -ge 1 ] || { echo "benchmark_steps must be >= 1"; exit 1; }
+ [ "${BENCHMARK_REPEATS}" -ge 1 ] || { echo "benchmark_repeats must be >= 1"; exit 1; }
+
+ - name: Install workflow dependencies
+ run: |
+ sudo apt-get update
+ sudo apt-get install -y jq openssh-client
+
+ - name: Configure SSH key
+ run: |
+ set -euo pipefail
+ mkdir -p ~/.ssh
+ printf '%s\n' '${{ secrets.RUNPOD_SSH_PRIVATE_KEY }}' > ~/.ssh/id_runpod
+ chmod 600 ~/.ssh/id_runpod
+
+ - name: Create RunPod on-demand pod
+ id: create-pod
+ run: |
+ set -euo pipefail
+
+ QUERY='mutation ($input: PodFindAndDeployOnDemandInput!) { podFindAndDeployOnDemand(input: $input) { id desiredStatus machine { podHostId } } }'
+
+ PAYLOAD=$(jq -n \
+ --arg query "$QUERY" \
+ --arg cloud "$RUNPOD_CLOUD_TYPE" \
+ --arg gpu "$RUNPOD_GPU_TYPE_ID" \
+ --arg image "$RUNPOD_IMAGE_NAME" \
+ --arg name "pytalises-gpu-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" \
+ --arg mount "$RUNPOD_VOLUME_MOUNT_PATH" \
+ --arg ports "22/tcp" \
+ --arg minVcpu "$RUNPOD_MIN_VCPU" \
+ --arg minMem "$RUNPOD_MIN_MEMORY_GB" \
+ --arg vol "$RUNPOD_VOLUME_GB" \
+ --arg disk "$RUNPOD_CONTAINER_DISK_GB" \
+ '{
+ query: $query,
+ variables: {
+ input: {
+ cloudType: $cloud,
+ gpuCount: 1,
+ gpuTypeId: $gpu,
+ name: $name,
+ imageName: $image,
+ startSsh: true,
+ supportPublicIp: true,
+ ports: $ports,
+ minVcpuCount: ($minVcpu | tonumber),
+ minMemoryInGb: ($minMem | tonumber),
+ volumeInGb: ($vol | tonumber),
+ containerDiskInGb: ($disk | tonumber),
+ volumeMountPath: $mount
+ }
+ }
+ }')
+
+ RESPONSE=$(curl --silent --show-error --fail \
+ --request POST \
+ --header 'content-type: application/json' \
+ --url "https://api.runpod.io/graphql?api_key=${RUNPOD_API_KEY}" \
+ --data "$PAYLOAD")
+
+ echo "$RESPONSE" | jq .
+
+ if [ "$(echo "$RESPONSE" | jq '.errors | length // 0')" != "0" ]; then
+ echo "RunPod API returned errors while creating pod"
+ exit 1
+ fi
+
+ POD_ID=$(echo "$RESPONSE" | jq -r '.data.podFindAndDeployOnDemand.id // empty')
+ POD_HOST_ID=$(echo "$RESPONSE" | jq -r '.data.podFindAndDeployOnDemand.machine.podHostId // empty')
+ if [ -z "$POD_ID" ]; then
+ echo "RunPod pod id missing in response"
+ exit 1
+ fi
+
+ echo "pod_id=$POD_ID" >> "$GITHUB_OUTPUT"
+ if [ -n "$POD_HOST_ID" ]; then
+ echo "pod_host_id=$POD_HOST_ID" >> "$GITHUB_OUTPUT"
+ fi
+
+ - name: Wait for SSH endpoint
+ id: wait-ssh
+ env:
+ POD_ID: ${{ steps.create-pod.outputs.pod_id }}
+ POD_HOST_ID: ${{ steps.create-pod.outputs.pod_host_id }}
+ RUN_BENCHMARKS: ${{ env.RUN_BENCHMARKS }}
+ run: |
+ set -euo pipefail
+
+ QUERY='query ($input: PodFilter) { pod(input: $input) { id desiredStatus runtime { ports { ip isIpPublic privatePort publicPort type } } } }'
+
+ for attempt in $(seq 1 30); do
+ PAYLOAD=$(jq -n \
+ --arg query "$QUERY" \
+ --arg podId "$POD_ID" \
+ '{query: $query, variables: {input: {podId: $podId}}}')
+
+ HTTP_BODY=$(mktemp)
+ HTTP_CODE=$(curl --silent --show-error \
+ --output "$HTTP_BODY" \
+ --write-out '%{http_code}' \
+ --request POST \
+ --header 'content-type: application/json' \
+ --url "https://api.runpod.io/graphql?api_key=${RUNPOD_API_KEY}" \
+ --data "$PAYLOAD")
+
+ RESPONSE=$(cat "$HTTP_BODY")
+ rm -f "$HTTP_BODY"
+
+ if [ "$HTTP_CODE" != "200" ]; then
+ echo "Attempt ${attempt}: RunPod returned HTTP ${HTTP_CODE}; retrying..."
+ sleep 10
+ continue
+ fi
+
+ ERRORS=$(echo "$RESPONSE" | jq '.errors | length // 0')
+ if [ "$ERRORS" != "0" ]; then
+ echo "Attempt ${attempt}: GraphQL errors while waiting for pod endpoint:"
+ echo "$RESPONSE" | jq '.errors'
+ sleep 10
+ continue
+ fi
+
+ STATUS=$(echo "$RESPONSE" | jq -r '.data.pod.desiredStatus // "UNKNOWN"')
+
+ HOST=$(echo "$RESPONSE" | jq -r '.data.pod.runtime.ports[]? | select((.privatePort | tostring) == "22" and .publicPort != null and .ip != null and (.isIpPublic == true or .isIpPublic == null)) | .ip' | head -n1)
+ PORT=$(echo "$RESPONSE" | jq -r '.data.pod.runtime.ports[]? | select((.privatePort | tostring) == "22" and .publicPort != null and .ip != null and (.isIpPublic == true or .isIpPublic == null)) | .publicPort' | head -n1)
+
+ if [ -n "${HOST}" ] && [ -n "${PORT}" ] && [ "${HOST}" != "null" ] && [ "${PORT}" != "null" ]; then
+ echo "ssh_host=$HOST" >> "$GITHUB_OUTPUT"
+ echo "ssh_port=$PORT" >> "$GITHUB_OUTPUT"
+ echo "ssh_user=root" >> "$GITHUB_OUTPUT"
+ echo "ssh_mode=public-ip" >> "$GITHUB_OUTPUT"
+ echo "Found public SSH endpoint: ${HOST}:${PORT}"
+ exit 0
+ fi
+
+ if [ "$STATUS" = "RUNNING" ] && [ -n "${POD_HOST_ID:-}" ]; then
+ if [ "${RUN_BENCHMARKS}" = "true" ] && [ "$attempt" -lt 18 ]; then
+ echo "Attempt ${attempt}: pod is running but public SSH endpoint is not ready yet (benchmarks prefer public IP for SCP artifact transfer)."
+ else
+ if [ "${RUN_BENCHMARKS}" = "true" ]; then
+ echo "Public SSH endpoint still unavailable after waiting; falling back to RunPod SSH proxy with log-based benchmark payload extraction."
+ else
+ echo "No public SSH endpoint yet; falling back to RunPod SSH proxy via ssh.runpod.io"
+ fi
+ echo "ssh_host=ssh.runpod.io" >> "$GITHUB_OUTPUT"
+ echo "ssh_port=22" >> "$GITHUB_OUTPUT"
+ echo "ssh_user=$POD_HOST_ID" >> "$GITHUB_OUTPUT"
+ echo "ssh_mode=proxy" >> "$GITHUB_OUTPUT"
+ exit 0
+ fi
+ fi
+
+ echo "Attempt ${attempt}: status=${STATUS}, waiting for SSH endpoint..."
+ sleep 10
+ done
+
+ echo "Timed out waiting for RunPod SSH endpoint"
+ exit 1
+
+ - name: Wait for SSH connectivity
+ env:
+ SSH_HOST: ${{ steps.wait-ssh.outputs.ssh_host }}
+ SSH_PORT: ${{ steps.wait-ssh.outputs.ssh_port }}
+ SSH_USER: ${{ steps.wait-ssh.outputs.ssh_user }}
+ SSH_MODE: ${{ steps.wait-ssh.outputs.ssh_mode }}
+ run: |
+ set -euo pipefail
+ echo "Using SSH mode: ${SSH_MODE} (${SSH_USER}@${SSH_HOST}:${SSH_PORT})"
+
+ SSH_TTY_FLAG=""
+ if [ "${SSH_MODE}" = "proxy" ]; then
+ echo "Proxy mode detected; skipping explicit connectivity probe."
+ exit 0
+ fi
+
+ for attempt in $(seq 1 30); do
+ set +e
+ OUT=$(timeout 20s ssh ${SSH_TTY_FLAG} \
+ -o StrictHostKeyChecking=no \
+ -o UserKnownHostsFile=/dev/null \
+ -o ConnectTimeout=8 \
+ -o ConnectionAttempts=1 \
+ -o BatchMode=yes \
+ -i ~/.ssh/id_runpod -p "$SSH_PORT" "$SSH_USER"@"$SSH_HOST" \
+ 'echo connected; exit' 2>&1)
+ SSH_RC=$?
+ set -e
+
+ echo "$OUT"
+ if [ "$SSH_RC" -eq 0 ] && echo "$OUT" | grep -q "connected"; then
+ exit 0
+ fi
+
+ echo "Attempt ${attempt}: waiting for SSH connectivity..."
+ sleep 10
+ done
+ echo "SSH never became reachable"
+ exit 1
+
+ - name: Upload source, run GPU tests, and collect benchmark output
+ env:
+ SSH_HOST: ${{ steps.wait-ssh.outputs.ssh_host }}
+ SSH_PORT: ${{ steps.wait-ssh.outputs.ssh_port }}
+ SSH_USER: ${{ steps.wait-ssh.outputs.ssh_user }}
+ SSH_MODE: ${{ steps.wait-ssh.outputs.ssh_mode }}
+ RUN_BENCHMARKS: ${{ env.RUN_BENCHMARKS }}
+ run: |
+ set -euo pipefail
+
+ SSH_TTY_FLAG=""
+ if [ "${SSH_MODE}" = "proxy" ]; then
+ SSH_TTY_FLAG="-tt"
+ fi
+
+ SSH_LOG=$(mktemp)
+ set +e
+ timeout 30m ssh ${SSH_TTY_FLAG} \
+ -o StrictHostKeyChecking=no \
+ -o UserKnownHostsFile=/dev/null \
+ -o ConnectTimeout=10 \
+ -o ConnectionAttempts=1 \
+ -o BatchMode=yes \
+ -i ~/.ssh/id_runpod -p "$SSH_PORT" "$SSH_USER"@"$SSH_HOST" \
+ "bash -se" \
+ <<'REMOTE' > "$SSH_LOG" 2>&1
+ set -euo pipefail
+
+ REPO='${{ github.repository }}'
+ SHA='${{ steps.resolve-sha.outputs.sha }}'
+ RUN_BENCH='${{ env.RUN_BENCHMARKS }}'
+ BENCH_SIZES='${{ env.BENCHMARK_SIZES }}'
+ BENCH_STEPS='${{ env.BENCHMARK_STEPS }}'
+ BENCH_REPEATS='${{ env.BENCHMARK_REPEATS }}'
+ BENCH_WARMUP='${{ env.BENCHMARK_WARMUP }}'
+ export REPO SHA RUN_BENCH BENCH_SIZES BENCH_STEPS BENCH_REPEATS BENCH_WARMUP
+
+ cat > /tmp/pytalises-ci.sh <<'SCRIPT'
+ #!/usr/bin/env bash
+ set -euo pipefail
+
+ if ! command -v python3 >/dev/null 2>&1; then
+ echo "python3 is required in the selected RunPod image."
+ exit 1
+ fi
+ if ! command -v git >/dev/null 2>&1; then
+ echo "git is required in the selected RunPod image."
+ exit 1
+ fi
+
+ rm -rf /workspace/pytalises
+ git clone --depth 1 https://github.com/${REPO}.git /workspace/pytalises
+ cd /workspace/pytalises
+ git fetch --depth 1 origin ${SHA}
+ git checkout --detach ${SHA}
+
+ python3 -m pip install --upgrade pip
+ python3 -m pip install -e ".[gpu]" pytest
+ pytest tests/backend_test.py tests/cupy_backend_test.py -q
+
+ if [ "${RUN_BENCH}" = "true" ]; then
+ python3 benchmarks/backend_benchmark.py \
+ --sizes "${BENCH_SIZES}" \
+ --workloads free,potential \
+ --steps "${BENCH_STEPS}" \
+ --repeats "${BENCH_REPEATS}" \
+ --warmup "${BENCH_WARMUP}" \
+ --json-out /workspace/backend-benchmark.json
+
+ python3 - <<'PY'
+ import base64
+ p = '/workspace/backend-benchmark.json'
+ with open(p, 'rb') as f:
+ data = f.read()
+ print('BENCHMARK_JSON_B64=' + base64.b64encode(data).decode())
+ PY
+ fi
+ SCRIPT
+
+ chmod +x /tmp/pytalises-ci.sh
+ set +e
+ /tmp/pytalises-ci.sh
+ EC=$?
+ set -e
+ echo "__RUNPOD_EXIT_CODE__:${EC}"
+ exit "$EC"
+ REMOTE
+ SSH_RC=$?
+ set -e
+
+ cat "$SSH_LOG"
+
+ REMOTE_EC=$(grep -o '__RUNPOD_EXIT_CODE__:[0-9]\+' "$SSH_LOG" | tail -n1 | cut -d: -f2)
+ if [ -z "${REMOTE_EC:-}" ]; then
+ echo "Remote exit marker missing"
+ exit 1
+ fi
+ if [ "$SSH_RC" -ne 0 ] || [ "$REMOTE_EC" -ne 0 ]; then
+ echo "Remote CI script failed (ssh_rc=$SSH_RC, remote_ec=$REMOTE_EC)"
+ exit 1
+ fi
+
+ if [ "${RUN_BENCHMARKS}" = "true" ]; then
+ if [ "${SSH_MODE}" = "public-ip" ]; then
+ timeout 5m scp \
+ -o StrictHostKeyChecking=no \
+ -o UserKnownHostsFile=/dev/null \
+ -o ConnectTimeout=10 \
+ -o ConnectionAttempts=1 \
+ -o BatchMode=yes \
+ -i ~/.ssh/id_runpod -P "$SSH_PORT" "$SSH_USER"@"$SSH_HOST":/workspace/backend-benchmark.json /tmp/backend-benchmark.json
+ else
+ echo "Public IP mode unavailable; falling back to benchmark payload extraction from SSH logs."
+ CLEAN_LOG=$(mktemp)
+ tr -d '\r' < "$SSH_LOG" \
+ | sed -r 's/\x1B\[[0-9;?]*[ -\/]*[@-~]//g' \
+ > "$CLEAN_LOG"
+
+ B64=$(grep -ao 'BENCHMARK_JSON_B64=[A-Za-z0-9+/=]\+' "$CLEAN_LOG" | tail -n1 | cut -d= -f2-)
+ if [ -z "$B64" ]; then
+ echo "Could not extract benchmark payload from SSH output"
+ rm -f "$CLEAN_LOG"
+ exit 1
+ fi
+ printf '%s' "$B64" | base64 -d > /tmp/backend-benchmark.json
+ rm -f "$CLEAN_LOG"
+ fi
+ fi
+
+ rm -f "$SSH_LOG"
+
+ - name: Add benchmark summary
+ if: ${{ inputs.run_benchmarks }}
+ run: |
+ python3 - <<'PY'
+ import json
+ import os
+ from pathlib import Path
+
+ path = Path('/tmp/backend-benchmark.json')
+ data = json.loads(path.read_text())
+
+ lines = [
+ '## Backend benchmark summary',
+ '',
+ '| Backend | Workload | Grid | Mean (s) | Std (s) |',
+ '|---|---|---:|---:|---:|',
+ ]
+
+ for row in data.get('results', []):
+ lines.append(
+ f"| {row['backend']} | {row['workload']} | {row['size']} | "
+ f"{row['mean_seconds']:.4f} | {row['std_seconds']:.4f} |"
+ )
+
+ speedups = data.get('speedups_numpy_over_cupy', {})
+ if speedups:
+ lines.append('')
+ lines.append('**Speedups (NumPy/CuPy):**')
+ for key, value in sorted(speedups.items()):
+ lines.append(f"- `{key}`: {value:.2f}x")
+
+ metadata = data.get('metadata', {})
+ if isinstance(metadata, dict):
+ lines.append('')
+ lines.append(
+ f"**Host:** CPU cores={metadata.get('cpu_count')}, "
+ f"backends={metadata.get('available_backends')}"
+ )
+ cuda = metadata.get('cuda')
+ if isinstance(cuda, dict):
+ lines.append(
+ f"**CUDA device:** {cuda.get('device_name')} "
+ f"(CuPy {cuda.get('cupy_version')})"
+ )
+
+ parity = data.get('parity')
+ if isinstance(parity, dict):
+ lines.append('')
+ lines.append(
+ '**Parity max abs diff:** '
+ f"density={parity['max_density_abs_diff']:.3e}, "
+ f"occupation={parity['max_occupation_abs_diff']:.3e}"
+ )
+
+ summary_path = Path(os.environ['GITHUB_STEP_SUMMARY'])
+ with open(summary_path, 'a', encoding='utf-8') as fh:
+ fh.write('\n'.join(lines) + '\n')
+ PY
+
+ - name: Upload benchmark artifact
+ if: ${{ inputs.run_benchmarks }}
+ uses: actions/upload-artifact@v4
+ with:
+ name: backend-benchmark-${{ github.run_id }}
+ path: /tmp/backend-benchmark.json
+
+ - name: Terminate RunPod pod
+ if: ${{ always() && steps.create-pod.outputs.pod_id != '' }}
+ env:
+ POD_ID: ${{ steps.create-pod.outputs.pod_id }}
+ run: |
+ set -euo pipefail
+
+ QUERY='mutation ($input: PodTerminateInput!) { podTerminate(input: $input) }'
+ PAYLOAD=$(jq -n \
+ --arg query "$QUERY" \
+ --arg podId "$POD_ID" \
+ '{query: $query, variables: {input: {podId: $podId}}}')
+
+ RESPONSE=$(curl --silent --show-error --fail \
+ --request POST \
+ --header 'content-type: application/json' \
+ --url "https://api.runpod.io/graphql?api_key=${RUNPOD_API_KEY}" \
+ --data "$PAYLOAD")
+
+ echo "$RESPONSE" | jq .
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
new file mode 100644
index 0000000..eb6cdd8
--- /dev/null
+++ b/.readthedocs.yaml
@@ -0,0 +1,18 @@
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html
+
+version: 2
+
+build:
+ os: ubuntu-22.04
+ tools:
+ python: "3.12"
+
+sphinx:
+ configuration: docs/source/conf.py
+
+python:
+ install:
+ - requirements: docs/requirements.txt
+ - method: pip
+ path: .
diff --git a/README.md b/README.md
index be7c2cc..6ae2b61 100644
--- a/README.md
+++ b/README.md
@@ -1,91 +1,155 @@
-[](https://pypi.org/project/pytalises/)
[](https://pypi.org/project/pytalises/)
[](https://anaconda.org/conda-forge/pytalises)
+[](https://pypi.org/project/pytalises/)
[](https://github.com/savowe/pytalises/actions/workflows/ci.yml)
-[](https://github.com/savowe/pytalises/actions/workflows/python-publish.yml)
-[](https://pytalises.readthedocs.io/en/latest/?badge=latest)
+[](https://pytalises.readthedocs.io/en/latest/)
+
-
+
# pyTALISES
-**pyTALISES** (This Ain't a LInear Schrödinger Equation Solver) is an easy-to-use Python implementation of the Split-Step Fourier Method, for numeric calculation of a wave function's time-propagation under the Schrödinger equation.
+**This Ain't a LInear Schrödinger Equation Solver**
+
+*Split-step Fourier method for quantum wavefunction propagation*
+
+
+
+
+
+## Quickstart
+
+```python
+import pytalises as pt
+
+# Define a 1D grid
+grid = pt.Grid(shape=(256,), extent=((-10, 10),))
+
+# Create a Gaussian wavepacket
+psi = pt.Wavefunction(initial="exp(-x**2)", grid=grid)
+
+# Propagate freely for 1000 steps
+psi.freely_propagate(steps=1000, dt=1e-4)
+```
+
+That's it. For coupled states, potentials, and GPU acceleration, see the [documentation](https://pytalises.readthedocs.io/).
+
+## Installation
+
+**pip**:
+```bash
+pip install pytalises
+```
+
+**With GPU support** (requires NVIDIA CUDA):
+```bash
+pip install pytalises[gpu]
+```
+
+## Why pyTALISES?
+
+pyTALISES excels at **position-space wavefunction dynamics**:
+
+- **Matter-wave propagation** — free expansion, wavepacket dynamics
+- **Atom optics** — Bragg diffraction, beam splitters, interferometry
+- **Cold atom physics** — BEC dynamics, nonlinear interactions
+- **Light-matter coupling** — Rabi oscillations, Raman transitions
+- **Multi-level systems** — arbitrary internal state structure
+
+If you need to propagate wavefunctions on spatial grids with time-dependent potentials, pyTALISES makes it simple.
+
+## Features
-### Features
-- Calculation of a wavefunction's time propagation under a (non)linear Schrödinger equation: 
-- the wave-function  may include an arbitrary number of internal and external degrees of freedom
-- simple implementation of Hamiltonians
-- speed of the [FFTW](https://pypi.org/project/pyFFTW/), [BLAS](https://www.netlib.org/blas/) and [numexpr](https://numexpr.readthedocs.io/en/latest/) libraries with multithreading
-- crucial functions are just-in-time compiled with [numba](https://numba.readthedocs.io/en/stable/)
+- **Multi-dimensional grids** — 1D, 2D, 3D spatial simulations
+- **Coupled internal states** — Two-level systems, Raman transitions, Bragg diffraction
+- **String-based potentials** — Define V(x,t) as human-readable expressions
+- **Performance** — FFTW, numba JIT, numexpr, multithreading
+- **GPU acceleration** — Optional CuPy backend for large grids
-### v2 API Quickstart
+## Examples
+
+
+Two-level Rabi oscillations
```python
import pytalises as pt
grid = pt.Grid(shape=(256,), extent=((-4, 4),))
psi = pt.Wavefunction(
- initial=["exp(-x**2)", "0"],
+ initial=["exp(-x**2)", "0"], # Start in ground state
grid=grid,
)
+# Coupling potential (off-diagonal drives transitions)
V = pt.HermitianPotential.from_lower_triangular([
- "0",
- "Omega*cos(t)",
- "Delta",
+ "0", # V_11: ground state energy
+ "Omega*cos(t)", # V_21: coupling
+ "Delta", # V_22: excited state detuning
])
psi.propagate(
potential=V,
steps=1000,
dt=1e-6,
- variables={"Omega": 2.0, "Delta": 1.0},
- options=pt.PropagationOptions(threads=4),
+ variables={"Omega": 2.0, "Delta": 0.0},
)
```
-For migration details from the pre-v2 API, see `docs/source/v2_migration.rst`.
+
-### Documentation
-Read the [documentation](https://pytalises.readthedocs.io/en/latest/) to learn more about pytalises' capabilities.
-The documentation features many examples, among others
-[2D harmonic potentials](https://pytalises.readthedocs.io/en/latest/examples.html#2D-harmonic-potential),
-[BEC scattering](https://pytalises.readthedocs.io/en/latest/examples.html#Nonlinear-interactions-between-internal-states),
-[three-level Raman transitions](https://pytalises.readthedocs.io/en/latest/additional_examples.html#Three-level-Raman-transitions),
-[single-Bragg diffraction](https://pytalises.readthedocs.io/en/latest/additional_examples.html#Single-Bragg-diffraction)
-and
-[atom interferometry](https://pytalises.readthedocs.io/en/latest/additional_examples.html#Light-pulse-atom-interferometry-with-single-Bragg-diffraction).
+
+GPU-accelerated propagation
+```python
+import pytalises as pt
+# Auto-detects GPU if available
+psi.propagate(
+ potential=V,
+ steps=10000,
+ dt=1e-6,
+ options=pt.PropagationOptions(backend="auto"),
+)
-Installing pytalises
-====================
-**We recommend installing pytalises via conda**
+# Or explicitly:
+options = pt.PropagationOptions(backend="cupy", dtype="complex128")
+```
-pytalises supports Python 3.9 through 3.13, and wheels are built and tested across those versions in CI.
+See the [GPU documentation](https://pytalises.readthedocs.io/en/latest/gpu_backend.html) for performance characteristics.
+
-#### Using conda
+## Documentation
-Installing `pytalises` from the `conda-forge` channel can be achieved by adding `conda-forge` to your channels with:
+**[Full documentation](https://pytalises.readthedocs.io/)**
-```
-conda config --add channels conda-forge
-```
+Includes:
+- [Usage examples](https://pytalises.readthedocs.io/en/latest/examples.html) — Gaussian wavepackets, harmonic potentials, BEC
+- [Advanced examples](https://pytalises.readthedocs.io/en/latest/additional_examples.html) — Raman transitions, Bragg diffraction, atom interferometry
+- [GPU backend guide](https://pytalises.readthedocs.io/en/latest/gpu_backend.html) — Installation, performance, troubleshooting
+- [Algorithm notes](https://pytalises.readthedocs.io/en/latest/notes.html) — Split-step Fourier method explained
-Once the `conda-forge` channel has been enabled, `pytalises` can be installed with:
+## Development
+```bash
+git clone https://github.com/savowe/pytalises.git
+cd pytalises
+pip install -e ".[dev]"
+pytest tests/
```
-conda install pytalises
-```
-
-#### Using pip
+## Citation
-pytalises is available on the Python Package Index and can be installed via
+If you use pyTALISES in academic work, please cite:
+```bibtex
+@software{pytalises,
+ author = {Vowe, Sascha},
+ title = {pyTALISES: Split-Step Fourier Method for the Schrödinger Equation},
+ url = {https://github.com/savowe/pytalises},
+}
```
-pip install pytalises
-```
-it has dependencies via `scipy` and `numba` on BLAS and LAPACK libraries that are not always found on windows systems. For linux they can usually be located.
+## License
+
+GPLv3 — see [LICENSE](LICENSE).
diff --git a/benchmarks/backend_benchmark.py b/benchmarks/backend_benchmark.py
new file mode 100644
index 0000000..fc8f616
--- /dev/null
+++ b/benchmarks/backend_benchmark.py
@@ -0,0 +1,467 @@
+#!/usr/bin/env python3
+"""Benchmark pyTALISES backend performance and parity metrics.
+
+This script is designed to compare backend behavior on the same machine,
+including GPU pods where both CPU and GPU backends are available.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+import platform
+import sys
+import time
+from dataclasses import asdict, dataclass
+
+import numpy as np
+
+# Support direct script execution from repo root:
+# python benchmarks/backend_benchmark.py ...
+if __package__ in (None, ""):
+ repo_root = Path(__file__).resolve().parents[1]
+ if str(repo_root) not in sys.path:
+ sys.path.insert(0, str(repo_root))
+
+import pytalises as pt
+
+
+@dataclass
+class BenchmarkResult:
+ backend: str
+ workload: str
+ size: int
+ steps: int
+ dt: float
+ repeats: int
+ mean_seconds: float
+ std_seconds: float
+
+
+def _sync_backend(backend: str) -> None:
+ if backend != "cupy":
+ return
+ import cupy as cp
+
+ cp.cuda.Stream.null.synchronize()
+
+
+def _make_problem(size: int):
+ grid = pt.Grid(shape=(size,), extent=((-4.0, 4.0),))
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+ return grid, initial, potential
+
+
+def _run_profiled_propagation(
+ *,
+ psi,
+ workload: str,
+ potential,
+ steps: int,
+ dt: float,
+ options,
+) -> dict[str, dict[str, float | int]]:
+ if workload == "free":
+ propagator = pt.Propagator(
+ psi,
+ potential=pt.DiagonalPotential(["0"] * psi.num_int_dim),
+ variables={},
+ options=options,
+ )
+ for _ in range(steps):
+ propagator.kinetic_prop(dt)
+ return propagator.stage_timings()
+
+ if workload == "potential":
+ propagator = pt.Propagator(
+ psi,
+ potential=potential,
+ variables={},
+ options=options,
+ )
+ propagator.kinetic_prop(dt / 2)
+ propagator.potential_prop(dt)
+ for _ in range(steps - 1):
+ propagator.kinetic_prop(dt)
+ propagator.potential_prop(dt)
+ propagator.kinetic_prop(dt / 2)
+ return propagator.stage_timings()
+
+ raise ValueError(f"Unknown workload '{workload}'")
+
+
+def _summarize_stage_runs(
+ stage_runs: list[dict[str, dict[str, float | int]]],
+ case_mean_seconds: float,
+) -> dict[str, dict[str, float]]:
+ if not stage_runs:
+ return {}
+
+ stage_names = sorted({name for run in stage_runs for name in run})
+ summary: dict[str, dict[str, float]] = {}
+
+ for name in stage_names:
+ seconds = np.asarray(
+ [float(run.get(name, {}).get("seconds", 0.0)) for run in stage_runs],
+ dtype=float,
+ )
+ calls = np.asarray(
+ [float(run.get(name, {}).get("calls", 0.0)) for run in stage_runs],
+ dtype=float,
+ )
+
+ mean_seconds = float(np.mean(seconds))
+ summary[name] = {
+ "mean_seconds": mean_seconds,
+ "std_seconds": float(np.std(seconds)),
+ "mean_calls": float(np.mean(calls)),
+ "share_of_case_runtime": (
+ mean_seconds / case_mean_seconds if case_mean_seconds > 0 else 0.0
+ ),
+ }
+
+ return summary
+
+
+def _run_case(
+ *,
+ backend: str,
+ workload: str,
+ size: int,
+ steps: int,
+ dt: float,
+ repeats: int,
+ warmup: int,
+ coupled_2x2_mode: str,
+ coupled_2x2_kernel: str,
+ dtype: str,
+) -> tuple[BenchmarkResult, dict[str, dict[str, float]]]:
+ timings: list[float] = []
+ stage_runs: list[dict[str, dict[str, float | int]]] = []
+ options = pt.PropagationOptions(
+ backend=backend,
+ profile_stages=True,
+ coupled_2x2_mode=coupled_2x2_mode,
+ coupled_2x2_kernel=coupled_2x2_kernel,
+ dtype=dtype,
+ )
+
+ for run_idx in range(warmup + repeats):
+ grid, initial, potential = _make_problem(size)
+ psi = pt.Wavefunction(initial, grid, normalize_const=1.0, backend=backend)
+
+ _sync_backend(backend)
+ t0 = time.perf_counter()
+ stage_profile = _run_profiled_propagation(
+ psi=psi,
+ workload=workload,
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=options,
+ )
+ _sync_backend(backend)
+ elapsed = time.perf_counter() - t0
+
+ if run_idx < warmup:
+ continue
+ timings.append(elapsed)
+ stage_runs.append(stage_profile)
+
+ if not timings:
+ raise RuntimeError("No benchmark timings collected; check repeats/warmup configuration.")
+
+ arr = np.asarray(timings, dtype=float)
+ result = BenchmarkResult(
+ backend=backend,
+ workload=workload,
+ size=size,
+ steps=steps,
+ dt=dt,
+ repeats=repeats,
+ mean_seconds=float(np.mean(arr)),
+ std_seconds=float(np.std(arr)),
+ )
+ stage_summary = _summarize_stage_runs(stage_runs, case_mean_seconds=result.mean_seconds)
+
+ return result, stage_summary
+
+
+def _parity_check(
+ size: int,
+ steps: int,
+ dt: float,
+ *,
+ coupled_2x2_mode: str,
+ coupled_2x2_kernel: str,
+ dtype: str,
+) -> dict[str, float] | None:
+ if not pt.has_cupy():
+ return None
+
+ import cupy as cp
+
+ grid, initial, potential = _make_problem(size)
+ psi_np = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="numpy")
+ psi_cp = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+
+ psi_np.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="numpy",
+ coupled_2x2_mode=coupled_2x2_mode,
+ coupled_2x2_kernel=coupled_2x2_kernel,
+ dtype=dtype,
+ ),
+ )
+ psi_cp.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="cupy",
+ coupled_2x2_mode=coupled_2x2_mode,
+ coupled_2x2_kernel=coupled_2x2_kernel,
+ dtype=dtype,
+ ),
+ )
+
+ amp_np = np.asarray(psi_np.amp)
+ amp_cp = cp.asnumpy(psi_cp.amp)
+
+ density_diff = np.abs(np.abs(amp_np) ** 2 - np.abs(amp_cp) ** 2)
+ occ_np = np.asarray(psi_np.state_occupation())
+ occ_cp = cp.asnumpy(psi_cp.state_occupation())
+
+ return {
+ "max_density_abs_diff": float(np.max(density_diff)),
+ "max_occupation_abs_diff": float(np.max(np.abs(occ_np - occ_cp))),
+ }
+
+
+def _collect_metadata() -> dict[str, object]:
+ meta: dict[str, object] = {
+ "python": sys.version.split()[0],
+ "platform": platform.platform(),
+ "cpu_count": os.cpu_count(),
+ "available_backends": list(pt.available_backends()),
+ }
+
+ if pt.has_cupy():
+ import cupy as cp
+
+ device = cp.cuda.Device()
+ props = cp.cuda.runtime.getDeviceProperties(device.id)
+ name_raw = props.get("name", b"unknown")
+ if isinstance(name_raw, bytes):
+ gpu_name = name_raw.decode("utf-8", errors="replace")
+ else:
+ gpu_name = str(name_raw)
+
+ meta["cuda"] = {
+ "device_id": int(device.id),
+ "device_name": gpu_name,
+ "cupy_version": cp.__version__,
+ }
+
+ return meta
+
+
+def _parse_csv(value: str) -> list[str]:
+ return [item.strip() for item in value.split(",") if item.strip()]
+
+
+def main() -> int:
+ parser = argparse.ArgumentParser(description="Benchmark pyTALISES backends")
+ parser.add_argument(
+ "--sizes",
+ default="192",
+ help="Comma-separated 1D grid sizes (e.g. 128,192,256)",
+ )
+ parser.add_argument("--steps", type=int, default=25, help="Propagation steps")
+ parser.add_argument("--dt", type=float, default=0.005, help="Time step")
+ parser.add_argument("--repeats", type=int, default=3, help="Benchmark repeats")
+ parser.add_argument("--warmup", type=int, default=1, help="Warmup runs")
+ parser.add_argument(
+ "--workloads",
+ default="free,potential",
+ help="Comma-separated workloads: free,potential",
+ )
+ parser.add_argument("--json-out", type=Path, default=None, help="Optional JSON output path")
+ parser.add_argument(
+ "--min-speedup",
+ type=float,
+ default=None,
+ help="Fail if CuPy speedup over NumPy is below this threshold",
+ )
+ parser.add_argument(
+ "--coupled-2x2-mode",
+ choices=("auto", "eigh"),
+ default="auto",
+ help="Potential-step strategy for 2x2 coupled systems",
+ )
+ parser.add_argument(
+ "--coupled-2x2-kernel",
+ choices=("vectorized", "fused"),
+ default="vectorized",
+ help="Analytic 2x2 kernel implementation mode",
+ )
+ parser.add_argument(
+ "--dtype",
+ choices=("complex128", "complex64"),
+ default="complex128",
+ help="Propagation precision mode",
+ )
+ args = parser.parse_args()
+
+ if args.repeats < 1:
+ parser.error("--repeats must be >= 1")
+ if args.warmup < 0:
+ parser.error("--warmup must be >= 0")
+
+ sizes = [int(s) for s in _parse_csv(args.sizes)]
+ workloads = _parse_csv(args.workloads)
+
+ backends = ["numpy"]
+ if pt.has_cupy():
+ backends.append("cupy")
+
+ results: list[BenchmarkResult] = []
+ stage_breakdown: list[dict[str, object]] = []
+ for workload in workloads:
+ for size in sizes:
+ for backend in backends:
+ result, stages = _run_case(
+ backend=backend,
+ workload=workload,
+ size=size,
+ steps=args.steps,
+ dt=args.dt,
+ repeats=args.repeats,
+ warmup=args.warmup,
+ coupled_2x2_mode=args.coupled_2x2_mode,
+ coupled_2x2_kernel=args.coupled_2x2_kernel,
+ dtype=args.dtype,
+ )
+ results.append(result)
+ stage_breakdown.append(
+ {
+ "backend": backend,
+ "workload": workload,
+ "size": size,
+ "steps": args.steps,
+ "dt": args.dt,
+ "coupled_2x2_mode": args.coupled_2x2_mode,
+ "coupled_2x2_kernel": args.coupled_2x2_kernel,
+ "dtype": args.dtype,
+ "stages": stages,
+ "profiled_stage_seconds_total": float(
+ sum(stage["mean_seconds"] for stage in stages.values())
+ ),
+ }
+ )
+
+ payload: dict[str, object] = {
+ "metadata": _collect_metadata(),
+ "config": {
+ "coupled_2x2_mode": args.coupled_2x2_mode,
+ "coupled_2x2_kernel": args.coupled_2x2_kernel,
+ "dtype": args.dtype,
+ },
+ "results": [asdict(r) for r in results],
+ "stage_breakdown": stage_breakdown,
+ "parity": _parity_check(
+ size=max(sizes),
+ steps=max(8, args.steps // 2),
+ dt=args.dt,
+ coupled_2x2_mode=args.coupled_2x2_mode,
+ coupled_2x2_kernel=args.coupled_2x2_kernel,
+ dtype=args.dtype,
+ ),
+ }
+
+ print("Backend benchmark results:")
+ print(f"- coupled 2x2 mode: {args.coupled_2x2_mode}")
+ print(f"- coupled 2x2 kernel: {args.coupled_2x2_kernel}")
+ print(f"- dtype: {args.dtype}")
+ for r in results:
+ print(
+ f"- {r.backend:>5} | {r.workload:>9} | n={r.size:>4}: "
+ f"{r.mean_seconds:.4f}s ± {r.std_seconds:.4f}s"
+ )
+
+ speedups: dict[str, float] = {}
+ if len(backends) >= 2:
+ for workload in workloads:
+ for size in sizes:
+ numpy_row = next(
+ (
+ r
+ for r in results
+ if r.backend == "numpy" and r.workload == workload and r.size == size
+ ),
+ None,
+ )
+ cupy_row = next(
+ (
+ r
+ for r in results
+ if r.backend == "cupy" and r.workload == workload and r.size == size
+ ),
+ None,
+ )
+ if numpy_row and cupy_row:
+ key = f"{workload}:n{size}"
+ speedup = numpy_row.mean_seconds / cupy_row.mean_seconds
+ speedups[key] = speedup
+ print(f"- speedup {key} (numpy/cupy): {speedup:.2f}x")
+
+ payload["speedups_numpy_over_cupy"] = speedups
+
+ parity = payload.get("parity")
+ if isinstance(parity, dict):
+ print(
+ "- parity max abs diff: density={:.3e}, occupation={:.3e}".format(
+ parity["max_density_abs_diff"],
+ parity["max_occupation_abs_diff"],
+ )
+ )
+
+ if args.json_out is not None:
+ args.json_out.parent.mkdir(parents=True, exist_ok=True)
+ args.json_out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+ if args.min_speedup is not None:
+ if "cupy" not in backends:
+ print(
+ "ERROR: --min-speedup requires CuPy benchmark results, "
+ "but CuPy is unavailable on this host."
+ )
+ return 1
+ if not speedups:
+ print("ERROR: --min-speedup requested but no NumPy/CuPy comparisons were produced.")
+ return 1
+ if min(speedups.values()) < args.min_speedup:
+ worst = min(speedups, key=speedups.get)
+ print(
+ f"ERROR: speedup {worst}={speedups[worst]:.2f}x is below required "
+ f"minimum {args.min_speedup:.2f}x"
+ )
+ return 1
+
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/docs/assets/diffraction_grating.gif b/docs/assets/diffraction_grating.gif
new file mode 100644
index 0000000..8b65c6f
Binary files /dev/null and b/docs/assets/diffraction_grating.gif differ
diff --git a/docs/assets/generate_diffraction_gif.py b/docs/assets/generate_diffraction_gif.py
new file mode 100644
index 0000000..4a5521f
--- /dev/null
+++ b/docs/assets/generate_diffraction_gif.py
@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""Generate a diffraction grating GIF for the README."""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+import numexpr as ne
+from pathlib import Path
+
+# Use legacy API for compatibility with notebook example
+import pytalises as pt
+
+# Create wavefunction - 2D Gaussian with momentum
+psi = pt.legacy.Wavefunction(
+ "exp(-((x-x0)/sigmax)**2)*exp(-((y-y0)/sigmay)**2)*exp(1j*ky*y)",
+ variables={'x0': 0, 'y0': -6, 'sigmax': 4, 'sigmay': 1, 'ky': 4},
+ number_of_grid_points=(128, 320),
+ spatial_ext=[(-10, 10), (-12, 12)],
+)
+
+# Periodic grating potential
+v = "where(y<.2, 1, 0)*where(y>-.2, 1, 0)*where(cos(3*x)<0, 1, 0)*1000"
+potential = ne.evaluate(v, local_dict=psi.default_var_dict)[:, :, 0]
+
+# Create figure with dark theme
+fig, ax = plt.subplots(figsize=(6, 6), dpi=80)
+fig.patch.set_facecolor('#0d1117')
+ax.set_facecolor('#0d1117')
+
+# Initial plot
+data = (np.abs(psi.amp**2) + potential * 0.0001).T
+im = ax.imshow(
+ data,
+ origin='lower',
+ extent=[-10, 10, -12, 12],
+ vmax=np.max(np.abs(psi.amp**2)),
+ cmap='inferno',
+ aspect='equal',
+)
+
+ax.set_xlabel('x', color='#c9d1d9', fontsize=11)
+ax.set_ylabel('y', color='#c9d1d9', fontsize=11)
+ax.tick_params(colors='#c9d1d9')
+for spine in ax.spines.values():
+ spine.set_color('#30363d')
+
+# Add subtle grating overlay
+grating_vis = np.where(potential > 0, 0.3, 0)
+ax.imshow(
+ grating_vis.T,
+ origin='lower',
+ extent=[-10, 10, -12, 12],
+ cmap='Greys',
+ alpha=0.5,
+ aspect='equal',
+)
+
+ax.set_title('Diffraction on periodic grating', color='#c9d1d9', fontsize=12, pad=10)
+
+frames_data = []
+
+# Pre-compute frames
+n_frames = 200
+for i in range(n_frames):
+ psi.propagate(v, num_time_steps=6, delta_t=0.003, diag=True)
+ frame = (np.abs(psi.amp**2) + potential * 0.00005).T.copy()
+ frames_data.append(frame)
+ if i % 30 == 0:
+ print(f"Computing frame {i}/{n_frames}")
+
+def animate(i):
+ im.set_data(frames_data[i])
+ return [im]
+
+anim = animation.FuncAnimation(
+ fig, animate, frames=n_frames, interval=40, blit=True
+)
+
+# Save as GIF
+output_path = Path(__file__).parent / 'diffraction_grating.gif'
+print("Saving animation...")
+anim.save(output_path, writer='pillow', fps=25)
+print(f"Saved animation to {output_path}")
+
+plt.close()
diff --git a/docs/assets/generate_readme_animation.py b/docs/assets/generate_readme_animation.py
new file mode 100644
index 0000000..04f646c
--- /dev/null
+++ b/docs/assets/generate_readme_animation.py
@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+"""Generate a simple GIF animation of a Gaussian wavepacket for the README."""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+from pathlib import Path
+
+# Simulation parameters
+N = 512
+x = np.linspace(-15, 15, N)
+dx = x[1] - x[0]
+k = np.fft.fftfreq(N, dx) * 2 * np.pi
+
+# Initial Gaussian wavepacket with momentum
+sigma = 1.0
+k0 = 3.0 # initial momentum
+psi = np.exp(-x**2 / (2 * sigma**2)) * np.exp(1j * k0 * x)
+psi = psi / np.sqrt(np.sum(np.abs(psi)**2) * dx)
+
+# Time evolution parameters
+dt = 0.02
+n_frames = 120
+
+# Kinetic propagator (split-step)
+kinetic_prop = np.exp(-0.5j * k**2 * dt)
+
+# Store frames
+frames = []
+psi_t = psi.copy()
+
+for _ in range(n_frames):
+ # Split-step propagation (free particle)
+ psi_k = np.fft.fft(psi_t)
+ psi_k *= kinetic_prop
+ psi_t = np.fft.ifft(psi_k)
+ frames.append(np.abs(psi_t)**2)
+
+# Create animation
+fig, ax = plt.subplots(figsize=(8, 3), dpi=100)
+fig.patch.set_facecolor('#0d1117') # GitHub dark mode
+ax.set_facecolor('#0d1117')
+
+line, = ax.plot([], [], color='#58a6ff', linewidth=2)
+fill = ax.fill_between([], [], alpha=0.3, color='#58a6ff')
+
+ax.set_xlim(-15, 15)
+ax.set_ylim(0, 0.5)
+ax.set_xlabel('Position', color='#c9d1d9', fontsize=12)
+ax.set_ylabel('|ψ|²', color='#c9d1d9', fontsize=12)
+ax.tick_params(colors='#c9d1d9')
+for spine in ax.spines.values():
+ spine.set_color('#30363d')
+
+# Title
+ax.set_title('Gaussian wavepacket propagation', color='#c9d1d9', fontsize=14, pad=10)
+
+def init():
+ line.set_data([], [])
+ return line,
+
+def animate(i):
+ global fill
+ fill.remove()
+ y = frames[i]
+ line.set_data(x, y)
+ fill = ax.fill_between(x, y, alpha=0.3, color='#58a6ff')
+ return line, fill
+
+anim = animation.FuncAnimation(
+ fig, animate, init_func=init,
+ frames=n_frames, interval=50, blit=False
+)
+
+# Save as GIF
+output_path = Path(__file__).parent / 'wavepacket_animation.gif'
+anim.save(output_path, writer='pillow', fps=20)
+print(f"Saved animation to {output_path}")
+
+plt.close()
diff --git a/docs/assets/wavepacket_animation.gif b/docs/assets/wavepacket_animation.gif
new file mode 100644
index 0000000..ef915a7
Binary files /dev/null and b/docs/assets/wavepacket_animation.gif differ
diff --git a/docs/requirements.txt b/docs/requirements.txt
index a387e6e..9d358a4 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,4 +1,9 @@
--c ../requirements-constraints.txt
+# Sphinx and extensions
+sphinx>=7.0
+sphinx-rtd-theme>=2.0
nbsphinx>=0.9.4
+pypandoc_binary>=1.13
+
+# For notebook rendering
ipython>=8.12
pygments>=2.10
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 8fe62eb..665c9b5 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -53,6 +53,9 @@
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = []
+# Suppress warnings for ambiguous cross-references (legacy module duplicates)
+suppress_warnings = ["ref.python"]
+
# -- Options for HTML output -------------------------------------------------
diff --git a/docs/source/gpu_backend.rst b/docs/source/gpu_backend.rst
new file mode 100644
index 0000000..cd8a4ce
--- /dev/null
+++ b/docs/source/gpu_backend.rst
@@ -0,0 +1,196 @@
+GPU Acceleration
+================
+
+pyTALISES supports GPU acceleration via `CuPy `_, providing
+significant speedups for large-grid simulations.
+
+Installation
+------------
+
+Install with GPU support using the optional ``gpu`` extra:
+
+.. code-block:: bash
+
+ pip install pytalises[gpu]
+
+This installs CuPy with CUDA 12.x support. For other CUDA versions, install
+CuPy separately:
+
+.. code-block:: bash
+
+ pip install pytalises
+ pip install cupy-cuda11x # For CUDA 11.x
+
+Requirements:
+
+- NVIDIA GPU with CUDA support
+- CUDA toolkit (version matching your CuPy install)
+- At least one visible CUDA device
+
+Backend Selection
+-----------------
+
+By default, pyTALISES auto-detects the best available backend:
+
+.. code-block:: python
+
+ import pytalises as pt
+
+ # Auto-detection: uses CuPy if available, otherwise NumPy
+ psi.propagate(potential=V, steps=1000, dt=1e-6)
+
+ # Check what's available
+ from pytalises.backends import available_backends, has_cupy
+ print(available_backends()) # ('cupy', 'numpy') or ('numpy',)
+ print(has_cupy()) # True if CuPy + GPU available
+
+Explicit backend selection via :class:`~pytalises.PropagationOptions`:
+
+.. code-block:: python
+
+ from pytalises import PropagationOptions
+
+ # Force NumPy (useful for debugging or comparison)
+ options = PropagationOptions(backend="numpy")
+ psi.propagate(potential=V, steps=1000, dt=1e-6, options=options)
+
+ # Force CuPy (raises error if unavailable)
+ options = PropagationOptions(backend="cupy")
+ psi.propagate(potential=V, steps=1000, dt=1e-6, options=options)
+
+Global default backend:
+
+.. code-block:: python
+
+ from pytalises.backends import set_default_backend
+
+ set_default_backend("cupy") # All subsequent propagations use CuPy
+
+Performance Characteristics
+---------------------------
+
+GPU acceleration is most beneficial for large grids. Typical crossover points
+on modern hardware (RTX A4500):
+
+================== ================ ==============
+Grid size Free propagation With potential
+================== ================ ==============
+n < 2048 NumPy faster ~Equal
+n = 4096 ~Equal CuPy ~1.2x
+n = 8192 CuPy ~1.1x CuPy ~1.3x
+n = 32768 CuPy ~6x CuPy ~3.5x
+n = 131072 CuPy ~12x CuPy ~14x
+================== ================ ==============
+
+.. note::
+
+ Exact crossover points depend on your specific GPU, CPU, and simulation
+ type. The potential workload benefits more at large grids due to the
+ analytic 2×2 coupled propagator optimization.
+
+Precision Control
+-----------------
+
+Control numerical precision via ``dtype``:
+
+.. code-block:: python
+
+ # Double precision (default) - best accuracy
+ options = PropagationOptions(dtype="complex128")
+
+ # Single precision - faster, uses less memory
+ options = PropagationOptions(dtype="complex64")
+
+Trade-offs:
+
+- **complex128** (default): Full double precision. Parity errors ~1e-14.
+- **complex64**: Single precision. ~1.3-2x faster, but parity errors ~1e-5.
+
+Use ``complex64`` when:
+
+- Speed is critical and moderate precision is acceptable
+- GPU memory is constrained
+- Running many simulations for parameter sweeps
+
+Advanced Options
+----------------
+
+The :class:`~pytalises.PropagationOptions` class provides fine-grained control:
+
+.. code-block:: python
+
+ options = PropagationOptions(
+ backend="auto", # "auto", "numpy", or "cupy"
+ dtype="complex128", # "complex128" or "complex64"
+ threads=4, # CPU threads (NumPy backend)
+ coupled_2x2_mode="auto", # "auto" or "eigh"
+ coupled_2x2_kernel="vectorized", # "vectorized" or "fused"
+ profile_stages=False, # Enable timing profiler
+ )
+
+**coupled_2x2_mode**:
+
+- ``auto``: Use analytic closed-form matrix exponential for 2×2 Hermitian
+ potentials (much faster, especially on GPU)
+- ``eigh``: Force eigendecomposition (slower but handles edge cases)
+
+**coupled_2x2_kernel**:
+
+- ``vectorized``: Standard vectorized operations
+- ``fused``: Fused kernel with reduced memory traffic (experimental)
+
+**profile_stages**: When ``True``, timing breakdowns are collected during
+propagation. Access via ``psi.stage_timings`` after propagation.
+
+Troubleshooting
+---------------
+
+**CuPy not detected:**
+
+.. code-block:: python
+
+ from pytalises.backends import has_cupy, available_backends
+ print(has_cupy()) # False
+ print(available_backends()) # ('numpy',)
+
+Check:
+
+1. CuPy is installed: ``pip show cupy-cuda12x``
+2. CUDA is visible: ``nvidia-smi``
+3. GPU is accessible from Python:
+
+ .. code-block:: python
+
+ import cupy as cp
+ print(cp.cuda.runtime.getDeviceCount())
+
+**Out of GPU memory:**
+
+- Reduce grid size
+- Use ``dtype="complex64"``
+- Free other GPU processes
+
+**Results differ between backends:**
+
+Small numerical differences (~1e-14 for complex128) are expected due to
+different floating-point ordering. Larger differences may indicate:
+
+- Single precision (complex64) accumulating error over many steps
+- Edge cases in potential evaluation
+
+For validation, compare both backends on a short run:
+
+.. code-block:: python
+
+ import numpy as np
+
+ psi_np = psi.copy()
+ psi_cp = psi.copy()
+
+ psi_np.propagate(potential=V, steps=100, dt=1e-6,
+ options=PropagationOptions(backend="numpy"))
+ psi_cp.propagate(potential=V, steps=100, dt=1e-6,
+ options=PropagationOptions(backend="cupy"))
+
+ diff = np.max(np.abs(psi_np.psi - psi_cp.psi.get()))
+ print(f"Max difference: {diff}") # Should be ~1e-14 for complex128
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 5b4bfd7..d886b8d 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -21,6 +21,7 @@ Welcome to pytalises's documentation!
examples
additional_examples
notes
+ gpu_backend
v2_migration
pytalises.rst
diff --git a/docs/source/v2_migration.rst b/docs/source/v2_migration.rst
index 6021903..db8ae85 100644
--- a/docs/source/v2_migration.rst
+++ b/docs/source/v2_migration.rst
@@ -13,6 +13,17 @@ Quick mapping
- ``num_time_steps`` -> ``steps``
- ``delta_t`` -> ``dt``
- ``num_of_threads`` / ``FFTWflags`` -> :class:`pytalises.PropagationOptions`
+- optional GPU execution -> ``PropagationOptions(backend="cupy")``
+ (see :doc:`gpu_backend` for details)
+
+Backend internals
+-----------------
+
+The backend/engine internals are private implementation details in v2.
+Only the public simulation API is considered stable:
+:class:`pytalises.Grid`, :class:`pytalises.Wavefunction`,
+:func:`pytalises.propagate`, :func:`pytalises.freely_propagate`, and
+structured potentials/options.
Example (before)
----------------
diff --git a/pyproject.toml b/pyproject.toml
index 860f353..da5cd81 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,11 @@ dependencies = [
"numexpr>=2.8.4",
]
+[project.optional-dependencies]
+gpu = [
+ "cupy-cuda12x>=13.0.0; platform_system=='Linux'",
+]
+
[project.urls]
Homepage = "http://github.com/savowe/pytalises"
Repository = "http://github.com/savowe/pytalises"
diff --git a/pytalises/__init__.py b/pytalises/__init__.py
index f674f29..b950295 100644
--- a/pytalises/__init__.py
+++ b/pytalises/__init__.py
@@ -1,4 +1,10 @@
-from .backends import get_backend, register_backend, set_default_backend # noqa
+from .backends import ( # noqa
+ available_backends,
+ get_backend,
+ has_cupy,
+ register_backend,
+ set_default_backend,
+)
from .exceptions import ( # noqa
BackendError,
GridMismatchError,
diff --git a/pytalises/backends/__init__.py b/pytalises/backends/__init__.py
index 3442e90..3805d05 100644
--- a/pytalises/backends/__init__.py
+++ b/pytalises/backends/__init__.py
@@ -11,6 +11,13 @@
"numpy": NumpyBackend,
}
+try: # optional dependency
+ from .cupy_backend import CupyBackend
+except ImportError: # pragma: no cover - depends on optional runtime
+ CupyBackend = None
+else:
+ _BACKENDS["cupy"] = CupyBackend
+
_default_backend: Optional[Backend] = None
@@ -19,8 +26,26 @@ def register_backend(name: str, backend_class):
_BACKENDS[name] = backend_class
+def available_backends() -> tuple[str, ...]:
+ """Return names of currently importable backends."""
+ return tuple(sorted(_BACKENDS.keys()))
+
+
+def has_cupy() -> bool:
+ """Return ``True`` if a usable CuPy backend is available."""
+ if "cupy" not in _BACKENDS:
+ return False
+ try:
+ import cupy as cp
+
+ return cp.cuda.runtime.getDeviceCount() > 0
+ except Exception:
+ return False
+
+
def _auto_detect_backend() -> str:
- # Future: GPU detection for CuPy/JAX backends.
+ if has_cupy():
+ return "cupy"
return "numpy"
@@ -40,6 +65,12 @@ def get_backend(name: Union[str, Backend, None] = None, **kwargs) -> Backend:
available = sorted(_BACKENDS.keys())
raise ValueError(f"Unknown backend '{name}'. Available: {available}")
+ if name == "cupy" and not has_cupy():
+ raise ValueError(
+ "CuPy backend requested but no usable CUDA device is available. "
+ "Install CuPy with CUDA support and ensure at least one visible GPU."
+ )
+
return _BACKENDS[name](**kwargs)
diff --git a/pytalises/backends/cupy_backend.py b/pytalises/backends/cupy_backend.py
new file mode 100644
index 0000000..4e8d69d
--- /dev/null
+++ b/pytalises/backends/cupy_backend.py
@@ -0,0 +1,143 @@
+"""CuPy backend implementation for pyTALISES.
+
+Notes
+-----
+This backend requires CuPy and a working CUDA runtime.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+try:
+ import cupy as cp
+except ImportError as exc: # pragma: no cover - exercised via conditional import
+ raise ImportError("CupyBackend requires CuPy to be installed.") from exc
+
+from .base import Backend, FFTPlan
+
+
+class CupyFFTPlan(FFTPlan):
+ """Simple FFT plan wrapper using CuPy FFT routines."""
+
+ def __init__(self, array: Any, axes: tuple[int, ...]):
+ self._array = array
+ self._axes = axes
+
+ def forward(self) -> None:
+ self._array[...] = cp.fft.fftn(self._array, axes=self._axes)
+
+ def backward(self) -> None:
+ self._array[...] = cp.fft.ifftn(self._array, axes=self._axes)
+
+
+class CupyBackend(Backend):
+ """CuPy backend for GPU-resident propagation."""
+
+ name = "cupy"
+
+ def __init__(self, num_threads: int = 1, device_id: int | None = None):
+ self.num_threads = num_threads
+ self.device_id = device_id
+ if device_id is not None:
+ cp.cuda.Device(device_id).use()
+ self._eval_globals = {
+ "cp": cp,
+ "np": cp,
+ "pi": cp.pi,
+ "exp": cp.exp,
+ "sin": cp.sin,
+ "cos": cp.cos,
+ "tan": cp.tan,
+ "arcsin": cp.arcsin,
+ "arccos": cp.arccos,
+ "arctan": cp.arctan,
+ "sinh": cp.sinh,
+ "cosh": cp.cosh,
+ "tanh": cp.tanh,
+ "sqrt": cp.sqrt,
+ "log": cp.log,
+ "log10": cp.log10,
+ "abs": cp.abs,
+ "real": cp.real,
+ "imag": cp.imag,
+ "conjugate": cp.conjugate,
+ "conj": cp.conj,
+ "where": cp.where,
+ "minimum": cp.minimum,
+ "maximum": cp.maximum,
+ "power": cp.power,
+ }
+
+ def set_num_threads(self, n: int) -> None:
+ # CuPy/CUDA thread/block scheduling is handled by CUDA runtime.
+ self.num_threads = n
+
+ def zeros(self, shape, dtype: Any = "complex128"):
+ return cp.zeros(shape, dtype=dtype)
+
+ def empty_aligned(self, shape, dtype: Any = "complex128"):
+ return cp.empty(shape, dtype=dtype)
+
+ def asarray(self, data, dtype: Any = None):
+ return cp.asarray(data, dtype=dtype)
+
+ def create_fft_plan(
+ self,
+ array,
+ axes,
+ num_threads: int = 1,
+ flags: tuple[str, ...] = ("FFTW_ESTIMATE", "FFTW_DESTROY_INPUT"),
+ ):
+ del num_threads, flags
+ return CupyFFTPlan(array, axes)
+
+ def exp(self, x):
+ return cp.exp(x)
+
+ def einsum(self, subscripts: str, *operands, out=None):
+ result = cp.einsum(subscripts, *operands, optimize="greedy")
+ if out is not None:
+ out[...] = result
+ return out
+ return result
+
+ def evaluate(self, expr: str, local_dict: dict, global_dict: dict | None = None):
+ scope = dict(self._eval_globals)
+ if global_dict:
+ scope.update(global_dict)
+ scope.update(local_dict)
+
+ # CuPy ufunc dispatch may rely on importing internals at runtime.
+ # Keep builtins locked down except for __import__ to avoid KeyError.
+ safe_builtins = {"__import__": __import__}
+ return eval(expr, {"__builtins__": safe_builtins}, scope)
+
+ def eigh(self, matrices):
+ eigvals, eigvecs = cp.linalg.eigh(matrices)
+ matrices[...] = eigvecs
+ return eigvals, matrices
+
+ def meshgrid(self, *arrays, indexing="ij"):
+ return cp.meshgrid(*arrays, indexing=indexing)
+
+ def linspace(self, start, stop, num):
+ return cp.linspace(start, stop, num)
+
+ def fftfreq(self, n):
+ return cp.fft.fftfreq(n)
+
+ def sum(self, x, axis=None):
+ return cp.sum(x, axis=axis)
+
+ def abs(self, x):
+ return cp.abs(x)
+
+ def sqrt(self, x):
+ return cp.sqrt(x)
+
+ def conjugate(self, x):
+ return cp.conjugate(x)
+
+ def power(self, x, n):
+ return cp.power(x, n)
diff --git a/pytalises/engine/__init__.py b/pytalises/engine/__init__.py
new file mode 100644
index 0000000..2ecda71
--- /dev/null
+++ b/pytalises/engine/__init__.py
@@ -0,0 +1,9 @@
+"""Private engine internals for pyTALISES propagation.
+
+The symbols in this package are intentionally unstable.
+"""
+
+from .evaluator import ExpressionEvaluator
+from .factory import create_engine
+
+__all__ = ["ExpressionEvaluator", "create_engine"]
diff --git a/pytalises/engine/base.py b/pytalises/engine/base.py
new file mode 100644
index 0000000..afada1f
--- /dev/null
+++ b/pytalises/engine/base.py
@@ -0,0 +1,151 @@
+"""Private propagation engine abstractions.
+
+These classes are internal implementation details and may change without notice.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Engine(ABC):
+ """Private engine contract used by propagation kernels."""
+
+ name: str
+ xp: Any
+
+ def __init__(self, backend: Any):
+ self.backend = backend
+
+ def zeros(self, shape: tuple[int, ...], dtype: Any = "complex128") -> Any:
+ return self.backend.zeros(shape, dtype=dtype)
+
+ def asarray(self, data: Any, dtype: Any = None) -> Any:
+ return self.backend.asarray(data, dtype=dtype)
+
+ def set_num_threads(self, n: int) -> None:
+ self.backend.set_num_threads(n)
+
+ def create_fft_plan(
+ self,
+ array: Any,
+ axes: tuple[int, ...],
+ num_threads: int = 1,
+ flags: tuple[str, ...] = ("FFTW_ESTIMATE", "FFTW_DESTROY_INPUT"),
+ ) -> Any:
+ return self.backend.create_fft_plan(
+ array,
+ axes=axes,
+ num_threads=num_threads,
+ flags=flags,
+ )
+
+ def apply_kinetic_phase(
+ self,
+ amp: Any,
+ *,
+ kmesh: tuple[Any, Any, Any],
+ alpha: float,
+ dt: float,
+ ) -> None:
+ """Apply kinetic split-step phase in reciprocal space in-place."""
+ kx, ky, kz = kmesh
+ imag_unit = self.xp.asarray(1j, dtype=amp.dtype)
+ phase = self.xp.exp(-imag_unit * alpha * dt * (kx**2 + ky**2 + kz**2))
+ amp *= phase[..., self.xp.newaxis]
+
+ def apply_diagonal_phase(self, amp: Any, *, diagonal: Any, dt: float) -> None:
+ """Apply diagonal potential phase in-place."""
+ imag_unit = self.xp.asarray(1j, dtype=amp.dtype)
+ amp *= self.xp.exp(-imag_unit * diagonal * dt)
+
+ def apply_coupled_phase(
+ self,
+ amp: Any,
+ *,
+ eigvals: Any,
+ eigvecs: Any,
+ dt: float,
+ ) -> None:
+ """Apply non-diagonal potential step via eigendecomposition in-place."""
+ # psi' = U @ diag(exp(-i*lambda*dt)) @ U† @ psi
+ u_dag = self.xp.swapaxes(self.xp.conjugate(eigvecs), -1, -2)
+ tmp = self.xp.matmul(u_dag, amp[..., self.xp.newaxis])[..., 0]
+ imag_unit = self.xp.asarray(1j, dtype=amp.dtype)
+ tmp *= self.xp.exp(-imag_unit * eigvals * dt)
+ amp[...] = self.xp.matmul(eigvecs, tmp[..., self.xp.newaxis])[..., 0]
+
+ def apply_coupled_phase_2x2(
+ self,
+ amp: Any,
+ *,
+ matrix: Any,
+ dt: float,
+ kernel: str = "vectorized",
+ ) -> None:
+ """Apply non-diagonal potential step for 2x2 Hermitian matrices.
+
+ Uses the closed-form matrix exponential of a Hermitian 2x2 block:
+ ``V = c I + [[z, b], [conj(b), -z]]``.
+ """
+ if amp.shape[-1] != 2 or matrix.shape[-2:] != (2, 2):
+ raise ValueError("apply_coupled_phase_2x2 requires 2x2 coupled amplitudes")
+ if kernel not in {"vectorized", "fused"}:
+ raise ValueError(
+ "apply_coupled_phase_2x2 kernel must be 'vectorized' or 'fused'"
+ )
+ # Base implementation is the vectorized path used by NumPy and as fallback.
+
+ xp = self.xp
+
+ a = matrix[..., 0, 0]
+ d = matrix[..., 1, 1]
+ b = matrix[..., 0, 1]
+
+ c = 0.5 * (a + d)
+ z = 0.5 * (a - d)
+ r = xp.sqrt((z * xp.conjugate(z)).real + (b * xp.conjugate(b)).real)
+
+ imag_unit = xp.asarray(1j, dtype=amp.dtype)
+ dt_scalar = xp.asarray(dt, dtype=r.dtype)
+
+ phase = xp.exp(-imag_unit * c * dt_scalar)
+ cos_term = xp.cos(r * dt_scalar)
+
+ eps = xp.asarray(1e-15, dtype=r.dtype)
+ safe_r = xp.where(r > eps, r, xp.asarray(1.0, dtype=r.dtype))
+ sin_over_r = xp.sin(r * dt_scalar) / safe_r
+ sin_over_r = xp.where(r > eps, sin_over_r, dt_scalar)
+
+ u00 = cos_term - imag_unit * sin_over_r * z
+ u11 = cos_term + imag_unit * sin_over_r * z
+ u01 = -imag_unit * sin_over_r * b
+ u10 = -imag_unit * sin_over_r * xp.conjugate(b)
+
+ psi0 = amp[..., 0]
+ psi1 = amp[..., 1]
+ new0 = phase * (u00 * psi0 + u01 * psi1)
+ new1 = phase * (u10 * psi0 + u11 * psi1)
+
+ amp[..., 0] = new0
+ amp[..., 1] = new1
+
+ def inner_product(self, lhs: Any, rhs: Any, *, volume_element: float) -> Any:
+ """Return including spatial volume element."""
+ return self.xp.sum(lhs * self.xp.conjugate(rhs)) * volume_element
+
+ def state_occupation(self, amp: Any, *, state_index: int, volume_element: float) -> Any:
+ """Return occupation of a single internal state."""
+ return (
+ self.xp.sum(self.xp.abs(amp[:, :, :, state_index]) ** 2)
+ * volume_element
+ )
+
+ def synchronize(self) -> None:
+ """Synchronize outstanding backend work when needed for profiling."""
+ return None
+
+ @abstractmethod
+ def eigendecompose_hermitian(self, matrices: Any) -> tuple[Any, Any]:
+ """Return eigenvalues and eigenvectors for Hermitian blocks."""
diff --git a/pytalises/engine/cpu.py b/pytalises/engine/cpu.py
new file mode 100644
index 0000000..7adc62b
--- /dev/null
+++ b/pytalises/engine/cpu.py
@@ -0,0 +1,17 @@
+"""CPU propagation engine."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .base import Engine
+
+
+class CpuEngine(Engine):
+ """CPU engine backed by NumPy/pyFFTW backend primitives."""
+
+ name = "cpu"
+ xp = np
+
+ def eigendecompose_hermitian(self, matrices):
+ return self.backend.eigh(matrices)
diff --git a/pytalises/engine/evaluator.py b/pytalises/engine/evaluator.py
new file mode 100644
index 0000000..a5df2ad
--- /dev/null
+++ b/pytalises/engine/evaluator.py
@@ -0,0 +1,77 @@
+"""Expression evaluation helpers for propagation internals.
+
+This module intentionally keeps expression handling separate from backend kernels.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numexpr as ne
+
+
+class ExpressionEvaluator:
+ """Evaluate symbolic expressions on CPU/GPU array backends."""
+
+ def __init__(self, *, backend_name: str):
+ self.backend_name = backend_name
+ self._gpu_scope = None
+
+ if backend_name == "cupy":
+ import cupy as cp # optional dependency
+
+ self._gpu_scope = {
+ "cp": cp,
+ "np": cp,
+ "pi": cp.pi,
+ "exp": cp.exp,
+ "sin": cp.sin,
+ "cos": cp.cos,
+ "tan": cp.tan,
+ "arcsin": cp.arcsin,
+ "arccos": cp.arccos,
+ "arctan": cp.arctan,
+ "sinh": cp.sinh,
+ "cosh": cp.cosh,
+ "tanh": cp.tanh,
+ "sqrt": cp.sqrt,
+ "log": cp.log,
+ "log10": cp.log10,
+ "abs": cp.abs,
+ "real": cp.real,
+ "imag": cp.imag,
+ "conjugate": cp.conjugate,
+ "conj": cp.conj,
+ "where": cp.where,
+ "minimum": cp.minimum,
+ "maximum": cp.maximum,
+ "power": cp.power,
+ }
+
+ def eval(
+ self,
+ expr: str,
+ *,
+ local_dict: dict[str, Any],
+ global_dict: dict[str, Any] | None = None,
+ ) -> Any:
+ """Evaluate expression against backend-native arrays."""
+ if self.backend_name == "numpy":
+ return ne.evaluate(
+ expr,
+ local_dict=local_dict,
+ global_dict=global_dict or {},
+ order="C",
+ )
+
+ if self.backend_name == "cupy":
+ scope = dict(self._gpu_scope or {})
+ if global_dict:
+ scope.update(global_dict)
+ scope.update(local_dict)
+
+ # CuPy ufunc dispatch may import internals dynamically.
+ safe_builtins = {"__import__": __import__}
+ return eval(expr, {"__builtins__": safe_builtins}, scope)
+
+ raise ValueError(f"Unsupported backend for expression evaluation: {self.backend_name}")
diff --git a/pytalises/engine/factory.py b/pytalises/engine/factory.py
new file mode 100644
index 0000000..f777685
--- /dev/null
+++ b/pytalises/engine/factory.py
@@ -0,0 +1,18 @@
+"""Factory for private propagation engines."""
+
+from __future__ import annotations
+
+from pytalises.engine.cpu import CpuEngine
+
+
+def create_engine(backend):
+ """Create engine implementation for a backend instance."""
+ if backend.name == "numpy":
+ return CpuEngine(backend)
+
+ if backend.name == "cupy":
+ from pytalises.engine.gpu import GpuEngine
+
+ return GpuEngine(backend)
+
+ raise ValueError(f"No engine registered for backend '{backend.name}'.")
diff --git a/pytalises/engine/gpu.py b/pytalises/engine/gpu.py
new file mode 100644
index 0000000..4358c3c
--- /dev/null
+++ b/pytalises/engine/gpu.py
@@ -0,0 +1,84 @@
+"""GPU propagation engine."""
+
+from __future__ import annotations
+
+try:
+ import cupy as cp
+except Exception as exc: # pragma: no cover - optional dependency
+ raise ImportError("GpuEngine requires CuPy.") from exc
+
+from .base import Engine
+
+
+@cp.fuse()
+def _analytic_2x2_step_fused(psi0, psi1, a, d, b, dt, eps, imag_unit):
+ c = 0.5 * (a + d)
+ z = 0.5 * (a - d)
+ r = cp.sqrt(cp.real(z * cp.conj(z) + b * cp.conj(b)))
+
+ phase = cp.exp(-imag_unit * c * dt)
+ cos_term = cp.cos(r * dt)
+
+ safe_r = cp.where(r > eps, r, (eps * 0) + 1)
+ sin_over_r = cp.sin(r * dt) / safe_r
+ sin_over_r = cp.where(r > eps, sin_over_r, dt)
+
+ u00 = cos_term - imag_unit * sin_over_r * z
+ u11 = cos_term + imag_unit * sin_over_r * z
+ u01 = -imag_unit * sin_over_r * b
+ u10 = -imag_unit * sin_over_r * cp.conj(b)
+
+ new0 = phase * (u00 * psi0 + u01 * psi1)
+ new1 = phase * (u10 * psi0 + u11 * psi1)
+ return new0, new1
+
+
+class GpuEngine(Engine):
+ """GPU engine backed by CuPy kernels and linear algebra."""
+
+ name = "gpu"
+ xp = cp
+
+ def __init__(self, backend):
+ super().__init__(backend)
+ device_count = cp.cuda.runtime.getDeviceCount()
+ if device_count < 1:
+ raise RuntimeError("CuPy backend selected but no CUDA device is available.")
+
+ def synchronize(self) -> None:
+ cp.cuda.Stream.null.synchronize()
+
+ def eigendecompose_hermitian(self, matrices):
+ eigvals, eigvecs = cp.linalg.eigh(matrices)
+ return eigvals, eigvecs
+
+ def apply_coupled_phase_2x2(
+ self,
+ amp,
+ *,
+ matrix,
+ dt: float,
+ kernel: str = "vectorized",
+ ) -> None:
+ if kernel != "fused":
+ super().apply_coupled_phase_2x2(amp, matrix=matrix, dt=dt, kernel=kernel)
+ return
+
+ if amp.shape[-1] != 2 or matrix.shape[-2:] != (2, 2):
+ raise ValueError("apply_coupled_phase_2x2 requires 2x2 coupled amplitudes")
+
+ a = matrix[..., 0, 0]
+ d = matrix[..., 1, 1]
+ b = matrix[..., 0, 1]
+ psi0 = amp[..., 0]
+ psi1 = amp[..., 1]
+
+ dt_scalar = cp.asarray(dt, dtype=a.real.dtype)
+ eps = cp.asarray(1e-15, dtype=a.real.dtype)
+ imag_unit = cp.asarray(1j, dtype=amp.dtype)
+ new0, new1 = _analytic_2x2_step_fused(
+ psi0, psi1, a, d, b, dt_scalar, eps, imag_unit
+ )
+
+ amp[..., 0] = new0
+ amp[..., 1] = new1
diff --git a/pytalises/options.py b/pytalises/options.py
index ededa72..042c612 100644
--- a/pytalises/options.py
+++ b/pytalises/options.py
@@ -23,7 +23,23 @@ class PropagationOptions:
"FFTW_ESTIMATE",
"FFTW_DESTROY_INPUT",
)
+ profile_stages: bool = False
+ coupled_2x2_mode: str = "auto"
+ coupled_2x2_kernel: str = "vectorized"
def __post_init__(self) -> None:
if self.threads < 1:
raise ValueError("PropagationOptions.threads must be >= 1")
+ if self.dtype not in {"complex64", "complex128"}:
+ raise ValueError(
+ "PropagationOptions.dtype must be one of: 'complex64', 'complex128'"
+ )
+ if self.coupled_2x2_mode not in {"auto", "eigh"}:
+ raise ValueError(
+ "PropagationOptions.coupled_2x2_mode must be one of: 'auto', 'eigh'"
+ )
+ if self.coupled_2x2_kernel not in {"vectorized", "fused"}:
+ raise ValueError(
+ "PropagationOptions.coupled_2x2_kernel must be one of: "
+ "'vectorized', 'fused'"
+ )
diff --git a/pytalises/propagator.py b/pytalises/propagator.py
index f6516ea..8346ffd 100644
--- a/pytalises/propagator.py
+++ b/pytalises/propagator.py
@@ -2,10 +2,14 @@
from __future__ import annotations
+from collections import defaultdict
from typing import TYPE_CHECKING, Any
+import time
import numexpr as ne
+import numpy as np
+from pytalises.engine import ExpressionEvaluator
from pytalises.options import PropagationOptions
from pytalises.potentials import BasePotential, zero_potential
@@ -80,25 +84,63 @@ def __init__(
)
self._backend = self.psi._backend
- self._backend.set_num_threads(self.options.threads)
+ self._engine = self.psi._engine
+ self._engine.set_num_threads(self.options.threads)
+ self._evaluator = ExpressionEvaluator(backend_name=self._backend.name)
+
+ self._complex_dtype = self.options.dtype
+ self._real_dtype = "float32" if self._complex_dtype == "complex64" else "float64"
+ self._alpha = (
+ float(np.float32(self.psi.alpha))
+ if self._complex_dtype == "complex64"
+ else float(self.psi.alpha)
+ )
+
+ if str(self.psi._amp.dtype) != self._complex_dtype:
+ self.psi._amp = self._engine.asarray(self.psi._amp, dtype=self._complex_dtype)
+ for i in range(self.psi.num_int_dim):
+ self.psi.default_var_dict[f"psi{i}"] = self.psi._amp[:, :, :, i]
+
+ self._kmesh = (
+ self._engine.asarray(self.psi.kmesh[0], dtype=self._real_dtype),
+ self._engine.asarray(self.psi.kmesh[1], dtype=self._real_dtype),
+ self._engine.asarray(self.psi.kmesh[2], dtype=self._real_dtype),
+ )
+
+ self._profile_stages = bool(self.options.profile_stages)
+ self._stage_seconds: dict[str, float] = defaultdict(float)
+ self._stage_calls: dict[str, int] = defaultdict(int)
self.v = self.Potential.from_potential(
potential=potential,
variables=variables or {},
num_int_dim=psi.num_int_dim,
)
+ self._use_analytic_2x2 = (
+ self.options.coupled_2x2_mode == "auto"
+ and self.v.diag is False
+ and self.psi.num_int_dim == 2
+ )
assert isinstance(psi, pytalises.wavefunction.Wavefunction)
assert self.v.num_int_dim == self.psi.num_int_dim
assert self.psi._amp.shape[-1] == self.psi.num_int_dim
- self.V_eval_array = self._backend.zeros(
+ self.V_eval_array = self._engine.zeros(
psi.number_of_grid_points + (psi.num_int_dim, psi.num_int_dim),
- dtype="complex128",
+ dtype=self._complex_dtype,
+ )
+ self.V_eval_diag_array = self._engine.zeros(
+ psi.number_of_grid_points + (psi.num_int_dim,),
+ dtype=self._complex_dtype,
)
- self.V_eval_eigval_array = self._backend.zeros(
+ self.V_eval_eigval_array = self._engine.zeros(
psi.number_of_grid_points + (psi.num_int_dim,),
- dtype="complex128",
+ dtype=self._complex_dtype,
+ )
+ self.V_eval_eigvec_array = self._engine.zeros(
+ psi.number_of_grid_points + (psi.num_int_dim, psi.num_int_dim),
+ dtype=self._complex_dtype,
)
self.psi.construct_FFT(self.options.threads, self.options.fftw_flags)
@@ -113,88 +155,139 @@ def __init__(
self.eval_diag_V()
else:
self.eval_V()
- self.V_eval_eigval_array, self.V_eval_array = self._backend.eigh(
- self.V_eval_array
- )
+ if not self._use_analytic_2x2:
+ self._refresh_eigendecomposition()
+
+ def _record_stage(self, name: str, elapsed: float) -> None:
+ if not self._profile_stages:
+ return
+ self._stage_seconds[name] += elapsed
+ self._stage_calls[name] += 1
+
+ def _start_stage_timer(self) -> float:
+ if not self._profile_stages:
+ return 0.0
+ self._engine.synchronize()
+ return time.perf_counter()
+
+ def _stop_stage_timer(self, name: str, start: float) -> None:
+ if not self._profile_stages:
+ return
+ self._engine.synchronize()
+ self._record_stage(name, time.perf_counter() - start)
+
+ def stage_timings(self) -> dict[str, dict[str, float | int]]:
+ """Return accumulated stage timings for this propagator instance."""
+ if not self._profile_stages:
+ return {}
+ return {
+ name: {
+ "seconds": float(self._stage_seconds[name]),
+ "calls": int(self._stage_calls[name]),
+ }
+ for name in sorted(self._stage_seconds)
+ }
def potential_prop(self, dt: float) -> None:
"""Apply potential propagator step."""
self.prop_method(dt)
+ def _refresh_eigendecomposition(self) -> None:
+ t0 = self._start_stage_timer()
+ eigvals, eigvecs = self._engine.eigendecompose_hermitian(self.V_eval_array)
+ self.V_eval_eigval_array[...] = eigvals
+ self.V_eval_eigvec_array[...] = eigvecs
+ self._stop_stage_timer("_refresh_eigendecomposition", t0)
+
def nondiag_potential_prop(self, dt: float) -> None:
"""Apply potential step for non-diagonal potentials."""
if self.v.static is False:
self.eval_V()
- self.V_eval_eigval_array, self.V_eval_array = self._backend.eigh(
- self.V_eval_array
+ if not self._use_analytic_2x2:
+ self._refresh_eigendecomposition()
+
+ if self._use_analytic_2x2:
+ t0 = self._start_stage_timer()
+ self._engine.apply_coupled_phase_2x2(
+ self.psi._amp,
+ matrix=self.V_eval_array,
+ dt=dt,
+ kernel=self.options.coupled_2x2_kernel,
)
- self._backend.einsum(
- "xyzij,xyzj,xyzkj,xyzk->xyzi",
- self.V_eval_array,
- self._backend.evaluate(
- "exp(-1j*eigval*dt)",
- local_dict={"eigval": self.V_eval_eigval_array, "dt": dt},
- ),
- self._backend.conjugate(self.V_eval_array),
+ self._stop_stage_timer("apply_coupled_phase_analytic_2x2", t0)
+ return
+
+ t0 = self._start_stage_timer()
+ self._engine.apply_coupled_phase(
self.psi._amp,
- out=self.psi._amp,
+ eigvals=self.V_eval_eigval_array,
+ eigvecs=self.V_eval_eigvec_array,
+ dt=dt,
)
+ self._stop_stage_timer("apply_coupled_phase", t0)
def diag_potential_prop(self, dt: float) -> None:
"""Apply potential step for diagonal potentials."""
if self.v.static is False:
self.eval_diag_V()
- self._backend.einsum(
- "xyzii,xyzi->xyzi",
- self._backend.evaluate(
- "exp(-1j*V*dt)",
- local_dict={"V": self.V_eval_array, "dt": dt},
- ),
+ t0 = self._start_stage_timer()
+ self._engine.apply_diagonal_phase(
self.psi._amp,
- out=self.psi._amp,
+ diagonal=self.V_eval_diag_array,
+ dt=dt,
)
+ self._stop_stage_timer("apply_diagonal_phase", t0)
def kinetic_prop(self, dt: float) -> None:
"""Perform kinetic propagation step in reciprocal space."""
+ t0 = self._start_stage_timer()
self.psi.fft()
- self._backend.einsum(
- "xyz,xyzi->xyzi",
- self._backend.evaluate(
- "exp(-1j*alpha*dt*(kx**2+ky**2+kz**2))",
- local_dict={
- "kx": self.psi.kmesh[0],
- "ky": self.psi.kmesh[1],
- "kz": self.psi.kmesh[2],
- "alpha": self.psi.alpha,
- "dt": dt,
- },
- ),
+ self._engine.apply_kinetic_phase(
self.psi._amp,
- out=self.psi._amp,
+ kmesh=self._kmesh,
+ alpha=self._alpha,
+ dt=dt,
)
self.psi.ifft()
self.psi.t += dt
+ self._stop_stage_timer("kinetic_prop", t0)
+
+ def _evaluation_scope(self) -> tuple[dict[str, Any], dict[str, Any]]:
+ local_scope = {**self.v.variables, **self.psi.default_var_dict}
+ global_scope = {"t": self.psi.t}
+ return local_scope, global_scope
def eval_V(self) -> None:
"""Evaluate full potential matrix on the complete spatial grid."""
+ t0 = self._start_stage_timer()
+ self.V_eval_array[...] = 0
+
+ local_scope, global_scope = self._evaluation_scope()
k = 0
for i in range(self.psi.num_int_dim):
for j in range(i, self.psi.num_int_dim):
- self.V_eval_array[:, :, :, j, i] = self._backend.evaluate(
+ eval_ji = self._evaluator.eval(
self.v.potential_strings[k],
- local_dict={**self.v.variables, **self.psi.default_var_dict},
- global_dict={"t": self.psi.t},
+ local_dict=local_scope,
+ global_dict=global_scope,
)
+ self.V_eval_array[:, :, :, j, i] = eval_ji
+ if i != j:
+ self.V_eval_array[:, :, :, i, j] = self._engine.xp.conjugate(eval_ji)
k += 1
+ self._stop_stage_timer("eval_V", t0)
def eval_diag_V(self) -> None:
"""Evaluate diagonal potential matrix elements."""
+ t0 = self._start_stage_timer()
+ local_scope, global_scope = self._evaluation_scope()
for i in range(self.psi.num_int_dim):
- self.V_eval_array[:, :, :, i, i] = self._backend.evaluate(
+ self.V_eval_diag_array[:, :, :, i] = self._evaluator.eval(
self.v.potential_strings[i],
- local_dict={**self.v.variables, **self.psi.default_var_dict},
- global_dict={"t": self.psi.t},
+ local_dict=local_scope,
+ global_dict=global_scope,
)
+ self._stop_stage_timer("eval_diag_V", t0)
class Potential:
"""Simple container for potential metadata."""
diff --git a/pytalises/wavefunction.py b/pytalises/wavefunction.py
index f963443..d37907d 100644
--- a/pytalises/wavefunction.py
+++ b/pytalises/wavefunction.py
@@ -10,6 +10,7 @@
import pytalises.propagator
from pytalises.backends import get_backend
from pytalises.backends.base import Backend
+from pytalises.engine import ExpressionEvaluator, create_engine
from pytalises.grid import Grid
from pytalises.options import PropagationOptions
from pytalises.potentials import BasePotential
@@ -63,6 +64,8 @@ def __init__(
raise TypeError("variables must be dict")
self._backend = get_backend(backend, num_threads=num_threads)
+ self._engine = create_engine(self._backend)
+ self._evaluator = ExpressionEvaluator(backend_name=self._backend.name)
self.grid = grid
self.initial = list(initial)
@@ -75,9 +78,9 @@ def __init__(
self.num_ext_dim = int(sum(self.axes))
self.nX, self.nY, self.nZ = self.number_of_grid_points
- r: list[NDArray[np.floating[Any]]] = []
+ r: list[Any] = []
Delta_r: list[float] = []
- k: list[NDArray[np.floating[Any]] | float] = []
+ k: list[Any] = []
delta_k: list[float] = []
for i, spatial_ext_tuple in enumerate(self.spatial_ext):
@@ -89,7 +92,9 @@ def __init__(
if self.axes[i] == 0:
Delta_r.append(np.nan)
delta_k.append(np.nan)
- k.append(0.0)
+ # Keep inactive reciprocal axes as 1D arrays for backend parity
+ # (CuPy meshgrid requires array-like inputs, not Python scalars).
+ k.append(self._backend.asarray([0.0]))
else:
drange = r_max - r_min
Delta_r.append(drange)
@@ -128,7 +133,7 @@ def __init__(
self.variables = variables
for i in range(self.num_int_dim):
- self._amp[:, :, :, i] = self._backend.evaluate(
+ self._amp[:, :, :, i] = self._evaluator.eval(
self.initial[i],
local_dict={**self.default_var_dict, **self.variables},
global_dict={"t": self.t},
@@ -148,10 +153,10 @@ def construct_FFT(
"FFTW_DESTROY_INPUT",
),
) -> None:
- """Construct FFT bindings through backend plan interface."""
+ """Construct FFT bindings through engine plan interface."""
axes = tuple(i for i, active in enumerate(self.axes) if active)
- self._backend.set_num_threads(num_threads)
- self._fft_plan = self._backend.create_fft_plan(
+ self._engine.set_num_threads(num_threads)
+ self._fft_plan = self._engine.create_fft_plan(
self._amp,
axes=axes,
num_threads=num_threads,
@@ -203,12 +208,13 @@ def exp_pos(self, axis: int | None = None) -> NDArray[np.floating[Any]]:
if not (0 <= axis < self.num_ext_dim):
raise ValueError(f"axis must be in [0, {self.num_ext_dim - 1}]")
+ xp = self._engine.xp
phys_axis = active_axes[axis]
axes_to_trace = [0, 1, 2]
axes_to_trace.pop(phys_axis)
- psi_sq_amp = self._backend.power(self._backend.abs(self._amp), 2)
- traced_out_psi = self._backend.sum(psi_sq_amp, axis=tuple(axes_to_trace))
- exp_pos = self._backend.einsum("ri,r->", traced_out_psi, self._r[phys_axis])
+ psi_sq_amp = xp.abs(self._amp) ** 2
+ traced_out_psi = xp.sum(psi_sq_amp, axis=tuple(axes_to_trace))
+ exp_pos = xp.sum(traced_out_psi * self._r[phys_axis][:, xp.newaxis])
exp_pos *= self._volume_element()
return exp_pos
@@ -224,24 +230,26 @@ def var_pos(self, axis: int | None = None) -> NDArray[np.floating[Any]]:
if not (0 <= axis < self.num_ext_dim):
raise ValueError(f"axis must be in [0, {self.num_ext_dim - 1}]")
+ xp = self._engine.xp
phys_axis = active_axes[axis]
axes_to_trace = [0, 1, 2]
axes_to_trace.pop(phys_axis)
- psi_sq_amp = self._backend.power(self._backend.abs(self._amp), 2)
- traced_out_psi = self._backend.sum(psi_sq_amp, axis=tuple(axes_to_trace))
- var_pos = self._backend.einsum(
- "ri,r->",
- traced_out_psi,
- (self._r[phys_axis] - self.exp_pos(axis)) ** 2,
- )
+ psi_sq_amp = xp.abs(self._amp) ** 2
+ traced_out_psi = xp.sum(psi_sq_amp, axis=tuple(axes_to_trace))
+ variance_axis = (self._r[phys_axis] - self.exp_pos(axis)) ** 2
+ var_pos = xp.sum(traced_out_psi * variance_axis[:, xp.newaxis])
var_pos *= self._volume_element()
return var_pos
def normalize_to(self, n_const: float) -> None:
"""Normalize total wavefunction occupation to ``n_const``."""
- s = self._backend.einsum("xyzi,xyzi->", self._amp, self._backend.conjugate(self._amp))
- s *= self._volume_element()
- self._amp *= self._backend.sqrt(n_const / s)
+ xp = self._engine.xp
+ s = self._engine.inner_product(
+ self._amp,
+ self._amp,
+ volume_element=self._volume_element(),
+ )
+ self._amp *= xp.sqrt(n_const / s)
def state_occupation(
self, nth_state: int | None = None
@@ -256,9 +264,10 @@ def state_occupation(
if not (0 <= nth_state < self.num_int_dim):
raise ValueError(f"nth_state must be in [0, {self.num_int_dim - 1}]")
- return (
- self._backend.sum(self._backend.abs(self._amp[:, :, :, nth_state]) ** 2)
- * self._volume_element()
+ return self._engine.state_occupation(
+ self._amp,
+ state_index=nth_state,
+ volume_element=self._volume_element(),
)
def freely_propagate(
diff --git a/tests/backend_parity_test.py b/tests/backend_parity_test.py
new file mode 100644
index 0000000..5a9dbf5
--- /dev/null
+++ b/tests/backend_parity_test.py
@@ -0,0 +1,76 @@
+import numpy as np
+import pytest
+
+import pytalises as pt
+
+
+def _to_numpy(x):
+ if hasattr(x, "get"):
+ return np.asarray(x.get())
+ return np.asarray(x)
+
+
+def _backend_runtime_available(name: str) -> bool:
+ if name == "cupy":
+ return pt.has_cupy()
+ return True
+
+
+def _run_reference_case(backend: str):
+ grid = pt.Grid(shape=(96,), extent=((-4.0, 4.0),))
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+
+ psi = pt.Wavefunction(initial, grid, normalize_const=1.0, backend=backend)
+ psi.propagate(
+ potential=potential,
+ steps=10,
+ dt=0.005,
+ options=pt.PropagationOptions(backend=backend),
+ )
+ return psi
+
+
+CANDIDATE_BACKENDS = [name for name in pt.available_backends() if name != "numpy"]
+if not CANDIDATE_BACKENDS:
+ CANDIDATE_BACKENDS = [
+ pytest.param(
+ "_none_",
+ marks=pytest.mark.skip(reason="No non-NumPy backends are importable"),
+ )
+ ]
+
+
+@pytest.mark.parametrize("backend_name", CANDIDATE_BACKENDS)
+def test_backend_matches_numpy_reference(backend_name: str):
+ if not _backend_runtime_available(backend_name):
+ pytest.skip(f"{backend_name} runtime is not available")
+
+ psi_np = _run_reference_case("numpy")
+ psi_other = _run_reference_case(backend_name)
+
+ amp_np = _to_numpy(psi_np.amp)
+ amp_other = _to_numpy(psi_other.amp)
+
+ np.testing.assert_allclose(
+ np.abs(amp_other) ** 2,
+ np.abs(amp_np) ** 2,
+ rtol=5e-5,
+ atol=5e-7,
+ )
+ np.testing.assert_allclose(
+ _to_numpy(psi_other.state_occupation()),
+ _to_numpy(psi_np.state_occupation()),
+ rtol=5e-5,
+ atol=5e-7,
+ )
+
+
+def test_numpy_backend_is_available():
+ assert "numpy" in pt.available_backends()
diff --git a/tests/backend_test.py b/tests/backend_test.py
index 61635aa..13ed2f5 100644
--- a/tests/backend_test.py
+++ b/tests/backend_test.py
@@ -1,9 +1,13 @@
+import numpy as np
+import pytest
+
import pytalises as pt
+import pytalises.backends as backends
-def test_default_backend_is_numpy():
+def test_default_backend_is_registered_backend():
backend = pt.get_backend()
- assert backend.name == "numpy"
+ assert backend.name in {"numpy", "cupy"}
def test_wavefunction_uses_backend_instance():
@@ -11,3 +15,161 @@ def test_wavefunction_uses_backend_instance():
grid = pt.Grid(shape=(32,), extent=((-1, 1),))
psi = pt.Wavefunction("exp(-x**2)", grid=grid, backend=backend)
assert psi._backend is backend
+
+
+def test_requesting_cupy_backend_matches_availability():
+ if pt.has_cupy():
+ backend = pt.get_backend("cupy")
+ assert backend.name == "cupy"
+ else:
+ with pytest.raises(ValueError):
+ pt.get_backend("cupy")
+
+
+@pytest.mark.skipif("cupy" not in pt.available_backends(), reason="CuPy backend module not importable")
+def test_requesting_cupy_backend_fails_fast_without_visible_device(monkeypatch):
+ monkeypatch.setattr(backends, "has_cupy", lambda: False)
+ with pytest.raises(ValueError, match="no usable CUDA device"):
+ backends.get_backend("cupy")
+
+
+@pytest.mark.skipif(not pt.has_cupy(), reason="CuPy backend not available")
+def test_wavefunction_can_run_on_cupy_backend():
+ cupy = pytest.importorskip("cupy")
+ grid = pt.Grid(shape=(16,), extent=((-1, 1),))
+ psi = pt.Wavefunction("exp(-x**2)", grid=grid, backend="cupy")
+ assert isinstance(psi.amp, cupy.ndarray)
+ psi.freely_propagate(steps=2, dt=0.01)
+ occ = psi.state_occupation()
+ assert occ.shape == (1,)
+
+
+def test_core_propagation_path_does_not_use_backend_einsum(monkeypatch):
+ grid = pt.Grid(shape=(32,), extent=((-2, 2),))
+ psi = pt.Wavefunction("exp(-x**2)", grid=grid, backend="numpy")
+
+ def _unexpected_einsum(*args, **kwargs):
+ raise AssertionError("Propagation should not call backend.einsum in v2 core path")
+
+ monkeypatch.setattr(psi._backend, "einsum", _unexpected_einsum)
+
+ psi.freely_propagate(steps=2, dt=0.01)
+ psi.propagate(
+ potential=pt.DiagonalPotential("0.1*x**2"),
+ steps=2,
+ dt=0.01,
+ options=pt.PropagationOptions(backend="numpy"),
+ )
+
+
+def test_coupled_2x2_mode_rejects_invalid_value():
+ with pytest.raises(ValueError, match="coupled_2x2_mode"):
+ pt.PropagationOptions(coupled_2x2_mode="invalid")
+
+
+def test_coupled_2x2_kernel_rejects_invalid_value():
+ with pytest.raises(ValueError, match="coupled_2x2_kernel"):
+ pt.PropagationOptions(coupled_2x2_kernel="invalid")
+
+
+def test_dtype_rejects_invalid_value():
+ with pytest.raises(ValueError, match="PropagationOptions.dtype"):
+ pt.PropagationOptions(dtype="float64")
+
+
+def test_analytic_2x2_path_matches_eigh_reference_numpy():
+ grid = pt.Grid(shape=(64,), extent=((-3, 3),))
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+
+ psi_auto = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="numpy")
+ psi_eigh = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="numpy")
+
+ steps = 12
+ dt = 0.005
+
+ psi_auto.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(backend="numpy", coupled_2x2_mode="auto"),
+ )
+ psi_eigh.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(backend="numpy", coupled_2x2_mode="eigh"),
+ )
+
+ np.testing.assert_allclose(
+ np.abs(np.asarray(psi_auto.amp)) ** 2,
+ np.abs(np.asarray(psi_eigh.amp)) ** 2,
+ rtol=1e-10,
+ atol=1e-12,
+ )
+ np.testing.assert_allclose(
+ np.asarray(psi_auto.state_occupation()),
+ np.asarray(psi_eigh.state_occupation()),
+ rtol=1e-10,
+ atol=1e-12,
+ )
+
+
+def test_complex64_mode_matches_complex128_reference_numpy():
+ grid = pt.Grid(shape=(64,), extent=((-3, 3),))
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+
+ psi64 = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="numpy")
+ psi128 = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="numpy")
+
+ steps = 16
+ dt = 0.005
+
+ psi64.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="numpy",
+ dtype="complex64",
+ coupled_2x2_mode="auto",
+ ),
+ )
+ psi128.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="numpy",
+ dtype="complex128",
+ coupled_2x2_mode="auto",
+ ),
+ )
+
+ assert np.asarray(psi64.amp).dtype == np.complex64
+
+ np.testing.assert_allclose(
+ np.abs(np.asarray(psi64.amp)) ** 2,
+ np.abs(np.asarray(psi128.amp)) ** 2,
+ rtol=5e-4,
+ atol=5e-6,
+ )
+ np.testing.assert_allclose(
+ np.asarray(psi64.state_occupation()),
+ np.asarray(psi128.state_occupation()),
+ rtol=5e-4,
+ atol=5e-6,
+ )
diff --git a/tests/benchmark_script_test.py b/tests/benchmark_script_test.py
new file mode 100644
index 0000000..b4608c6
--- /dev/null
+++ b/tests/benchmark_script_test.py
@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+
+import benchmarks.backend_benchmark as bench
+
+
+def test_benchmark_script_runs_directly_from_repo_root():
+ repo_root = Path(__file__).resolve().parents[1]
+ proc = subprocess.run(
+ [sys.executable, "benchmarks/backend_benchmark.py", "--help"],
+ cwd=repo_root,
+ capture_output=True,
+ text=True,
+ check=False,
+ )
+ assert proc.returncode == 0
+ assert "Benchmark pyTALISES backends" in proc.stdout
+
+
+def test_main_rejects_zero_repeats(monkeypatch):
+ monkeypatch.setattr(sys, "argv", ["backend_benchmark.py", "--repeats", "0"])
+ with pytest.raises(SystemExit):
+ bench.main()
+
+
+def test_min_speedup_fails_when_cupy_unavailable(monkeypatch):
+ monkeypatch.setattr(bench.pt, "has_cupy", lambda: False)
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "backend_benchmark.py",
+ "--sizes",
+ "16",
+ "--steps",
+ "1",
+ "--repeats",
+ "1",
+ "--warmup",
+ "0",
+ "--workloads",
+ "free",
+ "--min-speedup",
+ "1.1",
+ ],
+ )
+
+ assert bench.main() == 1
+
+
+def test_stage_breakdown_is_emitted_in_json(monkeypatch, tmp_path):
+ json_out = tmp_path / "benchmark.json"
+ monkeypatch.setattr(
+ sys,
+ "argv",
+ [
+ "backend_benchmark.py",
+ "--sizes",
+ "16",
+ "--steps",
+ "2",
+ "--repeats",
+ "1",
+ "--warmup",
+ "0",
+ "--workloads",
+ "free",
+ "--coupled-2x2-mode",
+ "eigh",
+ "--coupled-2x2-kernel",
+ "fused",
+ "--dtype",
+ "complex64",
+ "--json-out",
+ str(json_out),
+ ],
+ )
+
+ assert bench.main() == 0
+
+ payload = json.loads(json_out.read_text(encoding="utf-8"))
+ assert payload["config"]["coupled_2x2_mode"] == "eigh"
+ assert payload["config"]["coupled_2x2_kernel"] == "fused"
+ assert payload["config"]["dtype"] == "complex64"
+
+ breakdown = payload.get("stage_breakdown")
+ assert isinstance(breakdown, list)
+ assert breakdown
+
+ entry = breakdown[0]
+ assert entry["backend"] == "numpy"
+ assert entry["workload"] == "free"
+ assert entry["coupled_2x2_mode"] == "eigh"
+ assert entry["coupled_2x2_kernel"] == "fused"
+ assert entry["dtype"] == "complex64"
+ assert "kinetic_prop" in entry["stages"]
+ assert entry["stages"]["kinetic_prop"]["mean_seconds"] >= 0.0
diff --git a/tests/cupy_backend_test.py b/tests/cupy_backend_test.py
new file mode 100644
index 0000000..f3f0a33
--- /dev/null
+++ b/tests/cupy_backend_test.py
@@ -0,0 +1,283 @@
+import numpy as np
+import pytest
+
+import pytalises as pt
+
+
+pytestmark = pytest.mark.skipif(not pt.has_cupy(), reason="CuPy backend not available")
+
+
+def g1(n=64, a=-2.0, b=2.0):
+ return pt.Grid(shape=(n,), extent=((a, b),))
+
+
+def _to_numpy(x):
+ import cupy as cp
+
+ if isinstance(x, cp.ndarray):
+ return cp.asnumpy(x)
+ return np.asarray(x)
+
+
+def test_cupy_free_propagation_conserves_norm():
+ psi = pt.Wavefunction(
+ ["exp(-x**2)"],
+ g1(64, -3, 3),
+ normalize_const=1.0,
+ backend="cupy",
+ )
+ psi.freely_propagate(
+ steps=5,
+ dt=0.01,
+ options=pt.PropagationOptions(backend="cupy"),
+ )
+ np.testing.assert_allclose(np.sum(psi.state_occupation()), 1.0, atol=1e-5)
+
+
+def test_cupy_diagonal_time_dependent_potential():
+ psi = pt.Wavefunction(
+ ["exp(-x**2)", "0"],
+ g1(64, -3, 3),
+ normalize_const=1.0,
+ backend="cupy",
+ )
+ psi.propagate(
+ potential=pt.DiagonalPotential(["t*x**2", "t*x**2"]),
+ steps=4,
+ dt=0.01,
+ options=pt.PropagationOptions(backend="cupy"),
+ )
+ np.testing.assert_allclose(np.sum(psi.state_occupation()), 1.0, atol=1e-5)
+
+
+def test_cupy_nondiagonal_potential():
+ psi = pt.Wavefunction(
+ ["exp(-x**2)", "0"],
+ g1(32, -2, 2),
+ normalize_const=1.0,
+ backend="cupy",
+ )
+ psi.propagate(
+ potential=pt.HermitianPotential.from_lower_triangular(["0", "sin(t)", "0"]),
+ steps=3,
+ dt=0.01,
+ options=pt.PropagationOptions(backend="cupy"),
+ )
+ np.testing.assert_allclose(np.sum(psi.state_occupation()), 1.0, atol=1e-5)
+
+
+def test_cupy_matches_numpy_reference_dynamics():
+ grid = g1(96, -4, 4)
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+
+ psi_np = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="numpy")
+ psi_cp = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+
+ steps = 10
+ dt = 0.005
+ psi_np.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(backend="numpy"),
+ )
+ psi_cp.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(backend="cupy"),
+ )
+
+ amp_np = _to_numpy(psi_np.amp)
+ amp_cp = _to_numpy(psi_cp.amp)
+
+ np.testing.assert_allclose(
+ np.abs(amp_cp) ** 2,
+ np.abs(amp_np) ** 2,
+ rtol=5e-5,
+ atol=5e-7,
+ )
+ np.testing.assert_allclose(
+ _to_numpy(psi_cp.state_occupation()),
+ _to_numpy(psi_np.state_occupation()),
+ rtol=5e-5,
+ atol=5e-7,
+ )
+
+
+def test_cupy_backend_einsum_out_semantics():
+ cp = pytest.importorskip("cupy")
+ backend = pt.get_backend("cupy")
+
+ a = cp.arange(6, dtype=cp.float64).reshape(2, 3)
+ b = cp.arange(3, dtype=cp.float64)
+ out = cp.empty((2,), dtype=cp.float64)
+
+ result = backend.einsum("ij,j->i", a, b, out=out)
+
+ np.testing.assert_allclose(cp.asnumpy(out), cp.asnumpy(cp.einsum("ij,j->i", a, b)))
+ assert result is out
+
+
+def test_cupy_analytic_2x2_matches_eigh_reference():
+ grid = g1(96, -4, 4)
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+
+ psi_auto = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+ psi_eigh = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+
+ steps = 10
+ dt = 0.005
+ psi_auto.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(backend="cupy", coupled_2x2_mode="auto"),
+ )
+ psi_eigh.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(backend="cupy", coupled_2x2_mode="eigh"),
+ )
+
+ amp_auto = _to_numpy(psi_auto.amp)
+ amp_eigh = _to_numpy(psi_eigh.amp)
+
+ np.testing.assert_allclose(
+ np.abs(amp_auto) ** 2,
+ np.abs(amp_eigh) ** 2,
+ rtol=5e-5,
+ atol=5e-7,
+ )
+ np.testing.assert_allclose(
+ _to_numpy(psi_auto.state_occupation()),
+ _to_numpy(psi_eigh.state_occupation()),
+ rtol=5e-5,
+ atol=5e-7,
+ )
+
+
+def test_cupy_fused_analytic_2x2_matches_vectorized_reference():
+ grid = g1(96, -4, 4)
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+
+ psi_vec = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+ psi_fused = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+
+ steps = 10
+ dt = 0.005
+ psi_vec.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="cupy",
+ coupled_2x2_mode="auto",
+ coupled_2x2_kernel="vectorized",
+ ),
+ )
+ psi_fused.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="cupy",
+ coupled_2x2_mode="auto",
+ coupled_2x2_kernel="fused",
+ ),
+ )
+
+ amp_vec = _to_numpy(psi_vec.amp)
+ amp_fused = _to_numpy(psi_fused.amp)
+
+ np.testing.assert_allclose(
+ np.abs(amp_fused) ** 2,
+ np.abs(amp_vec) ** 2,
+ rtol=5e-5,
+ atol=5e-7,
+ )
+ np.testing.assert_allclose(
+ _to_numpy(psi_fused.state_occupation()),
+ _to_numpy(psi_vec.state_occupation()),
+ rtol=5e-5,
+ atol=5e-7,
+ )
+
+
+def test_cupy_complex64_mode_matches_complex128_reference():
+ grid = g1(96, -4, 4)
+ initial = ["exp(-x**2)", "0.2*exp(-(x-0.5)**2)"]
+ potential = pt.HermitianPotential.from_lower_triangular(
+ [
+ "0.2*x**2 + 0.1*t",
+ "0.03*cos(x) + 0.02*sin(t)",
+ "0.1*x**2 - 0.05*t",
+ ]
+ )
+
+ psi64 = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+ psi128 = pt.Wavefunction(initial, grid, normalize_const=1.0, backend="cupy")
+
+ steps = 10
+ dt = 0.005
+
+ psi64.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="cupy",
+ dtype="complex64",
+ coupled_2x2_mode="auto",
+ ),
+ )
+ psi128.propagate(
+ potential=potential,
+ steps=steps,
+ dt=dt,
+ options=pt.PropagationOptions(
+ backend="cupy",
+ dtype="complex128",
+ coupled_2x2_mode="auto",
+ ),
+ )
+
+ amp64 = _to_numpy(psi64.amp)
+ amp128 = _to_numpy(psi128.amp)
+
+ assert amp64.dtype == np.complex64
+
+ np.testing.assert_allclose(
+ np.abs(amp64) ** 2,
+ np.abs(amp128) ** 2,
+ rtol=2e-3,
+ atol=2e-5,
+ )
+ np.testing.assert_allclose(
+ _to_numpy(psi64.state_occupation()),
+ _to_numpy(psi128.state_occupation()),
+ rtol=2e-3,
+ atol=2e-5,
+ )
diff --git a/tests/engine_kernel_test.py b/tests/engine_kernel_test.py
new file mode 100644
index 0000000..e5fd22f
--- /dev/null
+++ b/tests/engine_kernel_test.py
@@ -0,0 +1,28 @@
+import numpy as np
+
+import pytalises as pt
+from pytalises.engine import create_engine
+
+
+def test_coupled_kernel_matches_manual_matrix_update():
+ backend = pt.get_backend("numpy")
+ engine = create_engine(backend)
+
+ amp = backend.asarray([[[[1.0 + 0.0j, 0.25 - 0.1j]]]], dtype="complex128")
+ eigvals = backend.asarray([[[[0.3, -0.2]]]], dtype="complex128")
+
+ s = 1 / np.sqrt(2)
+ eigvec = np.array([[s, s], [s, -s]], dtype=np.complex128)
+ eigvecs = backend.asarray([[[eigvec]]], dtype="complex128")
+
+ dt = 0.7
+
+ expected = np.asarray(amp).copy()
+ propagator = eigvec @ np.diag(np.exp(-1j * np.array([0.3, -0.2]) * dt)) @ np.conjugate(
+ eigvec.T
+ )
+ expected[0, 0, 0, :] = propagator @ expected[0, 0, 0, :]
+
+ engine.apply_coupled_phase(amp, eigvals=eigvals, eigvecs=eigvecs, dt=dt)
+
+ np.testing.assert_allclose(np.asarray(amp), expected, atol=1e-12, rtol=1e-12)