diff --git a/.github/actions/run-benchmarks/action.yml b/.github/actions/run-benchmarks/action.yml index b567c493f..e5d2342c3 100644 --- a/.github/actions/run-benchmarks/action.yml +++ b/.github/actions/run-benchmarks/action.yml @@ -31,7 +31,7 @@ runs: shell: bash run: | sudo apt-get update - sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev libhdf5-dev + sudo apt-get install -y build-essential cmake libnetcdf-dev libomp-dev libhdf5-dev - name: Mount bazel cache uses: actions/cache@v4 diff --git a/.github/constraints.txt b/.github/constraints.txt new file mode 100644 index 000000000..27937b037 --- /dev/null +++ b/.github/constraints.txt @@ -0,0 +1,3 @@ +# Block simsopt 1.11.0 wheels with unsupported AVX-512 instructions. Python 3.14 +# has no wheel and must use the corrected 1.11.0 source distribution. +simsopt!=1.11.0; python_version != "3.14" diff --git a/.github/workflows/benchmarks.yaml b/.github/workflows/benchmarks.yaml index 5bd06fa2a..f60fed849 100644 --- a/.github/workflows/benchmarks.yaml +++ b/.github/workflows/benchmarks.yaml @@ -8,6 +8,11 @@ permissions: contents: read pull-requests: write +# CI-only dependency workaround, see .github/constraints.txt +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + jobs: benchmark: runs-on: ubuntu-22.04 diff --git a/.github/workflows/clang_tidy.yaml b/.github/workflows/clang_tidy.yaml index 72d944cd7..f31b3ec15 100644 --- a/.github/workflows/clang_tidy.yaml +++ b/.github/workflows/clang_tidy.yaml @@ -28,7 +28,7 @@ jobs: id: review with: config_file: ".clang-tidy" - apt_packages: libhdf5-dev, liblapack-dev, libnetcdf-dev, gfortran, python3-dev, libomp-dev + apt_packages: zlib1g-dev, gfortran, python3-dev, libomp-dev build_dir: build cmake_command: cmake -B build && cmake -DCMAKE_EXPORT_COMPILE_COMMANDS:STRING=ON build # bazel-built test and benchmark files have no cmake compile_commands diff --git a/.github/workflows/copilot-setup-steps.yml b/.github/workflows/copilot-setup-steps.yml index df04747cd..4766f3400 100644 --- a/.github/workflows/copilot-setup-steps.yml +++ b/.github/workflows/copilot-setup-steps.yml @@ -11,6 +11,11 @@ on: paths: - .github/workflows/copilot-setup-steps.yml +# CI-only dependency workaround, see .github/constraints.txt +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + jobs: # The job MUST be called `copilot-setup-steps` or it will not be picked up by Copilot. copilot-setup-steps: @@ -37,10 +42,8 @@ jobs: sudo apt-get install -y \ build-essential \ cmake \ - libnetcdf-dev \ - liblapack-dev \ + zlib1g-dev \ libomp-dev \ - libhdf5-dev \ python3-dev - name: Install vmecpp with test dependencies diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index 7fc0352f7..c6d00590d 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -16,6 +16,11 @@ concurrency: group: "pages" cancel-in-progress: false +# CI-only dependency workaround, see .github/constraints.txt +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + jobs: benchmark: runs-on: ubuntu-22.04 @@ -48,7 +53,7 @@ jobs: - name: Build docs run: | sudo apt-get update - sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev libhdf5-dev + sudo apt-get install -y build-essential cmake zlib1g-dev libomp-dev # must actually install the package otherwise sphinx can't autogenerate the API reference python -m pip install -v .[docs] sphinx-apidoc --module-first --no-toc --force --separate -o docs/api src/vmecpp/ && sphinx-build docs html_docs diff --git a/.github/workflows/full_validation.yaml b/.github/workflows/full_validation.yaml index 7fb8179a3..324be2968 100644 --- a/.github/workflows/full_validation.yaml +++ b/.github/workflows/full_validation.yaml @@ -20,6 +20,11 @@ concurrency: group: ${{ github.workflow }}-${{ github.ref || github.run_id }} cancel-in-progress: ${{ github.ref_name != 'main' }} +# CI-only dependency workaround, see .github/constraints.txt +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + jobs: validation: name: Run VMEC++ validation @@ -28,6 +33,9 @@ jobs: # wout with Fortran VMEC2000 in Docker plus one with VMEC++: budget generously. # The short run finishes in minutes and won't come close to this limit. timeout-minutes: 360 + env: + PIP_CONSTRAINT: ${{ github.workspace }}/vmecpp/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/vmecpp/.github/constraints.txt steps: - name: Check out VMEC++ uses: actions/checkout@v4 @@ -49,8 +57,8 @@ jobs: - name: Install required system packages for Ubuntu run: | sudo apt-get update && sudo apt-get install -y \ - build-essential cmake libnetcdf-dev liblapack-dev liblapacke-dev libopenmpi-dev \ - libomp-dev libeigen3-dev nlohmann-json3-dev libhdf5-dev + build-essential cmake zlib1g-dev libopenmpi-dev \ + libomp-dev libeigen3-dev nlohmann-json3-dev - name: Install vmecpp-validation's Python requirements run: | @@ -61,7 +69,7 @@ jobs: python -m pip install -r requirements.no-vmecpp.txt - name: Install VMEC++ from this commit - run: python -m pip install ./vmecpp + run: python -m pip install "./vmecpp" "numpy<2.1" - name: Run short validation if: github.event_name == 'push' @@ -96,8 +104,8 @@ jobs: - name: Install required packages for Ubuntu run: | sudo apt-get update && sudo apt-get install -y \ - build-essential cmake libnetcdf-dev liblapack-dev liblapacke-dev libopenmpi-dev \ - libomp-dev libeigen3-dev nlohmann-json3-dev libhdf5-dev + build-essential cmake zlib1g-dev libopenmpi-dev \ + libomp-dev libeigen3-dev nlohmann-json3-dev - name: Install package run: python -m pip install -v .[test] - name: Run slow QUASR free-boundary tests diff --git a/.github/workflows/pypi_publish.yml b/.github/workflows/pypi_publish.yml index b5f268bb3..375c94318 100644 --- a/.github/workflows/pypi_publish.yml +++ b/.github/workflows/pypi_publish.yml @@ -4,9 +4,16 @@ on: # allows running workflows manually workflow_dispatch: + pull_request: + release: types: [published] +# CI-only dependency workaround, see .github/constraints.txt +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + jobs: main: name: Build wheels on ${{ matrix.os }} @@ -14,8 +21,7 @@ jobs: strategy: fail-fast: false matrix: - # ubuntu-24.04-arm not supported due to lack of netcdf-devel hdf5-devel - os: [ubuntu-24.04, macos-14, macos-15-intel] + os: [ubuntu-24.04, ubuntu-24.04-arm, macos-14, macos-15-intel] steps: - name: Checkout repository @@ -26,11 +32,15 @@ jobs: uses: pypa/cibuildwheel@v3.3.1 # All the settings for cibuildwheel are in pyproject.toml env: + # Linux builds run in a container, where PIP_CONSTRAINT above does not + # apply and the repo is mounted at /project. + CIBW_ENVIRONMENT_LINUX: > + PIP_CONSTRAINT=/project/.github/constraints.txt + UV_CONSTRAINT=/project/.github/constraints.txt CIBW_ENVIRONMENT_MACOS: > MACOSX_DEPLOYMENT_TARGET=${{ matrix.os == 'macos-14' && '14.0' || '15.0' }} DYLD_LIBRARY_PATH=/usr/local/opt/gcc/lib/gcc/current/:$DYLD_LIBRARY_PATH FC=gfortran-14 - HDF5_ROOT=$(brew --prefix hdf5) OpenMP_ROOT=$(brew --prefix)/opt/libomp with: package-dir: . @@ -76,7 +86,15 @@ jobs: # "debian:testing-20250224", # TODO(jurasic): re-enable when 3.14 support unblocked "fedora:41", ] - runs-on: ubuntu-24.04 + runner: [ubuntu-24.04] + include: + # Exercise the aarch64 wheels too, on a native arm64 runner. One + # distro is enough here -- this job is about testing the wheel + # actually works on a plain distro, and the multi-distro matrix + # above already covers distro-specific quirks on x86_64. + - distro: "ubuntu:24.04" + runner: ubuntu-24.04-arm + runs-on: ${{ matrix.runner }} container: ${{ matrix.distro }} steps: @@ -116,19 +134,17 @@ jobs: merge-multiple: true path: ./wheelhouse/ - name: Install wheel - # Install the appropriate file for vmecpp from the wheelhouse directory. - # --no-index prevents installing from PyPi, so we don't accidentally install - # a published version instead of the newly built ones. At the same time it prevents - # loading dependencies from PyPi, so the first installation fails. We only use the console output - # from this command to get the appropriate .whl file for this platform and environment. run: | ls wheelhouse/ $PYTHON -m venv venv . venv/bin/activate - wheel_filename=$(find wheelhouse/$($PYTHON --version | awk '{print $2}' | awk -F. '{printf "vmecpp*-cp%s%s-*manylinux*.whl", $1, $2}')) + # Let pip's own tag matching pick the single wheel compatible with this + # interpreter and architecture; --no-index keeps it from reaching PyPi. + pip download --no-index --no-deps --find-links=wheelhouse/ --dest selected-wheel/ vmecpp + wheel_filename=$(ls selected-wheel/*.whl) echo "Found matching wheel: $wheel_filename" - pip install $wheel_filename + pip install "$wheel_filename" pip install pytest python -m pytest tests/ diff --git a/.github/workflows/pyright.yaml b/.github/workflows/pyright.yaml index a62ddea7a..e02ec5d3b 100644 --- a/.github/workflows/pyright.yaml +++ b/.github/workflows/pyright.yaml @@ -24,7 +24,7 @@ jobs: - name: Install dependencies run: | sudo apt-get update - sudo apt-get install build-essential cmake libnetcdf-dev liblapack-dev python3-dev + sudo apt-get install build-essential cmake zlib1g-dev python3-dev # Parse the dependencies from the pyproject.toml file and install them, # without installing and building the entire vmecpp project - run: | diff --git a/.github/workflows/test_bazel.yaml b/.github/workflows/test_bazel.yaml index 215dd65bc..775984d82 100644 --- a/.github/workflows/test_bazel.yaml +++ b/.github/workflows/test_bazel.yaml @@ -25,7 +25,7 @@ jobs: lfs: true # for the mgrid test data - name: Install required system packages for Ubuntu run: | - sudo apt-get update && sudo apt-get install -y build-essential cmake liblapack-dev libomp-dev python3-dev python3-pip python-is-python3 + sudo apt-get update && sudo apt-get install -y build-essential cmake libomp-dev python3-dev python3-pip python-is-python3 - name: Mount bazel cache uses: actions/cache@v4 with: diff --git a/.github/workflows/test_conda_env.yaml b/.github/workflows/test_conda_env.yaml index 46679be68..7093e82c4 100644 --- a/.github/workflows/test_conda_env.yaml +++ b/.github/workflows/test_conda_env.yaml @@ -12,6 +12,11 @@ concurrency: group: ${{ github.workflow }}-${{ github.ref || github.run_id }} cancel-in-progress: ${{ github.ref_name != 'main' }} +# CI-only dependency workaround, see .github/constraints.txt +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + jobs: test-conda-env: name: Test conda env diff --git a/.github/workflows/test_enzyme.yaml b/.github/workflows/test_enzyme.yaml index 1e226087c..83aaeeefc 100644 --- a/.github/workflows/test_enzyme.yaml +++ b/.github/workflows/test_enzyme.yaml @@ -32,7 +32,8 @@ jobs: sudo ./llvm.sh "${LLVM_VERSION}" sudo apt-get install -y \ "llvm-${LLVM_VERSION}-dev" "clang-${LLVM_VERSION}" "libclang-${LLVM_VERSION}-dev" \ - cmake ninja-build + "libomp-${LLVM_VERSION}-dev" \ + cmake gfortran libhdf5-dev liblapack-dev libnetcdf-dev ninja-build python3-dev - name: Cache the ClangEnzyme plugin build id: enzyme-cache @@ -75,3 +76,28 @@ jobs: --copt="-fplugin=${ENZYME_PLUGIN}" \ --test_output=all \ //vmecpp/common/enzyme:enzyme_smoke_test + + - uses: astral-sh/setup-uv@v6 + + - name: Run Enzyme force-derivative integration tests + env: + CC: clang-${{ env.LLVM_VERSION }} + CXX: clang++-${{ env.LLVM_VERSION }} + CMAKE_ARGS: >- + -DVMECPP_ENABLE_ENZYME=ON + -DVMECPP_ENZYME_PLUGIN=${{ env.ENZYME_PLUGIN }} + run: | + uv sync --reinstall-package vmecpp + uv run --with pytest pytest -q \ + tests/test_hessian.py \ + tests/test_external_optimizers.py + + - name: Run exact-HVP benchmark + run: | + uv run python benchmarks/benchmark_exact_hvp.py | tee exact_hvp_benchmark.txt + { + echo "### Exact-HVP benchmark" + echo '```text' + cat exact_hvp_benchmark.txt + echo '```' + } >> "${GITHUB_STEP_SUMMARY}" diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 6607cb90b..31bc72573 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -12,6 +12,11 @@ concurrency: group: ${{ github.workflow }}-${{ github.ref || github.run_id }} cancel-in-progress: ${{ github.ref_name != 'main' }} +# CI-only dependency workaround, see .github/constraints.txt +env: + PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt + jobs: pre-commit: env: @@ -32,7 +37,7 @@ jobs: matrix: os: [ubuntu-22.04, macos-14] editable: [true, false] - python-version: ["3.10", "3.11", "3.12", "3.13"] + python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v4 with: @@ -44,12 +49,14 @@ jobs: - name: Install required packages for MacOS if: ${{ contains(matrix.os, 'macos') }} run: | - brew install ninja libomp netcdf-cxx eigen lapack git + brew install ninja libomp eigen git - name: Install required packages for Ubuntu if: ${{ contains(matrix.os, 'ubuntu') }} run: | - # install VMEC++ deps as well as VMEC2000 deps (we need to import VMEC2000 in a test) - sudo apt-get update && sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev + # install VMEC++ deps as well as VMEC2000 deps (we need to import VMEC2000 in a test). + # VMEC2000 itself gets libopenblas-dev/libscalapack-mpi-dev below, which cover its + # LAPACK needs; VMEC++ no longer links LAPACK at all. + sudo apt-get update && sudo apt-get install -y build-essential cmake zlib1g-dev libomp-dev - name: Cache the VMEC2000 wheel if: ${{ matrix.os == 'ubuntu-22.04' && matrix.python-version == '3.10' }} uses: actions/cache@v4 @@ -111,7 +118,7 @@ jobs: cache: 'pip' - name: Install required packages for Ubuntu run: | - sudo apt-get update && sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev + sudo apt-get update && sudo apt-get install -y build-essential cmake zlib1g-dev libomp-dev - name: Install package run: python -m pip install .[test] - name: Test example scripts diff --git a/CMakeLists.txt b/CMakeLists.txt index 9dbf919b2..2e1543a4d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -12,6 +12,20 @@ set(CMAKE_CXX_FLAGS "-fPIC -Wall -Wextra") set(CMAKE_CXX_FLAGS_DEBUG "-O0 -g") set(CMAKE_CXX_FLAGS_RELEASE "-O3 -DNDEBUG -fno-math-errno") +# LTO: cross-TU inlining and dead-code elimination, ~1/3 smaller binaries. +include(CheckIPOSupported) +check_ipo_supported(RESULT VMECPP_IPO_SUPPORTED OUTPUT VMECPP_IPO_ERROR) +if(VMECPP_IPO_SUPPORTED) + set(CMAKE_INTERPROCEDURAL_OPTIMIZATION_RELEASE TRUE) +else() + message(STATUS "LTO unavailable, building without it: ${VMECPP_IPO_ERROR}") +endif() + +# -march raises Eigen's EIGEN_MAX_ALIGN_BYTES from 16 to 32, which changes +# Eigen::aligned_allocator and is ABI-affecting. Pin it so cores built with +# different -march settings stay interchangeable. +add_compile_definitions(EIGEN_MAX_ALIGN_BYTES=32 EIGEN_MAX_STATIC_ALIGN_BYTES=32) + # use ccache if available find_program(CCACHE_COMMAND NAMES ccache ccache-swig) if(EXISTS ${CCACHE_COMMAND}) @@ -24,23 +38,105 @@ endif() # First check if required libraries are installed locally find_package(OpenMP REQUIRED) -find_package(HDF5 REQUIRED COMPONENTS C CXX) -include_directories(${HDF5_INCLUDE_DIRS} ${HDF5_CXX_INCLUDE_DIRS}) - -find_package(netCDF) -if(NOT netCDF_FOUND) - # Finds the netCDF installation using CMake's PkgConfig - set(netCDF_PC_FILE netcdf) - find_package(PkgConfig REQUIRED) - set(PKG_CONFIG_USE_CMAKE_PREFIX_PATH TRUE) - pkg_check_modules(netCDF REQUIRED ${netCDF_PC_FILE}>=4.3.0 IMPORTED_TARGET) - pkg_get_variable(netCDF_PREFIX ${netCDF_PC_FILE} prefix) - message(STATUS "netCDF prefix: ${netCDF_PREFIX}") - message(STATUS "netCDF include dirs: ${netCDF_INCLUDE_DIRS}") - message(STATUS "netCDF libraries: ${netCDF_LIBRARIES}") +# HDF5 and netCDF-C are built from source, as static libraries, rather than +# relying on whatever happens to be installed on the system. This is to +# ensure reproducibility and minimize binary size by turning off netcdf +# features we don't use. Versions must stay in sync with the Bazel build's +# src/vmecpp/cpp/third_party/non_module_deps.bzl. +include(ExternalProject) +set(_native_deps_prefix "${CMAKE_BINARY_DIR}/native-deps") +set(_native_deps_lib "${_native_deps_prefix}/lib") +set(_native_deps_include "${_native_deps_prefix}/include") +# INTERFACE_INCLUDE_DIRECTORIES below must exist already at generate time. +file(MAKE_DIRECTORY "${_native_deps_include}") + +# HDF5 is built with zlib support (required by netCDF-4), but as a static +# lib it doesn't bundle zlib's own object code, so it must be linked in +# explicitly wherever hdf5_c is used. +find_package(ZLIB REQUIRED) + +set(_native_deps_cmake_args + -DCMAKE_INSTALL_PREFIX=${_native_deps_prefix} + -DCMAKE_INSTALL_LIBDIR=lib + -DCMAKE_PREFIX_PATH=${_native_deps_prefix} + -DCMAKE_BUILD_TYPE=Release + -DCMAKE_MAKE_PROGRAM=${CMAKE_MAKE_PROGRAM} + -DBUILD_SHARED_LIBS=OFF + -DCMAKE_POSITION_INDEPENDENT_CODE=ON + -DBUILD_TESTING=OFF) + +if(APPLE) + # CC=brew's gcc-14 (set by some CI jobs) chokes on Apple's SDK math.h for + # plain C. -D always wins over the CC/CXX env vars, unlike omitting it. + list(APPEND _native_deps_cmake_args + -DCMAKE_C_COMPILER=/usr/bin/cc + -DCMAKE_CXX_COMPILER=/usr/bin/c++) endif() -include_directories(${netCDF_INCLUDE_DIRS}) +ExternalProject_Add(hdf5_native + URL "https://github.com/HDFGroup/hdf5/archive/refs/tags/hdf5-1_14_3.tar.gz" + URL_HASH SHA256=df5ee33c74d5efb59738075ef96f4201588e1f1eeb233f047ac7fd1072dee1f6 + CMAKE_ARGS + ${_native_deps_cmake_args} + -DHDF5_BUILD_CPP_LIB=ON + -DHDF5_ENABLE_Z_LIB_SUPPORT=ON + -DHDF5_ENABLE_SZIP_SUPPORT=OFF + -DHDF5_ENABLE_SZIP_ENCODING=OFF + -DHDF5_BUILD_EXAMPLES=OFF + -DHDF5_BUILD_TOOLS=OFF + BUILD_BYPRODUCTS + "${_native_deps_lib}/libhdf5.a" + "${_native_deps_lib}/libhdf5_cpp.a" + "${_native_deps_lib}/libhdf5_hl.a") + +ExternalProject_Add(netcdf_native + DEPENDS hdf5_native + URL "https://github.com/Unidata/netcdf-c/archive/refs/tags/v4.9.3.tar.gz" + URL_HASH SHA256=990f46d49525d6ab5dc4249f8684c6deeaf54de6fec63a187e9fb382cc0ffdff + CMAKE_ARGS + ${_native_deps_cmake_args} + -DNETCDF_ENABLE_DAP=OFF + -DNETCDF_ENABLE_DAP2=OFF + -DNETCDF_ENABLE_DAP4=OFF + -DNETCDF_ENABLE_NCZARR=OFF + -DNETCDF_ENABLE_NCZARR_ZIP=OFF + -DNETCDF_ENABLE_TESTS=OFF + -DNETCDF_BUILD_UTILITIES=OFF + BUILD_BYPRODUCTS + "${_native_deps_lib}/libnetcdf.a") + +# Imported targets pointing at the static libraries the ExternalProjects +# above install at build time +add_library(hdf5_c STATIC IMPORTED GLOBAL) +set_target_properties(hdf5_c PROPERTIES + IMPORTED_LOCATION "${_native_deps_lib}/libhdf5.a" + INTERFACE_INCLUDE_DIRECTORIES "${_native_deps_include}" + INTERFACE_LINK_LIBRARIES ZLIB::ZLIB) +add_dependencies(hdf5_c hdf5_native) + +add_library(hdf5_hl STATIC IMPORTED GLOBAL) +set_target_properties(hdf5_hl PROPERTIES + IMPORTED_LOCATION "${_native_deps_lib}/libhdf5_hl.a" + INTERFACE_LINK_LIBRARIES hdf5_c) +add_dependencies(hdf5_hl hdf5_native) + +add_library(hdf5_cpp STATIC IMPORTED GLOBAL) +set_target_properties(hdf5_cpp PROPERTIES + IMPORTED_LOCATION "${_native_deps_lib}/libhdf5_cpp.a" + INTERFACE_LINK_LIBRARIES hdf5_c) +add_dependencies(hdf5_cpp hdf5_native) + +# netCDF's static lib references HDF5 HL symbols internally, so link it in +# too even though vmecpp's own code never calls it directly. +add_library(netcdf STATIC IMPORTED GLOBAL) +set_target_properties(netcdf PROPERTIES + IMPORTED_LOCATION "${_native_deps_lib}/libnetcdf.a" + INTERFACE_INCLUDE_DIRECTORIES "${_native_deps_include}" + INTERFACE_LINK_LIBRARIES "hdf5_hl;hdf5_c") +add_dependencies(netcdf netcdf_native) + +add_library(vmecpp_hdf5_netcdf INTERFACE) +target_link_libraries(vmecpp_hdf5_netcdf INTERFACE netcdf hdf5_cpp hdf5_hl hdf5_c) # Fetch all the remote dependencies include(FetchContent) @@ -56,13 +152,15 @@ FetchContent_Declare( ) FetchContent_MakeAvailable(eigen) include_directories(${eigen_SOURCE_DIR}) +# The only dense linear solve in this project (LaplaceSolver::DecomposeMatrix/ +# SolveForPotential) uses Eigen::PartialPivLU, which is pure C++ with no +# external LAPACK/BLAS dependency as long as EIGEN_USE_LAPACKE/EIGEN_USE_BLAS/ +# EIGEN_USE_MKL are never defined. Do not define those macros without +# reintroducing an explicit LAPACK/BLAS dependency here. FetchContent_Declare(nlohmann_json URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz) FetchContent_MakeAvailable(nlohmann_json) -find_package(LAPACK REQUIRED) - - FetchContent_Declare( abseil-cpp GIT_REPOSITORY https://github.com/abseil/abseil-cpp.git @@ -109,12 +207,47 @@ include_directories(${PROJECT_SOURCE_DIR}/src/vmecpp/cpp) set(vmecpp_sources ${abscab_sources}) add_subdirectory(src) -# Define a static library for the actual computation core of VMEC++. -add_library(vmecpp_core STATIC ${vmecpp_sources}) -target_link_libraries(vmecpp_core PRIVATE ${HDF5_CXX_LIBRARIES} ${HDF5_LIBRARIES}) -target_link_libraries(vmecpp_core PRIVATE ${netCDF_LIBRARIES}) -target_link_libraries(vmecpp_core PRIVATE nlohmann_json::nlohmann_json) -target_link_libraries(vmecpp_core PRIVATE LAPACK::LAPACK) +# The computation core, built once per target ISA. Variants share one SONAME and +# are told apart by directory, so the loader can substitute them (see the +# glibc-hwcaps install rules below). +function(vmecpp_add_core target) + cmake_parse_arguments(CORE "" "HWCAPS_SUBDIR" "COMPILE_OPTIONS" ${ARGN}) + + add_library(${target} SHARED ${vmecpp_sources}) + set_target_properties(${target} PROPERTIES OUTPUT_NAME vmecpp_core) + if(CORE_HWCAPS_SUBDIR) + set_target_properties(${target} PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/glibc-hwcaps/${CORE_HWCAPS_SUBDIR}") + endif() + target_compile_options(${target} PRIVATE ${CORE_COMPILE_OPTIONS}) + + target_link_libraries(${target} PUBLIC vmecpp_hdf5_netcdf) + target_link_libraries(${target} PUBLIC nlohmann_json::nlohmann_json) + target_link_libraries(${target} PUBLIC absl::algorithm absl::base + absl::synchronization absl::strings absl::str_format absl::log + absl::string_view absl::check absl::status absl::statusor) + + if(VMECPP_USE_FFTX) + # Built per variant: the codelets must match the ISA of the core linking them. + add_library(${target}_fftx STATIC ${_fftx_iprdft_srcs} ${_fftx_prdft_srcs}) + target_include_directories(${target}_fftx PUBLIC + "${_fftx_dir}" # for + "${_fftx_dir}/include" # for fftx_minimal.hpp + "${_fftx_dir}/lib_fftx_iprdftbat_cpu_srcs" + "${_fftx_dir}/lib_fftx_prdftbat_cpu_srcs") + target_compile_options(${target}_fftx PRIVATE -fPIC ${CORE_COMPILE_OPTIONS}) + target_link_libraries(${target} PRIVATE ${target}_fftx) + target_compile_definitions(${target} PUBLIC VMECPP_USE_FFTX) + endif() + + if(OpenMP_CXX_FOUND) + target_link_libraries(${target} PUBLIC OpenMP::OpenMP_CXX) + endif() + + # We multithread at the outer loop level, not for the individual Eigen operations + # https://libeigen.gitlab.io/eigen/docs-3.3/TopicMultiThreading.html + target_compile_definitions(${target} PRIVATE EIGEN_DONT_PARALLELIZE) +endfunction() # FFTX (SPIRAL-generated batched IPRDFT/PRDFT) for the toroidal FFT hot path. # The codelets are vendored under src/vmecpp/cpp/third_party/fftx_codelets/; @@ -127,35 +260,41 @@ if(VMECPP_USE_FFTX) "${_fftx_dir}/lib_fftx_iprdftbat_cpu_srcs/*.cpp") file(GLOB _fftx_prdft_srcs CONFIGURE_DEPENDS "${_fftx_dir}/lib_fftx_prdftbat_cpu_srcs/*.cpp") - add_library(fftx_codelets STATIC ${_fftx_iprdft_srcs} ${_fftx_prdft_srcs}) - target_include_directories(fftx_codelets PUBLIC - "${_fftx_dir}" # for - "${_fftx_dir}/include" # for fftx_minimal.hpp - "${_fftx_dir}/lib_fftx_iprdftbat_cpu_srcs" - "${_fftx_dir}/lib_fftx_prdftbat_cpu_srcs") - # Match vmecpp's release flags; codelets are >30k lines and benefit from -O3. - target_compile_options(fftx_codelets PRIVATE -fPIC) - target_link_libraries(vmecpp_core PRIVATE fftx_codelets) - target_compile_definitions(vmecpp_core PUBLIC VMECPP_USE_FFTX) list(LENGTH _fftx_iprdft_srcs _n_iprdft) list(LENGTH _fftx_prdft_srcs _n_prdft) message(STATUS "FFTX vendored codelets: ${_n_iprdft} iprdft + ${_n_prdft} prdft sources") endif() -target_link_libraries(vmecpp_core PRIVATE absl::algorithm absl::base absl::synchronization absl::strings absl::str_format absl::log absl::string_view absl::check absl::status absl::statusor) -if(OpenMP_CXX_FOUND) - target_link_libraries(vmecpp_core PRIVATE OpenMP::OpenMP_CXX) -endif() +vmecpp_add_core(vmecpp_core) -# We multithread at the outer loop level, not for the individual Eigen operations -# https://libeigen.gitlab.io/eigen/docs-3.3/TopicMultiThreading.html -target_compile_definitions(vmecpp_core PRIVATE EIGEN_DONT_PARALLELIZE) +set(VMECPP_HWCAPS_DISPATCH_DEFAULT OFF) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND + CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64)$") + set(VMECPP_HWCAPS_DISPATCH_DEFAULT ON) +endif() +option(VMECPP_HWCAPS_DISPATCH + "Also build x86-64-v3 core variant selected at load time via glibc-hwcaps" + ${VMECPP_HWCAPS_DISPATCH_DEFAULT}) +if(VMECPP_HWCAPS_DISPATCH) + if(NOT VMECPP_HWCAPS_DISPATCH_DEFAULT) + message(FATAL_ERROR + "VMECPP_HWCAPS_DISPATCH requires Linux on an x86-64 processor") + endif() + vmecpp_add_core(vmecpp_core_v3 HWCAPS_SUBDIR x86-64-v3 + COMPILE_OPTIONS -march=x86-64-v3) +endif() # Now also add the vmec_standalone executable. add_executable(vmec_standalone ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/vmec_standalone/vmec_standalone.cc) target_link_libraries(vmec_standalone vmecpp_core) -target_link_libraries(vmecpp_core PRIVATE absl::strings) +if(APPLE) + set(_vmecpp_loader_path "@loader_path") +else() + set(_vmecpp_loader_path "$ORIGIN") +endif() +set_target_properties(vmec_standalone PROPERTIES + BUILD_RPATH "${_vmecpp_loader_path}") # Now add the pybind11 module for VMEC++. FetchContent_MakeAvailable(pybind11) @@ -172,7 +311,16 @@ if(OpenMP_CXX_FOUND) target_link_libraries(_vmecpp PRIVATE OpenMP::OpenMP_CXX) endif() +set_target_properties(_vmecpp PROPERTIES + INSTALL_RPATH "${_vmecpp_loader_path}") install(TARGETS _vmecpp LIBRARY DESTINATION vmecpp/cpp/.) +# glibc >= 2.33 prefers glibc-hwcaps/ subdirectories; older loaders ignore +# them and pick up the baseline next to the extension module. +install(TARGETS vmecpp_core LIBRARY DESTINATION vmecpp/cpp/.) +if(VMECPP_HWCAPS_DISPATCH) + install(TARGETS vmecpp_core_v3 + LIBRARY DESTINATION vmecpp/cpp/glibc-hwcaps/x86-64-v3/.) +endif() install(TARGETS indata2json DESTINATION vmecpp/cpp/third_party/indata2json/) # Optional Enzyme automatic-differentiation target: exact autodiff (forward and @@ -185,26 +333,121 @@ install(TARGETS indata2json DESTINATION vmecpp/cpp/third_party/indata2json/) # -DVMECPP_ENABLE_ENZYME=ON -DVMECPP_ENZYME_PLUGIN=/path/to/ClangEnzyme-NN.so option(VMECPP_ENABLE_ENZYME "Build Enzyme autodiff targets" OFF) if(VMECPP_ENABLE_ENZYME) - if(NOT CMAKE_CXX_COMPILER_ID MATCHES "Clang") + if(NOT CMAKE_CXX_COMPILER_ID MATCHES "Clang|IntelLLVM") message(FATAL_ERROR - "VMECPP_ENABLE_ENZYME requires a Clang compiler (got " - "${CMAKE_CXX_COMPILER_ID}); Enzyme attaches as a Clang plugin.") + "VMECPP_ENABLE_ENZYME requires a Clang-based compiler (got " + "${CMAKE_CXX_COMPILER_ID}); Enzyme attaches as a Clang plugin. " + "Intel oneAPI icx/icpx (IntelLLVM) is Clang-based and supported.") endif() set(VMECPP_ENZYME_PLUGIN "" CACHE FILEPATH "Path to ClangEnzyme-NN.so") - if(NOT VMECPP_ENZYME_PLUGIN OR NOT EXISTS "${VMECPP_ENZYME_PLUGIN}") - message(FATAL_ERROR - "VMECPP_ENABLE_ENZYME=ON requires " - "-DVMECPP_ENZYME_PLUGIN=/path/to/ClangEnzyme-NN.so") + set(VMECPP_ENZYME_CORE_TARGETS vmecpp_core) + if(VMECPP_HWCAPS_DISPATCH) + list(APPEND VMECPP_ENZYME_CORE_TARGETS vmecpp_core_v3) + endif() + # The Enzyme pass is an LLVM-version-specific plugin. icx/icpx (IntelLLVM) + # cannot load a plugin built against a different upstream LLVM. Enzyme's output + # is an ordinary object file, though, so for such compilers the two Enzyme TUs + # can be precompiled with a matching ClangEnzyme and linked in, while the rest + # is built with icx. + set(VMECPP_ENZYME_JVP_OBJECT "" CACHE FILEPATH + "Prebuilt exact_force_jvp object (compiled with a matching ClangEnzyme)") + set(VMECPP_ENZYME_VJP_OBJECT "" CACHE FILEPATH + "Prebuilt exact_force_vjp object (compiled with a matching ClangEnzyme)") + # Backward-compatible name for the original JVP-only input. Keep its meaning + # explicit: the cotangent implementation must still be supplied separately. + set(VMECPP_ENZYME_OBJECT "" CACHE FILEPATH + "Deprecated JVP-only alias for VMECPP_ENZYME_JVP_OBJECT") + set(VMECPP_ENZYME_EFFECTIVE_JVP_OBJECT "${VMECPP_ENZYME_JVP_OBJECT}") + if(VMECPP_ENZYME_OBJECT) + if(VMECPP_ENZYME_JVP_OBJECT AND + NOT "${VMECPP_ENZYME_JVP_OBJECT}" STREQUAL "${VMECPP_ENZYME_OBJECT}") + message(FATAL_ERROR + "VMECPP_ENZYME_OBJECT and VMECPP_ENZYME_JVP_OBJECT name different " + "files; remove the deprecated VMECPP_ENZYME_OBJECT setting") + endif() + message(DEPRECATION + "VMECPP_ENZYME_OBJECT supplies only the JVP and is deprecated; use " + "VMECPP_ENZYME_JVP_OBJECT together with VMECPP_ENZYME_VJP_OBJECT") + set(VMECPP_ENZYME_EFFECTIVE_JVP_OBJECT "${VMECPP_ENZYME_OBJECT}") + endif() + if(VMECPP_ENZYME_EFFECTIVE_JVP_OBJECT OR VMECPP_ENZYME_VJP_OBJECT) + if(NOT VMECPP_ENZYME_EFFECTIVE_JVP_OBJECT) + message(FATAL_ERROR + "Prebuilt Enzyme mode requires VMECPP_ENZYME_JVP_OBJECT") + endif() + if(NOT VMECPP_ENZYME_VJP_OBJECT) + message(FATAL_ERROR + "Prebuilt Enzyme mode requires VMECPP_ENZYME_VJP_OBJECT; the " + "deprecated VMECPP_ENZYME_OBJECT supplies only the JVP") + endif() + foreach(VMECPP_ENZYME_PREBUILT_OBJECT + IN ITEMS "${VMECPP_ENZYME_EFFECTIVE_JVP_OBJECT}" + "${VMECPP_ENZYME_VJP_OBJECT}") + if(NOT EXISTS "${VMECPP_ENZYME_PREBUILT_OBJECT}") + message(FATAL_ERROR + "Prebuilt Enzyme object not found: ${VMECPP_ENZYME_PREBUILT_OBJECT}") + endif() + endforeach() + set(VMECPP_ENZYME_PREBUILT_OBJECTS + "${VMECPP_ENZYME_EFFECTIVE_JVP_OBJECT}" + "${VMECPP_ENZYME_VJP_OBJECT}") + message(STATUS + "Enzyme: linking prebuilt JVP object ${VMECPP_ENZYME_EFFECTIVE_JVP_OBJECT}") + message(STATUS + "Enzyme: linking prebuilt VJP object ${VMECPP_ENZYME_VJP_OBJECT}") + set_source_files_properties(${VMECPP_ENZYME_PREBUILT_OBJECTS} PROPERTIES + EXTERNAL_OBJECT TRUE GENERATED TRUE) + foreach(VMECPP_ENZYME_CORE_TARGET IN LISTS VMECPP_ENZYME_CORE_TARGETS) + target_sources(${VMECPP_ENZYME_CORE_TARGET} PRIVATE + ${VMECPP_ENZYME_PREBUILT_OBJECTS}) + target_compile_definitions(${VMECPP_ENZYME_CORE_TARGET} + PUBLIC VMECPP_ENABLE_ENZYME) + endforeach() + else() + if(NOT VMECPP_ENZYME_PLUGIN OR NOT EXISTS "${VMECPP_ENZYME_PLUGIN}") + message(FATAL_ERROR + "VMECPP_ENABLE_ENZYME=ON requires " + "-DVMECPP_ENZYME_PLUGIN=/path/to/ClangEnzyme-NN.so " + "(or both VMECPP_ENZYME_JVP_OBJECT and VMECPP_ENZYME_VJP_OBJECT " + "for prebuilt objects)") + endif() + message(STATUS "Enzyme plugin: ${VMECPP_ENZYME_PLUGIN}") + enable_testing() + # Exact autodiff (forward and reverse) of a real VMEC nonlinear kernel: the + # half-grid Jacobian. Enzyme runs as an optimization-time pass, so it needs + # -O2 and the plugin attached. + add_executable(jacobian_kernel_autodiff_test + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/common/enzyme/jacobian_kernel_autodiff_test.cc) + target_compile_options(jacobian_kernel_autodiff_test PRIVATE + -O2 -fplugin=${VMECPP_ENZYME_PLUGIN}) + add_test(NAME jacobian_kernel_autodiff COMMAND jacobian_kernel_autodiff_test) + # Exact Hessian of the composed local force map (all six force-chain kernels): + # forward/reverse Jacobian vs finite differences, and JVP cost vs FD-HVP. + add_executable(local_force_hessian_test + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc) + target_compile_options(local_force_hessian_test PRIVATE + -O2 -fplugin=${VMECPP_ENZYME_PLUGIN}) + add_test(NAME local_force_hessian COMMAND local_force_hessian_test) + + # Wire the exact force Hessian-vector product into the core library and the + # Python extension. Every runtime-selectable core variant must export these + # functions. The Enzyme translation units are compiled with the plugin; the + # rest of each core stays normally compiled. + foreach(VMECPP_ENZYME_CORE_TARGET IN LISTS VMECPP_ENZYME_CORE_TARGETS) + target_sources(${VMECPP_ENZYME_CORE_TARGET} PRIVATE + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.cc) + target_compile_definitions(${VMECPP_ENZYME_CORE_TARGET} + PUBLIC VMECPP_ENABLE_ENZYME) + endforeach() + # Enzyme does not support the AVX mask intrinsics emitted for x86-64-v3. + # Compile these two translation units for the baseline ISA in every core; + # the rest of vmecpp_core_v3 remains optimized for x86-64-v3. + set_source_files_properties( + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc + ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.cc + PROPERTIES COMPILE_OPTIONS + "-march=x86-64;-O2;-fplugin=${VMECPP_ENZYME_PLUGIN}") endif() - message(STATUS "Enzyme plugin: ${VMECPP_ENZYME_PLUGIN}") - enable_testing() - # Self-contained over flat buffers; checks Jv and J^T u against finite - # differences and against each other (adjoint identity). Enzyme runs as an - # optimization-time pass, so it needs -O2 and the plugin attached. -O2 also - # guards against a Debug (-O0) configuration where the AD pass would not run. - add_executable(jacobian_kernel_autodiff_test - ${PROJECT_SOURCE_DIR}/src/vmecpp/cpp/vmecpp/common/enzyme/jacobian_kernel_autodiff_test.cc) - target_compile_options(jacobian_kernel_autodiff_test PRIVATE - -O2 -fplugin=${VMECPP_ENZYME_PLUGIN}) - add_test(NAME jacobian_kernel_autodiff COMMAND jacobian_kernel_autodiff_test) endif() diff --git a/README.md b/README.md index a16840c8c..d86d214a9 100644 --- a/README.md +++ b/README.md @@ -152,7 +152,7 @@ Ubuntu 22.04 and 24.04, as well as Debian 12 are officially supported. 1. Install required system packages: ```shell -sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev libhdf5-dev python3-dev +sudo apt-get install -y build-essential cmake gfortran libnetcdf-dev libomp-dev libhdf5-dev python3-dev ``` 2. Install VMEC++ as a Python package (possibly after creating a dedicated virtual environment): @@ -172,7 +172,7 @@ Otherwise the Ubuntu package `python-is-python3` provides the `python` alias. 1. Install required system packages: ```shell -pacman -Sy --noconfirm python-pip gcc gcc-fortran openmp hdf5 netcdf lapack +pacman -Sy --noconfirm python-pip gcc gcc-fortran openmp hdf5 netcdf ``` 2. Install VMEC++ as a Python package (possibly after creating a virtual environment): @@ -188,7 +188,7 @@ python -m pip install git+https://github.com/proximafusion/vmecpp 1. Install required system packages: ```shell -dnf install -y python3.10-devel cmake g++ gfortran libomp-devel hdf5-devel netcdf-devel lapack-devel +dnf install -y python3.10-devel cmake g++ gfortran libomp-devel hdf5-devel netcdf-devel ``` 2. Install VMEC++ as a Python package (possibly after creating a virtual environment): @@ -232,7 +232,7 @@ python -m pip install -e .[test] ``` The shell provides Python 3.13 together with the native build dependencies needed -to build and test VMEC++, including CMake, GCC, GFortran, HDF5, NetCDF, LAPACK, +to build and test VMEC++, including CMake, GCC, GFortran, HDF5, NetCDF, OpenMPI, and Git LFS. ### As part of a conda environment @@ -279,24 +279,24 @@ VMEC++ on a configuration that is very similar to the converged equilibrium. ```python import vmecpp -input = vmecpp.VmecInput.from_file("w7x.json") +vmec_input = vmecpp.VmecInput.from_file("w7x.json") # Base run -output = vmecpp.run(input) +vmec_output = vmecpp.run(vmec_input) # Now let's perturb the plasma boundary a little bit... -input.rbc[0, 0] *= 0.8 -input.rbc[1, 0] *= 1.2 +vmec_input.rbc[0, 0] *= 0.8 +vmec_input.rbc[1, 0] *= 1.2 # ...and fix up the multigrid steps: hot-restarted runs only allow a single step -input.ns_array = input.ns_array[-1:] -input.ftol_array = input.ftol_array[-1:] -input.niter_array = input.niter_array[-1:] +vmec_input.ns_array = vmec_input.ns_array[-1:] +vmec_input.ftol_array = vmec_input.ftol_array[-1:] +vmec_input.niter_array = vmec_input.niter_array[-1:] # We can now run with hot restart: -# passing the previously obtained output ensures that +# passing the previously obtained vmec_output ensures that # the run starts already close to the equilibrium, so it will take # very few iterations to converge this time! -hot_restarted_output = vmecpp.run(input, restart_from=output) +hot_restarted_output = vmecpp.run(vmec_input, restart_from=vmec_output) ``` ## Full tests and validation against the reference Fortran VMEC v8.52 @@ -326,16 +326,16 @@ VMEC++: - reports issues via standard Python exceptions and has a zero crash policy - allows hot-restarting a run from a previous converged state (see [Hot restart](#hot-restart)) - supports inputs in the classic INDATA format as well as simpler-to-parse JSON files; it is also simple to construct input objects programmatically in Python -- employs the same parallelization strategy as Fortran VMEC, but VMEC++ leverages OpenMP for a multi-thread implementation rather than Fortran VMEC's MPI parallelization: as a consequence it cannot parallelize over multiple nodes +- employs the same parallelization strategy as Fortran VMEC, but VMEC++ leverages OpenMP for a multi-thread implementation rather than Fortran VMEC's MPI parallelization: as a consequence it cannot parallelize over multiple nodes, but can utilize shared caches between threads +- Reduces the force spikes between multi-grid stages in free-boundary, which should lead to faster and more robust free-boundary convergence (for details see https://github.com/proximafusion/vmecpp/releases/tag/v0.7.0) +- Stable recurrence for Neumann kernel integrals enables convergence of free-boundary solves at high mpol, ntor (see https://github.com/proximafusion/vmecpp/releases/tag/v0.5.3) - Uses FFT kernels optimized for small mode numbers [generated using FFTX](https://github.com/spiral-software/fftx) instead of DFT for supported resolutions. They give a 10-20% speedup relative to the DFT counterparts. - implements the iteration algorithm of Fortran VMEC 8.52, which sometimes has different convergence behavior from (PAR)VMEC 9.0: some configurations might converge with VMEC++ and not with (PAR)VMEC 9.0, and vice versa. One deliberate exception: at multigrid grid transitions, the rollback backup of the state vector is taken *after* the radial interpolation of the coarse-grid solution (matching PARVMEC/VMEC2000 since 2017-01-24, "SPH 012417"), not before it as in VMEC 8.52 -- with the 8.52 ordering, the first restart of a stage silently discards the interpolated state and the finer stages effectively re-solve from a cold start ### Limitations with respect to the Fortran implementations -- non-stellarator-symmetric terms (`lasym == true`) are not supported yet - free-boundary works only for `ntor > 0` - axisymmetric (`ntor = 0`) free-boundary runs don't work yet - `lgiveup`/`fgiveup` logic for early termination of a multi-grid sequence is not implemented yet - `lbsubs` logic in computing outputs is not implemented yet -- `lforbal` logic for non-variational forces near the magnetic axis is not implemented yet - `lrfp` flag is available for wout compatibility, but RFP-specific physics is not implemented yet - only stellarators/Tokamaks for now - several profile parameterizations are not fully implemented yet: * `gauss_trunc` @@ -369,12 +369,12 @@ Some of the things we are planning for VMEC++'s future: - [x] free-boundary hot-restart in Python - [X] open-sourcing the full VMEC++ test suite (including the Verification&Validation part that compares `wout` contents) - [x] open-sourcing the source code to reproduce VMEC++'s performance benchmarks -- [ ] VMEC++ usable as a C++ bazel module +- [x] VMEC++ usable as a C++ bazel module Some items we do not plan to work on, but where community ownership is welcome: - [ ] packaging VMEC++ for platforms or package managers other than pip (e.g. conda, homebrew, ...) - [ ] native Windows support -- [ ] ARM support +- [x] ARM support - [ ] 2D preconditioner using [`bcyclic_plus_plus`](https://code.ornl.gov/m4c/bcyclic_plus_plus) ## Related repositories diff --git a/benchmarks/test_benchmarks.py b/benchmarks/test_benchmarks.py index 445175bed..2101c89ef 100644 --- a/benchmarks/test_benchmarks.py +++ b/benchmarks/test_benchmarks.py @@ -10,7 +10,10 @@ import subprocess import sys +from collections.abc import Callable +from importlib import import_module from pathlib import Path +from typing import Any, cast import pytest @@ -19,6 +22,7 @@ REPO_ROOT = Path(__file__).parent.parent TEST_DATA_DIR = REPO_ROOT / "src" / "vmecpp" / "cpp" / "vmecpp" / "test_data" EXAMPLES_DATA_DIR = REPO_ROOT / "examples" / "data" +sys.path.insert(0, str(EXAMPLES_DATA_DIR.parent)) # --------------------------------------------------------------------------- @@ -84,20 +88,6 @@ def test_bench_cli_startup(benchmark): assert result.returncode == 0 -def test_bench_cli_invalid_input(benchmark): - """Benchmark CLI error path via `vmecpp invalid_input`.""" - - def run_invalid(): - return subprocess.run( - [sys.executable, "-m", "vmecpp", "invalid_input"], - capture_output=True, - check=False, - ) - - result = benchmark(run_invalid) - assert result.returncode != 0 - - # --------------------------------------------------------------------------- # Fixed-boundary solver benchmarks # --------------------------------------------------------------------------- @@ -165,3 +155,37 @@ def test_bench_free_boundary(benchmark, free_boundary_input, response_table): warmup_rounds=1, ) assert result.wout.volume == pytest.approx(0.3075, rel=1e-3) + + +# --------------------------------------------------------------------------- +# SIMSOPT adjoint gradient benchmarks +# --------------------------------------------------------------------------- + + +def _gradient_cost(*, analytic): + module = import_module("simsopt_vmec_gradient") + gradient_cost = cast(Callable[..., Any], module.gradient_cost) + + return gradient_cost(analytic=analytic) + + +def test_bench_simsopt_adjoint_gradient(benchmark): + """Track the analytic adjoint gradient cost without timing assertions.""" + result = benchmark.pedantic( + _gradient_cost, + kwargs={"analytic": True}, + rounds=1, + warmup_rounds=0, + ) + assert result.force_evals > 0 + + +def test_bench_simsopt_finite_difference_gradient(benchmark): + """Track the finite-difference reference cost for comparison.""" + result = benchmark.pedantic( + _gradient_cost, + kwargs={"analytic": False}, + rounds=1, + warmup_rounds=0, + ) + assert result.force_evals > 0 diff --git a/docker/Dockerfile b/docker/Dockerfile index 15c332a8b..fb99bb4af 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -7,7 +7,6 @@ RUN apt-get -q update && \ build-essential \ cmake \ libnetcdf-dev \ - liblapack-dev \ libomp-dev \ libhdf5-dev \ python3-pip \ diff --git a/environment.yml b/environment.yml index 9ba51a4af..d6d4b93c9 100644 --- a/environment.yml +++ b/environment.yml @@ -50,7 +50,6 @@ dependencies: - libgomp=14.2.0=h77fa898_1 - libiconv=1.17=hd590300_2 - libjpeg-turbo=3.0.0=hd590300_1 - - liblapack=3.9.0=6_ha36c22a_netlib - liblzma=5.6.3=hb9d3cd8_1 - libnetcdf=4.9.2=nompi_h00e09a9_116 - libnghttp2=1.64.0=h161d5f1_0 diff --git a/examples/simsopt_vmec_gradient.py b/examples/simsopt_vmec_gradient.py new file mode 100644 index 000000000..e88219141 --- /dev/null +++ b/examples/simsopt_vmec_gradient.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Use VMEC++ as a gradient-providing equilibrium component for SIMSOPT. + +A SIMSOPT optimization over a plasma boundary normally differentiates VMEC by +finite differences: one equilibrium re-solve per boundary degree of freedom and +per outer iteration. VMEC++ instead provides the boundary gradient analytically +through the implicit-function adjoint (``vmecpp_adjoint.boundary_gradient``): one +extra Hessian solve, independent of the number of boundary DOFs. + +``VmecEnergy`` wraps that as a SIMSOPT ``Optimizable`` whose objective is the MHD +energy of the converged equilibrium and whose ``dJ`` is the adjoint gradient. +``optimize_to_target`` runs a gradient-based optimization of the boundary toward +a target energy, with analytic gradient or with finite differences, and +reports the cost (forward-model evaluations counted inside VMEC++, outer +iterations, wall time) so the two can be compared. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from simsopt._core import Optimizable +from simsopt._core.derivative import Derivative, OptimizableDefaultDict, derivative_dec +from vmecpp_adjoint import ( + DEFAULT_INPUT, + boundary_gradient, + finite_difference_boundary_gradient, + make_model, + mhd_energy, + partition, + solve_interior, +) + + +class VmecBoundaryProblem: + """Equilibrium energy and its boundary gradient, cached per boundary state.""" + + def __init__(self, input_path: Path = DEFAULT_INPUT, ns: int = 11): + self.model = make_model(input_path, ns) + self.model.solve() + self.ns = ns + self.interior, self.boundary = partition(self.model, ns) + self._x_full = np.asarray(self.model.get_state(), float).copy() + self._cached_p = self._x_full[self.boundary].copy() + + @property + def x0(self): + return self._x_full[self.boundary].copy() + + def _resolve(self, p): + p = np.asarray(p, float) + if not np.array_equal(p, self._cached_p): + self._x_full = solve_interior( + self.model, self._x_full, self.interior, self.boundary, p + ) + self._cached_p = p.copy() + + def value(self, p): + self._resolve(p) + self.model.set_state(np.ascontiguousarray(self._x_full)) + self.model.evaluate(2, 2, False) + return self.model.mhd_energy + + def gradient(self, p): + self._resolve(p) + return boundary_gradient( + self.model, self._x_full, self.interior, self.boundary, mhd_energy + ) + + +def make_simsopt_optimizable(problem: VmecBoundaryProblem): + """Wrap the problem as a SIMSOPT Optimizable exposing an analytic gradient.""" + + class VmecEnergy(Optimizable): + def __init__(self): + x0 = problem.x0 + super().__init__(x0=x0, names=[f"boundary{i}" for i in range(x0.size)]) + + def J(self): + return problem.value(self.local_full_x) + + @derivative_dec + def dJ(self): + data = OptimizableDefaultDict({self: problem.gradient(self.local_full_x)}) + return Derivative(data) + + return VmecEnergy() + + +@dataclass +class GradResult: + method: str + force_evals: int + seconds: float + gradient: np.ndarray + + +def gradient_cost(input_path: Path = DEFAULT_INPUT, ns: int = 11, analytic=True): + """Cost of one full boundary gradient at the converged equilibrium. + + This is what an external optimizer pays per iteration. The analytic adjoint needs + one Hessian solve regardless of the number of boundary DOFs; finite differences re- + converge the equilibrium twice per boundary DOF. + """ + problem = VmecBoundaryProblem(input_path, ns) + x_star = problem._x_full.copy() + interior, boundary = problem.interior, problem.boundary + problem.model.reset_force_eval_count() + t0 = time.perf_counter() + if analytic: + g_dict = None + g = boundary_gradient(problem.model, x_star, interior, boundary, mhd_energy) + else: + g_dict = finite_difference_boundary_gradient( + problem.model, + x_star, + interior, + boundary, + mhd_energy, + range(boundary.size), + ) + g = np.array([g_dict[j] for j in range(boundary.size)]) + return GradResult( + "analytic adjoint" if analytic else "finite differences", + problem.model.force_eval_count, + time.perf_counter() - t0, + g, + ) + + +def main(): + analytic = gradient_cost(analytic=True) + fd = gradient_cost(analytic=False) + rel = np.linalg.norm(analytic.gradient - fd.gradient) / np.linalg.norm(fd.gradient) + n_boundary = analytic.gradient.size + print(f"boundary gradient cost ({n_boundary} boundary DOFs, solovev ns=11)\n") + print(f"{'method':20s} {'F-evals':>9s} {'time[s]':>8s}") + for r in (fd, analytic): + print(f"{r.method:20s} {r.force_evals:9d} {r.seconds:8.2f}") + print( + f"\nspeedup (force evals): {fd.force_evals / max(analytic.force_evals, 1):.1f}x" + f" gradient agreement: {rel:.1e}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/vmecpp_adjoint.py b/examples/vmecpp_adjoint.py new file mode 100644 index 000000000..281459f38 --- /dev/null +++ b/examples/vmecpp_adjoint.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Adjoint sensitivity of a converged VMEC++ equilibrium to its boundary. + +A fixed-boundary equilibrium satisfies the interior force balance F_I(x) = 0, +where x is the decomposed internal-basis state and F is the gradient of VMEC's +augmented functional. The outermost flux surface (the boundary) is the last +radial block of the state and is held fixed during the solve. For a scalar +objective J(x), the sensitivity to the boundary degrees of freedom follows from +the implicit function theorem: + + dJ/dx_B = dJ/dx_B|_x - (dF_I/dx_B)^T lambda, H_II lambda = dJ/dx_I, + +with H = dF/dx the (symmetric) Hessian of the augmented functional. Every +operator is matrix-free and already exposed by VmecModel: the Hessian-vector +product (``hessian_vector_product``) and the preconditioner +(``apply_preconditioner``), used to solve the adjoint system. Only one Hessian +solve is needed for the full boundary gradient, versus one equilibrium re-solve +per boundary degree of freedom for finite differences. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from scipy.optimize import root +from scipy.sparse.linalg import LinearOperator, gmres + +from vmecpp.cpp import _vmecpp # type: ignore + +DEFAULT_INPUT = ( + Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" +) + + +def make_model(input_path: Path = DEFAULT_INPUT, ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(input_path)), ns) + + +def partition(model, ns: int): + """Indices of the interior (free) and boundary (LCFS) state components.""" + k = model.mpol * (model.ntor + 1) + n = np.asarray(model.get_state()).size + per_span = ns * k + n_span = n // per_span + boundary = [] + for s in range(n_span): + boundary.extend(range(s * per_span + (ns - 1) * k, s * per_span + ns * k)) + boundary = np.array(sorted(boundary)) + interior = np.setdiff1d(np.arange(n), boundary) + return interior, boundary + + +def _raw_force(model, x): + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), float) + + +def _interior_operators(model, x, interior): + # The caller must set the base state to x and assemble the preconditioner + # (evaluate(2, 2, True)) before using these. hessian_vector_product uses the + # current state as its base point and restores it, so no per-matvec state + # update is needed: that keeps each Hessian matvec at two force evaluations. + n = x.size + ni = interior.size + + def hii(vi): + v = np.zeros(n) + v[interior] = vi + return np.asarray(model.hessian_vector_product(np.ascontiguousarray(v)), float)[ + interior + ] + + def mii(bi): + v = np.zeros(n) + v[interior] = bi + return np.asarray(model.apply_preconditioner(np.ascontiguousarray(v)), float)[ + interior + ] + + return ( + LinearOperator((ni, ni), matvec=hii), # type: ignore[call-overload] + LinearOperator((ni, ni), matvec=mii), # type: ignore[call-overload] + ) + + +class _VmecPreconditioner(LinearOperator): + """Adaptive VMEC preconditioner for SciPy's Newton-Krylov solver.""" + + def __init__(self, model, x_template, interior): + self._model = model + self._x_template = x_template + self._interior = interior + self._x = x_template.copy() + super().__init__(dtype=float, shape=(interior.size, interior.size)) + + def update(self, x, _f): + self._x = self._x_template.copy() + self._x[self._interior] = x + self._model.set_state(np.ascontiguousarray(self._x)) + self._model.evaluate(2, 2, True) + + def _matvec(self, x): + v = np.zeros_like(self._x) + v[self._interior] = x + self._model.set_state(np.ascontiguousarray(self._x)) + return np.asarray(self._model.apply_preconditioner(v), float)[self._interior] + + +def solve_interior(model, x0, interior, boundary, x_boundary, tol=1e-10, max_newton=80): + """Converge the interior to force balance with the boundary held fixed. + + Use SciPy's Newton-Krylov implementation with VMEC's adaptive preconditioner. + + SciPy owns the nonlinear iteration, inner Krylov solve, and line search. VMEC supplies + the state-dependent preconditioner through SciPy's public ``inner_M`` interface. + """ + x_template = np.asarray(x0, float).copy() + x_template[boundary] = x_boundary + + def residual(xi): + x = x_template.copy() + x[interior] = xi + return _raw_force(model, x)[interior] + + preconditioner = _VmecPreconditioner(model, x_template, interior) + preconditioner.update(x_template[interior], np.zeros(interior.size)) + solution = root( + residual, + x_template[interior], + method="krylov", + options={ + "fatol": tol, + "line_search": "armijo", + "maxiter": max_newton, + "jac_options": { + "method": "lgmres", + "inner_M": preconditioner, + "inner_rtol": 1e-4, + "inner_maxiter": 300, + }, + }, + ) + + x = x_template.copy() + x[interior] = solution.x + if not solution.success: + error_message = ( + f"Interior solve failed: {solution.message}; " + f"max|F|={np.max(np.abs(solution.fun)):.3e}" + ) + raise RuntimeError(error_message) + return x + + +def objective_state_gradient(model, x, objective, h=1e-6): + """Partial derivative dJ/dx at fixed state, by central finite differences.""" + n = x.size + g = np.zeros(n) + for i in range(n): + xp = x.copy() + xp[i] += h + model.set_state(np.ascontiguousarray(xp)) + model.evaluate(2, 2, False) + jp = objective(model) + xm = x.copy() + xm[i] -= h + model.set_state(np.ascontiguousarray(xm)) + model.evaluate(2, 2, False) + jm = objective(model) + g[i] = (jp - jm) / (2 * h) + return g + + +def boundary_gradient(model, x_star, interior, boundary, objective, h=1e-6): + """Adjoint gradient dJ/dx_B at the converged equilibrium x_star.""" + n = x_star.size + dj = objective_state_gradient(model, x_star, objective, h) + model.set_state(np.ascontiguousarray(x_star)) + model.evaluate(2, 2, True) # assemble preconditioner + set base state to x_star + h_op, m_op = _interior_operators(model, x_star, interior) + lam, _ = gmres(h_op, dj[interior], M=m_op, rtol=1e-6, restart=100, maxiter=30) + embedded = np.zeros(n) + embedded[interior] = lam + model.set_state(np.ascontiguousarray(x_star)) + model.evaluate(2, 2, False) + coupling = np.asarray( + model.hessian_vector_product(np.ascontiguousarray(embedded)), float + )[boundary] + return dj[boundary] - coupling + + +def finite_difference_boundary_gradient( + model, x_star, interior, boundary, objective, dofs, h=1e-5 +): + """Reference gradient: re-solve the interior for each perturbed boundary DOF.""" + g = {} + for j in dofs: + xbp = x_star[boundary].copy() + xbp[j] += h + xp = solve_interior(model, x_star, interior, boundary, xbp) + model.set_state(np.ascontiguousarray(xp)) + model.evaluate(2, 2, False) + jp = objective(model) + xbm = x_star[boundary].copy() + xbm[j] -= h + xm = solve_interior(model, x_star, interior, boundary, xbm) + model.set_state(np.ascontiguousarray(xm)) + model.evaluate(2, 2, False) + jm = objective(model) + g[j] = (jp - jm) / (2 * h) + return g + + +def mhd_energy(model): + return model.mhd_energy diff --git a/pyproject.toml b/pyproject.toml index 4c850d5cc..e9181a72f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "vmecpp" dynamic = ["version"] description = "Proxima Fusion's reimplementation of the Variational Moments Equilibrium Code (VMEC), a free-boundary ideal-MHD equilibrium solver for stellarators and Tokamaks." readme = "README.md" -requires-python = ">=3.10,<3.14" +requires-python = ">=3.10" license = "MIT" license-files = ["LICENSE.txt"] keywords = [] @@ -25,13 +25,13 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Programming Language :: Python :: Implementation :: CPython", ] dependencies = [ "netCDF4", "numpy", "pydantic>=2.11", # Introduced support for by_alias validation - "pydantic<2.13", # Breaking change currently under investigation "jaxtyping>=0.2.6", "jax!=0.5.2", # To prevent a circular import bug in jax on arch Linux "jax!=0.5.1", # To prevent a circular import bug in jax on arch Linux @@ -105,16 +105,16 @@ version_file = "src/vmecpp/__about__.py" [tool.cibuildwheel] archs = ["native"] skip = ["*-musllinux*"] -build-frontend = "build[uv]" # More modern package manager than pip, with faster builds test-skip = ["cp314-*", "cp314t-*"] # Skip tests on Python 3.14 until all our dependencies start supporting it. The wheels are still valid, and become installable the moment simsopt releases 3.14 support. +build-frontend = "build[uv]" # More modern package manager than pip, with faster builds test-requires = "pytest" test-command = "pytest {package}/tests/test_simsopt_compat.py {package}/tests/test_free_boundary.py" [tool.cibuildwheel.linux] -before-build = "yum install -y lapack-devel netcdf-devel hdf5-devel" +before-build = "yum install -y zlib-devel" [tool.cibuildwheel.macos] -before-build = "brew install gcc lapack netcdf hdf5 libomp" +before-build = "brew install gcc libomp" [tool.coverage.run] source_pkgs = ["vmecpp", "tests"] diff --git a/src/vmecpp/__init__.py b/src/vmecpp/__init__.py index 409569e38..dd21c5c33 100644 --- a/src/vmecpp/__init__.py +++ b/src/vmecpp/__init__.py @@ -34,7 +34,12 @@ solve_equilibrium, solve_multigrid, ) -from vmecpp._pydantic_numpy import BaseModelWithNumpy +from vmecpp._pydantic_numpy import ( + _DAPPER_TYPE_FIELD, + BaseModelWithNumpy, + own_model_fields, +) +from vmecpp._rescale import rescale from vmecpp.cpp import _vmecpp # type: ignore # bindings to the C++ core logger = logging.getLogger(__name__) @@ -112,6 +117,14 @@ def _final_resolution(value: int | np.ndarray) -> int: _ArrayType, pydantic.BeforeValidator(lambda x: _util.right_pad(x, ndfmax, -1.0)), ] +# `VmecWOut.am`/`ac`/`ai` share one axis ("_preset") and so must all agree on +# length; right-pad to the classic fixed VMEC size (like AuxFType/AuxSType +# already do for the aux profile arrays) instead of leaving them at whatever +# length the user's profile happened to be. +ProfileCoeffType = typing.Annotated[ + _ArrayType, + pydantic.BeforeValidator(lambda x: _util.right_pad(x, preset, 0.0)), +] MgridModeType: typing.TypeAlias = typing.Annotated[ typing.Literal["R", "S", ""], pydantic.Field(max_length=1) @@ -681,7 +694,7 @@ def _from_cpp_vmecindata( # VmecInput does _not_ have any default values. vmec_input_dict = { attr_name: getattr(vmecindata, attr_name) - for attr_name in VmecInput.model_fields + for attr_name in own_model_fields(VmecInput) } vmec_input_dict["ns_array"] = vmec_input_dict["ns_array"].astype(np.int64) vmec_input_dict["niter_array"] = vmec_input_dict["niter_array"].astype(np.int64) @@ -713,7 +726,7 @@ def _to_cpp_vmecindata(self) -> _vmecpp.VmecINDATA: "zbc", } - for attr in VmecInput.model_fields: + for attr in own_model_fields(VmecInput): if attr in readonly_attrs or attr in ( "free_boundary_method", "iteration_style", @@ -795,10 +808,10 @@ class VmecWOut(BaseModelWithNumpy): """ # We use alias names to map to the wout keys, when they differ from the variable - # names in Python (e.g. lasym__logical__ instead of lasym). By default, we want - # to use the nicer Python names and explicitly opt in to use the wout names. + # names in Python (e.g. lasym__logical__ instead of lasym). Accept both the + # Python names and the wout aliases when validating. model_config = pydantic.ConfigDict( - validate_by_alias=False, + validate_by_alias=True, validate_by_name=True, serialize_by_alias=False, # Allow for variables in the wout file even if VMEC++ doesn't use them. @@ -1211,13 +1224,13 @@ def reason(self) -> str: piota_type: ProfileType """Parametrization of iota profile (copied from input).""" - am: jt.Float[np.ndarray, "_preset"] + am: ProfileCoeffType[jt.Float[np.ndarray, "_preset"]] """Mass/pressure profile coefficients (copied from input).""" - ac: jt.Float[np.ndarray, "_preset"] + ac: ProfileCoeffType[jt.Float[np.ndarray, "_preset"]] """Enclosed toroidal current profile coefficients (copied from input).""" - ai: jt.Float[np.ndarray, "_preset"] + ai: ProfileCoeffType[jt.Float[np.ndarray, "_preset"]] """Iota profile coefficients (copied from input).""" am_aux_s: AuxSType[jt.Float[np.ndarray, "_ndfmax"]] @@ -1449,7 +1462,7 @@ def _save_to_netcdf3(self, out_path: str | Path) -> None: # Convert VmeWOut to its NetCDF3 compatible representation # (wout compatible names and datatypes) - dumped_fields = self.model_dump(by_alias=True) + dumped_fields = self.model_dump(by_alias=True, exclude={_DAPPER_TYPE_FIELD}) # Make a dictionary of alias names to field info, from # model_fields (dictionary of non-alias names) @@ -1457,7 +1470,7 @@ def _save_to_netcdf3(self, out_path: str | Path) -> None: ( field_info.alias if field_info.alias is not None else field ): field_info - for field, field_info in VmecWOut.model_fields.items() + for field, field_info in own_model_fields(VmecWOut).items() } # jaxtyping does not expose a stable public API for dimension marker # types. Older versions expose `_AnonymousDim`, while newer versions @@ -1612,7 +1625,7 @@ def _from_cpp_wout(cpp_wout: _vmecpp.VmecppWOut) -> VmecWOut: attrs = {} # These attributes are the same in VMEC++ and in Fortran VMEC - for field in VmecWOut.model_fields: + for field in own_model_fields(VmecWOut): if field not in VmecWOut._CPP_WOUT_SPECIAL_HANDLING: attrs[field] = getattr(cpp_wout, field) @@ -1645,7 +1658,7 @@ def _to_cpp_wout(self) -> _vmecpp.WOutFileContents: cpp_wout = _vmecpp.WOutFileContents() # These attributes are the same in VMEC++ and in Fortran VMEC - for field in VmecWOut.model_fields: + for field in own_model_fields(VmecWOut): if field not in VmecWOut._CPP_WOUT_SPECIAL_HANDLING: setattr(cpp_wout, field, getattr(self, field)) @@ -1782,7 +1795,7 @@ def _from_cpp_threed1volumetrics( threed1volumetrics = Threed1Volumetrics( **{ attr: getattr(cpp_threed1volumetrics, attr) - for attr in Threed1Volumetrics.model_fields + for attr in own_model_fields(Threed1Volumetrics) } ) @@ -1847,7 +1860,7 @@ def _from_cpp_threed1_first_table( return Threed1FirstTable( **{ attr: getattr(cpp_threed1_first_table, attr) - for attr in Threed1FirstTable.model_fields + for attr in own_model_fields(Threed1FirstTable) } ) @@ -1958,20 +1971,25 @@ class Threed1GeometricAndMagneticQuantities(BaseModelWithNumpy): psi: jt.Float[np.ndarray, "num_full"] """Poloidal magnetic flux on the full grid.""" - ygeo: jt.Float[np.ndarray, "num_full"] - """Geometric minor radius profile.""" + ygeo: jt.Float[np.ndarray, "two_num_full"] + """Geometric minor radius profile, concatenated for the phi=0 and phi=pi symmetry + planes.""" - yinden: jt.Float[np.ndarray, "num_full"] - """Geometric indentation profile.""" + yinden: jt.Float[np.ndarray, "two_num_full"] + """Geometric indentation profile, concatenated for the phi=0 and phi=pi symmetry + planes.""" - yellip: jt.Float[np.ndarray, "num_full"] - """Geometric ellipticity profile.""" + yellip: jt.Float[np.ndarray, "two_num_full"] + """Geometric ellipticity profile, concatenated for the phi=0 and phi=pi symmetry + planes.""" - ytrian: jt.Float[np.ndarray, "num_full"] - """Geometric triangularity profile.""" + ytrian: jt.Float[np.ndarray, "two_num_full"] + """Geometric triangularity profile, concatenated for the phi=0 and phi=pi symmetry + planes.""" - yshift: jt.Float[np.ndarray, "num_full"] - """Geometric shift measured from the magnetic axis.""" + yshift: jt.Float[np.ndarray, "two_num_full"] + """Geometric shift measured from the magnetic axis, concatenated for the phi=0 and + phi=pi symmetry planes.""" loc_jpar_perp: jt.Float[np.ndarray, "num_full"] """Local parallel/perpendicular current ratio profile.""" @@ -1986,7 +2004,7 @@ def _from_cpp_threed1_geometric_and_magnetic_quantities( return Threed1GeometricAndMagneticQuantities( **{ attr: getattr(cpp_threed1_geometric_and_magnetic, attr) - for attr in Threed1GeometricAndMagneticQuantities.model_fields + for attr in own_model_fields(Threed1GeometricAndMagneticQuantities) } ) @@ -2004,22 +2022,35 @@ class Threed1AxisGeometry(BaseModelWithNumpy): zaxis_symm: jt.Float[np.ndarray, "ntor_plus_1"] """Stellarator-symmetric axis coefficients `Z * sin(n * zeta)`.""" - raxis_asym: jt.Float[np.ndarray, "ntor_plus_1"] - """Non-stellarator-symmetric axis coefficients `R * sin(n * zeta)`.""" + raxis_asym: jt.Float[np.ndarray, "ntor_plus_1"] | None = None + """Non-stellarator-symmetric axis coefficients `R * sin(n * zeta)`. - zaxis_asym: jt.Float[np.ndarray, "ntor_plus_1"] - """Non-stellarator-symmetric axis coefficients `Z * cos(n * zeta)`.""" + `None` when `lasym=False`. + """ + + zaxis_asym: jt.Float[np.ndarray, "ntor_plus_1"] | None = None + """Non-stellarator-symmetric axis coefficients `Z * cos(n * zeta)`. + + `None` when `lasym=False`. + """ @staticmethod def _from_cpp_threed1_axis_geometry( cpp_threed1_axis: _vmecpp.Threed1AxisGeometry, ) -> Threed1AxisGeometry: - return Threed1AxisGeometry( - **{ - attr: getattr(cpp_threed1_axis, attr) - for attr in Threed1AxisGeometry.model_fields - } - ) + attrs = { + attr: getattr(cpp_threed1_axis, attr) + for attr in own_model_fields(Threed1AxisGeometry) + if attr not in ("raxis_asym", "zaxis_asym") + } + # raxis_asym/zaxis_asym are only populated (non-empty) when lasym=True; + # they keep their `None` default otherwise (C++ leaves them as + # default-constructed, empty vectors, which don't fit the "ntor_plus_1" + # shape). + if cpp_threed1_axis.raxis_asym.size > 0: + attrs["raxis_asym"] = cpp_threed1_axis.raxis_asym + attrs["zaxis_asym"] = cpp_threed1_axis.zaxis_asym + return Threed1AxisGeometry(**attrs) class Threed1Betas(BaseModelWithNumpy): @@ -2052,7 +2083,7 @@ def _from_cpp_threed1_betas( return Threed1Betas( **{ attr: getattr(cpp_threed1_betas, attr) - for attr in Threed1Betas.model_fields + for attr in own_model_fields(Threed1Betas) } ) @@ -2121,7 +2152,7 @@ def _from_cpp_threed1_shafranov_integrals( ) -> Threed1ShafranovIntegrals: # `lambda_` maps to the C++ member `lambda` (a Python keyword). values = {} - for field_name in Threed1ShafranovIntegrals.model_fields: + for field_name in own_model_fields(Threed1ShafranovIntegrals): cpp_name = "lambda" if field_name == "lambda_" else field_name values[field_name] = getattr(cpp_threed1_shafranov_integrals, cpp_name) return Threed1ShafranovIntegrals(**values) @@ -2178,7 +2209,7 @@ class Mercier(BaseModelWithNumpy): @staticmethod def _from_cpp_mercier(cpp_mercier: _vmecpp.Mercier) -> Mercier: mercier = Mercier( - **{attr: getattr(cpp_mercier, attr) for attr in Mercier.model_fields} + **{attr: getattr(cpp_mercier, attr) for attr in own_model_fields(Mercier)} ) return mercier @@ -2187,13 +2218,13 @@ def _from_cpp_mercier(cpp_mercier: _vmecpp.Mercier) -> Mercier: class JxBOut(BaseModelWithNumpy): model_config = pydantic.ConfigDict(extra="forbid") - itheta: jt.Float[np.ndarray, "num_half nZnT"] + itheta: jt.Float[np.ndarray, "num_full nZnT"] r"""Poloidal surface current. :math:`itheta = (\frac{\partial B_s}{\partial \Phi} - \frac{\partial B_\phi}{\partial s}) / \mu_0` """ - izeta: jt.Float[np.ndarray, "num_half nZnT"] + izeta: jt.Float[np.ndarray, "num_full nZnT"] r"""Toroidal surface current. :math:`izeta = (-\frac{\partial B_s}{\partial \Theta} + \frac{\partial @@ -2277,7 +2308,7 @@ class JxBOut(BaseModelWithNumpy): @staticmethod def _from_cpp_jxbout(cpp_jxbout: _vmecpp.JxBOutFileContents) -> JxBOut: jxbout = JxBOut( - **{attr: getattr(cpp_jxbout, attr) for attr in JxBOut.model_fields} + **{attr: getattr(cpp_jxbout, attr) for attr in own_model_fields(JxBOut)} ) return jxbout @@ -2643,6 +2674,7 @@ def set_profile( __all__ = [ # noqa: RUF022 "run", "interpolate_solution", + "rescale", "VmecInput", "VmecOutput", "VmecWOut", diff --git a/src/vmecpp/__main__.py b/src/vmecpp/__main__.py index 8041ad111..f2946c036 100644 --- a/src/vmecpp/__main__.py +++ b/src/vmecpp/__main__.py @@ -43,7 +43,7 @@ def parse_arguments(argv: list[str] | None = None) -> argparse.Namespace: p.add_argument( "-t", "--max-threads", - help="Maximum number of threads that VMEC++ should spawn. The actual number might still be lower that this in case there are too few flux surfaces to keep these many threads busy.", + help="Maximum number of threads that VMEC++ should spawn. The actual number might still be lower than this in case there are too few flux surfaces to keep these many threads busy.", type=int, ) p.add_argument( diff --git a/src/vmecpp/_free_boundary.py b/src/vmecpp/_free_boundary.py index 690554c78..202f8d297 100644 --- a/src/vmecpp/_free_boundary.py +++ b/src/vmecpp/_free_boundary.py @@ -9,7 +9,7 @@ import numpy as np import pydantic -from vmecpp._pydantic_numpy import BaseModelWithNumpy +from vmecpp._pydantic_numpy import BaseModelWithNumpy, own_model_fields from vmecpp.cpp import _vmecpp # type: ignore @@ -48,13 +48,16 @@ def _from_cpp_makegrid_parameters( cpp_obj: _vmecpp.MakegridParameters, ) -> MakegridParameters: makegrid_parameters = MakegridParameters( - **{attr: getattr(cpp_obj, attr) for attr in MakegridParameters.model_fields} + **{ + attr: getattr(cpp_obj, attr) + for attr in own_model_fields(MakegridParameters) + } ) return makegrid_parameters def _to_cpp_makegrid_parameters(self) -> _vmecpp.MakegridParameters: return _vmecpp.MakegridParameters( - *[getattr(self, attr) for attr in MakegridParameters.model_fields] + *[getattr(self, attr) for attr in own_model_fields(MakegridParameters)] ) @staticmethod @@ -122,8 +125,8 @@ def _to_cpp_magnetic_field_response_table( """Convert the Pydantic object to a C++ MagneticFieldResponseTable object, avoiding a copy if possible.""" # If vmecpp.MagneticFieldResponseTable was created from a C++ object, the - # arrays be views to the memory of the C++ object. We don't need to create - # a new object and just return the underling one. + # arrays are views to the memory of the C++ object. We don't need to create + # a new object and just return the underlying one. underlying = self.b_r.base if ( isinstance(underlying, _vmecpp.MagneticFieldResponseTable) diff --git a/src/vmecpp/_pydantic_numpy.py b/src/vmecpp/_pydantic_numpy.py index 5be497bc3..7a4ac33b8 100644 --- a/src/vmecpp/_pydantic_numpy.py +++ b/src/vmecpp/_pydantic_numpy.py @@ -10,67 +10,99 @@ import numpy as np import pydantic +try: + import dapper # pyright: ignore[reportMissingImports] +except ImportError: + dapper = None + +if dapper is not None: + # BaseModelWithNumpy is a small subset of an internal Proxima library, dapper. + # Use it directly when available instead of redefining the same behavior. + BaseModelWithNumpy = dapper.DapperData +else: + + class BaseModelWithNumpy(pydantic.BaseModel): + """A minimal layer on top of pydantic to help with serialization and de- + serialization of classes with numpy arrays, annotated with jaxtyping. + + Does checks on the shape and type of arrays, may do casting if needed. + """ + + model_config = pydantic.ConfigDict( + arbitrary_types_allowed=True, + # Serialize NaN and infinite floats as strings in JSON output. + # Due to a bug in Pydantic, this setting is ignored by + # model_dump(mode="json"), so we override it below to also convert values + # to string there. + ser_json_inf_nan="strings", + ) -class BaseModelWithNumpy(pydantic.BaseModel): - """A minimal layer on top of pydantic to help with serialization and de- - serialization of classes with numpy arrays, annotated with jaxtyping. + @pydantic.field_serializer("*", mode="wrap", when_used="json") + def _serialize_field( + self, + value: typing.Any, + default_handler: pydantic.SerializerFunctionWrapHandler, + info: pydantic.FieldSerializationInfo, + # Note: Do NOT annotate return type with "-> Any" here, since that removes + # types from the JSON schema of all fields (in serialization mode). + ): + value = serialize_special_field(type(self), info.field_name, value) + return default_handler(value) + + @pydantic.field_validator("*", mode="wrap") + @classmethod + def _validate_field( + cls, + value: typing.Any, + default_handler: pydantic.ValidatorFunctionWrapHandler, + info: pydantic.ValidationInfo, + ) -> typing.Any: + assert info.field_name is not None + # We consciously *ignore* info.mode_is_json() here to allow validating from + # JSON-like dicts without having to go through strings. This is consistent + # with Pydantic's default behavior: while serializers retain Python objects + # in model_dump, but convert them to JSON values in model_dump_json, + # validators are lenient and accept JSON values in both modes. + value = deserialize_special_field(cls, info.field_name, value) + return default_handler(value) + + # This override is necessary to make also `model_dump(mode="json")` respect the + # ser_json_inf_nan="strings" setting in model_config. Without this fix, Pydantic + # would keep returning NaN/Inf as Python floats from this function, which leads + # to invalid JSON, e.g. when the output is used with `json.dumps`. + # This bug will probably only be fixed in Pydantic V3, since they consider it a + # breaking change. + # See: https://github.com/pydantic/pydantic/issues/10037#issuecomment-2314751795 + # Hide the override from the type system to "pass through" docstring and + # parameter types and defaults. + if not typing.TYPE_CHECKING: + + def model_dump( + self, *, mode: str = "python", **kwargs: Any + ) -> dict[str, Any]: + output_dict = super().model_dump(mode=mode, **kwargs) + if mode == "json": + sanitize_floats_in_container(output_dict) + return output_dict + + +# Outside Proxima, BaseModelWithNumpy has no extra fields and this is a no-op. +# When dapper is installed, BaseModelWithNumpy gains a "dapper_type" field that +# has no meaning to VMEC++ itself, so it must never be forwarded to/from the +# C++ bindings or written to non-dapper file formats (e.g. NetCDF wout files). +_DAPPER_TYPE_FIELD = "dapper_type" + + +def own_model_fields( + cls: type[pydantic.BaseModel], +) -> dict[str, pydantic.fields.FieldInfo]: + """Like `cls.model_fields`, but excludes `dapper_type`. - Does checks on the shape and type of arrays, may do casting if needed. + Use this instead of `cls.model_fields` wherever fields are bridged to the C++ + bindings or to a non-dapper file format. For `model.model_dump(...)` calls + with the same need, pass `exclude={_DAPPER_TYPE_FIELD}` directly instead. """ - - model_config = pydantic.ConfigDict( - arbitrary_types_allowed=True, - # Serialize NaN and infinite floats as strings in JSON output. - # Due to a bug in Pydantic, this setting is ignored by model_dump(mode="json"), - # so we override it below to also convert values to string there. - ser_json_inf_nan="strings", - ) - - @pydantic.field_serializer("*", mode="wrap", when_used="json") - def _serialize_field( - self, - value: typing.Any, - default_handler: pydantic.SerializerFunctionWrapHandler, - info: pydantic.FieldSerializationInfo, - # Note: Do NOT annotate return type with "-> Any" here, since that removes types - # from the JSON schema of all fields (in serialization mode). - ): - value = serialize_special_field(type(self), info.field_name, value) - return default_handler(value) - - @pydantic.field_validator("*", mode="wrap") - @classmethod - def _validate_field( - cls, - value: typing.Any, - default_handler: pydantic.ValidatorFunctionWrapHandler, - info: pydantic.ValidationInfo, - ) -> typing.Any: - assert info.field_name is not None - # We consciously *ignore* info.mode_is_json() here to allow validating from - # JSON-like dicts without having to go through strings. This is consistent with - # Pydantic's default behavior: while serializers retain Python objects in - # model_dump, but convert them to JSON values in model_dump_json, validators - # are lenient and accept JSON values in both modes. - value = deserialize_special_field(cls, info.field_name, value) - return default_handler(value) - - # This override is necessary to make also `model_dump(mode="json")` respect the - # ser_json_inf_nan="strings" setting in model_config. Without this fix, Pydantic - # would keep returning NaN/Inf as Python floats from this function, which leads to - # invalid JSON, e.g. when the output is used with `json.dumps`. - # This bug will probably only be fixed in Pydantic V3, since they consider it a - # breaking change. - # See: https://github.com/pydantic/pydantic/issues/10037#issuecomment-2314751795 - # Hide the override from the type system to "pass through" docstring and parameter - # types and defaults. - if not typing.TYPE_CHECKING: - - def model_dump(self, *, mode: str = "python", **kwargs: Any) -> dict[str, Any]: - output_dict = super().model_dump(mode=mode, **kwargs) - if mode == "json": - sanitize_floats_in_container(output_dict) - return output_dict + return {k: v for k, v in cls.model_fields.items() if k != _DAPPER_TYPE_FIELD} """This module handles special cases for serialization of BaseModelWithNumpy types. @@ -101,8 +133,7 @@ def model_dump(self, *, mode: str = "python", **kwargs: Any) -> dict[str, Any]: The mechanism for serialization is to convert "special" fields to JSON-serializable values and then let Pydantic take care of the rest. In practice: -- If the field is "special", (e.g. `np.ndarray`) it is converted to either a primitive - list. +- If the field is "special", (e.g. `np.ndarray`) it is converted to a primitive list. - If the field is a generic/composed type such as a list, an optional or a union, recurse and do the same for the inner types. - Otherwise just return the field as is. diff --git a/src/vmecpp/_rescale.py b/src/vmecpp/_rescale.py new file mode 100644 index 000000000..d4f028c11 --- /dev/null +++ b/src/vmecpp/_rescale.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import typing + +import numpy as np + +if typing.TYPE_CHECKING: + from vmecpp import VmecOutput + + +def rescale( + output: "VmecOutput", + b_scale: float, + r_scale: float, + scale_pressure: bool = True, +) -> "VmecOutput": + from . import run # noqa: PLC0415 + + """Rescale the equilibrium state. + + This scales the underlying geometry and inputs to represent an equilibrium + with a scaled major radius R -> r_scale * R and magnetic field B -> b_scale * B. + + Args: + output: a converged VmecOutput instance. + b_scale: factor to scale the magnetic field by. + r_scale: factor to scale the major radius by. + scale_pressure: whether to scale pressure to maintain force balance (default: True). + + Returns: + A new VmecOutput object with all derived parameters properly rescaled. + """ + # Scale INDATA parameters + scaled_input = output.input.model_copy(deep=True) + scaled_input.phiedge *= b_scale * (r_scale**2) + + if scale_pressure: + scaled_input.pres_scale *= b_scale**2 + + scaled_input.curtor *= b_scale * r_scale + + scaled_input.rbc *= r_scale + scaled_input.zbs *= r_scale + if scaled_input.rbs is not None: + scaled_input.rbs *= r_scale + if scaled_input.zbc is not None: + scaled_input.zbc *= r_scale + + scaled_input.raxis_c *= r_scale + scaled_input.zaxis_s *= r_scale + if scaled_input.raxis_s is not None: + scaled_input.raxis_s *= r_scale + if scaled_input.zaxis_c is not None: + scaled_input.zaxis_c *= r_scale + + # Force exactly 1 iteration with high tolerance so it exits immediately + # (niter_array > 0 is strictly required by the underlying C++ solver) + scaled_input.ns_array = np.array([scaled_input.ns_array[-1]]) + scaled_input.niter_array = np.array([1]) + scaled_input.ftol_array = np.array([1.0]) + + # Scale WOUT geometry + scaled_wout = output.wout.model_copy(deep=True) + scaled_wout.rmnc *= r_scale + scaled_wout.zmns *= r_scale + if scaled_wout.rmns is not None: + scaled_wout.rmns *= r_scale + if scaled_wout.zmnc is not None: + scaled_wout.zmnc *= r_scale + + # Create intermediate VmecOutput for restart_from + # (The C++ side only uses wout and indata for HotRestartState, so it's fine if the rest is unscaled) + intermediate_output = output.model_copy(deep=False) + intermediate_output.input = scaled_input + intermediate_output.wout = scaled_wout + + # Call run with 0 iterations + return run(scaled_input, restart_from=intermediate_output) diff --git a/src/vmecpp/_util.py b/src/vmecpp/_util.py index f5518fc9c..8e63a9abf 100644 --- a/src/vmecpp/_util.py +++ b/src/vmecpp/_util.py @@ -32,16 +32,16 @@ def distribution_root() -> Path: package_root, but doesn't have to be. The two differ in editable installations, where package_root() will point to the - source files, and distribution will point the /site-packages/vmecpp folder of your - python environment. It is the correct path to use for accessing shared libraries and - executables that come with vmecpp. + source files, and distribution will point to the /site-packages/vmecpp folder of + your python environment. It is the correct path to use for accessing shared + libraries and executables that come with vmecpp. """ return Path(importlib.metadata.distribution("vmecpp").locate_file("vmecpp")) # type: ignore @contextlib.contextmanager def change_working_directory_to(path: Path) -> Generator[None, None, None]: - """Changes the working director within a context manager. + """Changes the working directory within a context manager. Args: path: The path to change the working directory to. @@ -170,7 +170,7 @@ def indata_to_json( # adapted from https://github.com/jonathanschilling/indata2json/blob/4274976/json2indata def vmecpp_json_to_indata(vmecpp_json: dict[str, Any]) -> str: """Convert a dictionary with the contents of a VMEC++ JSON input file to the - corresponding conents of a VMEC2000 INDATA file.""" + corresponding contents of a VMEC2000 INDATA file.""" indata: str = "&INDATA\n" diff --git a/src/vmecpp/cpp/docker/tsan/Dockerfile b/src/vmecpp/cpp/docker/tsan/Dockerfile index 9de1cdd30..cb293e1d2 100644 --- a/src/vmecpp/cpp/docker/tsan/Dockerfile +++ b/src/vmecpp/cpp/docker/tsan/Dockerfile @@ -41,8 +41,7 @@ RUN apt-get install -y curl && curl -L "https://github.com/bazelbuild/bazelisk/r # Install packages required for building VMEC++. RUN apt-get install -y \ - libnetcdf-dev \ - liblapacke-dev + libnetcdf-dev # Install packages for in-docker development RUN apt-get install -y vim diff --git a/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc b/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc index 302f0eb93..425d933f3 100644 --- a/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc +++ b/src/vmecpp/cpp/vmecpp/common/enzyme/local_force_hessian_test.cc @@ -28,13 +28,8 @@ #include #include "vmecpp/common/enzyme/enzyme.h" -#include "vmecpp/vmec/ideal_mhd_model/bco_kernel.h" -#include "vmecpp/vmec/ideal_mhd_model/bcontra_kernel.h" -#include "vmecpp/vmec/ideal_mhd_model/jacobian_kernel.h" -#include "vmecpp/vmec/ideal_mhd_model/lambda_force_kernel.h" -#include "vmecpp/vmec/ideal_mhd_model/metric_kernel.h" -#include "vmecpp/vmec/ideal_mhd_model/mhdforce_kernel.h" -#include "vmecpp/vmec/ideal_mhd_model/pressure_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/exact_force_jvp.h" +#include "vmecpp/vmec/ideal_mhd_model/local_force_composition.h" // Problem dimensions and the constant (non-differentiated) context. struct Ctx { @@ -58,173 +53,26 @@ enum { kGeomBlocks = 16, kForceBlocks = 16 }; // g: geometry -> force density, composing the MHD and lambda-force kernels. __attribute__((noinline)) void LocalForce(const double* geom, double* work, double* force, const Ctx* c) { - const int nF = c->nFull; - const int nH = c->nHalf; - const int nZnT = c->nZnT, nsH = c->nsH; - // geometry blocks - const double* r1e = geom + 0 * nF; - const double* r1o = geom + 1 * nF; - const double* z1e = geom + 2 * nF; - const double* z1o = geom + 3 * nF; - const double* rue = geom + 4 * nF; - const double* ruo = geom + 5 * nF; - const double* zue = geom + 6 * nF; - const double* zuo = geom + 7 * nF; - const double* rve = geom + 8 * nF; - const double* rvo = geom + 9 * nF; - const double* zve = geom + 10 * nF; - const double* zvo = geom + 11 * nF; - const double* lue = geom + 12 * nF; - const double* luo = geom + 13 * nF; - const double* lve = geom + 14 * nF; - const double* lvo = geom + 15 * nF; - // half-grid intermediates from work - double* p = work; - double* r12 = p; - p += nH; - double* ru12 = p; - p += nH; - double* zu12 = p; - p += nH; - double* rs = p; - p += nH; - double* zs = p; - p += nH; - double* tau = p; - p += nH; - double* gsqrt = p; - p += nH; - double* guu = p; - p += nH; - double* guv = p; - p += nH; - double* gvv = p; - p += nH; - double* bsupu = p; - p += nH; - double* bsupv = p; - p += nH; - double* bsubu = p; - p += nH; - double* bsubv = p; - p += nH; - double* tp = p; - p += nH; - // per-nZnT scratch for the force kernel (26 blocks) - double* sc = p; // 26 * nZnT - - vmecpp::ComputeHalfGridJacobian( - r1e, r1o, z1e, z1o, rue, ruo, zue, zuo, c->sqrtSH, c->deltaS, - /*dSHalfDsInterp=*/0.25, nZnT, 0, 0, nsH, r12, ru12, zu12, rs, zs, tau); - vmecpp::ComputeMetricElements(r1e, r1o, rue, ruo, zue, zuo, rve, rvo, zve, - zvo, tau, r12, c->sqrtSF, c->sqrtSH, c->lthreed, - nZnT, 0, 0, nsH, gsqrt, guu, guv, gvv); - vmecpp::ComputeBsupContra(lue, luo, lve, lvo, gsqrt, c->sqrtSH, c->lthreed, - nZnT, 0, 0, nsH, bsupu, bsupv); - for (int jH = 0; jH < nsH; ++jH) { - for (int kl = 0; kl < nZnT; ++kl) { - const int ih = jH * nZnT + kl; - bsupu[ih] += c->chipH[jH] / gsqrt[ih]; - } - } - vmecpp::ComputeBCo(guu, guv, gvv, bsupu, bsupv, c->lthreed, nH, bsubu, bsubv); - vmecpp::ComputeMagneticPressure(bsupu, bsubu, bsupv, bsubv, nH, tp); - for (int jH = 0; jH < nsH; ++jH) { - for (int kl = 0; kl < nZnT; ++kl) tp[jH * nZnT + kl] += c->presH[jH]; - } - double* s = sc; - double* P_i = s; - s += nZnT; - double* rup_i = s; - s += nZnT; - double* zup_i = s; - s += nZnT; - double* rsp_i = s; - s += nZnT; - double* zsp_i = s; - s += nZnT; - double* taup_i = s; - s += nZnT; - double* gbubu_i = s; - s += nZnT; - double* gbubv_i = s; - s += nZnT; - double* gbvbv_i = s; - s += nZnT; - double* P_o = s; - s += nZnT; - double* rup_o = s; - s += nZnT; - double* zup_o = s; - s += nZnT; - double* rsp_o = s; - s += nZnT; - double* zsp_o = s; - s += nZnT; - double* taup_o = s; - s += nZnT; - double* gbubu_o = s; - s += nZnT; - double* gbubv_o = s; - s += nZnT; - double* gbvbv_o = s; - s += nZnT; - double* P_avg = s; - s += nZnT; - double* P_wavg = s; - s += nZnT; - double* gbubu_avg = s; - s += nZnT; - double* gbubu_wavg = s; - s += nZnT; - double* gbvbv_avg = s; - s += nZnT; - double* gbvbv_wavg = s; - s += nZnT; - double* gbubv_avg = s; - s += nZnT; - double* gbubv_wavg = s; - s += nZnT; - // lambda-force radial-sweep scratch (carried inside half-grid point) - double* bsubu_i = s; - s += nZnT; - double* bsubv_i = s; - s += nZnT; - double* gvv_gsqrt_i = s; - s += nZnT; - double* guv_bsupu_i = s; - s += nZnT; - double* armn_e = force + 0 * nH; - double* armn_o = force + 1 * nH; - double* azmn_e = force + 2 * nH; - double* azmn_o = force + 3 * nH; - double* brmn_e = force + 4 * nH; - double* brmn_o = force + 5 * nH; - double* bzmn_e = force + 6 * nH; - double* bzmn_o = force + 7 * nH; - double* crmn_e = force + 8 * nH; - double* crmn_o = force + 9 * nH; - double* czmn_e = force + 10 * nH; - double* czmn_o = force + 11 * nH; - vmecpp::ComputeMHDForceDensity( - r1e, r1o, rue, ruo, zue, zuo, z1o, rve, rvo, zve, zvo, r12, ru12, zu12, - rs, zs, tau, tp, gsqrt, bsupu, bsupv, c->sqrtSF, c->sqrtSH, P_i, rup_i, - zup_i, rsp_i, zsp_i, taup_i, gbubu_i, gbubv_i, gbvbv_i, P_o, rup_o, zup_o, - rsp_o, zsp_o, taup_o, gbubu_o, gbubv_o, gbvbv_o, P_avg, P_wavg, gbubu_avg, - gbubu_wavg, gbvbv_avg, gbvbv_wavg, gbubv_avg, gbubv_wavg, c->deltaS, nZnT, - /*nsMinF=*/0, 0, 0, nsH, /*jMaxRZ=*/nsH, c->lthreed, armn_e, armn_o, - azmn_e, azmn_o, brmn_e, brmn_o, bzmn_e, bzmn_o, crmn_e, crmn_o, czmn_e, - czmn_o); - // lambda force (blmn_e/o, clmn_e/o) from the shared kernel - double* blmn_e = force + 12 * nH; - double* blmn_o = force + 13 * nH; - double* clmn_e = force + 14 * nH; - double* clmn_o = force + 15 * nH; - vmecpp::ComputeHybridLambdaForce( - bsubu, bsubv, gvv, gsqrt, guv, bsupu, lue, luo, c->sqrtSH, c->sqrtSF, - c->radialBlending, c->lamscale, c->lthreed, nZnT, /*nsMinF=*/0, - /*nsMinF1=*/0, /*nsMinH=*/0, nsH, /*nsMaxFIncludingLcfs=*/nsH, bsubu_i, - bsubv_i, gvv_gsqrt_i, guv_bsupu_i, blmn_e, blmn_o, clmn_e, clmn_o); + vmecpp::LocalForceComposition lc; + lc.nZnT = c->nZnT; + lc.geom_stride = c->nFull; + lc.force_stride = c->nHalf; + lc.nsMinF = 0; + lc.nsMinF1 = 0; + lc.nsMinH = 0; + lc.nsMaxH = c->nsH; + lc.jMaxRZ = c->nsH; + lc.nsMaxFIncludingLcfs = c->nsH; + lc.sqrtSF = c->sqrtSF; + lc.sqrtSH = c->sqrtSH; + lc.chipH = c->chipH; + lc.presH = c->presH; + lc.radialBlending = c->radialBlending; + lc.deltaS = c->deltaS; + lc.dSHalfDsInterp = 0.25; + lc.lamscale = c->lamscale; + lc.lthreed = c->lthreed; + vmecpp::ComputeLocalForceDensity(geom, work, force, &lc); } // Scalar objective L = 0.5 ||force||^2; work and force are caller-owned @@ -307,9 +155,46 @@ int main() { printf(" forward / reverse agreement : %.2e\n", std::fabs(dfwd - drev) / (std::fabs(drev) + 1e-300)); + // Validate the standalone plugin JVP wrapper (exact_force_jvp.cc, the same + // entry point the exact Hessian-vector product calls) against a finite + // difference of the composition's force-density output. + vmecpp::LocalForceComposition lc; + lc.nZnT = c.nZnT; + lc.geom_stride = c.nFull; + lc.force_stride = c.nHalf; + lc.nsMinF = 0; + lc.nsMinF1 = 0; + lc.nsMinH = 0; + lc.nsMaxH = c.nsH; + lc.jMaxRZ = c.nsH; + lc.nsMaxFIncludingLcfs = c.nsH; + lc.sqrtSF = c.sqrtSF; + lc.sqrtSH = c.sqrtSH; + lc.chipH = c.chipH; + lc.presH = c.presH; + lc.radialBlending = c.radialBlending; + lc.deltaS = c.deltaS; + lc.dSHalfDsInterp = 0.25; + lc.lamscale = c.lamscale; + lc.lthreed = c.lthreed; + std::vector jf(nFc, 0.0), jdf(nFc, 0.0), jw(nW, 0.0), jdw(nW, 0.0); + vmecpp::ExactForceDensityJvp(geom.data(), v.data(), jw.data(), jdw.data(), + jf.data(), jdf.data(), &lc); + std::vector fpf(nFc, 0.0), fmf(nFc, 0.0), fw(nW, 0.0); + vmecpp::ComputeLocalForceDensity(gp.data(), fw.data(), fpf.data(), &lc); + vmecpp::ComputeLocalForceDensity(gm.data(), fw.data(), fmf.data(), &lc); + double jvp_err = 0.0, jvp_scale = 1e-300; + for (int i = 0; i < nFc; ++i) { + const double fd = (fpf[i] - fmf[i]) / (2 * h); + jvp_err = std::max(jvp_err, std::fabs(jdf[i] - fd)); + jvp_scale = std::max(jvp_scale, std::fabs(fd)); + } + printf(" plugin JVP wrapper vs finite-diff : %.2e\n", jvp_err / jvp_scale); + const bool ok = std::fabs(drev - dfd) < 1e-5 * scale && std::fabs(dfwd - dfd) < 1e-5 * scale && - std::fabs(dfwd - drev) < 1e-9 * (std::fabs(drev) + 1e-300); + std::fabs(dfwd - drev) < 1e-9 * (std::fabs(drev) + 1e-300) && + jvp_err < 1e-5 * jvp_scale; printf("%s\n", ok ? "PASS" : "FAIL"); return ok ? 0 : 1; } diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.cc b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.cc index 85d3b6551..000d3ff6e 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.cc +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.cc @@ -395,4 +395,79 @@ void FourierBasis::computeConversionIndices(Eigen::VectorXi& m_xm, template class FourierBasis; template class FourierBasis; +double SpectralWidth(const SurfaceFourierGeometry& geometry, const Sizes& sizes, + std::span mscale, + std::span nscale, const int p, const int q) { + double spectral_width_numerator = 0.0; + double spectral_width_denominator = 0.0; + + // note that we exclude m = 0 + for (int m = 1; m < sizes.mpol; ++m) { + for (int n = 0; n < sizes.ntor + 1; ++n) { + const int fourier_index = m * (sizes.ntor + 1) + n; + + const double basis_norm = mscale[m] * nscale[n]; + + // Use Eigen for vectorized norm computation + Eigen::Vector4d r_coefficients = Eigen::Vector4d::Zero(); + Eigen::Vector4d z_coefficients = Eigen::Vector4d::Zero(); + int basis_dimension = 0; + + r_coefficients[basis_dimension] = geometry.rmncc[fourier_index]; + z_coefficients[basis_dimension] = geometry.zmnsc[fourier_index]; + basis_dimension++; + + // CONVERT FROM INTERNAL XC REPRESENTATION FOR m=1 MODES, + // R+(at rsc) = .5(rsc + zcc), + // R-(at zcc) = .5(rsc - zcc), + // TO REQUIRED rsc, zcc FORMS + if (sizes.lthreed) { + if (m == 1) { + const double r_plus = geometry.rmnss[fourier_index]; + const double r_minus = geometry.zmncs[fourier_index]; + // rmnss + r_coefficients[basis_dimension] = r_plus + r_minus; + // zmncs + z_coefficients[basis_dimension] = r_plus - r_minus; + } else { + r_coefficients[basis_dimension] = geometry.rmnss[fourier_index]; + z_coefficients[basis_dimension] = geometry.zmncs[fourier_index]; + } + basis_dimension++; + } + if (sizes.lasym) { + if (m == 1) { + const double r_plus = geometry.rmnsc[fourier_index]; + const double r_minus = geometry.zmncc[fourier_index]; + // rmnsc + r_coefficients[basis_dimension] = r_plus + r_minus; + // zmncc + z_coefficients[basis_dimension] = r_plus - r_minus; + } else { + r_coefficients[basis_dimension] = geometry.rmnsc[fourier_index]; + z_coefficients[basis_dimension] = geometry.zmncc[fourier_index]; + } + basis_dimension++; + } + + if (sizes.lasym && sizes.lthreed) { + r_coefficients[basis_dimension] = geometry.rmncs[fourier_index]; + z_coefficients[basis_dimension] = geometry.zmnss[fourier_index]; + basis_dimension++; + } + + // Vectorized squared norm computation + double coefficient_norm = + r_coefficients.head(basis_dimension).squaredNorm() + + z_coefficients.head(basis_dimension).squaredNorm(); + coefficient_norm *= basis_norm * basis_norm; + + spectral_width_numerator += coefficient_norm * std::pow(m, p + q); + spectral_width_denominator += coefficient_norm * std::pow(m, p); + } // n + } // m + + return spectral_width_numerator / spectral_width_denominator; +} // SpectralWidth + } // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.h b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.h index f4148e007..1caf2da28 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.h +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis/fourier_basis.h @@ -359,6 +359,52 @@ class FourierBasis { using FourierBasisFastPoloidal = FourierBasis; using FourierBasisFastToroidal = FourierBasis; +// R and Z Fourier coefficients of a single flux surface in the VMEC-internal +// product basis, indexed as [m * (ntor + 1) + n]. The components excluded by +// the symmetry flags are left empty and are never read. +struct SurfaceFourierGeometry { + // contrib to R ~ cos(m * theta) * cos(n * zeta) + std::span rmncc; + + // contrib to R ~ sin(m * theta) * sin(n * zeta) + std::span rmnss; + + // contrib to R ~ sin(m * theta) * cos(n * zeta) + std::span rmnsc; + + // contrib to R ~ cos(m * theta) * sin(n * zeta) + std::span rmncs; + + // contrib to Z ~ sin(m * theta) * cos(n * zeta) + std::span zmnsc; + + // contrib to Z ~ cos(m * theta) * sin(n * zeta) + std::span zmncs; + + // contrib to Z ~ cos(m * theta) * cos(n * zeta) + std::span zmncc; + + // contrib to Z ~ sin(m * theta) * sin(n * zeta) + std::span zmnss; +}; + +// Spectral width of a single flux surface, +// +// = \sum_{m \ge 1, n} |c_{mn}|^2 m^{p+q} / \sum_{m \ge 1, n} |c_{mn}|^2 +// m^p , +// +// where |c_{mn}|^2 is the summed squared R and Z amplitude of mode (m, n). +// m = 0 is excluded because it carries the average position of the surface +// rather than its poloidal shape, and the m = 1 coefficients are unpacked from +// the internal R+/R- representation before they enter the sums. +// +// `mscale` and `nscale` undo the normalization of the internal Fourier basis: +// pass FourierBasis::mscale and ::nscale for coefficients held in that basis, +// and arrays of ones for plain Fourier amplitudes. +double SpectralWidth(const SurfaceFourierGeometry& geometry, const Sizes& sizes, + std::span mscale, + std::span nscale, int p = 4, int q = 1); + } // namespace vmecpp #endif // VMECPP_COMMON_FOURIER_BASIS_FOURIER_BASIS_H_ diff --git a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal_test.cc b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal_test.cc index 722b663da..513898b0f 100644 --- a/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal_test.cc +++ b/src/vmecpp/cpp/vmecpp/common/fourier_basis_fast_poloidal/fourier_basis_fast_poloidal_test.cc @@ -999,4 +999,137 @@ TEST(TestFourierBasisFastPoloidal, CheckInternally) { } } // CheckInternally +namespace { +constexpr double kTolerance = 1.0e-12; + +SurfaceFourierGeometry StellaratorSymmetricSurface( + const std::vector& rmncc, const std::vector& zmnsc) { + SurfaceFourierGeometry geometry; + geometry.rmncc = rmncc; + geometry.zmnsc = zmnsc; + return geometry; +} +} // namespace + +TEST(TestSpectralWidth, SingleModeGivesItsOwnPoloidalModeNumber) { + // With only one poloidal mode present, both sums carry the same coefficient + // and collapses to that mode number for any p and q. + const Sizes sizes(/*lasym=*/false, /*nfp=*/1, /*mpol=*/8, /*ntor=*/0, + /*ntheta=*/0, /*nzeta=*/0); + const std::vector unit_scale(sizes.mpol, 1.0); + + for (int m = 1; m < sizes.mpol; ++m) { + std::vector rmncc(sizes.mpol, 0.0); + std::vector zmnsc(sizes.mpol, 0.0); + rmncc[m] = 0.3; + zmnsc[m] = -0.2; + + const double spectral_width = + SpectralWidth(StellaratorSymmetricSurface(rmncc, zmnsc), sizes, + unit_scale, unit_scale); + EXPECT_NEAR(spectral_width, m, kTolerance); + } +} + +TEST(TestSpectralWidth, TwoModesAreWeightedByTheirEnergyAndExponents) { + const Sizes sizes(/*lasym=*/false, /*nfp=*/1, /*mpol=*/8, /*ntor=*/0, + /*ntheta=*/0, /*nzeta=*/0); + const std::vector unit_scale(sizes.mpol, 1.0); + + const double low_amplitude = 0.5; + const double high_amplitude = 0.1; + std::vector rmncc(sizes.mpol, 0.0); + const std::vector zmnsc(sizes.mpol, 0.0); + rmncc[2] = low_amplitude; + rmncc[6] = high_amplitude; + + const int p = 4; + const int q = 1; + const double expected = + (low_amplitude * low_amplitude * std::pow(2, p + q) + + high_amplitude * high_amplitude * std::pow(6, p + q)) / + (low_amplitude * low_amplitude * std::pow(2, p) + + high_amplitude * high_amplitude * std::pow(6, p)); + + EXPECT_NEAR(SpectralWidth(StellaratorSymmetricSurface(rmncc, zmnsc), sizes, + unit_scale, unit_scale, p, q), + expected, kTolerance); +} + +TEST(TestSpectralWidth, BasisScalesWeightTheToroidalModes) { + // The internal basis carries a factor nscale[n] on each coefficient, so + // passing the basis arrays weights an n != 0 mode twice as strongly as the + // same coefficient at n = 0. + const Sizes sizes(/*lasym=*/false, /*nfp=*/5, /*mpol=*/4, /*ntor=*/2, + /*ntheta=*/0, /*nzeta=*/0); + const FourierBasisFastPoloidal fourier_basis(&sizes); + + const int coefficients_per_surface = sizes.mpol * (sizes.ntor + 1); + std::vector rmncc(coefficients_per_surface, 0.0); + std::vector rmnss(coefficients_per_surface, 0.0); + const std::vector zeros(coefficients_per_surface, 0.0); + + // m = 2 at n = 0, and m = 3 at n = 1. + rmncc[2 * (sizes.ntor + 1) + 0] = 0.4; + rmncc[3 * (sizes.ntor + 1) + 1] = 0.4; + + SurfaceFourierGeometry geometry; + geometry.rmncc = rmncc; + geometry.rmnss = rmnss; + geometry.zmnsc = zeros; + geometry.zmncs = zeros; + + const std::vector unit_scale(std::max(sizes.mpol, sizes.ntor + 1) + 1, + 1.0); + const double unweighted = + SpectralWidth(geometry, sizes, unit_scale, unit_scale); + + const std::span mscale(fourier_basis.mscale.data(), + fourier_basis.mscale.size()); + const std::span nscale(fourier_basis.nscale.data(), + fourier_basis.nscale.size()); + const double weighted = SpectralWidth(geometry, sizes, mscale, nscale); + + // nscale[1]^2 = 2 doubles the weight of the m = 3 mode, which pulls up. + EXPECT_GT(weighted, unweighted); + + const double amplitude = 0.4 * 0.4; + const int p = 4; + const int q = 1; + const double expected = + (amplitude * std::pow(2, p + q) + 2.0 * amplitude * std::pow(3, p + q)) / + (amplitude * std::pow(2, p) + 2.0 * amplitude * std::pow(3, p)); + // The constant mscale[m >= 1]^2 cancels between numerator and denominator. + EXPECT_NEAR(weighted, expected, kTolerance); +} + +TEST(TestSpectralWidth, ModeOneIsUnpackedFromTheInternalRepresentation) { + // m = 1 is stored as R+ = (rmnss + zmncs) / 2 and R- = (rmnss - zmncs) / 2, + // so the sums have to see the original rmnss and zmncs again. + const Sizes sizes(/*lasym=*/false, /*nfp=*/5, /*mpol=*/3, /*ntor=*/1, + /*ntheta=*/0, /*nzeta=*/0); + const int coefficients_per_surface = sizes.mpol * (sizes.ntor + 1); + const std::vector unit_scale(coefficients_per_surface, 1.0); + + const double rmnss_value = 0.6; + const double zmncs_value = 0.2; + const int index = 1 * (sizes.ntor + 1) + 1; + + std::vector rmnss(coefficients_per_surface, 0.0); + std::vector zmncs(coefficients_per_surface, 0.0); + const std::vector zeros(coefficients_per_surface, 0.0); + rmnss[index] = 0.5 * (rmnss_value + zmncs_value); + zmncs[index] = 0.5 * (rmnss_value - zmncs_value); + + SurfaceFourierGeometry geometry; + geometry.rmncc = zeros; + geometry.rmnss = rmnss; + geometry.zmnsc = zeros; + geometry.zmncs = zmncs; + + // Only m = 1 carries any amplitude, so has to come out as exactly 1. + EXPECT_NEAR(SpectralWidth(geometry, sizes, unit_scale, unit_scale), 1.0, + kTolerance); +} + } // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/BUILD.bazel b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/BUILD.bazel index d45bfdad1..08bc91738 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/BUILD.bazel @@ -13,9 +13,6 @@ cc_library( "//vmecpp/common/fourier_basis_fast_toroidal", "//vmecpp/free_boundary/tangential_partitioning:tangential_partitioning", ], - linkopts = [ - "-llapack", - ], ) cc_test( @@ -41,5 +38,4 @@ cc_binary( "//vmecpp/free_boundary/tangential_partitioning:tangential_partitioning", "@google_benchmark//:benchmark_main", ], - linkopts = ["-llapack"], ) diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.cc b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.cc index 16aa7bda2..f4fd61e10 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.cc +++ b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.cc @@ -5,33 +5,26 @@ #include "vmecpp/free_boundary/laplace_solver/laplace_solver.h" #include -#include #include #include "absl/algorithm/container.h" #include "absl/log/check.h" -/// LU factorization of a general M-by-N matrix A -extern "C" void dgetrf_(int* m, int* n, double* a, int* lda, int* ipiv, - int* info); -/// Solves a system of linear equations using the LU factorization. -extern "C" void dgetrs_(char* transpose, int* num_rows, int* num_columns, - double* matrix, int* leading_dim, int* pivot, double* y, - int* y_leading_dim, int* info); - namespace vmecpp { -LaplaceSolver::LaplaceSolver(const Sizes* s, const FourierBasisFastToroidal* fb, - const TangentialPartitioning* tp, int nf, int mf, - std::span matrixShare, std::span iPiv, - std::span bvecShare) +LaplaceSolver::LaplaceSolver( + const Sizes* s, const FourierBasisFastToroidal* fb, + const TangentialPartitioning* tp, int nf, int mf, + std::span matrixShare, + Eigen::PartialPivLU* lu_decomposition, + std::span bvecShare) : s_(*s), fb_(*fb), tp_(*tp), nf(nf), mf(mf), matrixShare(matrixShare), - iPiv(iPiv), + lu_decomposition_(lu_decomposition), bvecShare(bvecShare) { // thread-local tangential grid point range numLocal = tp_.ztMax - tp_.ztMin; @@ -578,28 +571,17 @@ void LaplaceSolver::BuildMatrix() { void LaplaceSolver::DecomposeMatrix() { const int mnpd = (mf + 1) * (2 * nf + 1); - int mnpd_dim = s_.lasym ? 2 * mnpd : mnpd; - - // NOTE: - // As soon as LAPACK starts working on `matrixShare`, - // it is not consistent with the value on entry anymore - // and thus cannot be used for testing anymore. + const int mnpd_dim = s_.lasym ? 2 * mnpd : mnpd; // perform LU factorization of the matrix // (only needed when matrix is updated --> every nvacskip iterations) - int info; - dgetrf_(&mnpd_dim, &mnpd_dim, matrixShare.data(), &mnpd_dim, iPiv.data(), - &info); - - if (info < 0) { - std::cout << -info << "-th argument to dgetrf is wrong\n"; - } else if (info > 0) { - std::cout << absl::StrFormat( - "U(%d,%d) is exactly zero in dgetrf --> singular matrix!\n", info, - info); - } + Eigen::Map matrix_map(matrixShare.data(), mnpd_dim, + mnpd_dim); + lu_decomposition_->compute(matrix_map); - CHECK_EQ(info, 0) << "dgetrf error"; + // A zero diagonal entry in the U factor means an exactly singular matrix. + CHECK((lu_decomposition_->matrixLU().diagonal().array() != 0.0).all()) + << "singular matrix in LaplaceSolver::DecomposeMatrix"; } // DecomposeMatrix void LaplaceSolver::SolveForPotential( @@ -653,19 +635,13 @@ void LaplaceSolver::SolveForPotential( } } - // solve for given RHS - int one = 1; - int info; - char no_transpose = 'N'; - int n = mnpd_dim; - dgetrs_(&no_transpose, &n, &one, matrixShare.data(), &n, iPiv.data(), - bvecShare.data(), &n, &info); - - if (info < 0) { - std::cout << -info << "-th argument to dgetrs wrong\n"; - } - - CHECK_EQ(info, 0) << "dgetrs error"; + // solve for given RHS using the factorization computed in + // DecomposeMatrix(). Use a temporary for the result: bvecShare is both + // the right-hand side and the destination, and aliasing it directly + // against the input of PartialPivLU::solve() is not guaranteed safe. + Eigen::Map rhs(bvecShare.data(), mnpd_dim); + const Eigen::VectorXd solution = lu_decomposition_->solve(rhs); + Eigen::Map(bvecShare.data(), mnpd_dim) = solution; } #ifdef _OPENMP #pragma omp barrier diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.h b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.h index d319e1d09..1621854b3 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.h +++ b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver.h @@ -17,9 +17,18 @@ namespace vmecpp { class LaplaceSolver { public: + // lu_decomposition is a non-owning pointer to a LU decomposition object + // that is shared across all per-thread LaplaceSolver instances of a given + // vacuum solve, exactly like matrixShare/bvecShare are spans into shared + // backing storage: DecomposeMatrix() (called by exactly one thread, guarded + // by an outer `#pragma omp single`) may run on a different LaplaceSolver + // object than the one whose SolveForPotential() (also single-threaded + // internally) later reads the factorization, so the decomposition itself + // must live outside any individual LaplaceSolver instance. LaplaceSolver(const Sizes* s, const FourierBasisFastToroidal* fb, const TangentialPartitioning* tp, int nf, int mf, - std::span matrixShare, std::span iPiv, + std::span matrixShare, + Eigen::PartialPivLU* lu_decomposition, std::span bvecShare); void TransformGreensFunctionDerivative(const Eigen::VectorXd& greenp); @@ -96,10 +105,13 @@ class LaplaceSolver { int nf; int mf; - // needed for LAPACK's dgetrf + // needed for the dense linear solve (BuildMatrix/DecomposeMatrix/ + // SolveForPotential) // non-owning pointers std::span matrixShare; - std::span iPiv; + // Shared LU decomposition object; see the constructor comment for why this + // cannot be a plain member of LaplaceSolver. + Eigen::PartialPivLU* lu_decomposition_; std::span bvecShare; // ---------------- diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_bench.cc b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_bench.cc index 752970dd8..7d80cea8d 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_bench.cc +++ b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_bench.cc @@ -7,8 +7,8 @@ // // The dense linear system assembled and factorized here is the dominant // free-boundary cost when NESTOR performs a full update: an -// mnpd x mnpd (mnpd = (mf+1)*(2*nf+1)) matrix is built, LU-factorized via -// LAPACK dgetrf, and back-substituted via dgetrs. We benchmark: +// mnpd x mnpd (mnpd = (mf+1)*(2*nf+1)) matrix is built, LU-factorized, and +// back-substituted via Eigen::PartialPivLU. We benchmark: // // * TransformGreensFunctionDerivative -- the largest of the Fourier // transforms feeding the source term (uses the manufactured single-mode @@ -53,8 +53,8 @@ constexpr ResParams kResolutions[] = { // ---------------------------------------------------------------------------- // Fixture matching laplace_solver_test.cc's minimal single-threaded setup: -// a non-mocked LaplaceSolver with flat matrixShare / iPiv / bvecShare backing -// storage. +// a non-mocked LaplaceSolver with flat matrixShare / bvecShare backing +// storage and a shared LU decomposition object. // ---------------------------------------------------------------------------- struct BenchFixture { Sizes s; @@ -66,8 +66,8 @@ struct BenchFixture { int mnpd; std::vector matrixShare; - std::vector iPiv; std::vector bvecShare; + Eigen::PartialPivLU lu_decomposition; std::unique_ptr ls; @@ -87,16 +87,16 @@ struct BenchFixture { mf(mpol + 1), mnpd((mf + 1) * (2 * nf + 1)), matrixShare(mnpd * mnpd, 0.0), - iPiv(mnpd, 0), bvecShare(mnpd, 0.0) { ls = std::make_unique( - &s, &fb, &tp, nf, mf, std::span(matrixShare), - std::span(iPiv), std::span(bvecShare)); + &s, &fb, &tp, nf, mf, std::span(matrixShare), &lu_decomposition, + std::span(bvecShare)); std::mt19937 rng(42); std::uniform_real_distribution dist(-1.0, 1.0); - // Diagonally-dominant matrix so dgetrf never hits a singular pivot. + // Diagonally-dominant matrix so the LU factorization never hits a + // singular pivot. amat_seed.resize(mnpd * mnpd); for (int col = 0; col < mnpd; ++col) { for (int row = 0; row < mnpd; ++row) { @@ -129,9 +129,8 @@ struct BenchFixture { } } - // Re-seed the LaplaceSolver's system state to the well-conditioned baseline. - // dgetrf destroys matrixShare in place, so this must run before each - // decompose+solve iteration. + // Re-seed the LaplaceSolver's system state to the well-conditioned + // baseline before each decompose+solve iteration. void ResetSystem() { ls->amat_sin_sin = amat_seed; ls->bvec_sin = Eigen::VectorXd::Zero(mnpd); @@ -156,7 +155,7 @@ void BM_LaplaceSolve(benchmark::State& state) { state.SetLabel(kResolutions[kIdx].label); } -// Just the LU factorization (dgetrf), the O(mnpd^3) hotspot. +// Just the LU factorization, the O(mnpd^3) hotspot. template void BM_LaplaceDecompose(benchmark::State& state) { static BenchFixture fx(kResolutions[kIdx].nfp, kResolutions[kIdx].mpol, diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_test.cc b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_test.cc index bb9c8fbfa..4e0e4c09e 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_test.cc +++ b/src/vmecpp/cpp/vmecpp/free_boundary/laplace_solver/laplace_solver_test.cc @@ -58,11 +58,11 @@ TEST_P(TransformGreensFunctionDerivativeTest, SingleModeRoundTrip) { // Dummy shared arrays (single-threaded, no cross-thread accumulation needed) std::vector matrixShare(mnpd * mnpd, 0.0); - std::vector iPiv(mnpd, 0); std::vector bvecShare(mnpd, 0.0); + Eigen::PartialPivLU lu_decomposition; LaplaceSolver ls(&s, &fb, &tp, nf, mf, std::span(matrixShare), - std::span(iPiv), std::span(bvecShare)); + &lu_decomposition, std::span(bvecShare)); // Build greenp[klpRel * nThetaEven * nZeta + l * nZeta + k]. // For each klp, inject exactly one mode: @@ -219,11 +219,11 @@ TEST(LaplaceSolverTest, ZeroKernelSolveGivesHalfInverse) { const int numLocal = tp.ztMax - tp.ztMin; std::vector matrixShare(mnpd * mnpd, 0.0); - std::vector iPiv(mnpd, 0); std::vector bvecShare(mnpd, 0.0); + Eigen::PartialPivLU lu_decomposition; LaplaceSolver ls(&s, &fb, &tp, nf, mf, std::span(matrixShare), - std::span(iPiv), std::span(bvecShare)); + &lu_decomposition, std::span(bvecShare)); // Zero greenp: no double-layer kernel contribution. Eigen::VectorXd greenp = diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.cc b/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.cc index 26afc08eb..60954d291 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.cc +++ b/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.cc @@ -11,7 +11,8 @@ namespace vmecpp { Nestor::Nestor(const Sizes* s, const TangentialPartitioning* tp, const MGridProvider* mgrid, std::span matrixShare, std::span bvecShare, std::span bSqVacShare, - std::span iPiv, std::span vacuum_b_r_share, + Eigen::PartialPivLU* lu_decomposition, + std::span vacuum_b_r_share, std::span vacuum_b_phi_share, std::span vacuum_b_z_share) : FreeBoundaryBase(s, tp, mgrid, bSqVacShare, vacuum_b_r_share, @@ -20,7 +21,7 @@ Nestor::Nestor(const Sizes* s, const TangentialPartitioning* tp, mf(s_.mpol + 1), si_(s, &fb_, tp, &sg_, nf, mf), ri_(s, tp, &sg_), - ls_(s, &fb_, tp, nf, mf, matrixShare, iPiv, bvecShare), + ls_(s, &fb_, tp, nf, mf, matrixShare, lu_decomposition, bvecShare), bvecShare(bvecShare) { int numLocal = tp_.ztMax - tp_.ztMin; diff --git a/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.h b/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.h index 6fd9198ec..1d6fef4af 100644 --- a/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.h +++ b/src/vmecpp/cpp/vmecpp/free_boundary/nestor/nestor.h @@ -26,7 +26,8 @@ class Nestor : public FreeBoundaryBase { Nestor(const Sizes* s, const TangentialPartitioning* tp, const MGridProvider* mgrid, std::span matrixShare, std::span bvecShare, std::span bSqVacShare, - std::span iPiv, std::span vacuum_b_r_share, + Eigen::PartialPivLU* lu_decomposition, + std::span vacuum_b_r_share, std::span vacuum_b_phi_share, std::span vacuum_b_z_share); diff --git a/src/vmecpp/cpp/vmecpp/test_data/BUILD.bazel b/src/vmecpp/cpp/vmecpp/test_data/BUILD.bazel index eff6b6a6e..55fab7339 100644 --- a/src/vmecpp/cpp/vmecpp/test_data/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/test_data/BUILD.bazel @@ -106,6 +106,16 @@ filegroup( ], ) +filegroup( + name = "solovev_free_bdy_lforbal", + visibility = ["//visibility:public"], + srcs = [ + "solovev_free_bdy_lforbal.json", + "mgrid_solovev.nc", + "wout_solovev_free_bdy_T.nc", + ], +) + filegroup( name = "cth_like_free_bdy_multigrid", visibility = ["//visibility:public"], diff --git a/src/vmecpp/cpp/vmecpp/test_data/solovev_free_bdy_lforbal.json b/src/vmecpp/cpp/vmecpp/test_data/solovev_free_bdy_lforbal.json new file mode 100644 index 000000000..1ad6e22a5 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/test_data/solovev_free_bdy_lforbal.json @@ -0,0 +1,63 @@ +{ +"lasym":false, +"nfp":1, +"mpol":12, +"ntor":0, +"ntheta":0, +"nzeta":1, +"ns_array":[ +16, +32], +"ftol_array":[ +1.0000000000000000E-010, +1.0000000000000000E-014], +"niter_array":[ +10000, +20000], +"delt":0.90000000000000002, +"tcon0":1.0000000000000000, +"aphi":[ +1.0000000000000000], +"phiedge":1.0000000000000000, +"nstep":100, +"pmass_type":"power_series", +"am":[ +0.12500000000000000, +-0.12500000000000000], +"pres_scale":1.0000000000000000, +"gamma":0.0000000000000000, +"spres_ped":1.0000000000000000, +"ncurr":0, +"piota_type":"power_series", +"ai":[ +0.48999999999999999, +-0.29999999999999999], +"lfreeb":true, +"mgrid_file":"vmecpp/test_data/mgrid_solovev.nc", +"extcur":[ +3884526.4098763089, +-293557.71237379522, +-17348.518536770429, +60021.370169731599, +60025.409404908867, +-17349.931031838169, +-293553.15363085101, +-356063.91087172751, +-65884.347192830843, +-11543.877747129871, +-11535.465107552191, +-65883.008583646064, +-356058.93884688552], +"nvacskip":6, +"lforbal":true, +"raxis_c":[ +3.9990000000000001], +"zaxis_s":[ +0.0000000000000000], +"rbc":[ +{"n": 0,"m": 0,"value": 0.399900000000000011E+01}, +{"n": 0,"m": 1,"value": 0.102600000000000002E+01}, +{"n": 0,"m": 2,"value": 0.680000000000000049E-01}], +"zbs":[ +{"n": 0,"m": 1,"value": 0.158000000000000007E+01}, +{"n": 0,"m": 2,"value":-0.100000000000000002E-01}]} diff --git a/src/vmecpp/cpp/vmecpp/test_data/wout_cma.nc b/src/vmecpp/cpp/vmecpp/test_data/wout_cma.nc index 2854b9c5f..cb3a59fd3 100644 Binary files a/src/vmecpp/cpp/vmecpp/test_data/wout_cma.nc and b/src/vmecpp/cpp/vmecpp/test_data/wout_cma.nc differ diff --git a/src/vmecpp/cpp/vmecpp/test_data/wout_cth_like_free_bdy.nc b/src/vmecpp/cpp/vmecpp/test_data/wout_cth_like_free_bdy.nc index 3ce51dea5..eb908a523 100644 Binary files a/src/vmecpp/cpp/vmecpp/test_data/wout_cth_like_free_bdy.nc and b/src/vmecpp/cpp/vmecpp/test_data/wout_cth_like_free_bdy.nc differ diff --git a/src/vmecpp/cpp/vmecpp/test_data/wout_solovev_free_bdy_T.nc b/src/vmecpp/cpp/vmecpp/test_data/wout_solovev_free_bdy_T.nc new file mode 100644 index 000000000..f73e888c5 Binary files /dev/null and b/src/vmecpp/cpp/vmecpp/test_data/wout_solovev_free_bdy_T.nc differ diff --git a/src/vmecpp/cpp/vmecpp/vmec/boundaries/BUILD.bazel b/src/vmecpp/cpp/vmecpp/vmec/boundaries/BUILD.bazel index 5a149edf1..1ec09d175 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/boundaries/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/vmec/boundaries/BUILD.bazel @@ -23,5 +23,23 @@ cc_library( "//vmecpp/common/fourier_basis_fast_poloidal", "//vmecpp/common/vmec_indata:vmec_indata", "@abseil-cpp//absl/algorithm:container", + "@abseil-cpp//absl/log:log", + "@abseil-cpp//absl/strings:str_format", ], ) + +cc_test( + name = "boundaries_test", + srcs = ["boundaries_test.cc"], + data = [ + "//vmecpp/test_data:cma", + "//vmecpp/test_data:cth_like_fixed_bdy", + "//vmecpp/test_data:solovev", + ], + deps = [ + ":boundaries", + "//vmecpp/vmec/vmec_constants:vmec_constants", + "@googletest//:gtest_main", + ], + size = "small", +) diff --git a/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.cc b/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.cc index f443ebcf5..ffa2a5c9f 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.cc @@ -5,12 +5,21 @@ #include "vmecpp/vmec/boundaries/boundaries.h" #include +#include #include "absl/algorithm/container.h" +#include "absl/log/log.h" +#include "absl/strings/str_format.h" #include "vmecpp/vmec/boundaries/guess_magnetic_axis.h" namespace vmecpp { +namespace { +std::span AsSpan(const Eigen::VectorXd& coefficients) { + return {coefficients.data(), static_cast(coefficients.size())}; +} +} // namespace + Boundaries::Boundaries(const Sizes* s, const FourierBasisFastPoloidal* t, const int sign_of_jacobian) : s_(*s), t_(*t), sign_of_jacobian_(sign_of_jacobian) { @@ -51,9 +60,44 @@ bool Boundaries::setupFromIndata(const VmecINDATA& id, bool verbose) { // activate m=1-constraint ensureM1Constrained(0.5); + // A boundary carrying a lot of high-poloidal-mode content is still a valid + // input, so this is a warning rather than an error. A free-boundary run + // relaxes the initial boundary, so only fixed-boundary inputs are flagged. + if (verbose && !id.lfreeb) { + const double spectral_width = ComputeSpectralWidth(); + if (spectral_width > kSpectrallyDenseBoundaryThreshold) { + LOG(WARNING) << absl::StrFormat( + "Input boundary is spectrally dense (spectral width = %.1f). " + "Runs with high spectral content may converge poorly or incorrectly " + "resolve the physics. Consider spectral condensation of the boundary " + "first, e.g. using simsopt.geo.surfacerzfourier.condense_spectrum", + spectral_width); + } + } + return haveToFlipTheta; } +double Boundaries::ComputeSpectralWidth(const int p, const int q) const { + // The boundary coefficients are plain Fourier amplitudes, whereas + // FourierGeometry holds them divided by mscale * nscale; unit scales make + // SpectralWidth weight both representations the same way. + const Eigen::VectorXd unit_mscale = Eigen::VectorXd::Ones(s_.mpol); + const Eigen::VectorXd unit_nscale = Eigen::VectorXd::Ones(s_.ntor + 1); + + const SurfaceFourierGeometry boundary = {.rmncc = AsSpan(rbcc), + .rmnss = AsSpan(rbss), + .rmnsc = AsSpan(rbsc), + .rmncs = AsSpan(rbcs), + .zmnsc = AsSpan(zbsc), + .zmncs = AsSpan(zbcs), + .zmncc = AsSpan(zbcc), + .zmnss = AsSpan(zbss)}; + + return SpectralWidth(boundary, s_, AsSpan(unit_mscale), AsSpan(unit_nscale), + p, q); +} // ComputeSpectralWidth + void Boundaries::parseToInternalArrays(const VmecINDATA& id, bool verbose) { // copy over axis from INDATA to this class for (int n = 0; n < s_.ntor + 1; ++n) { diff --git a/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.h b/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.h index 51ee1b83d..639652080 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.h +++ b/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries.h @@ -14,6 +14,11 @@ namespace vmecpp { +// Boundary spectral width above which setupFromIndata warns that a +// fixed-boundary run may converge poorly or resolve the physics incorrectly. +// Every input bundled with VMEC++ stays below this. +inline constexpr double kSpectrallyDenseBoundaryThreshold = 2.0; + /** radial boundaries: magnetic axis and last closed flux surface */ class Boundaries { public: @@ -23,6 +28,11 @@ class Boundaries { bool setupFromIndata(const VmecINDATA& id, bool verbose = true); void ensureM1Constrained(double scaling_factor); + // Spectral width of the boundary, the same quantity VMEC++ reports for + // each flux surface in the column of the iteration table. Only meaningful + // once setupFromIndata has filled the coefficients below. + double ComputeSpectralWidth(int p = 4, int q = 1) const; + // This object is initialized with an initial guess for the magnetic axis // geometry via setupFromIndata() that is provided by the user. This method // replaces the axis geometry in this object with a new guess for the magnetic diff --git a/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries_test.cc b/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries_test.cc new file mode 100644 index 000000000..961fb535c --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/boundaries/boundaries_test.cc @@ -0,0 +1,69 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#include "vmecpp/vmec/boundaries/boundaries.h" + +#include +#include + +#include "gtest/gtest.h" +#include "vmecpp/vmec/vmec_constants/vmec_algorithm_constants.h" + +namespace vmecpp { + +namespace { +double BoundarySpectralWidth(const VmecINDATA& indata) { + const Sizes sizes(indata); + const FourierBasisFastPoloidal fourier_basis(&sizes); + Boundaries boundaries(&sizes, &fourier_basis, + vmec_algorithm_constants::kSignOfJacobian); + boundaries.setupFromIndata(indata, /*verbose=*/false); + return boundaries.ComputeSpectralWidth(); +} +} // namespace + +TEST(TestBoundaries, SpectralWidthOfSmoothBoundaryIsLow) { + const VmecINDATA indata = + VmecINDATA::FromFile("vmecpp/test_data/solovev.json"); + + // solovev has a smooth boundary (R at m=0,1,2; Z at m=1,2), so its spectral + // width is close to 1 and well below the spectrally-dense warning threshold. + const double spectral_width = BoundarySpectralWidth(indata); + EXPECT_GE(spectral_width, 1.0); + EXPECT_LT(spectral_width, kSpectrallyDenseBoundaryThreshold); +} + +TEST(TestBoundaries, SpectralWidthOfDenseBoundaryExceedsThreshold) { + VmecINDATA indata = VmecINDATA::FromFile("vmecpp/test_data/solovev.json"); + + // Make room for high-m modes and inject a high-poloidal-mode ripple, which + // drives the boundary spectral width above the warning threshold. + indata.SetMpolNtor(/*new_mpol=*/12, /*new_ntor=*/0); + for (int m = 8; m <= 10; ++m) { + indata.rbc(m, 0) = 0.05; + indata.zbs(m, 0) = 0.05; + } + EXPECT_GT(BoundarySpectralWidth(indata), kSpectrallyDenseBoundaryThreshold); + + // A spectrally dense boundary is still a valid input: it is warned about, not + // rejected. + EXPECT_TRUE(IsConsistent(indata, /*enable_info_messages=*/false).ok()); +} + +TEST(TestBoundaries, BundledInputsAreNotSpectrallyDense) { + // The threshold has to sit above every input VMEC++ ships, or a routine run + // would warn about a boundary that is known to be fine. cma is the densest of + // them; cth_like_fixed_bdy is the densest three-dimensional one. + const std::vector input_files = { + "vmecpp/test_data/solovev.json", "vmecpp/test_data/cma.json", + "vmecpp/test_data/cth_like_fixed_bdy.json"}; + + for (const std::string& input_file : input_files) { + const VmecINDATA indata = VmecINDATA::FromFile(input_file); + EXPECT_LT(BoundarySpectralWidth(indata), kSpectrallyDenseBoundaryThreshold) + << "for " << input_file; + } +} + +} // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/BUILD.bazel b/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/BUILD.bazel index b7da681e9..4e9a2779a 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/BUILD.bazel @@ -9,6 +9,7 @@ cc_library( deps = [ "//vmecpp/common/util:util", "//vmecpp/common/sizes:sizes", + "//vmecpp/common/fourier_basis_fast_poloidal", "//vmecpp/vmec/radial_partitioning:radial_partitioning", "//vmecpp/vmec/handover_storage:handover_storage", "//vmecpp/vmec/boundaries:boundaries", diff --git a/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.cc b/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.cc index 02c8ebf7c..8c8741f70 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.cc @@ -12,6 +12,19 @@ namespace vmecpp { +namespace { +// The coefficients of one flux surface, or an empty span for a component the +// symmetry flags exclude and which therefore was never allocated. +std::span SurfaceOf(std::span coefficients, + int surface_offset, + int coefficients_per_surface) { + if (coefficients.empty()) { + return {}; + } + return coefficients.subspan(surface_offset, coefficients_per_surface); +} +} // namespace + FourierGeometry::FourierGeometry(const Sizes* s, const RadialPartitioning* r, int ns) : FourierCoeffs(s, r, r->nsMinF1, r->nsMaxF1, ns), @@ -351,6 +364,47 @@ void FourierGeometry::extrapolateTowardsAxis() { } // n } +void FourierGeometry::extrapolateTowardsAxisTranspose() { + if (nsMin_ > 0) { + return; + } + int axis = 0; + int firstSurface = 1; + for (int n = 0; n < s_.ntor + 1; ++n) { + int m0 = 0; + int m1 = 1; + int axis0 = (axis * s_.mpol + m0) * (s_.ntor + 1) + n; + int axis1 = (axis * s_.mpol + m1) * (s_.ntor + 1) + n; + int firstSurface0 = (firstSurface * s_.mpol + m0) * (s_.ntor + 1) + n; + int firstSurface1 = (firstSurface * s_.mpol + m1) * (s_.ntor + 1) + n; + + auto fold = [](std::span c, int axisIdx, int firstIdx) { + c[firstIdx] += c[axisIdx]; + c[axisIdx] = 0.0; + }; + fold(rmncc, axis1, firstSurface1); + fold(zmnsc, axis1, firstSurface1); + fold(lmnsc, axis1, firstSurface1); + if (s_.lthreed) { + fold(rmnss, axis1, firstSurface1); + fold(zmncs, axis1, firstSurface1); + fold(lmncs, axis1, firstSurface1); + fold(lmncs, axis0, firstSurface0); + } + if (s_.lasym) { + fold(rmnsc, axis1, firstSurface1); + fold(zmncc, axis1, firstSurface1); + fold(lmncc, axis1, firstSurface1); + fold(lmncc, axis0, firstSurface0); + if (s_.lthreed) { + fold(rmncs, axis1, firstSurface1); + fold(zmnss, axis1, firstSurface1); + fold(lmnss, axis1, firstSurface1); + } + } + } // n +} + void FourierGeometry::ComputeSpectralWidth( const FourierBasisFastPoloidal& fourier_basis, RadialProfiles& m_radial_profiles, const int p, const int q) const { @@ -362,80 +416,29 @@ void FourierGeometry::ComputeSpectralWidth( m_radial_profiles.spectral_width[nsMin_ - r_.nsMinF1] = 1.0; } - // compute only on unique full-grid points - for (int jF = minimum_j; jF < nsMax_; ++jF) { - double spectral_width_numerator = 0.0; - double spectral_width_denominator = 0.0; + const int coefficients_per_surface = s_.mpol * (s_.ntor + 1); - // note that we exclude m = 0 - for (int m = 1; m < s_.mpol; ++m) { - for (int n = 0; n < s_.ntor + 1; ++n) { - int fourier_index = ((jF - nsMin_) * s_.mpol + m) * (s_.ntor + 1) + n; - - const double basis_norm = - fourier_basis.mscale[m] * fourier_basis.nscale[n]; - - // Use Eigen for vectorized norm computation - Eigen::Vector4d r_coefficients = Eigen::Vector4d::Zero(); - Eigen::Vector4d z_coefficients = Eigen::Vector4d::Zero(); - int basis_dimension = 0; + const std::span mscale(fourier_basis.mscale.data(), + fourier_basis.mscale.size()); + const std::span nscale(fourier_basis.nscale.data(), + fourier_basis.nscale.size()); - r_coefficients[basis_dimension] = rmncc[fourier_index]; - z_coefficients[basis_dimension] = zmnsc[fourier_index]; - basis_dimension++; - - // CONVERT FROM INTERNAL XC REPRESENTATION FOR m=1 MODES, - // R+(at rsc) = .5(rsc + zcc), - // R-(at zcc) = .5(rsc - zcc), - // TO REQUIRED rsc, zcc FORMS - if (s_.lthreed) { - if (m == 1) { - const double r_plus = rmnss[fourier_index]; - const double r_minus = zmncs[fourier_index]; - // rmnss - r_coefficients[basis_dimension] = r_plus + r_minus; - // zmncs - z_coefficients[basis_dimension] = r_plus - r_minus; - } else { - r_coefficients[basis_dimension] = rmnss[fourier_index]; - z_coefficients[basis_dimension] = zmncs[fourier_index]; - } - basis_dimension++; - } - if (s_.lasym) { - if (m == 1) { - const double r_plus = rmnsc[fourier_index]; - const double r_minus = zmncc[fourier_index]; - // rmnsc - r_coefficients[basis_dimension] = r_plus + r_minus; - // zmncc - z_coefficients[basis_dimension] = r_plus - r_minus; - } else { - r_coefficients[basis_dimension] = rmnsc[fourier_index]; - z_coefficients[basis_dimension] = zmncc[fourier_index]; - } - basis_dimension++; - } - - if (s_.lasym && s_.lthreed) { - r_coefficients[basis_dimension] = rmncs[fourier_index]; - z_coefficients[basis_dimension] = zmnss[fourier_index]; - basis_dimension++; - } - - // Vectorized squared norm computation - double coefficient_norm = - r_coefficients.head(basis_dimension).squaredNorm() + - z_coefficients.head(basis_dimension).squaredNorm(); - coefficient_norm *= basis_norm * basis_norm; - - spectral_width_numerator += coefficient_norm * std::pow(m, p + q); - spectral_width_denominator += coefficient_norm * std::pow(m, p); - } // m - } // n + // compute only on unique full-grid points + for (int jF = minimum_j; jF < nsMax_; ++jF) { + const int surface_offset = (jF - nsMin_) * coefficients_per_surface; + + const SurfaceFourierGeometry surface = { + .rmncc = SurfaceOf(rmncc, surface_offset, coefficients_per_surface), + .rmnss = SurfaceOf(rmnss, surface_offset, coefficients_per_surface), + .rmnsc = SurfaceOf(rmnsc, surface_offset, coefficients_per_surface), + .rmncs = SurfaceOf(rmncs, surface_offset, coefficients_per_surface), + .zmnsc = SurfaceOf(zmnsc, surface_offset, coefficients_per_surface), + .zmncs = SurfaceOf(zmncs, surface_offset, coefficients_per_surface), + .zmncc = SurfaceOf(zmncc, surface_offset, coefficients_per_surface), + .zmnss = SurfaceOf(zmnss, surface_offset, coefficients_per_surface)}; m_radial_profiles.spectral_width[jF - r_.nsMinF1] = - spectral_width_numerator / spectral_width_denominator; + SpectralWidth(surface, s_, mscale, nscale, p, q); } // jF } // ComputeSpectralWidth diff --git a/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.h b/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.h index fa65a020f..1e6059b4a 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.h +++ b/src/vmecpp/cpp/vmecpp/vmec/fourier_geometry/fourier_geometry.h @@ -17,14 +17,14 @@ namespace vmecpp { class FourierGeometry : public FourierCoeffs { public: - FourierGeometry(const Sizes *s, const RadialPartitioning *r, int ns); - FourierGeometry(const FourierGeometry &other); - FourierGeometry &operator=(const FourierGeometry &other); - FourierGeometry(FourierGeometry &&other) noexcept; - FourierGeometry &operator=(FourierGeometry &&other) noexcept; + FourierGeometry(const Sizes* s, const RadialPartitioning* r, int ns); + FourierGeometry(const FourierGeometry& other); + FourierGeometry& operator=(const FourierGeometry& other); + FourierGeometry(FourierGeometry&& other) noexcept; + FourierGeometry& operator=(FourierGeometry&& other) noexcept; - void interpFromBoundaryAndAxis(const FourierBasisFastPoloidal &t, - const Boundaries &b, const RadialProfiles &p); + void interpFromBoundaryAndAxis(const FourierBasisFastPoloidal& t, + const Boundaries& b, const RadialProfiles& p); // Initialize the state of this FourierGeometry with the given Fourier // coefficients. If a Boundaries object is specified (defaults to nullptr; in @@ -33,18 +33,23 @@ class FourierGeometry : public FourierCoeffs { // instead of from the Fourier coefficient matrices. // This latter use case applies to fixed-boundary hot-restart operation of // VMEC++. - void InitFromState(const FourierBasisFastPoloidal &fb, - const RowMatrixXd &rmnc, const RowMatrixXd &zmns, - const RowMatrixXd &lmns_full, const RadialProfiles &p, - const VmecConstants &constants, - const Boundaries *b = nullptr); + void InitFromState(const FourierBasisFastPoloidal& fb, + const RowMatrixXd& rmnc, const RowMatrixXd& zmns, + const RowMatrixXd& lmns_full, const RadialProfiles& p, + const VmecConstants& constants, + const Boundaries* b = nullptr); void extrapolateTowardsAxis(); + // Transpose of extrapolateTowardsAxis: the forward copies surface-1 m=1 (and + // m=0 lambda) coefficients onto the axis, so the adjoint folds the axis + // cotangent back into surface 1 and zeroes the axis. + void extrapolateTowardsAxisTranspose(); + // Compute the spectral width of the R and Z Fourier coefficients // and write it into the spectral_width vector in the given RadialProfiles. - void ComputeSpectralWidth(const FourierBasisFastPoloidal &fourier_basis, - RadialProfiles &m_radial_profiles, int p = 4, + void ComputeSpectralWidth(const FourierBasisFastPoloidal& fourier_basis, + RadialProfiles& m_radial_profiles, int p = 4, int q = 1) const; // N.B. all raw pointers below are non-owning pointers to data in diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel index 7f5c618d6..132a8f590 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/BUILD.bazel @@ -82,7 +82,7 @@ cc_test( cc_library( name = "ideal_mhd_model", srcs = ["ideal_mhd_model.cc"], -# The shared force-chain kernels and the autodiff composition/JVP header are + # The shared force-chain kernels and the autodiff composition/JVP header are # header-only and included by ideal_mhd_model.cc; declare whatever exists on # this branch so the bazel sandbox can see them. The Enzyme JVP entry point # (exact_force_jvp.cc) is built only by the CMake VMECPP_ENABLE_ENZYME path, @@ -91,6 +91,7 @@ cc_library( "*_kernel.h", "local_force_composition.h", "exact_force_jvp.h", + "exact_force_vjp.h", ]), visibility = ["//visibility:public"], deps = [ diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h index b5078a6f1..56fb0a538 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h @@ -63,6 +63,148 @@ inline void AddConstraintForces( } // jF } +// Fourier-space bandpass of the constraint force: forward transform gConEff to +// the (m, n) coefficients gsc/gcs scaled by tcon, then inverse transform back +// to real space scaled by faccon[m]. Bandpass keeps m in [1, mpol-1). The axis +// surface has no poloidal angle and is skipped. Allocation-free over flat +// buffers with explicit reductions (no Eigen temporaries), so it differentiates +// under Enzyme; the free function vmecpp::deAliasConstraintForce wraps this +// with the partition/basis structs. Basis layout: sinmui/cosmui/sinmu/cosmu +// indexed m*nThetaReduced + l; cosnv/sinnv indexed k*(nnyq2+1) + n; +// gConEff/gCon indexed +// ((jF-nsMinF)*nZeta + k)*nThetaEff + l; all Fourier scratch arrays have +// length ntor+1. In the non-symmetric case, gConAsym has nZeta*nThetaEff +// entries for one surface and refl has nThetaReduced entries. +inline void ComputeDeAliasConstraintForce( + const double* gConEff, const double* faccon, const double* tcon, + const double* sinmui, const double* cosmui, const double* cosnv, + const double* sinnv, const double* sinmu, const double* cosmu, int nsMinF, + int nsMaxF, int nZeta, int nThetaEff, int nThetaReduced, int nThetaEven, + int mpol, int ntor, int nnyq2, bool lasym, double* m_gsc, double* m_gcs, + double* m_gcc, double* m_gss, double* m_gConAsym, double* m_refl, + double* m_gCon) { + for (int i = 0; i < (nsMaxF - nsMinF) * nZeta * nThetaEff; ++i) { + m_gCon[i] = 0.0; + } + const int jMin = (nsMinF == 0) ? 1 : 0; + for (int jF = (jMin > nsMinF ? jMin : nsMinF); jF < nsMaxF; ++jF) { + if (lasym) { + for (int i = 0; i < nZeta * nThetaEff; ++i) { + m_gConAsym[i] = 0.0; + } + } + for (int m = 1; m < mpol - 1; ++m) { + for (int n = 0; n < ntor + 1; ++n) { + m_gsc[n] = 0.0; + m_gcs[n] = 0.0; + if (lasym) { + m_gcc[n] = 0.0; + m_gss[n] = 0.0; + } + } + for (int k = 0; k < nZeta; ++k) { + const int kl_base = ((jF - nsMinF) * nZeta + k) * nThetaEff; + const int ml_base = m * nThetaReduced; + double w0 = 0.0; + double w1 = 0.0; + for (int l = 0; l < nThetaReduced; ++l) { + w0 += gConEff[kl_base + l] * sinmui[ml_base + l]; + w1 += gConEff[kl_base + l] * cosmui[ml_base + l]; + } + const double tc = tcon[jF - nsMinF]; + if (!lasym) { + for (int n = 0; n < ntor + 1; ++n) { + const int idx_kn = k * (nnyq2 + 1) + n; + m_gsc[n] += cosnv[idx_kn] * w0 * tc; + m_gcs[n] += sinnv[idx_kn] * w1 * tc; + } + } else { + const int k_rev = (nZeta - k) % nZeta; + const int refl_base = ((jF - nsMinF) * nZeta + k_rev) * nThetaEff; + for (int l = 0; l < nThetaReduced; ++l) { + const int l_rev = (nThetaEven - l) % nThetaEven; + m_refl[l] = gConEff[refl_base + l_rev]; + } + double w3 = 0.0; + double w4 = 0.0; + for (int l = 0; l < nThetaReduced; ++l) { + w3 += m_refl[l] * cosmui[ml_base + l]; + w4 += m_refl[l] * sinmui[ml_base + l]; + } + for (int n = 0; n < ntor + 1; ++n) { + const int idx_kn = k * (nnyq2 + 1) + n; + const double cosnv_value = cosnv[idx_kn]; + const double sinnv_value = sinnv[idx_kn]; + m_gcs[n] += 0.5 * tc * sinnv_value * (w1 - w3); + m_gsc[n] += 0.5 * tc * cosnv_value * (w0 - w4); + m_gss[n] += 0.5 * tc * sinnv_value * (w0 + w4); + m_gcc[n] += 0.5 * tc * cosnv_value * (w1 + w3); + } + } + } + + for (int k = 0; k < nZeta; ++k) { + const int kn_base = k * (nnyq2 + 1); + double w0 = 0.0; + double w1 = 0.0; + double a0 = 0.0; + double a1 = 0.0; + for (int n = 0; n < ntor + 1; ++n) { + w0 += m_gsc[n] * cosnv[kn_base + n]; + w1 += m_gcs[n] * sinnv[kn_base + n]; + if (lasym) { + a0 += m_gcc[n] * cosnv[kn_base + n]; + a1 += m_gss[n] * sinnv[kn_base + n]; + } + } + for (int l = 0; l < nThetaReduced; ++l) { + const int idx_kl = ((jF - nsMinF) * nZeta + k) * nThetaEff + l; + const int idx_ml = m * nThetaReduced + l; + m_gCon[idx_kl] += + faccon[m] * (w0 * sinmu[idx_ml] + w1 * cosmu[idx_ml]); + if (lasym) { + const int within = k * nThetaEff + l; + m_gConAsym[within] += + faccon[m] * (a0 * cosmu[idx_ml] + a1 * sinmu[idx_ml]); + } + } + } + } + if (lasym) { + const int surface_base = (jF - nsMinF) * nZeta * nThetaEff; + for (int k = 0; k < nZeta; ++k) { + const int k_rev = (nZeta - k) % nZeta; + for (int l = nThetaReduced; l < nThetaEven; ++l) { + const int target = surface_base + k * nThetaEff + l; + const int reflected = k_rev * nThetaEff + (nThetaEven - l); + m_gCon[target] = + -m_gCon[surface_base + reflected] + m_gConAsym[reflected]; + } + } + for (int k = 0; k < nZeta; ++k) { + for (int l = 0; l < nThetaReduced; ++l) { + const int within = k * nThetaEff + l; + m_gCon[surface_base + within] += m_gConAsym[within]; + } + } + } + } +} + +// Symmetric convenience overload used by standalone tests and benchmarks. +inline void ComputeDeAliasConstraintForce( + const double* gConEff, const double* faccon, const double* tcon, + const double* sinmui, const double* cosmui, const double* cosnv, + const double* sinnv, const double* sinmu, const double* cosmu, int nsMinF, + int nsMaxF, int nZeta, int nThetaEff, int nThetaReduced, int mpol, int ntor, + int nnyq2, double* m_gsc, double* m_gcs, double* m_gCon) { + ComputeDeAliasConstraintForce( + gConEff, faccon, tcon, sinmui, cosmui, cosnv, sinnv, sinmu, cosmu, nsMinF, + nsMaxF, nZeta, nThetaEff, nThetaReduced, 2 * (nThetaReduced - 1), mpol, + ntor, nnyq2, false, m_gsc, m_gcs, nullptr, nullptr, nullptr, nullptr, + m_gCon); +} + } // namespace vmecpp #endif // VMECPP_VMEC_IDEAL_MHD_MODEL_CONSTRAINT_FORCE_KERNEL_H_ diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc new file mode 100644 index 000000000..2880969d1 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.cc @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT + +// Forward-mode Jacobian-vector product of the local force-density composition. +// This translation unit is compiled with the Clang/Enzyme plugin so that the +// __enzyme_fwddiff call differentiates ComputeLocalForceDensity. The rest of +// VMEC++ is compiled normally; the exact Hessian-vector product calls into here +// for the single nonlinear pass and wraps it with the linear spectral +// transforms. + +#include "vmecpp/vmec/ideal_mhd_model/exact_force_jvp.h" + +#include "vmecpp/common/enzyme/enzyme.h" +#include "vmecpp/vmec/ideal_mhd_model/local_force_composition.h" + +namespace vmecpp { + +void ExactForceDensityJvp(const double* geom, const double* dgeom, double* work, + double* dwork, double* force, double* dforce, + const LocalForceComposition* c) { + __enzyme_fwddiff(reinterpret_cast(ComputeLocalForceDensity), + enzyme_dup, geom, dgeom, enzyme_dup, work, dwork, + enzyme_dup, force, dforce, enzyme_const, c); +} + +} // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.h new file mode 100644 index 000000000..66e3c46be --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_jvp.h @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#ifndef VMECPP_VMEC_IDEAL_MHD_MODEL_EXACT_FORCE_JVP_H_ +#define VMECPP_VMEC_IDEAL_MHD_MODEL_EXACT_FORCE_JVP_H_ + +#include "vmecpp/vmec/ideal_mhd_model/local_force_composition.h" + +namespace vmecpp { + +// Forward-mode Jacobian-vector product of ComputeLocalForceDensity: given the +// geometry primal and tangent, returns the force-density tangent in dforce in +// one Enzyme forward pass. work/dwork and force/dforce are caller-owned scratch +// sized as ComputeLocalForceDensity requires. Defined in exact_force_jvp.cc, +// which is compiled with the Clang/Enzyme plugin. +void ExactForceDensityJvp(const double* geom, const double* dgeom, double* work, + double* dwork, double* force, double* dforce, + const LocalForceComposition* c); + +} // namespace vmecpp + +#endif // VMECPP_VMEC_IDEAL_MHD_MODEL_EXACT_FORCE_JVP_H_ diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.cc new file mode 100644 index 000000000..8d217a26c --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.cc @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT + +// Reverse-mode vector-Jacobian product of the local force-density composition. +// Compiled with the Clang/Enzyme plugin so that __enzyme_autodiff +// differentiates ComputeLocalForceDensity in reverse. This is the transpose +// J_g^T of the forward pass in exact_force_jvp.cc; the transposed exact +// Hessian-vector product wraps it with the transposed linear spectral +// transforms. + +#include "vmecpp/vmec/ideal_mhd_model/exact_force_vjp.h" + +#include "vmecpp/common/enzyme/enzyme.h" +#include "vmecpp/vmec/ideal_mhd_model/local_force_composition.h" + +namespace vmecpp { + +void ExactForceDensityVjp(const double* geom, double* geom_bar, double* work, + double* work_bar, double* force, double* force_bar, + const LocalForceComposition* c) { + // force_bar carries the output cotangent (seed); geom_bar accumulates + // J_g^T force_bar. geom_bar and work_bar are zeroed by the caller. + __enzyme_autodiff(reinterpret_cast(ComputeLocalForceDensity), + enzyme_dup, geom, geom_bar, enzyme_dup, work, work_bar, + enzyme_dup, force, force_bar, enzyme_const, c); +} + +} // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.h new file mode 100644 index 000000000..97b07df47 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/exact_force_vjp.h @@ -0,0 +1,26 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#ifndef VMECPP_VMEC_IDEAL_MHD_MODEL_EXACT_FORCE_VJP_H_ +#define VMECPP_VMEC_IDEAL_MHD_MODEL_EXACT_FORCE_VJP_H_ + +#include "vmecpp/vmec/ideal_mhd_model/local_force_composition.h" + +namespace vmecpp { + +// Reverse-mode vector-Jacobian product of ComputeLocalForceDensity: given the +// geometry primal and a force-density cotangent in force_bar, accumulates the +// geometry cotangent J_g^T force_bar into geom_bar in one Enzyme reverse pass. +// geom_bar and work_bar must be zeroed by the caller; work/work_bar/force are +// caller-owned scratch sized as ComputeLocalForceDensity requires. This is the +// transpose of ExactForceDensityJvp and is the nonlinear factor of the +// transposed exact Hessian-vector product. Defined in exact_force_vjp.cc, which +// is compiled with the Clang/Enzyme plugin. +void ExactForceDensityVjp(const double* geom, double* geom_bar, double* work, + double* work_bar, double* force, double* force_bar, + const LocalForceComposition* c); + +} // namespace vmecpp + +#endif // VMECPP_VMEC_IDEAL_MHD_MODEL_EXACT_FORCE_VJP_H_ diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc index bebacbd60..fdf17746a 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc @@ -23,8 +23,11 @@ #include "vmecpp/vmec/ideal_mhd_model/bco_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/bcontra_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/exact_force_jvp.h" +#include "vmecpp/vmec/ideal_mhd_model/exact_force_vjp.h" #include "vmecpp/vmec/ideal_mhd_model/jacobian_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/lambda_force_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/local_force_composition.h" #include "vmecpp/vmec/ideal_mhd_model/metric_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/mhdforce_kernel.h" #include "vmecpp/vmec/ideal_mhd_model/pressure_kernel.h" @@ -96,161 +99,44 @@ void vmecpp::deAliasConstraintForce( const Eigen::VectorXd& faccon, const Eigen::VectorXd& tcon, const Eigen::VectorXd& gConEff, Eigen::VectorXd& m_gsc, Eigen::VectorXd& m_gcs, Eigen::VectorXd& m_gCon) { - absl::c_fill_n(m_gCon, (rp.nsMaxF - rp.nsMinF) * s_.nZnT, 0); - - // For non-stellarator-symmetric runs the spectral-condensation constraint - // force carries an antisymmetric parity as well. Following educational_VMEC - // alias.f90, accumulate the antisymmetric toroidal coefficients (gcc, gss) - // and the antisymmetric real-space force (gcona) per flux surface, then - // extend the odd-parity constraint force onto the full poloidal interval - // (symrzl convention). The stellarator-symmetric path is left unchanged. - const bool lasym = s_.lasym; - Eigen::VectorXd gcc, gss, gcona, refl; - if (lasym) { + if (!s_.lasym) { + ComputeDeAliasConstraintForce( + gConEff.data(), faccon.data(), tcon.data(), fb.sinmui.data(), + fb.cosmui.data(), fb.cosnv.data(), fb.sinnv.data(), fb.sinmu.data(), + fb.cosmu.data(), rp.nsMinF, rp.nsMaxF, s_.nZeta, s_.nThetaEff, + s_.nThetaReduced, s_.mpol, s_.ntor, s_.nnyq2, m_gsc.data(), + m_gcs.data(), m_gCon.data()); + return; + } + Eigen::VectorXd gcc; + Eigen::VectorXd gss; + Eigen::VectorXd gConAsym; + Eigen::VectorXd refl; + if (s_.lasym) { gcc.setZero(s_.ntor + 1); gss.setZero(s_.ntor + 1); - gcona.setZero(s_.nZnT); + gConAsym.setZero(s_.nZnT); refl.setZero(s_.nThetaReduced); } + deAliasConstraintForce(rp, fb, s_, faccon, tcon, gConEff, m_gsc, m_gcs, gcc, + gss, gConAsym, refl, m_gCon); +} - // no constraint on axis --> has no poloidal angle - int jMin = 0; - if (rp.nsMinF == 0) { - jMin = 1; - } - - for (int jF = std::max(jMin, rp.nsMinF); jF < rp.nsMaxF; ++jF) { - if (lasym) { - absl::c_fill_n(gcona, s_.nZnT, 0); - } - for (int m = 1; m < s_.mpol - 1; ++m) { - absl::c_fill_n(m_gsc, s_.ntor + 1, 0); - absl::c_fill_n(m_gcs, s_.ntor + 1, 0); - if (lasym) { - absl::c_fill_n(gcc, s_.ntor + 1, 0); - absl::c_fill_n(gss, s_.ntor + 1, 0); - } - - for (int k = 0; k < s_.nZeta; ++k) { - // fwd transform in poloidal direction - // integrate poloidally to get m-th poloidal Fourier coefficient - const int kl_base = ((jF - rp.nsMinF) * s_.nZeta + k) * s_.nThetaEff; - const int ml_base = m * s_.nThetaReduced; - - auto gConEff_seg = Eigen::Map( - gConEff.data() + kl_base, s_.nThetaReduced); - auto sinmui_seg = fb.sinmui.segment(ml_base, s_.nThetaReduced); - auto cosmui_seg = fb.cosmui.segment(ml_base, s_.nThetaReduced); - - double w0 = gConEff_seg.dot(sinmui_seg); - double w1 = gConEff_seg.dot(cosmui_seg); - - const double tc = tcon[jF - rp.nsMinF]; - - if (!lasym) { - // forward Fourier transform in toroidal direction for full set of - // mode numbers (n = 0, 1, ..., ntor) - for (int n = 0; n < s_.ntor + 1; ++n) { - int idx_kn = k * (s_.nnyq2 + 1) + n; - - // NOTE: `tcon` comes into play here - m_gsc[n] += fb.cosnv[idx_kn] * w0 * tc; - m_gcs[n] += fb.sinnv[idx_kn] * w1 * tc; - } - } else { - // effective force at the reflected point (theta -> 2pi - theta, - // zeta -> 2pi - zeta), sampled over the reduced poloidal interval - const int kRev = (s_.nZeta - k) % s_.nZeta; - const int refl_base = - ((jF - rp.nsMinF) * s_.nZeta + kRev) * s_.nThetaEff; - for (int l = 0; l < s_.nThetaReduced; ++l) { - const int lRev = (s_.nThetaEven - l) % s_.nThetaEven; - refl[l] = gConEff[refl_base + lRev]; - } - double w3 = refl.dot(cosmui_seg); - double w4 = refl.dot(sinmui_seg); - for (int n = 0; n < s_.ntor + 1; ++n) { - int idx_kn = k * (s_.nnyq2 + 1) + n; - const double cosnv = fb.cosnv[idx_kn]; - const double sinnv = fb.sinnv[idx_kn]; - m_gcs[n] += 0.5 * tc * sinnv * (w1 - w3); - m_gsc[n] += 0.5 * tc * cosnv * (w0 - w4); - gss[n] += 0.5 * tc * sinnv * (w0 + w4); - gcc[n] += 0.5 * tc * cosnv * (w1 + w3); - } - } - } // k - - // ------------------------------------------ - // need to "wait" (= finish k loop) here - // to get Fourier coefficients fully defined! - // ------------------------------------------ - - // inverse Fourier-transform from reduced set of mode numbers - for (int k = 0; k < s_.nZeta; ++k) { - // collect contribution to current grid point from n-th toroidal mode - const int kn_base = k * (s_.nnyq2 + 1); - auto cosnv_seg = fb.cosnv.segment(kn_base, s_.ntor + 1); - auto sinnv_seg = fb.sinnv.segment(kn_base, s_.ntor + 1); - - auto m_gsc_seg = - Eigen::Map(m_gsc.data(), s_.ntor + 1); - auto m_gcs_seg = - Eigen::Map(m_gcs.data(), s_.ntor + 1); - - double w0 = m_gsc_seg.dot(cosnv_seg); - double w1 = m_gcs_seg.dot(sinnv_seg); - - double a0 = 0.0; - double a1 = 0.0; - if (lasym) { - auto gcc_seg = - Eigen::Map(gcc.data(), s_.ntor + 1); - auto gss_seg = - Eigen::Map(gss.data(), s_.ntor + 1); - a0 = gcc_seg.dot(cosnv_seg); - a1 = gss_seg.dot(sinnv_seg); - } - - // inv transform in poloidal direction - for (int l = 0; l < s_.nThetaReduced; ++l) { - int idx_kl = ((jF - rp.nsMinF) * s_.nZeta + k) * s_.nThetaEff + l; - const int idx_ml = m * s_.nThetaReduced + l; - - // NOTE: `faccon` comes into play here - m_gCon[idx_kl] += - faccon[m] * (w0 * fb.sinmu[idx_ml] + w1 * fb.cosmu[idx_ml]); - if (lasym) { - const int within = k * s_.nThetaEff + l; - gcona[within] += - faccon[m] * (a0 * fb.cosmu[idx_ml] + a1 * fb.sinmu[idx_ml]); - } - } // l - } // k - } // m - - if (lasym) { - // Extend the odd-parity constraint force onto theta in [pi, 2pi[ as the - // parity-signed reflection (-sym + asym), then add the antisymmetric - // piece on [0, pi] (alias.f90 / symrzl). The extension reads the reduced - // interval, which still holds the pure symmetric values here. - const int surf_base = (jF - rp.nsMinF) * s_.nZnT; - for (int k = 0; k < s_.nZeta; ++k) { - const int kRev = (s_.nZeta - k) % s_.nZeta; - for (int l = s_.nThetaReduced; l < s_.nThetaEven; ++l) { - const int jl = surf_base + k * s_.nThetaEff + l; - const int within_rev = kRev * s_.nThetaEff + (s_.nThetaEven - l); - m_gCon[jl] = -m_gCon[surf_base + within_rev] + gcona[within_rev]; - } - } - for (int k = 0; k < s_.nZeta; ++k) { - for (int l = 0; l < s_.nThetaReduced; ++l) { - const int within = k * s_.nThetaEff + l; - m_gCon[surf_base + within] += gcona[within]; - } - } - } - } +void vmecpp::deAliasConstraintForce( + const vmecpp::RadialPartitioning& rp, + const vmecpp::FourierBasisFastPoloidal& fb, const vmecpp::Sizes& s_, + const Eigen::VectorXd& faccon, const Eigen::VectorXd& tcon, + const Eigen::VectorXd& gConEff, Eigen::VectorXd& m_gsc, + Eigen::VectorXd& m_gcs, Eigen::VectorXd& m_gcc, Eigen::VectorXd& m_gss, + Eigen::VectorXd& m_gConAsym, Eigen::VectorXd& m_refl, + Eigen::VectorXd& m_gCon) { + ComputeDeAliasConstraintForce( + gConEff.data(), faccon.data(), tcon.data(), fb.sinmui.data(), + fb.cosmui.data(), fb.cosnv.data(), fb.sinnv.data(), fb.sinmu.data(), + fb.cosmu.data(), rp.nsMinF, rp.nsMaxF, s_.nZeta, s_.nThetaEff, + s_.nThetaReduced, s_.nThetaEven, s_.mpol, s_.ntor, s_.nnyq2, s_.lasym, + m_gsc.data(), m_gcs.data(), m_gcc.data(), m_gss.data(), m_gConAsym.data(), + m_refl.data(), m_gCon.data()); } namespace vmecpp { @@ -420,6 +306,12 @@ IdealMhdModel::IdealMhdModel( gConEff.setZero(nrztIncludingBoundary); gsc.setZero(s_.ntor + 1); gcs.setZero(s_.ntor + 1); + if (s_.lasym) { + gcc.setZero(s_.ntor + 1); + gss.setZero(s_.ntor + 1); + gConAsym.setZero(s_.nZnT); + refl.setZero(s_.nThetaReduced); + } gCon.setZero(nrztIncludingBoundary); frcon_e.setZero(nrzt); @@ -479,10 +371,33 @@ IdealMhdModel::IdealMhdModel( } void IdealMhdModel::setFromINDATA(int ncurr, double adiabaticIndex, - double tcon0) { + double tcon0, bool lforbal) { this->ncurr = ncurr; this->adiabaticIndex = adiabaticIndex; this->tcon0 = tcon0; + // The m=1 trig weights below are built on the reduced poloidal grid, so the + // force-balance modification is restricted to the stellarator-symmetric case. + this->lforbal = lforbal && !s_.lasym; + + if (this->lforbal) { + // m=1,n=0 force-balance factors (full grid) and the m=1 trig weights on the + // (theta, zeta) grid. cos01 = cos(u) * mscale(1), sin01 = -sin(u) * + // mscale(1), i.e. educational_VMEC fixaray cos01/sin01 at m=1, n=0. + const int num_full = r_.nsMaxF - r_.nsMinF; + rzu_fac.setZero(num_full); + rru_fac.setZero(num_full); + frcc_fac.setZero(num_full); + fzsc_fac.setZero(num_full); + cos01.setZero(s_.nZnT); + sin01.setZero(s_.nZnT); + for (int kl = 0; kl < s_.nZnT; ++kl) { + const int l = kl % s_.nThetaEff; + // cosmu/sinmu are laid out as [m * nThetaReduced + l] and already include + // mscale(m); the m=1 row is the cos01/sin01 weight (symmetric grid). + cos01[kl] = t_.cosmu[s_.nThetaReduced + l]; + sin01[kl] = -t_.sinmu[s_.nThetaReduced + l]; + } + } } void IdealMhdModel::evalFResInvar(const Eigen::Vector3d& localFResInvar) { @@ -2171,16 +2086,32 @@ void IdealMhdModel::updateRadialPreconditioner() { updateLambdaPreconditioner(); // compute preconditioning matrix for R - // TODO(jons): also cos01, rzu_fac for lforbal computePreconditioningMatrix(zs, zu12, zu_e, zu_o, z1_o, arm, ard, brm, brd, - cxd); + cxd, cos01, rzu_fac); // compute preconditioning matrix for Z - // TODO(jons): also sin01, rru_fac for lforbal computePreconditioningMatrix(rs, ru12, ru_e, ru_o, r1_o, azm, azd, bzm, bzd, - cxd); - - // (compute stuff for lforbal: scaleEqFactor --> later) + cxd, sin01, rru_fac); + + if (lforbal) { + // Form the m=1,n=0 force-balance factors from the R/Z preconditioner + // diagonals (educational_VMEC bcovar): scale by sqrt(s), take the + // reciprocals frcc_fac/fzsc_fac, then halve rzu_fac/rru_fac. Interior + // full-grid surfaces only. + for (int jF = r_.nsMinF; jF < r_.nsMaxF; ++jF) { + if (jF == 0 || jF >= m_fc_.ns - 1) { + continue; + } + const int i = jF - r_.nsMinF; + const double sj = m_p_.sqrtSF[jF - r_.nsMinF1]; + rzu_fac[i] *= sj; + rru_fac[i] *= sj; + frcc_fac[i] = 1.0 / rzu_fac[i]; + rzu_fac[i] *= 0.5; + fzsc_fac[i] = -1.0 / rru_fac[i]; + rru_fac[i] *= 0.5; + } + } } void IdealMhdModel::updateLambdaPreconditioner() { @@ -2297,10 +2228,20 @@ void IdealMhdModel::computePreconditioningMatrix( const Eigen::VectorXd& xs, const Eigen::VectorXd& xu12, const Eigen::VectorXd& xu_e, const Eigen::VectorXd& xu_o, const Eigen::VectorXd& x1_o, Eigen::VectorXd& m_axm, Eigen::VectorXd& m_axd, - Eigen::VectorXd& m_bxm, Eigen::VectorXd& m_bxd, Eigen::VectorXd& m_cxd) { + Eigen::VectorXd& m_bxm, Eigen::VectorXd& m_bxd, Eigen::VectorXd& m_cxd, + const Eigen::VectorXd& trigmult, Eigen::VectorXd& m_eqfactor) { // zs, zu12, zu, z1 --> arm, ard, brm, brd, cxd // rs, ru12, ru, r1 --> azm, azd, bzm, bzd, cxd + // lforbal: when m_eqfactor is sized, accumulate the flux-averaged + // force-balance weight (temp_h, half-grid) using the m=1 trig weights + // trigmult, and assemble the force-balance scale factor below. + const bool do_eqfactor = m_eqfactor.size() > 0; + Eigen::VectorXd temp_h; + if (do_eqfactor) { + temp_h.setZero(r_.nsMaxH - r_.nsMinH); + } + // restored in v8.51 // TODO(jons): what is this? double pFactor = -4.0; @@ -2372,6 +2313,14 @@ void IdealMhdModel::computePreconditioningMatrix( // --> essentially, 0.25 * pFactor simply introduces a (-1) here! cx[jH - r_.nsMinH] += 0.25 * pFactor * bsupv[iHalf] * bsupv[iHalf] * gsqrt[iHalf] * s_.wInt[l]; + + if (do_eqfactor) { + // Fortran precondn: temp(js) += (pfactor*r12*bsq*wint)*trigmult*xu12. + // pTau = pFactor*r12*totalPressure/tau*wInt, so pTau*tau equals that + // weight (r0scale == 1, as noted for cx above). + temp_h[jH - r_.nsMinH] += + pTau * tau[iHalf] * trigmult[kl] * xu12[iHalf]; + } } // kl } // jH @@ -2418,6 +2367,25 @@ void IdealMhdModel::computePreconditioningMatrix( m_cxd[jF - r_.nsMinF] = (jF > 0 ? cx[jH_i] : 0.0) + (jF < m_fc_.ns - 1 ? cx[jH_o] : 0.0); } + + if (do_eqfactor) { + // Flux-averaged force-balance scale factor (educational_VMEC precondn): + // temp /= vp (half grid), couple js and js+1 onto the full grid (signgs), + // then eqfactor = axd(m=1) * hs^2 / temp on interior full-grid surfaces. + const double hs2 = m_fc_.deltaS * m_fc_.deltaS; + for (int jF = r_.nsMinF; jF < r_.nsMaxF; ++jF) { + const int jH_i = jF - 1 - r_.nsMinH; + const int jH_o = jF - r_.nsMinH; + double temp_f = + (jF > 0 ? temp_h[jH_i] / m_p_.dVdsH[jH_i] : 0.0) + + (jF < m_fc_.ns - 1 ? temp_h[jH_o] / m_p_.dVdsH[jH_o] : 0.0); + temp_f *= signOfJacobian; + const double axd_m1 = m_axd[(jF - r_.nsMinF) * 2 + kOddParity]; + m_eqfactor[jF - r_.nsMinF] = + (jF > 0 && jF < m_fc_.ns - 1 && temp_f != 0.0) ? axd_m1 * hs2 / temp_f + : 0.0; + } + } } /** @@ -2425,6 +2393,12 @@ void IdealMhdModel::computePreconditioningMatrix( * Note that this needs to have the radial preconditioner updated. */ absl::Status IdealMhdModel::constraintForceMultiplier() { + // Freeze: reuse the existing tcon so the raw force is a function of the state + // alone, matching the exact HVP (which freezes tcon). Requires a prior + // unfrozen evaluation to have populated tcon. + if (freeze_constraint_multiplier_) { + return absl::OkStatus(); + } // tcon // TODO(jons): some parabola in ns, @@ -2500,7 +2474,7 @@ void IdealMhdModel::effectiveConstraintForce() { // and apply scaling (tcon[j]) and preconditioning (faccon[m]) void IdealMhdModel::deAliasConstraintForce() { vmecpp::deAliasConstraintForce(r_, t_, s_, faccon, tcon, gConEff, gsc, gcs, - gCon); + gcc, gss, gConAsym, refl, gCon); } // add constraint force to MHD force @@ -2535,6 +2509,702 @@ void IdealMhdModel::assembleTotalForces() { fzcon_e.data(), fzcon_o.data()); } +#ifdef VMECPP_ENABLE_ENZYME +void IdealMhdModel::packGeometry(FourierGeometry& m_decomposed, + FourierGeometry& m_physical_scratch, + double* out, int gS, bool primal) { + // Linear pre-chain decomposed -> real-space geometry, identical to the head + // of update(). Applied to a state it yields the geometry; applied to a + // tangent it yields the exact geometry tangent (the chain is linear), so no + // finite difference is needed. + m_decomposed.decomposeInto(m_physical_scratch, m_p_.scalxc); + m_physical_scratch.m1Constraint(1.0); + m_physical_scratch.extrapolateTowardsAxis(); + geometryFromFourier(m_physical_scratch); + + auto blk = [&](int b, const Eigen::VectorXd& src) { + const int n = static_cast(src.size()); + for (int i = 0; i < n; ++i) out[b * gS + i] = src[i]; + }; + blk(0, r1_e); + blk(1, r1_o); + blk(2, z1_e); + blk(3, z1_o); + blk(4, ru_e); + blk(5, ru_o); + blk(6, zu_e); + blk(7, zu_o); + blk(8, rv_e); + blk(9, rv_o); + blk(10, zv_e); + blk(11, zv_o); + // lambda carries the computeBContra normalization: *lamscale, and for the + // primal also + phipF on lu_e (a constant, so it drops from the tangent). + auto blk_lam = [&](int b, const Eigen::VectorXd& src) { + const int n = static_cast(src.size()); + for (int i = 0; i < n; ++i) out[b * gS + i] = constants_.lamscale * src[i]; + }; + blk_lam(12, lu_e); + blk_lam(13, lu_o); + blk_lam(14, lv_e); + blk_lam(15, lv_o); + if (primal) { + const int nFullSurf = static_cast(lu_e.size()) / s_.nZnT; + for (int jF = 0; jF < nFullSurf; ++jF) { + const double phip = m_p_.phipF[jF]; + for (int kl = 0; kl < s_.nZnT; ++kl) { + out[12 * gS + jF * s_.nZnT + kl] += phip; + } + } + } + blk(16, rCon); + blk(17, zCon); + blk(18, ruFull); + blk(19, zuFull); +} + +LocalForceComposition IdealMhdModel::makeLocalForceComposition( + int geom_stride) { + LocalForceComposition comp; + comp.nZnT = s_.nZnT; + comp.geom_stride = geom_stride; + comp.force_stride = (r_.nsMaxFIncludingLcfs - r_.nsMinF) * s_.nZnT; + comp.nsMinF = r_.nsMinF; + comp.nsMinF1 = r_.nsMinF1; + comp.nsMinH = r_.nsMinH; + comp.nsMaxH = r_.nsMaxH; + comp.jMaxRZ = std::min(r_.nsMaxF, m_fc_.ns - 1); + comp.nsMaxFIncludingLcfs = r_.nsMaxFIncludingLcfs; + comp.sqrtSF = m_p_.sqrtSF.data(); + comp.sqrtSH = m_p_.sqrtSH.data(); + comp.chipH = m_p_.chipH.data(); + comp.presH = m_p_.presH.data(); + comp.radialBlending = m_p_.radialBlending.data(); + comp.deltaS = m_fc_.deltaS; + comp.dSHalfDsInterp = dSHalfDsInterp; + comp.lamscale = constants_.lamscale; + comp.lthreed = s_.lthreed; + comp.with_constraint = true; + comp.lasym = s_.lasym; + comp.nsMaxF = r_.nsMaxF; + comp.nZeta = s_.nZeta; + comp.nThetaEff = s_.nThetaEff; + comp.ncurr = ncurr; + comp.currH = m_p_.currH.data(); + comp.wInt = s_.wInt.data(); + comp.nThetaEven = s_.nThetaEven; + comp.nThetaReduced = s_.nThetaReduced; + comp.mpol = s_.mpol; + comp.ntor = s_.ntor; + comp.nnyq2 = s_.nnyq2; + comp.rCon0 = rCon0.data(); + comp.zCon0 = zCon0.data(); + comp.faccon = faccon.data(); + comp.tcon = tcon.data(); + comp.sinmui = t_.sinmui.data(); + comp.cosmui = t_.cosmui.data(); + comp.cosnv = t_.cosnv.data(); + comp.sinnv = t_.sinnv.data(); + comp.sinmu = t_.sinmu.data(); + comp.cosmu = t_.cosmu.data(); + return comp; +} + +void IdealMhdModel::applyExactForceJacobian(const double* geomP, + const double* dgeom, + int geom_stride, + FourierForces& m_physical_f, + FourierForces& m_decomposed_hv, + bool fix_m1_gauge) { + LocalForceComposition comp = makeLocalForceComposition(geom_stride); + const int nForce = comp.force_stride; + + const int nH = (r_.nsMaxH - r_.nsMinH) * s_.nZnT; + // work holds the half-grid and per-point scratch plus the constraint scratch. + const int nWork = 15 * nH + 30 * s_.nZnT + 4 * nForce + 4 * (s_.ntor + 1) + + s_.nZnT + s_.nThetaReduced; + std::vector work(nWork, 0.0); + std::vector dwork(nWork, 0.0); + std::vector force(20 * nForce, 0.0); + std::vector dforce(20 * nForce, 0.0); + + // single nonlinear forward pass: J_g . (T v) + ExactForceDensityJvp(geomP, dgeom, work.data(), dwork.data(), force.data(), + dforce.data(), &comp); + + // scatter the force-density tangent into the real-space force members + auto put = [&](int block, Eigen::VectorXd& dst) { + const int n = static_cast(dst.size()); + for (int i = 0; i < n; ++i) { + dst[i] = dforce[block * nForce + i]; + } + }; + put(0, armn_e); + put(1, armn_o); + put(2, azmn_e); + put(3, azmn_o); + put(4, brmn_e); + put(5, brmn_o); + put(6, bzmn_e); + put(7, bzmn_o); + put(12, blmn_e); + put(13, blmn_o); + if (s_.lthreed) { + put(8, crmn_e); + put(9, crmn_o); + put(10, czmn_e); + put(11, czmn_o); + put(14, clmn_e); + put(15, clmn_o); + } + put(16, frcon_e); + put(17, frcon_o); + put(18, fzcon_e); + put(19, fzcon_o); + + // linear forward transform and preconditioner decomposition, mirroring the + // tail of update() + forcesToFourier(m_physical_f); + m_physical_f.decomposeInto(m_decomposed_hv, m_p_.scalxc); + m_decomposed_hv.m1Constraint(1.0 / std::numbers::sqrt2); + if (fix_m1_gauge) { + m_decomposed_hv.zeroZForceForM1(); + } +} + +// Raw force-density tangent (20 blocks of (nsMaxFIncludingLcfs-nsMinF)*nZnT) +// from one Enzyme forward pass, with no transform applied. For isolating the +// JVP from the spectral-transform wrapping. +void IdealMhdModel::exactForceDensityTangent(const double* geomP, + const double* dgeom, + int geom_stride, + double* dforce_out) { + LocalForceComposition comp = makeLocalForceComposition(geom_stride); + const int nForce = comp.force_stride; + const int nH = (r_.nsMaxH - r_.nsMinH) * s_.nZnT; + const int nWork = 15 * nH + 30 * s_.nZnT + 4 * nForce + 4 * (s_.ntor + 1) + + s_.nZnT + s_.nThetaReduced; + std::vector work(nWork, 0.0); + std::vector dwork(nWork, 0.0); + std::vector force(20 * nForce, 0.0); + ExactForceDensityJvp(geomP, dgeom, work.data(), dwork.data(), force.data(), + dforce_out, &comp); +} + +void IdealMhdModel::exactForceDensityCotangent(const double* geomP, + const double* force_bar, + int geom_stride, + double* geom_bar_out) { + LocalForceComposition comp = makeLocalForceComposition(geom_stride); + const int nForce = comp.force_stride; + const int nH = (r_.nsMaxH - r_.nsMinH) * s_.nZnT; + const int nWork = 15 * nH + 30 * s_.nZnT + 4 * nForce + 4 * (s_.ntor + 1) + + s_.nZnT + s_.nThetaReduced; + std::vector work(nWork, 0.0); + std::vector work_bar(nWork, 0.0); + std::vector force(20 * nForce, 0.0); + // force_bar is the output cotangent seed; Enzyme consumes (and may clobber) + // the shadow, so pass a private copy. geom_bar_out is zeroed by the caller. + std::vector fbar(force_bar, force_bar + 20 * nForce); + ExactForceDensityVjp(geomP, geom_bar_out, work.data(), work_bar.data(), + force.data(), fbar.data(), &comp); +} + +#endif // VMECPP_ENABLE_ENZYME + +// The four transform transposes below are also useful as standalone linear +// operators, so keep them available in non-Enzyme builds for direct adjoint +// identity tests. The nonlinear force-chain wrappers resume below. + +// (forcesToFourier)^T for the 2D case: scatter a decomposed-force Fourier +// cotangent (frcc, fzsc, flsc) back to the real-space force-density members +// through the same weighted basis the forward projection uses. Transpose of +// dft_ForcesToFourier_2d_symm. +void IdealMhdModel::dft_ForcesToFourierTranspose_2d_symm( + const FourierForces& m_coeff_bar) { + for (auto* v : + {&armn_e, &armn_o, &brmn_e, &brmn_o, &azmn_e, &azmn_o, &bzmn_e, &bzmn_o, + &frcon_e, &frcon_o, &fzcon_e, &fzcon_o, &blmn_e, &blmn_o}) { + v->setZero(); + } + int jMaxRZ = std::min(r_.nsMaxF, m_fc_.ns - 1); + if (m_fc_.lfreeb && + (m_vacuum_pressure_state_ == VacuumPressureState::kInitialized || + m_vacuum_pressure_state_ == VacuumPressureState::kActive)) { + jMaxRZ = std::min(r_.nsMaxF, m_fc_.ns); + } + for (int jF = r_.nsMinF; jF < jMaxRZ; ++jF) { + const int num_m = (jF == 0) ? 1 : s_.mpol; + for (int m = 0; m < num_m; ++m) { + const bool m_even = m % 2 == 0; + auto& armn = m_even ? armn_e : armn_o; + auto& brmn = m_even ? brmn_e : brmn_o; + auto& azmn = m_even ? azmn_e : azmn_o; + auto& bzmn = m_even ? bzmn_e : bzmn_o; + auto& frcon = m_even ? frcon_e : frcon_o; + auto& fzcon = m_even ? fzcon_e : fzcon_o; + const int idx_jm = (jF - r_.nsMinF) * s_.mpol + m; + const double fr = m_coeff_bar.frcc[idx_jm]; + const double fz = m_coeff_bar.fzsc[idx_jm]; + for (int l = 0; l < s_.nThetaReduced; ++l) { + const int idx_jl = (jF - r_.nsMinF) * s_.nThetaEff + l; + const int idx_ml = m * s_.nThetaReduced + l; + const double cosmui = t_.cosmui[idx_ml]; + const double sinmumi = t_.sinmumi[idx_ml]; + const double sinmui = t_.sinmui[idx_ml]; + const double cosmumi = t_.cosmumi[idx_ml]; + armn[idx_jl] += fr * cosmui; + brmn[idx_jl] += fr * sinmumi; + frcon[idx_jl] += fr * xmpq[m] * cosmui; + azmn[idx_jl] += fz * sinmui; + bzmn[idx_jl] += fz * cosmumi; + fzcon[idx_jl] += fz * xmpq[m] * sinmui; + } // l + } // m + } // jF + for (int jF = std::max(1, r_.nsMinF); jF < r_.nsMaxFIncludingLcfs; ++jF) { + for (int m = 0; m < s_.mpol; ++m) { + const bool m_even = m % 2 == 0; + auto& blmn = m_even ? blmn_e : blmn_o; + const int idx_jm = (jF - r_.nsMinF) * s_.mpol + m; + const double fl = m_coeff_bar.flsc[idx_jm]; + for (int l = 0; l < s_.nThetaReduced; ++l) { + const int idx_jl = (jF - r_.nsMinF) * s_.nThetaEff + l; + const double cosmumi = t_.cosmumi[m * s_.nThetaReduced + l]; + blmn[idx_jl] += fl * cosmumi; + } // l + } // m + } // jF +} + +// (geometryFromFourier)^T for the 2D case: project the real-space geometry +// member cotangents (r1_e .. zCon) back to Fourier coefficient cotangents +// through the unweighted basis. Transpose of dft_FourierToReal_2d_symm. +void IdealMhdModel::dft_FourierToRealTranspose_2d_symm( + FourierGeometry& m_coeff_bar_out) { + m_coeff_bar_out.setZero(); + for (int jF = r_.nsMinF1; jF < r_.nsMaxF1; ++jF) { + double* dst_rcc = &(m_coeff_bar_out.rmncc[(jF - r_.nsMinF1) * s_.mnsize]); + double* dst_zsc = &(m_coeff_bar_out.zmnsc[(jF - r_.nsMinF1) * s_.mnsize]); + double* dst_lsc = &(m_coeff_bar_out.lmnsc[(jF - r_.nsMinF1) * s_.mnsize]); + for (int l = 0; l < s_.nThetaReduced; ++l) { + const int idx_jl = (jF - r_.nsMinF1) * s_.nThetaEff + l; + const double r1eb = r1_e[idx_jl], rueb = ru_e[idx_jl], + z1eb = z1_e[idx_jl], zueb = zu_e[idx_jl], + lueb = lu_e[idx_jl]; + const double r1ob = r1_o[idx_jl], ruob = ru_o[idx_jl], + z1ob = z1_o[idx_jl], zuob = zu_o[idx_jl], + luob = lu_o[idx_jl]; + const int num_m = (jF == 0) ? 2 : s_.mpol; + for (int m = 0; m < num_m; ++m) { + const int p = m % 2; + const int idx_ml = m * s_.nThetaReduced + l; + const double cosmu = t_.cosmu[idx_ml]; + const double sinmum = t_.sinmum[idx_ml]; + const double sinmu = t_.sinmu[idx_ml]; + const double cosmum = t_.cosmum[idx_ml]; + dst_rcc[m] += (p ? r1ob : r1eb) * cosmu + (p ? ruob : rueb) * sinmum; + dst_zsc[m] += (p ? z1ob : z1eb) * sinmu + (p ? zuob : zueb) * cosmum; + dst_lsc[m] += (p ? luob : lueb) * cosmum; + } // m + } // l + } // jF + for (int jF = r_.nsMinF; jF < r_.nsMaxFIncludingLcfs; ++jF) { + double* dst_rcc = &(m_coeff_bar_out.rmncc[(jF - r_.nsMinF1) * s_.mnsize]); + double* dst_zsc = &(m_coeff_bar_out.zmnsc[(jF - r_.nsMinF1) * s_.mnsize]); + const int num_m = (jF == 0) ? 2 : s_.mpol; + for (int m = 0; m < num_m; ++m) { + const int p = m % 2; + const double scale = xmpq[m] * (1 - p + p * m_p_.sqrtSF[jF - r_.nsMinF1]); + for (int l = 0; l < s_.nThetaReduced; ++l) { + const int idx_ml = m * s_.nThetaReduced + l; + const int idx_con = (jF - r_.nsMinF) * s_.nThetaEff + l; + dst_rcc[m] += rCon[idx_con] * t_.cosmu[idx_ml] * scale; + dst_zsc[m] += zCon[idx_con] * t_.sinmu[idx_ml] * scale; + } // l + } // m + } // jF +} + +// (forcesToFourier)^T for the 3D case. Transpose of +// ForcesToFourier3DSymmFastPoloidal: undo the toroidal scatter, then the +// poloidal projection, back onto the real-space force-density members. +void IdealMhdModel::dft_ForcesToFourierTranspose_3d_symm( + const FourierForces& m_coeff_bar) { + for (auto* v : + {&armn_e, &armn_o, &azmn_e, &azmn_o, &blmn_e, &blmn_o, &brmn_e, + &brmn_o, &bzmn_e, &bzmn_o, &clmn_e, &clmn_o, &crmn_e, &crmn_o, + &czmn_e, &czmn_o, &frcon_e, &frcon_o, &fzcon_e, &fzcon_o}) { + v->setZero(); + } + const int nThR = s_.nThetaReduced; + const int ntorp1 = s_.ntor + 1; + int jMaxRZ = std::min(r_.nsMaxF, m_fc_.ns - 1); + if (m_fc_.lfreeb && + (m_vacuum_pressure_state_ == VacuumPressureState::kInitialized || + m_vacuum_pressure_state_ == VacuumPressureState::kActive)) { + jMaxRZ = std::min(r_.nsMaxF, m_fc_.ns); + } + const int jMinL = 1; + for (int jF = r_.nsMinF; jF < jMaxRZ; ++jF) { + const int mmax = (jF == 0) ? 1 : s_.mpol; + for (int m = 0; m < mmax; ++m) { + const bool m_even = m % 2 == 0; + auto& armn = m_even ? armn_e : armn_o; + auto& azmn = m_even ? azmn_e : azmn_o; + auto& blmn = m_even ? blmn_e : blmn_o; + auto& brmn = m_even ? brmn_e : brmn_o; + auto& bzmn = m_even ? bzmn_e : bzmn_o; + auto& clmn = m_even ? clmn_e : clmn_o; + auto& crmn = m_even ? crmn_e : crmn_o; + auto& czmn = m_even ? czmn_e : czmn_o; + auto& frcon = m_even ? frcon_e : frcon_o; + auto& fzcon = m_even ? fzcon_e : fzcon_o; + const int idx_ml_base = m * nThR; + for (int k = 0; k < s_.nZeta; ++k) { + const int idx_kl_base = + ((jF - r_.nsMinF) * s_.nZeta + k) * s_.nThetaEff; + const int idx_kn_base = k * (s_.nnyq2 + 1); + const int idx_mn_base = ((jF - r_.nsMinF) * s_.mpol + m) * ntorp1; + double rmkcc = 0, rmkss = 0, zmksc = 0, zmkcs = 0, rmkcc_n = 0, + zmkcs_n = 0, rmkss_n = 0, zmksc_n = 0, lmksc = 0, lmkcs = 0, + lmkcs_n = 0, lmksc_n = 0; + for (int nn = 0; nn < ntorp1; ++nn) { + const int kn = idx_kn_base + nn; + const int mn = idx_mn_base + nn; + const double cosnv = t_.cosnv[kn], sinnv = t_.sinnv[kn], + cosnvn = t_.cosnvn[kn], sinnvn = t_.sinnvn[kn]; + const double frcc = m_coeff_bar.frcc[mn], frss = m_coeff_bar.frss[mn], + fzsc = m_coeff_bar.fzsc[mn], fzcs = m_coeff_bar.fzcs[mn]; + rmkcc += frcc * cosnv; + rmkcc_n += frcc * sinnvn; + rmkss += frss * sinnv; + rmkss_n += frss * cosnvn; + zmksc += fzsc * cosnv; + zmksc_n += fzsc * sinnvn; + zmkcs += fzcs * sinnv; + zmkcs_n += fzcs * cosnvn; + if (jMinL <= jF) { + const double flsc = m_coeff_bar.flsc[mn], + flcs = m_coeff_bar.flcs[mn]; + lmksc += flsc * cosnv; + lmksc_n += flsc * sinnvn; + lmkcs += flcs * sinnv; + lmkcs_n += flcs * cosnvn; + } + } + for (int l = 0; l < nThR; ++l) { + const int im = idx_ml_base + l; + const int kl = idx_kl_base + l; + const double cosmui = t_.cosmui[im], sinmui = t_.sinmui[im], + cosmumi = t_.cosmumi[im], sinmumi = t_.sinmumi[im]; + const double tR = rmkcc * cosmui + rmkss * sinmui; + armn[kl] += tR; + frcon[kl] += xmpq[m] * tR; + brmn[kl] += rmkcc * sinmumi + rmkss * cosmumi; + const double tZ = zmksc * sinmui + zmkcs * cosmui; + azmn[kl] += tZ; + fzcon[kl] += xmpq[m] * tZ; + bzmn[kl] += zmksc * cosmumi + zmkcs * sinmumi; + crmn[kl] += -(rmkcc_n * cosmui + rmkss_n * sinmui); + czmn[kl] += -(zmkcs_n * cosmui + zmksc_n * sinmui); + blmn[kl] += lmksc * cosmumi + lmkcs * sinmumi; + clmn[kl] += -(lmkcs_n * cosmui + lmksc_n * sinmui); + } // l + } // k + } // m + } // jF + for (int jF = jMaxRZ; jF < r_.nsMaxFIncludingLcfs; ++jF) { + for (int m = 0; m < s_.mpol; ++m) { + const bool m_even = m % 2 == 0; + auto& blmn = m_even ? blmn_e : blmn_o; + auto& clmn = m_even ? clmn_e : clmn_o; + const int idx_ml_base = m * nThR; + for (int k = 0; k < s_.nZeta; ++k) { + const int idx_kl_base = + ((jF - r_.nsMinF) * s_.nZeta + k) * s_.nThetaEff; + const int idx_kn_base = k * (s_.nnyq2 + 1); + const int idx_mn_base = ((jF - r_.nsMinF) * s_.mpol + m) * ntorp1; + double lmksc = 0, lmkcs = 0, lmkcs_n = 0, lmksc_n = 0; + for (int nn = 0; nn < ntorp1; ++nn) { + const int kn = idx_kn_base + nn; + const int mn = idx_mn_base + nn; + lmksc += m_coeff_bar.flsc[mn] * t_.cosnv[kn]; + lmksc_n += m_coeff_bar.flsc[mn] * t_.sinnvn[kn]; + lmkcs += m_coeff_bar.flcs[mn] * t_.sinnv[kn]; + lmkcs_n += m_coeff_bar.flcs[mn] * t_.cosnvn[kn]; + } + for (int l = 0; l < nThR; ++l) { + const int im = idx_ml_base + l; + const int kl = idx_kl_base + l; + blmn[kl] += lmksc * t_.cosmumi[im] + lmkcs * t_.sinmumi[im]; + clmn[kl] += -(lmkcs_n * t_.cosmui[im] + lmksc_n * t_.sinmui[im]); + } // l + } // k + } // m + } // jF +} + +// (geometryFromFourier)^T for the 3D case. Transpose of +// FourierToReal3DSymmFastPoloidal: undo the poloidal evaluation, then the +// toroidal evaluation, back onto the Fourier coefficient cotangents. +void IdealMhdModel::dft_FourierToRealTranspose_3d_symm( + FourierGeometry& m_coeff_bar_out) { + m_coeff_bar_out.setZero(); + const int nThR = s_.nThetaReduced; + const int ntorp1 = s_.ntor + 1; + for (int jF = r_.nsMinF1; jF < r_.nsMaxF1; ++jF) { + for (int m = 0; m < s_.mpol; ++m) { + const bool m_even = m % 2 == 0; + const double con_factor = + m_even ? xmpq[m] : xmpq[m] * m_p_.sqrtSF[jF - r_.nsMinF1]; + auto& r1 = m_even ? r1_e : r1_o; + auto& ru = m_even ? ru_e : ru_o; + auto& rv = m_even ? rv_e : rv_o; + auto& z1 = m_even ? z1_e : z1_o; + auto& zu = m_even ? zu_e : zu_o; + auto& zv = m_even ? zv_e : zv_o; + auto& lu = m_even ? lu_e : lu_o; + auto& lv = m_even ? lv_e : lv_o; + const int jMin = (m == 0 || m == 1) ? 0 : 1; + if (jF < jMin) { + continue; + } + const int idx_ml_base = m * nThR; + for (int k = 0; k < s_.nZeta; ++k) { + const int idx_kl_base = + ((jF - r_.nsMinF1) * s_.nZeta + k) * s_.nThetaEff; + const bool con_in_range = + (r_.nsMinF <= jF && jF < r_.nsMaxFIncludingLcfs); + const int idx_con_base = + ((jF - r_.nsMinF) * s_.nZeta + k) * s_.nThetaEff; + double rmkcc = 0, rmkss = 0, rmkcc_n = 0, rmkss_n = 0, zmksc = 0, + zmkcs = 0, zmksc_n = 0, zmkcs_n = 0, lmksc = 0, lmkcs = 0, + lmksc_n = 0, lmkcs_n = 0; + for (int l = 0; l < nThR; ++l) { + const int im = idx_ml_base + l; + const int kl = idx_kl_base + l; + const double cosmu = t_.cosmu[im], sinmu = t_.sinmu[im], + sinmum = t_.sinmum[im], cosmum = t_.cosmum[im]; + const double r1b = r1[kl], rub = ru[kl], rvb = rv[kl], z1b = z1[kl], + zub = zu[kl], zvb = zv[kl], lub = lu[kl], lvb = lv[kl]; + double rConb = 0, zConb = 0; + if (con_in_range) { + rConb = rCon[idx_con_base + l]; + zConb = zCon[idx_con_base + l]; + } + rmkcc += r1b * cosmu + rub * sinmum + rConb * cosmu * con_factor; + rmkss += r1b * sinmu + rub * cosmum + rConb * sinmu * con_factor; + rmkcc_n += rvb * cosmu; + rmkss_n += rvb * sinmu; + zmksc += z1b * sinmu + zub * cosmum + zConb * sinmu * con_factor; + zmkcs += z1b * cosmu + zub * sinmum + zConb * cosmu * con_factor; + zmksc_n += zvb * sinmu; + zmkcs_n += zvb * cosmu; + lmksc += lub * cosmum; + lmkcs += lub * sinmum; + lmksc_n += -lvb * sinmu; + lmkcs_n += -lvb * cosmu; + } // l + const int idx_kn_base = k * (s_.nnyq2 + 1); + const int idx_mn_base = ((jF - r_.nsMinF1) * s_.mpol + m) * ntorp1; + for (int nn = 0; nn < ntorp1; ++nn) { + const int kn = idx_kn_base + nn; + const int mn = idx_mn_base + nn; + const double cosnv = t_.cosnv[kn], sinnv = t_.sinnv[kn], + cosnvn = t_.cosnvn[kn], sinnvn = t_.sinnvn[kn]; + m_coeff_bar_out.rmncc[mn] += rmkcc * cosnv + rmkcc_n * sinnvn; + m_coeff_bar_out.rmnss[mn] += rmkss * sinnv + rmkss_n * cosnvn; + m_coeff_bar_out.zmnsc[mn] += zmksc * cosnv + zmksc_n * sinnvn; + m_coeff_bar_out.zmncs[mn] += zmkcs * sinnv + zmkcs_n * cosnvn; + m_coeff_bar_out.lmnsc[mn] += lmksc * cosnv + lmksc_n * sinnvn; + m_coeff_bar_out.lmncs[mn] += lmkcs * sinnv + lmkcs_n * cosnvn; + } // nn + } // k + } // m + } // jF +} + +#ifdef VMECPP_ENABLE_ENZYME +void IdealMhdModel::applyExactForceJacobianTranspose( + const double* geomP, int geom_stride, FourierForces& m_decomposed_in, + FourierForces& m_physical_f, FourierGeometry& m_physical_scratch, + FourierGeometry& m_decomposed_out, bool fix_m1_gauge) { + const int gS = geom_stride; + const int nForce = (r_.nsMaxFIncludingLcfs - r_.nsMinF) * s_.nZnT; + + // C^T: transpose of [scatter -> forcesToFourier -> decompose -> m1 -> zeroZ]. + if (fix_m1_gauge) { + m_decomposed_in.zeroZForceForM1(); + } + m_decomposed_in.m1Constraint(1.0 / std::numbers::sqrt2); + m_decomposed_in.decomposeInto(m_physical_f, m_p_.scalxc); + if (s_.lthreed) { + dft_ForcesToFourierTranspose_3d_symm(m_physical_f); + } else { + dft_ForcesToFourierTranspose_2d_symm(m_physical_f); + } + + // Gather the force-density member cotangents into the 20-block flat layout. + std::vector force_bar(20 * nForce, 0.0); + auto gather = [&](int b, const Eigen::VectorXd& src) { + for (int i = 0; i < nForce; ++i) force_bar[b * nForce + i] = src[i]; + }; + gather(0, armn_e); + gather(1, armn_o); + gather(2, azmn_e); + gather(3, azmn_o); + gather(4, brmn_e); + gather(5, brmn_o); + gather(6, bzmn_e); + gather(7, bzmn_o); + gather(12, blmn_e); + gather(13, blmn_o); + gather(16, frcon_e); + gather(17, frcon_o); + gather(18, fzcon_e); + gather(19, fzcon_o); + if (s_.lthreed) { + gather(8, crmn_e); + gather(9, crmn_o); + gather(10, czmn_e); + gather(11, czmn_o); + gather(14, clmn_e); + gather(15, clmn_o); + } + + // J_g^T: reverse-mode force-density kernel. + std::vector geom_bar(20 * gS, 0.0); + exactForceDensityCotangent(geomP, force_bar.data(), gS, geom_bar.data()); + + // B^T: transpose of packGeometry [decompose -> m1 -> extrapolate -> + // geometryFromFourier -> block pack with lamscale + ruFull/zuFull]. + auto scat = [&](int b, Eigen::VectorXd& dst) { + const int sz = std::min(gS, static_cast(dst.size())); + for (int i = 0; i < sz; ++i) dst[i] = geom_bar[b * gS + i]; + }; + scat(0, r1_e); + scat(1, r1_o); + scat(2, z1_e); + scat(3, z1_o); + scat(4, ru_e); + scat(5, ru_o); + scat(6, zu_e); + scat(7, zu_o); + for (int i = 0; i < gS; ++i) { + lu_e[i] = constants_.lamscale * geom_bar[12 * gS + i]; + lu_o[i] = constants_.lamscale * geom_bar[13 * gS + i]; + } + if (s_.lthreed) { + scat(8, rv_e); + scat(9, rv_o); + scat(10, zv_e); + scat(11, zv_o); + for (int i = 0; i < gS; ++i) { + lv_e[i] = constants_.lamscale * geom_bar[14 * gS + i]; + lv_o[i] = constants_.lamscale * geom_bar[15 * gS + i]; + } + } + scat(16, rCon); + scat(17, zCon); + // ruFull/zuFull (blocks 18,19): forward ruFull = ru_e + sqrtSF*ru_o, so the + // adjoint folds the full-grid cotangent back into the even/odd ru, zu. + for (int jF = r_.nsMinF; jF < r_.nsMaxFIncludingLcfs; ++jF) { + const double sf = m_p_.sqrtSF[jF - r_.nsMinF1]; + for (int kl = 0; kl < s_.nZnT; ++kl) { + const int idx_kl1 = (jF - r_.nsMinF1) * s_.nZnT + kl; + const int idx_kl = (jF - r_.nsMinF) * s_.nZnT + kl; + ru_e[idx_kl1] += geom_bar[18 * gS + idx_kl]; + ru_o[idx_kl1] += sf * geom_bar[18 * gS + idx_kl]; + zu_e[idx_kl1] += geom_bar[19 * gS + idx_kl]; + zu_o[idx_kl1] += sf * geom_bar[19 * gS + idx_kl]; + } + } + if (s_.lthreed) { + dft_FourierToRealTranspose_3d_symm(m_physical_scratch); + } else { + dft_FourierToRealTranspose_2d_symm(m_physical_scratch); + } + m_physical_scratch.extrapolateTowardsAxisTranspose(); + m_physical_scratch.m1Constraint(1.0); + m_physical_scratch.decomposeInto(m_decomposed_out, m_p_.scalxc); +} + +// Diagnostic: max |composed force density - production force density| at the +// current state, to isolate composition bugs from the transform/tangent path. +double IdealMhdModel::composedForceResidual(const double* geomP, + int geom_stride) { + LocalForceComposition comp; + comp.nZnT = s_.nZnT; + comp.geom_stride = geom_stride; + const int nForce = (r_.nsMaxFIncludingLcfs - r_.nsMinF) * s_.nZnT; + comp.force_stride = nForce; + comp.nsMinF = r_.nsMinF; + comp.nsMinF1 = r_.nsMinF1; + comp.nsMinH = r_.nsMinH; + comp.nsMaxH = r_.nsMaxH; + comp.jMaxRZ = std::min(r_.nsMaxF, m_fc_.ns - 1); + comp.nsMaxFIncludingLcfs = r_.nsMaxFIncludingLcfs; + comp.sqrtSF = m_p_.sqrtSF.data(); + comp.sqrtSH = m_p_.sqrtSH.data(); + comp.chipH = m_p_.chipH.data(); + comp.presH = m_p_.presH.data(); + comp.radialBlending = m_p_.radialBlending.data(); + comp.deltaS = m_fc_.deltaS; + comp.dSHalfDsInterp = dSHalfDsInterp; + comp.lamscale = constants_.lamscale; + comp.lthreed = s_.lthreed; + comp.with_constraint = true; + comp.lasym = s_.lasym; + comp.nsMaxF = r_.nsMaxF; + comp.nZeta = s_.nZeta; + comp.nThetaEff = s_.nThetaEff; + comp.ncurr = ncurr; + comp.currH = m_p_.currH.data(); + comp.wInt = s_.wInt.data(); + comp.nThetaEven = s_.nThetaEven; + comp.nThetaReduced = s_.nThetaReduced; + comp.mpol = s_.mpol; + comp.ntor = s_.ntor; + comp.nnyq2 = s_.nnyq2; + comp.rCon0 = rCon0.data(); + comp.zCon0 = zCon0.data(); + comp.faccon = faccon.data(); + comp.tcon = tcon.data(); + comp.sinmui = t_.sinmui.data(); + comp.cosmui = t_.cosmui.data(); + comp.cosnv = t_.cosnv.data(); + comp.sinnv = t_.sinnv.data(); + comp.sinmu = t_.sinmu.data(); + comp.cosmu = t_.cosmu.data(); + + const int nH = (r_.nsMaxH - r_.nsMinH) * s_.nZnT; + const int nWork = 15 * nH + 30 * s_.nZnT + 4 * nForce + 4 * (s_.ntor + 1) + + s_.nZnT + s_.nThetaReduced; + std::vector work(nWork, 0.0); + std::vector force(20 * nForce, 0.0); + ComputeLocalForceDensity(geomP, work.data(), force.data(), &comp); + + double maxd = 0.0; + auto cmp = [&](int block, const Eigen::VectorXd& prod) { + for (int i = 0; i < static_cast(prod.size()); ++i) { + maxd = std::max(maxd, std::fabs(force[block * nForce + i] - prod[i])); + } + }; + cmp(0, armn_e); + cmp(1, armn_o); + cmp(2, azmn_e); + cmp(3, azmn_o); + cmp(4, brmn_e); + cmp(5, brmn_o); + cmp(6, bzmn_e); + cmp(7, bzmn_o); + cmp(12, blmn_e); + cmp(13, blmn_o); + return maxd; +} +#endif // VMECPP_ENABLE_ENZYME + void IdealMhdModel::forcesToFourier(FourierForces& m_physical_f) { if (s_.lasym) { // Split the real-space forces into their standard- and reversed-parity @@ -2555,6 +3225,27 @@ void IdealMhdModel::forcesToFourier(FourierForces& m_physical_f) { } else { dft_ForcesToFourier_2d_asymm(m_physical_f); } + } // lasym + + if (lforbal) { + // lforbal (educational_VMEC tomnsps): replace the m=1, n=0 R,Z forces with + // the flux-averaged radial force balance. equiF lives on the interior full + // grid [nsMinFi, nsMaxFi); the force-balance factors live on [nsMinF, + // nsMaxF). r0scale == 1, so the EQUIF weight is c = nscale(0). + const double c = t_.nscale[0]; + for (int jF = r_.nsMinFi; jF < r_.nsMaxFi; ++jF) { + if (jF == 0 || jF >= m_fc_.ns - 1) { + continue; + } + const int i = jF - r_.nsMinF; + const int idx_mn = ((jF - r_.nsMinF) * s_.mpol + 1) * (s_.ntor + 1); + const double equif = m_p_.equiF[jF - r_.nsMinFi]; + const double frcc = m_physical_f.frcc[idx_mn]; + const double fzsc = m_physical_f.fzsc[idx_mn]; + const double work1 = frcc_fac[i] * frcc + fzsc_fac[i] * fzsc; + m_physical_f.frcc[idx_mn] = rzu_fac[i] * (c * equif + work1); + m_physical_f.fzsc[idx_mn] = rru_fac[i] * (c * equif - work1); + } } } diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h index 8bd957ca9..fe80189b7 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h @@ -26,6 +26,7 @@ #include "vmecpp/vmec/fourier_geometry/fourier_geometry.h" #include "vmecpp/vmec/handover_storage/handover_storage.h" #include "vmecpp/vmec/ideal_mhd_model/dft_toroidal.h" +#include "vmecpp/vmec/ideal_mhd_model/local_force_composition.h" #ifdef VMECPP_USE_FFTX #include "vmecpp/vmec/ideal_mhd_model/fft_toroidal.h" #endif @@ -44,6 +45,15 @@ void deAliasConstraintForce(const RadialPartitioning& rp, const Eigen::VectorXd& gConEff, Eigen::VectorXd& m_gsc, Eigen::VectorXd& m_gcs, Eigen::VectorXd& m_gCon); +void deAliasConstraintForce(const RadialPartitioning& rp, + const FourierBasisFastPoloidal& fb, const Sizes& s_, + const Eigen::VectorXd& faccon, + const Eigen::VectorXd& tcon, + const Eigen::VectorXd& gConEff, + Eigen::VectorXd& m_gsc, Eigen::VectorXd& m_gcs, + Eigen::VectorXd& m_gcc, Eigen::VectorXd& m_gss, + Eigen::VectorXd& m_gConAsym, + Eigen::VectorXd& m_refl, Eigen::VectorXd& m_gCon); class IdealMhdModel { public: @@ -55,7 +65,8 @@ class IdealMhdModel { int vac_num_threads, int signOfJacobian, int nvacskip, VacuumPressureState* m_vacuum_pressure_state); - void setFromINDATA(int ncurr, double adiabaticIndex, double tCon0); + void setFromINDATA(int ncurr, double adiabaticIndex, double tCon0, + bool lforbal); // Compute the invariant (i.e., not preconditioned yet) force residuals. // Will put them into the provided array as { fsqr, fsqz, fsql }. @@ -75,7 +86,6 @@ class IdealMhdModel { const int iter2, const VmecCheckpoint& checkpoint = VmecCheckpoint::NONE, const int iterations_before_checkpointing = INT_MAX, bool verbose = true, bool always_fix_m1_gauge = false); - std::int64_t forceEvaluationCount() const { return force_evaluation_count_; } void resetForceEvaluationCount() { force_evaluation_count_ = 0; } @@ -176,6 +186,86 @@ class IdealMhdModel { // Fourier space. void forcesToFourier(FourierForces& m_physical_f); + // Exact Hessian-vector product of the local force chain. Given the packed + // real-space geometry primal geomP and a geometry tangent dgeom (each + // geom_stride doubles per block, see local_force_composition.h), + // differentiate the force density (MHD, hybrid lambda, and spectral- + // condensation constraint force) by one Enzyme forward pass, then apply the + // linear forward transform and preconditioner decomposition to obtain the + // decomposed force tangent in m_decomposed_hv. The constraint multiplier tcon + // is held frozen; freeze it in the raw force too + // (freeze_constraint_multiplier_) for an exactly consistent Jacobian. Used by + // the exact internal Newton-Krylov Hessian-vector product. This low-level + // kernel does not differentiate the state-dependent LFORBAL replacement; + // public callers must reject lforbal=true. + void applyExactForceJacobian(const double* geomP, const double* dgeom, + int geom_stride, FourierForces& m_physical_f, + FourierForces& m_decomposed_hv, + bool fix_m1_gauge); + + // Linear pre-chain decomposed -> real-space geometry (decomposeInto, + // m1Constraint, extrapolate, geometryFromFourier) packed into the 20-block + // layout of local_force_composition.h, with the computeBContra lambda + // normalization. Applied to a state (primal=true, adds phipF on lu_e) it + // gives the geometry; applied to a tangent (primal=false) it gives the exact + // geometry tangent, no finite difference. Uses m_physical_scratch as scratch. + void packGeometry(FourierGeometry& m_decomposed, + FourierGeometry& m_physical_scratch, double* out, int gS, + bool primal); + + // Diagnostic: max |composed force density - production force density| at the + // current state, to isolate composition bugs from the transform/tangent path. + double composedForceResidual(const double* geomP, int geom_stride); + + // Raw force-density tangent (20 blocks of (nsMaxFIncludingLcfs-nsMinF)*nZnT) + // from one Enzyme forward pass, no transform. For isolating the JVP from the + // spectral-transform wrapping. + void exactForceDensityTangent(const double* geomP, const double* dgeom, + int geom_stride, double* dforce_out); + + // Freeze/unfreeze the constraint-force multiplier tcon (see the member). + void setFreezeConstraintMultiplier(bool freeze) { + freeze_constraint_multiplier_ = freeze; + } + + // Reverse-mode force-density cotangent: J_g^T applied to the force-density + // cotangent force_bar (20 blocks of (nsMaxFIncludingLcfs-nsMinF)*nZnT), + // accumulated into geom_bar_out (20 blocks of geom_stride, zeroed by caller), + // by one Enzyme reverse pass. The transpose of exactForceDensityTangent. + void exactForceDensityCotangent(const double* geomP, const double* force_bar, + int geom_stride, double* geom_bar_out); + + // Fill the local force-density composition descriptor shared by the exact + // forward/reverse force-density passes (geom_stride sized blocks). + LocalForceComposition makeLocalForceComposition(int geom_stride); + + // Transpose of applyExactForceJacobian: H^T w. Given a decomposed-force + // cotangent (the space applyExactForceJacobian writes), apply C^T (transpose + // of the output transform), the reverse-mode force-density kernel J_g^T, and + // B^T (transpose of packGeometry's pre-chain), yielding the state cotangent + // in m_decomposed_out. The two spectral transforms are reused as each other's + // adjoint with the poloidal integration weight; the rest of the linear chain + // (decomposeInto, m1Constraint, zeroZForceForM1, extrapolateTowardsAxis, + // ruFull/zuFull, lamscale) is transposed analytically. This low-level kernel + // does not transpose the state-dependent LFORBAL replacement; public callers + // must reject lforbal=true. + void applyExactForceJacobianTranspose(const double* geomP, int geom_stride, + FourierForces& m_decomposed_in, + FourierForces& m_physical_f, + FourierGeometry& m_physical_scratch, + FourierGeometry& m_decomposed_out, + bool fix_m1_gauge); + + // Transposes of the spectral transforms, for the transposed exact Hessian. + // dft_ForcesToFourierTranspose: (forcesToFourier)^T, decomposed-force coeff + // cotangent -> real-space force-density member cotangents (armn_e .. + // fzcon_o). dft_FourierToRealTranspose: (geometryFromFourier)^T, real-space + // geometry member cotangents (r1_e .. zCon) -> Fourier coeff cotangent. + void dft_ForcesToFourierTranspose_2d_symm(const FourierForces& m_coeff_bar); + void dft_FourierToRealTranspose_2d_symm(FourierGeometry& m_coeff_bar_out); + void dft_ForcesToFourierTranspose_3d_symm(const FourierForces& m_coeff_bar); + void dft_FourierToRealTranspose_3d_symm(FourierGeometry& m_coeff_bar_out); + // Computes the forward-DFT of forces for the 3D (Stellarator) case. // Dispatching dft_ForcesToFourier_3d_symm void dft_ForcesToFourier_3d_symm(FourierForces& m_physical_f); @@ -215,7 +305,8 @@ class IdealMhdModel { const Eigen::VectorXd& xu_e, const Eigen::VectorXd& xu_o, const Eigen::VectorXd& x1_o, Eigen::VectorXd& m_axm, Eigen::VectorXd& m_axd, Eigen::VectorXd& m_bxm, Eigen::VectorXd& m_bxd, - Eigen::VectorXd& m_cxd); + Eigen::VectorXd& m_cxd, const Eigen::VectorXd& trigmult, + Eigen::VectorXd& m_eqfactor); // Applies the radial preconditioner for the m=1 Fourier coefficients of R and // Z. @@ -426,6 +517,18 @@ class IdealMhdModel { // crd == czd --> cxd Eigen::VectorXd cxd; + // lforbal: when set, the flux-averaged radial force balance evolves the + // m=1,n=0 R,Z components (non-variational). cos01/sin01 are the m=1 trig + // weights; rzu_fac/rru_fac/frcc_fac/fzsc_fac are the force-balance factors + // derived from the R,Z preconditioner diagonals. All unused when lforbal off. + bool lforbal = false; + Eigen::VectorXd cos01; + Eigen::VectorXd sin01; + Eigen::VectorXd rzu_fac; + Eigen::VectorXd rru_fac; + Eigen::VectorXd frcc_fac; + Eigen::VectorXd fzsc_fac; + Eigen::VectorXd ar; Eigen::VectorXd dr; Eigen::VectorXd br; @@ -448,6 +551,10 @@ class IdealMhdModel { // Fourier coefficients of constraint force - used during de-aliasing Eigen::VectorXd gsc; Eigen::VectorXd gcs; + Eigen::VectorXd gcc; + Eigen::VectorXd gss; + Eigen::VectorXd gConAsym; + Eigen::VectorXd refl; // de-aliased constraint force - what enters the Fourier coefficients of the // forces @@ -507,6 +614,14 @@ class IdealMhdModel { // 1 (default) -- full spectral condensation constraint force double tcon0; + // When true, constraintForceMultiplier reuses the existing tcon instead of + // recomputing it from the geometry. The exact Hessian-vector product freezes + // tcon (it depends on the preconditioner diagonal, not just the geometry), so + // freezing it in the raw force too keeps the force and its exact HVP a + // consistent function of the state -- the residual a Newton solver drives and + // the Jacobian it linearizes with then match. + bool freeze_constraint_multiplier_ = false; + // [mnsize] minimum flux surface index for which to apply radial // preconditioner for R and Z Eigen::VectorXi jMin; diff --git a/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/local_force_composition.h b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/local_force_composition.h new file mode 100644 index 000000000..7ee453f45 --- /dev/null +++ b/src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/local_force_composition.h @@ -0,0 +1,339 @@ +// SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +// +// +// SPDX-License-Identifier: MIT +#ifndef VMECPP_VMEC_IDEAL_MHD_MODEL_LOCAL_FORCE_COMPOSITION_H_ +#define VMECPP_VMEC_IDEAL_MHD_MODEL_LOCAL_FORCE_COMPOSITION_H_ + +#include "vmecpp/vmec/ideal_mhd_model/bco_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/bcontra_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/constraint_force_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/jacobian_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/lambda_force_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/metric_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/mhdforce_kernel.h" +#include "vmecpp/vmec/ideal_mhd_model/pressure_kernel.h" + +namespace vmecpp { + +// Composition of the local force-density chain as a single allocation-free map +// g: real-space geometry -> real-space force density. This is the nonlinear +// core of VMEC's force; the spectral transforms around it are linear and +// applied separately. Shared between the Enzyme autodiff validation and the +// exact Hessian-vector product. Covers the MHD force and the hybrid lambda +// force; when with_constraint is set it also computes the spectral-condensation +// constraint force (effective force, Fourier bandpass, assembly into the R/Z +// force), holding the multiplier tcon frozen. +// +// Geometry layout (each block GeomStride doubles, index (jF-nsMinF1)*nZnT): +// r1_e r1_o z1_e z1_o ru_e ru_o zu_e zu_o rv_e rv_o zv_e zv_o lu_e lu_o lv_e +// lv_o +// Force layout (each block ForceStride doubles): the 12 MHD densities then +// blmn_e blmn_o clmn_e clmn_o. +struct LocalForceComposition { + int nZnT; + int geom_stride; // doubles per geometry block (>= (nsMaxF1-nsMinF1)*nZnT) + int force_stride; // doubles per force block (>= + // (nsMaxFIncludingLcfs-nsMinF)*nZnT) + int nsMinF, nsMinF1, nsMinH, nsMaxH; + int jMaxRZ; // MHD force surfaces: [nsMinF, jMaxRZ) + int nsMaxFIncludingLcfs; // lambda force surfaces: [nsMinF, + // nsMaxFIncludingLcfs) + const double* sqrtSF; // index jF-nsMinF1 + const double* sqrtSH; // index jH-nsMinH + const double* chipH; // index jH-nsMinH (frozen for ncurr==0) + const double* presH; // index jH-nsMinH + const double* radialBlending; // index jF-nsMinF1 + double deltaS; + double dSHalfDsInterp; + double lamscale; + bool lthreed; + + // Constrained-current profile (ncurr==1): chi' is a function of geometry, + // recomputed each step, so it is differentiated in place. ncurr==0 uses the + // frozen chipH above. + int ncurr = 0; + int nThetaEff = 0; + const double* currH = nullptr; // index jH-nsMinH + const double* wInt = nullptr; // index kl % nThetaEff + + // Spectral-condensation constraint force. Enabled only when with_constraint + // is set; then geometry blocks 16-19 hold rCon, zCon, ruFull, zuFull and + // force blocks 16-19 receive frcon_e/o, fzcon_e/o. The bandpass uses the + // Fourier basis arrays and the tcon/faccon profiles. rCon0/zCon0 are + // recomputed in place from the live geometry (so they are differentiated); + // tcon is held frozen (see freeze_constraint_multiplier_). + bool with_constraint = false; + bool lasym = false; + int nsMaxF = 0; // constraint RZ range upper bound + int nZeta = 0, nThetaEven = 0, nThetaReduced = 0, mpol = 0, ntor = 0, + nnyq2 = 0; + const double* rCon0 = nullptr; + const double* zCon0 = nullptr; + const double* faccon = nullptr; + const double* tcon = nullptr; + const double* sinmui = nullptr; + const double* cosmui = nullptr; + const double* cosnv = nullptr; + const double* sinnv = nullptr; + const double* sinmu = nullptr; + const double* cosmu = nullptr; +}; + +// work must hold 15*nHalf + 30*nZnT plus the constraint scratch described +// below, where nHalf=(nsMaxH-nsMinH)*nZnT. +inline void ComputeLocalForceDensity(const double* geom, double* work, + double* force, + const LocalForceComposition* c) { + const int nZnT = c->nZnT; + const int gS = c->geom_stride; + const int fS = c->force_stride; + const int nH = (c->nsMaxH - c->nsMinH) * nZnT; + const double* r1e = geom + 0 * gS; + const double* r1o = geom + 1 * gS; + const double* z1e = geom + 2 * gS; + const double* z1o = geom + 3 * gS; + const double* rue = geom + 4 * gS; + const double* ruo = geom + 5 * gS; + const double* zue = geom + 6 * gS; + const double* zuo = geom + 7 * gS; + const double* rve = geom + 8 * gS; + const double* rvo = geom + 9 * gS; + const double* zve = geom + 10 * gS; + const double* zvo = geom + 11 * gS; + const double* lue = geom + 12 * gS; + const double* luo = geom + 13 * gS; + const double* lve = geom + 14 * gS; + const double* lvo = geom + 15 * gS; + + double* p = work; + double* r12 = p; + p += nH; + double* ru12 = p; + p += nH; + double* zu12 = p; + p += nH; + double* rs = p; + p += nH; + double* zs = p; + p += nH; + double* tau = p; + p += nH; + double* gsqrt = p; + p += nH; + double* guu = p; + p += nH; + double* guv = p; + p += nH; + double* gvv = p; + p += nH; + double* bsupu = p; + p += nH; + double* bsupv = p; + p += nH; + double* bsubu = p; + p += nH; + double* bsubv = p; + p += nH; + double* tp = p; + p += nH; + double* s = p; // 30 * nZnT + + ComputeHalfGridJacobian(r1e, r1o, z1e, z1o, rue, ruo, zue, zuo, c->sqrtSH, + c->deltaS, c->dSHalfDsInterp, nZnT, c->nsMinF1, + c->nsMinH, c->nsMaxH, r12, ru12, zu12, rs, zs, tau); + ComputeMetricElements(r1e, r1o, rue, ruo, zue, zuo, rve, rvo, zve, zvo, tau, + r12, c->sqrtSF, c->sqrtSH, c->lthreed, nZnT, c->nsMinF1, + c->nsMinH, c->nsMaxH, gsqrt, guu, guv, gvv); + ComputeBsupContra(lue, luo, lve, lvo, gsqrt, c->sqrtSH, c->lthreed, nZnT, + c->nsMinF1, c->nsMinH, c->nsMaxH, bsupu, bsupv); + for (int jH = c->nsMinH; jH < c->nsMaxH; ++jH) { + // For a prescribed-current profile (ncurr==1), chi' is recomputed from the + // geometry each step (constrained toroidal current), so differentiate it + // here rather than freezing it. For ncurr==0 chi' = iota*phi' is a fixed + // profile, so use the frozen c->chipH. + double chip = c->chipH[jH - c->nsMinH]; + if (c->ncurr == 1) { + double jvPlasma = 0.0; + double avg_guu_gsqrt = 0.0; + for (int kl = 0; kl < nZnT; ++kl) { + const int ih = (jH - c->nsMinH) * nZnT + kl; + const int l = kl % c->nThetaEff; + if (c->lthreed) { + jvPlasma += (guu[ih] * bsupu[ih] + guv[ih] * bsupv[ih]) * c->wInt[l]; + } else { + jvPlasma += guu[ih] * bsupu[ih] * c->wInt[l]; + } + avg_guu_gsqrt += guu[ih] / gsqrt[ih] * c->wInt[l]; + } + if (avg_guu_gsqrt != 0.0) { + chip = (c->currH[jH - c->nsMinH] - jvPlasma) / avg_guu_gsqrt; + } + } + for (int kl = 0; kl < nZnT; ++kl) { + const int ih = (jH - c->nsMinH) * nZnT + kl; + bsupu[ih] += chip / gsqrt[ih]; + } + } + ComputeBCo(guu, guv, gvv, bsupu, bsupv, c->lthreed, nH, bsubu, bsubv); + ComputeMagneticPressure(bsupu, bsubu, bsupv, bsubv, nH, tp); + for (int jH = c->nsMinH; jH < c->nsMaxH; ++jH) { + for (int kl = 0; kl < nZnT; ++kl) + tp[(jH - c->nsMinH) * nZnT + kl] += c->presH[jH - c->nsMinH]; + } + + double* P_i = s; + s += nZnT; + double* rup_i = s; + s += nZnT; + double* zup_i = s; + s += nZnT; + double* rsp_i = s; + s += nZnT; + double* zsp_i = s; + s += nZnT; + double* taup_i = s; + s += nZnT; + double* gbubu_i = s; + s += nZnT; + double* gbubv_i = s; + s += nZnT; + double* gbvbv_i = s; + s += nZnT; + double* P_o = s; + s += nZnT; + double* rup_o = s; + s += nZnT; + double* zup_o = s; + s += nZnT; + double* rsp_o = s; + s += nZnT; + double* zsp_o = s; + s += nZnT; + double* taup_o = s; + s += nZnT; + double* gbubu_o = s; + s += nZnT; + double* gbubv_o = s; + s += nZnT; + double* gbvbv_o = s; + s += nZnT; + double* P_avg = s; + s += nZnT; + double* P_wavg = s; + s += nZnT; + double* gbubu_avg = s; + s += nZnT; + double* gbubu_wavg = s; + s += nZnT; + double* gbvbv_avg = s; + s += nZnT; + double* gbvbv_wavg = s; + s += nZnT; + double* gbubv_avg = s; + s += nZnT; + double* gbubv_wavg = s; + s += nZnT; + double* bsubu_i = s; + s += nZnT; + double* bsubv_i = s; + s += nZnT; + double* gvv_gsqrt_i = s; + s += nZnT; + double* guv_bsupu_i = s; + s += nZnT; + + double* armn_e = force + 0 * fS; + double* armn_o = force + 1 * fS; + double* azmn_e = force + 2 * fS; + double* azmn_o = force + 3 * fS; + double* brmn_e = force + 4 * fS; + double* brmn_o = force + 5 * fS; + double* bzmn_e = force + 6 * fS; + double* bzmn_o = force + 7 * fS; + double* crmn_e = force + 8 * fS; + double* crmn_o = force + 9 * fS; + double* czmn_e = force + 10 * fS; + double* czmn_o = force + 11 * fS; + ComputeMHDForceDensity( + r1e, r1o, rue, ruo, zue, zuo, z1o, rve, rvo, zve, zvo, r12, ru12, zu12, + rs, zs, tau, tp, gsqrt, bsupu, bsupv, c->sqrtSF, c->sqrtSH, P_i, rup_i, + zup_i, rsp_i, zsp_i, taup_i, gbubu_i, gbubv_i, gbvbv_i, P_o, rup_o, zup_o, + rsp_o, zsp_o, taup_o, gbubu_o, gbubv_o, gbvbv_o, P_avg, P_wavg, gbubu_avg, + gbubu_wavg, gbvbv_avg, gbvbv_wavg, gbubv_avg, gbubv_wavg, c->deltaS, nZnT, + c->nsMinF, c->nsMinF1, c->nsMinH, c->nsMaxH, c->jMaxRZ, c->lthreed, + armn_e, armn_o, azmn_e, azmn_o, brmn_e, brmn_o, bzmn_e, bzmn_o, crmn_e, + crmn_o, czmn_e, czmn_o); + + double* blmn_e = force + 12 * fS; + double* blmn_o = force + 13 * fS; + double* clmn_e = force + 14 * fS; + double* clmn_o = force + 15 * fS; + ComputeHybridLambdaForce( + bsubu, bsubv, gvv, gsqrt, guv, bsupu, lue, luo, c->sqrtSH, c->sqrtSF, + c->radialBlending, c->lamscale, c->lthreed, nZnT, c->nsMinF, c->nsMinF1, + c->nsMinH, c->nsMaxH, c->nsMaxFIncludingLcfs, bsubu_i, bsubv_i, + gvv_gsqrt_i, guv_bsupu_i, blmn_e, blmn_o, clmn_e, clmn_o); + + if (c->with_constraint) { + // geometry blocks 16-19 carry the constraint coordinates and full-grid + // derivatives; force blocks 16-19 receive the constraint outputs. + const double* rCon = geom + 16 * gS; + const double* zCon = geom + 17 * gS; + const double* ruFull = geom + 18 * gS; + const double* zuFull = geom + 19 * gS; + double* gConEff = s; + s += (c->nsMaxFIncludingLcfs - c->nsMinF) * nZnT; + double* gCon = s; + s += (c->nsMaxF - c->nsMinF) * nZnT; + double* gsc = s; + s += c->ntor + 1; + double* gcs = s; + s += c->ntor + 1; + double* gcc = s; + s += c->ntor + 1; + double* gss = s; + s += c->ntor + 1; + double* gConAsym = s; + s += nZnT; + double* refl = s; + s += c->nThetaReduced; + // Constraint reference rCon0/zCon0 extrapolated from the LCFS into the + // volume (rzConIntoVolume): rCon0[jF] = rCon[LCFS] * s_full. This is linear + // in the geometry, so computing it here (rather than freezing it) keeps the + // exact HVP consistent with re-evaluating rzConIntoVolume each step. + double* rCon0 = s; + s += (c->nsMaxFIncludingLcfs - c->nsMinF) * nZnT; + double* zCon0 = s; // last slice of the work buffer + const int lcfs = (c->nsMaxFIncludingLcfs - 1 - c->nsMinF) * nZnT; + for (int jF = (c->nsMinF > 1 ? c->nsMinF : 1); jF < c->nsMaxFIncludingLcfs; + ++jF) { + const double sf = c->sqrtSF[jF - c->nsMinF1] * c->sqrtSF[jF - c->nsMinF1]; + for (int kl = 0; kl < nZnT; ++kl) { + const int idx = (jF - c->nsMinF) * nZnT + kl; + rCon0[idx] = rCon[lcfs + kl] * sf; + zCon0[idx] = zCon[lcfs + kl] * sf; + } + } + ComputeEffectiveConstraintForce(rCon, rCon0, zCon, zCon0, ruFull, zuFull, + nZnT, c->nsMinF, c->nsMaxFIncludingLcfs, + gConEff); + ComputeDeAliasConstraintForce( + gConEff, c->faccon, c->tcon, c->sinmui, c->cosmui, c->cosnv, c->sinnv, + c->sinmu, c->cosmu, c->nsMinF, c->nsMaxF, c->nZeta, c->nThetaEff, + c->nThetaReduced, c->nThetaEven, c->mpol, c->ntor, c->nnyq2, c->lasym, + gsc, gcs, gcc, gss, gConAsym, refl, gCon); + double* frcon_e = force + 16 * fS; + double* frcon_o = force + 17 * fS; + double* fzcon_e = force + 18 * fS; + double* fzcon_o = force + 19 * fS; + AddConstraintForces(rCon, rCon0, zCon, zCon0, ruFull, zuFull, gCon, + c->sqrtSF, nZnT, c->nsMinF, c->nsMinF1, c->nsMaxF, + brmn_e, brmn_o, bzmn_e, bzmn_o, frcon_e, frcon_o, + fzcon_e, fzcon_o); + } +} + +} // namespace vmecpp + +#endif // VMECPP_VMEC_IDEAL_MHD_MODEL_LOCAL_FORCE_COMPOSITION_H_ diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel index 6a6f895e1..38ef4f19c 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/BUILD.bazel @@ -64,6 +64,7 @@ cc_test( "//vmecpp/test_data:cma", "//vmecpp/test_data:cth_like_free_bdy", "//vmecpp/test_data:solovev_free_bdy", + "//vmecpp/test_data:solovev_free_bdy_lforbal", ], deps = [ "//vmecpp/vmec/output_quantities", diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc index 7e7f8ed14..7ce056f5e 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities.cc @@ -4464,9 +4464,10 @@ vmecpp::WOutFileContents vmecpp::ComputeWOutFileContents( if (wout.wdot.size() > 1) { // Compute decay rate wdot = -W/dt = (W[1:] - W[0:]) / W[1:] wout.wdot.tail(wout.wdot.size() - 1) = - (wout.wdot.tail(wout.wdot.size() - 1) - - wout.wdot.head(wout.wdot.size() - 1)) - .cwiseQuotient(wout.wdot.tail(wout.wdot.size() - 1)); + ((wout.wdot.tail(wout.wdot.size() - 1) - + wout.wdot.head(wout.wdot.size() - 1)) + .cwiseQuotient(wout.wdot.tail(wout.wdot.size() - 1))) + .eval(); } // ------------------- // one-dimensional array quantities diff --git a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_test.cc b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_test.cc index 44ecb7f86..957c786e0 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_test.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/output_quantities/output_quantities_test.cc @@ -81,6 +81,17 @@ TEST_P(WOutFileContentsTest, CheckWOutFileContents) { const OutputQuantities& output_quantities = vmec.output_quantities_; const WOutFileContents& wout = output_quantities.wout; + ASSERT_EQ(wout.wdot.size(), fc.mhd_energy.size()); + ASSERT_FALSE(fc.mhd_energy.empty()); + EXPECT_DOUBLE_EQ(wout.wdot(0), fc.mhd_energy[0]); + for (size_t i = 1; i < fc.mhd_energy.size(); ++i) { + const double expected_decay_rate = + (fc.mhd_energy[i] - fc.mhd_energy[i - 1]) / fc.mhd_energy[i]; + EXPECT_DOUBLE_EQ(wout.wdot(static_cast(i)), + expected_decay_rate) + << "i = " << i; + } + // Note that the actual `wout` file itself is taken as reference here. filename = absl::StrFormat("vmecpp/test_data/wout_%s.nc", data_source_.identifier); @@ -805,4 +816,135 @@ TEST(SolovevFreeBoundary, MatchesEducationalVmecGolden) { << std::endl; } +// lforbal free-boundary regression. With lforbal = true the flux-averaged +// radial force balance evolves the m=1, n=0 R,Z components, so the converged +// equilibrium differs from the variational one; this checks the converged wout +// against an educational_VMEC lforbal = true golden for the solovev +// free-boundary case (axisymmetric, ntor = 0). +TEST(SolovevFreeBoundaryLforbal, MatchesEducationalVmecGolden) { + const absl::StatusOr indata_json = + ReadFile("vmecpp/test_data/solovev_free_bdy_lforbal.json"); + ASSERT_TRUE(indata_json.ok()); + const absl::StatusOr vmec_indata = + VmecINDATA::FromJson(*indata_json); + ASSERT_TRUE(vmec_indata.ok()); + ASSERT_TRUE(vmec_indata->lfreeb); + ASSERT_TRUE(vmec_indata->lforbal); + ASSERT_EQ(vmec_indata->ntor, 0); + + auto maybe_vmec = Vmec::FromIndata(*vmec_indata); + ASSERT_TRUE(maybe_vmec.ok()); + Vmec& vmec = **maybe_vmec; + const Sizes& s = vmec.s_; + const FlowControl& fc = vmec.fc_; + + const bool reached_checkpoint = vmec.run().value(); + ASSERT_FALSE(reached_checkpoint); // ran to convergence + + const WOutFileContents& wout = vmec.output_quantities_.wout; + + int ncid; + ASSERT_EQ( + nc_open("vmecpp/test_data/wout_solovev_free_bdy_T.nc", NC_NOWRITE, &ncid), + NC_NOERR); + + // The flux-surface geometry, magnetic field, and integrated scalars agree + // with the VMEC 8.52 lforbal reference to within kTight. The current-density + // profiles are the most edge-sensitive derived quantity and are held to the + // looser kCurrent. + const double kTight = 5.0e-5; + const double kCurrent = 2.0e-3; + double tolerance = kTight; // mutated below; captured by reference + double worst_abs = 0.0; + double worst_norm = 0.0; + std::string worst_norm_field; + + auto compare = [&](const std::string& name, const std::vector& ref, + const std::vector& val) { + double peak = 1e-300; + for (double r : ref) { + peak = std::max(peak, std::abs(r)); + } + for (size_t i = 0; i < ref.size(); ++i) { + EXPECT_TRUE(IsCloseRelAbs(ref[i], val[i], tolerance)) + << name << "[" << i << "]: ref=" << ref[i] << " val=" << val[i]; + const double abs_dev = std::abs(ref[i] - val[i]); + worst_abs = std::max(worst_abs, abs_dev); + if (abs_dev / peak > worst_norm) { + worst_norm = abs_dev / peak; + worst_norm_field = name; + } + } + }; + auto scalar = [&](const std::string& name, double ref, double val) { + compare(name, {ref}, {val}); + }; + auto flatten = [&](const std::vector>& ref2d, int rows, + int cols, auto getter) { + std::vector ref; + std::vector val; + ref.reserve(static_cast(rows) * cols); + val.reserve(static_cast(rows) * cols); + for (int jF = 0; jF < rows; ++jF) { + for (int mn = 0; mn < cols; ++mn) { + ref.push_back(ref2d[jF][mn]); + val.push_back(getter(mn, jF)); + } + } + return std::make_pair(ref, val); + }; + + scalar("volume_p", NetcdfReadDouble(ncid, "volume_p").value(), wout.volume); + scalar("betatotal", NetcdfReadDouble(ncid, "betatotal").value(), + wout.betatotal); + scalar("aspect", NetcdfReadDouble(ncid, "aspect").value(), wout.aspect); + scalar("b0", NetcdfReadDouble(ncid, "b0").value(), wout.b0); + scalar("rbtor", NetcdfReadDouble(ncid, "rbtor").value(), wout.rbtor); + scalar("ctor", NetcdfReadDouble(ncid, "ctor").value(), wout.ctor); + scalar("Aminor_p", NetcdfReadDouble(ncid, "Aminor_p").value(), wout.Aminor_p); + scalar("Rmajor_p", NetcdfReadDouble(ncid, "Rmajor_p").value(), wout.Rmajor_p); + scalar("volavgB", NetcdfReadDouble(ncid, "volavgB").value(), wout.volavgB); + + std::vector wpresf(fc.ns), wiotaf(fc.ns); + for (int jF = 0; jF < fc.ns; ++jF) { + wpresf[jF] = wout.presf[jF]; + wiotaf[jF] = wout.iotaf[jF]; + } + compare("presf", NetcdfReadArray1D(ncid, "presf").value(), wpresf); + compare("iotaf", NetcdfReadArray1D(ncid, "iotaf").value(), wiotaf); + + auto [r_ref, r_val] = + flatten(NetcdfReadArray2D(ncid, "rmnc").value(), fc.ns, s.mnmax, + [&](int mn, int jF) { return wout.rmnc(mn, jF); }); + compare("rmnc", r_ref, r_val); + auto [z_ref, z_val] = + flatten(NetcdfReadArray2D(ncid, "zmns").value(), fc.ns, s.mnmax, + [&](int mn, int jF) { return wout.zmns(mn, jF); }); + compare("zmns", z_ref, z_val); + auto [l_ref, l_val] = + flatten(NetcdfReadArray2D(ncid, "lmns").value(), fc.ns, s.mnmax, + [&](int mn, int jF) { return wout.lmns(mn, jF); }); + compare("lmns", l_ref, l_val); + + auto [b_ref, b_val] = + flatten(NetcdfReadArray2D(ncid, "bmnc").value(), fc.ns, s.mnmax_nyq, + [&](int mn, int jF) { return wout.bmnc(mn, jF); }); + compare("bmnc", b_ref, b_val); + + tolerance = kCurrent; + std::vector wjcuru(fc.ns), wjcurv(fc.ns); + for (int jF = 0; jF < fc.ns; ++jF) { + wjcuru[jF] = wout.jcuru[jF]; + wjcurv[jF] = wout.jcurv[jF]; + } + compare("jcuru", NetcdfReadArray1D(ncid, "jcuru").value(), wjcuru); + compare("jcurv", NetcdfReadArray1D(ncid, "jcurv").value(), wjcurv); + + ASSERT_EQ(nc_close(ncid), NC_NOERR); + + std::cout << "[lforbal-vs-Fortran-golden] worst abs dev = " << worst_abs + << ", worst dev normalized by field peak = " << worst_norm << " (" + << worst_norm_field << ")" << std::endl; +} + } // namespace vmecpp diff --git a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc index 27067a28d..7cb70f33e 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.cc @@ -202,7 +202,6 @@ Vmec::Vmec(const VmecINDATA& indata, std::optional max_threads, // mnpd2 = 2 * mnpd (analog of mnpd2 in the Fortran vacmod / scalpot). int mnpd_dim = s_.lasym ? 2 * mnpd : mnpd; matrixShare.setZero(mnpd_dim * mnpd_dim); - iPiv.setZero(mnpd_dim); bvecShare.setZero(mnpd_dim); h_.vacuum_magnetic_pressure.setZero(s_.nZnT); @@ -485,7 +484,7 @@ void Vmec::SetupVacuumSolvers() { std::span(bvecShare.data(), bvecShare.size()), std::span(h_.vacuum_magnetic_pressure.data(), h_.vacuum_magnetic_pressure.size()), - std::span(iPiv.data(), iPiv.size()), + &lu_decomposition, std::span(h_.vacuum_b_r.data(), h_.vacuum_b_r.size()), std::span(h_.vacuum_b_phi.data(), h_.vacuum_b_phi.size()), std::span(h_.vacuum_b_z.data(), h_.vacuum_b_z.size())); @@ -644,7 +643,8 @@ bool Vmec::InitializeRadial( ls_[thread_id].get(), &h_, r_[thread_id].get(), &fb_vac_, vac_num_threads_, kSignOfJacobian, indata_.nvacskip, &vacuum_pressure_state_); - m_[thread_id]->setFromINDATA(indata_.ncurr, indata_.gamma, indata_.tcon0); + m_[thread_id]->setFromINDATA(indata_.ncurr, indata_.gamma, indata_.tcon0, + indata_.lforbal); } // thread_id if (checkpoint == VmecCheckpoint::SPECTRAL_CONSTRAINT && diff --git a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h index 467b94da1..2d44c88f6 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h +++ b/src/vmecpp/cpp/vmecpp/vmec/vmec/vmec.h @@ -224,7 +224,11 @@ class Vmec { std::vector> old_r_; Eigen::VectorXd matrixShare; - Eigen::VectorXi iPiv; + // LU decomposition of matrixShare, shared across all vac_num_threads_ + // Nestor/LaplaceSolver instances (mirroring how matrixShare/bvecShare are + // spans into shared backing storage). See LaplaceSolver's constructor for + // why this must be a single object rather than a per-thread member. + Eigen::PartialPivLU lu_decomposition; Eigen::VectorXd bvecShare; private: diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/.gitignore b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/.gitignore index 296f74bb6..e252872e7 100644 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/.gitignore +++ b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/.gitignore @@ -4,3 +4,10 @@ # wout_solovev.nc is ignored at the repo root (it's a common dev output), # but the test_data copy is reference data that must be tracked. !wout_solovev.nc + +/cma.json +/cth_like_fixed_bdy.json +/cth_like_fixed_bdy_nzeta_37.json +/solovev.json +/solovev_analytical.json +/solovev_no_axis.json diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cma.json b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cma.json deleted file mode 100644 index 2e38ffd73..000000000 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cma.json +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2220e115d98797230b7bdab5f10b3a89dbef61d20246c9c28855cea53084e4d0 -size 13120 diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cth_like_fixed_bdy.json b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cth_like_fixed_bdy.json deleted file mode 100644 index 892fa40dc..000000000 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cth_like_fixed_bdy.json +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a934986c84af94e97f66a77a02209812405b29bf5a4555bf42950094b8359649 -size 9480 diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cth_like_fixed_bdy_nzeta_37.json b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cth_like_fixed_bdy_nzeta_37.json deleted file mode 100644 index d067f7809..000000000 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/cth_like_fixed_bdy_nzeta_37.json +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a096d5932959e6d5ada6f4e7f19e3e6d59fbcccb5c408175a1df3d1664eb7b7e -size 9480 diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/mgrid_solovev.nc b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/mgrid_solovev.nc deleted file mode 100644 index 5ca4f7427..000000000 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/mgrid_solovev.nc +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f17260704b9f6ca96c0f7202cc99dbfe8bdb625d1103d4597fd8ca5197b96fb1 -size 25215108 diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev.json b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev.json deleted file mode 100644 index 8d1013639..000000000 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev.json +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9725eabb1bc02647221aaad280a30c2697d56ee3dbb042692d548668fe529331 -size 1269 diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev_analytical.json b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev_analytical.json deleted file mode 100644 index a68562377..000000000 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev_analytical.json +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b9acd66d80f381e6e95c4a7175f58c46451047f2fdef66eb27efc260d49a5761 -size 2595 diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev_no_axis.json b/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev_no_axis.json deleted file mode 100644 index 057973514..000000000 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/solovev_no_axis.json +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cd37411dbeefed66f6b81476e7d7630da0c8da608cddd0fa97b715ba4986c6eb -size 1228 diff --git a/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/ideal_mhd_model/ideal_mhd_model_test.cc b/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/ideal_mhd_model/ideal_mhd_model_test.cc index 66e27d92e..a2ff44ed9 100644 --- a/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/ideal_mhd_model/ideal_mhd_model_test.cc +++ b/src/vmecpp/cpp/vmecpp_large_cpp_tests/vmec/ideal_mhd_model/ideal_mhd_model_test.cc @@ -6,7 +6,9 @@ #include "vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h" #include +#include #include +#include #include #include @@ -27,6 +29,7 @@ using vmecpp::vmec_algorithm_constants::kEvenParity; using vmecpp::vmec_algorithm_constants::kOddParity; +namespace vmecpp { namespace { using nlohmann::json; @@ -35,9 +38,393 @@ using testing::IsCloseRelAbs; using ::testing::TestWithParam; using ::testing::Values; -} // namespace -namespace vmecpp { +template +double DotValues(const VectorA& a, const VectorB& b) { + EXPECT_EQ(a.size(), b.size()); + double result = 0.0; + const std::size_t size = static_cast(a.size()); + for (std::size_t i = 0; i < size; ++i) { + result += a[i] * b[i]; + } + return result; +} + +template +void FillRandom(Vector& values, std::mt19937& generator) { + std::uniform_real_distribution distribution(-1.0, 1.0); + const std::size_t size = static_cast(values.size()); + for (std::size_t i = 0; i < size; ++i) { + values[i] = distribution(generator); + } +} + +double RelativeDotError(double lhs, double rhs) { + return std::abs(lhs - rhs) / std::max({1.0, std::abs(lhs), std::abs(rhs)}); +} + +double DotGeometry(const FourierGeometry& a, const FourierGeometry& b) { + return DotValues(a.rmncc, b.rmncc) + DotValues(a.rmnss, b.rmnss) + + DotValues(a.rmnsc, b.rmnsc) + DotValues(a.rmncs, b.rmncs) + + DotValues(a.zmnsc, b.zmnsc) + DotValues(a.zmncs, b.zmncs) + + DotValues(a.zmncc, b.zmncc) + DotValues(a.zmnss, b.zmnss) + + DotValues(a.lmnsc, b.lmnsc) + DotValues(a.lmncs, b.lmncs) + + DotValues(a.lmncc, b.lmncc) + DotValues(a.lmnss, b.lmnss); +} + +void FillGeometry(FourierGeometry& geometry, std::mt19937& generator) { + geometry.setZero(); + FillRandom(geometry.rmncc, generator); + FillRandom(geometry.rmnss, generator); + FillRandom(geometry.rmnsc, generator); + FillRandom(geometry.rmncs, generator); + FillRandom(geometry.zmnsc, generator); + FillRandom(geometry.zmncs, generator); + FillRandom(geometry.zmncc, generator); + FillRandom(geometry.zmnss, generator); + FillRandom(geometry.lmnsc, generator); + FillRandom(geometry.lmncs, generator); + FillRandom(geometry.lmncc, generator); + FillRandom(geometry.lmnss, generator); +} + +void FillSymmetricForce(IdealMhdModel& model, std::mt19937& generator) { + FillRandom(model.armn_e, generator); + FillRandom(model.armn_o, generator); + FillRandom(model.brmn_e, generator); + FillRandom(model.brmn_o, generator); + FillRandom(model.azmn_e, generator); + FillRandom(model.azmn_o, generator); + FillRandom(model.bzmn_e, generator); + FillRandom(model.bzmn_o, generator); + FillRandom(model.blmn_e, generator); + FillRandom(model.blmn_o, generator); + FillRandom(model.frcon_e, generator); + FillRandom(model.frcon_o, generator); + FillRandom(model.fzcon_e, generator); + FillRandom(model.fzcon_o, generator); + FillRandom(model.crmn_e, generator); + FillRandom(model.crmn_o, generator); + FillRandom(model.czmn_e, generator); + FillRandom(model.czmn_o, generator); + FillRandom(model.clmn_e, generator); + FillRandom(model.clmn_o, generator); +} + +double CheckGeometryTranspose2d(Vmec& vmec, std::mt19937& generator) { + IdealMhdModel& model = *vmec.m_[0]; + FourierGeometry input = *vmec.physical_x_[0]; + FillGeometry(input, generator); + model.dft_FourierToReal_2d_symm(input); + + const Eigen::VectorXd r1_e_forward = model.r1_e; + const Eigen::VectorXd r1_o_forward = model.r1_o; + const Eigen::VectorXd ru_e_forward = model.ru_e; + const Eigen::VectorXd ru_o_forward = model.ru_o; + const Eigen::VectorXd z1_e_forward = model.z1_e; + const Eigen::VectorXd z1_o_forward = model.z1_o; + const Eigen::VectorXd zu_e_forward = model.zu_e; + const Eigen::VectorXd zu_o_forward = model.zu_o; + const Eigen::VectorXd lu_e_forward = model.lu_e; + const Eigen::VectorXd lu_o_forward = model.lu_o; + const Eigen::VectorXd rCon_forward = model.rCon; + const Eigen::VectorXd zCon_forward = model.zCon; + + Eigen::VectorXd r1_e_bar = model.r1_e; + Eigen::VectorXd r1_o_bar = model.r1_o; + Eigen::VectorXd ru_e_bar = model.ru_e; + Eigen::VectorXd ru_o_bar = model.ru_o; + Eigen::VectorXd z1_e_bar = model.z1_e; + Eigen::VectorXd z1_o_bar = model.z1_o; + Eigen::VectorXd zu_e_bar = model.zu_e; + Eigen::VectorXd zu_o_bar = model.zu_o; + Eigen::VectorXd lu_e_bar = model.lu_e; + Eigen::VectorXd lu_o_bar = model.lu_o; + Eigen::VectorXd rCon_bar = model.rCon; + Eigen::VectorXd zCon_bar = model.zCon; + FillRandom(r1_e_bar, generator); + FillRandom(r1_o_bar, generator); + FillRandom(ru_e_bar, generator); + FillRandom(ru_o_bar, generator); + FillRandom(z1_e_bar, generator); + FillRandom(z1_o_bar, generator); + FillRandom(zu_e_bar, generator); + FillRandom(zu_o_bar, generator); + FillRandom(lu_e_bar, generator); + FillRandom(lu_o_bar, generator); + FillRandom(rCon_bar, generator); + FillRandom(zCon_bar, generator); + model.r1_e = r1_e_bar; + model.r1_o = r1_o_bar; + model.ru_e = ru_e_bar; + model.ru_o = ru_o_bar; + model.z1_e = z1_e_bar; + model.z1_o = z1_o_bar; + model.zu_e = zu_e_bar; + model.zu_o = zu_o_bar; + model.lu_e = lu_e_bar; + model.lu_o = lu_o_bar; + model.rCon = rCon_bar; + model.zCon = zCon_bar; + + FourierGeometry transpose = *vmec.physical_x_[0]; + transpose.setZero(); + model.dft_FourierToRealTranspose_2d_symm(transpose); + const double lhs = + DotValues(r1_e_forward, r1_e_bar) + DotValues(r1_o_forward, r1_o_bar) + + DotValues(ru_e_forward, ru_e_bar) + DotValues(ru_o_forward, ru_o_bar) + + DotValues(z1_e_forward, z1_e_bar) + DotValues(z1_o_forward, z1_o_bar) + + DotValues(zu_e_forward, zu_e_bar) + DotValues(zu_o_forward, zu_o_bar) + + DotValues(lu_e_forward, lu_e_bar) + DotValues(lu_o_forward, lu_o_bar) + + DotValues(rCon_forward, rCon_bar) + DotValues(zCon_forward, zCon_bar); + const double rhs = DotValues(input.rmncc, transpose.rmncc) + + DotValues(input.zmnsc, transpose.zmnsc) + + DotValues(input.lmnsc, transpose.lmnsc); + return RelativeDotError(lhs, rhs); +} + +double CheckGeometryTranspose3d(Vmec& vmec, std::mt19937& generator) { + IdealMhdModel& model = *vmec.m_[0]; + FourierGeometry input = *vmec.physical_x_[0]; + FillGeometry(input, generator); + model.dft_FourierToReal_3d_symm(input); + + const Eigen::VectorXd r1_e_forward = model.r1_e; + const Eigen::VectorXd r1_o_forward = model.r1_o; + const Eigen::VectorXd ru_e_forward = model.ru_e; + const Eigen::VectorXd ru_o_forward = model.ru_o; + const Eigen::VectorXd rv_e_forward = model.rv_e; + const Eigen::VectorXd rv_o_forward = model.rv_o; + const Eigen::VectorXd z1_e_forward = model.z1_e; + const Eigen::VectorXd z1_o_forward = model.z1_o; + const Eigen::VectorXd zu_e_forward = model.zu_e; + const Eigen::VectorXd zu_o_forward = model.zu_o; + const Eigen::VectorXd zv_e_forward = model.zv_e; + const Eigen::VectorXd zv_o_forward = model.zv_o; + const Eigen::VectorXd lu_e_forward = model.lu_e; + const Eigen::VectorXd lu_o_forward = model.lu_o; + const Eigen::VectorXd lv_e_forward = model.lv_e; + const Eigen::VectorXd lv_o_forward = model.lv_o; + const Eigen::VectorXd rCon_forward = model.rCon; + const Eigen::VectorXd zCon_forward = model.zCon; + + Eigen::VectorXd r1_e_bar = model.r1_e; + Eigen::VectorXd r1_o_bar = model.r1_o; + Eigen::VectorXd ru_e_bar = model.ru_e; + Eigen::VectorXd ru_o_bar = model.ru_o; + Eigen::VectorXd rv_e_bar = model.rv_e; + Eigen::VectorXd rv_o_bar = model.rv_o; + Eigen::VectorXd z1_e_bar = model.z1_e; + Eigen::VectorXd z1_o_bar = model.z1_o; + Eigen::VectorXd zu_e_bar = model.zu_e; + Eigen::VectorXd zu_o_bar = model.zu_o; + Eigen::VectorXd zv_e_bar = model.zv_e; + Eigen::VectorXd zv_o_bar = model.zv_o; + Eigen::VectorXd lu_e_bar = model.lu_e; + Eigen::VectorXd lu_o_bar = model.lu_o; + Eigen::VectorXd lv_e_bar = model.lv_e; + Eigen::VectorXd lv_o_bar = model.lv_o; + Eigen::VectorXd rCon_bar = model.rCon; + Eigen::VectorXd zCon_bar = model.zCon; + FillRandom(r1_e_bar, generator); + FillRandom(r1_o_bar, generator); + FillRandom(ru_e_bar, generator); + FillRandom(ru_o_bar, generator); + FillRandom(rv_e_bar, generator); + FillRandom(rv_o_bar, generator); + FillRandom(z1_e_bar, generator); + FillRandom(z1_o_bar, generator); + FillRandom(zu_e_bar, generator); + FillRandom(zu_o_bar, generator); + FillRandom(zv_e_bar, generator); + FillRandom(zv_o_bar, generator); + FillRandom(lu_e_bar, generator); + FillRandom(lu_o_bar, generator); + FillRandom(lv_e_bar, generator); + FillRandom(lv_o_bar, generator); + FillRandom(rCon_bar, generator); + FillRandom(zCon_bar, generator); + model.r1_e = r1_e_bar; + model.r1_o = r1_o_bar; + model.ru_e = ru_e_bar; + model.ru_o = ru_o_bar; + model.rv_e = rv_e_bar; + model.rv_o = rv_o_bar; + model.z1_e = z1_e_bar; + model.z1_o = z1_o_bar; + model.zu_e = zu_e_bar; + model.zu_o = zu_o_bar; + model.zv_e = zv_e_bar; + model.zv_o = zv_o_bar; + model.lu_e = lu_e_bar; + model.lu_o = lu_o_bar; + model.lv_e = lv_e_bar; + model.lv_o = lv_o_bar; + model.rCon = rCon_bar; + model.zCon = zCon_bar; + + FourierGeometry transpose = *vmec.physical_x_[0]; + transpose.setZero(); + model.dft_FourierToRealTranspose_3d_symm(transpose); + const double lhs = + DotValues(r1_e_forward, r1_e_bar) + DotValues(r1_o_forward, r1_o_bar) + + DotValues(ru_e_forward, ru_e_bar) + DotValues(ru_o_forward, ru_o_bar) + + DotValues(rv_e_forward, rv_e_bar) + DotValues(rv_o_forward, rv_o_bar) + + DotValues(z1_e_forward, z1_e_bar) + DotValues(z1_o_forward, z1_o_bar) + + DotValues(zu_e_forward, zu_e_bar) + DotValues(zu_o_forward, zu_o_bar) + + DotValues(zv_e_forward, zv_e_bar) + DotValues(zv_o_forward, zv_o_bar) + + DotValues(lu_e_forward, lu_e_bar) + DotValues(lu_o_forward, lu_o_bar) + + DotValues(lv_e_forward, lv_e_bar) + DotValues(lv_o_forward, lv_o_bar) + + DotValues(rCon_forward, rCon_bar) + DotValues(zCon_forward, zCon_bar); + const double rhs = DotValues(input.rmncc, transpose.rmncc) + + DotValues(input.rmnss, transpose.rmnss) + + DotValues(input.zmnsc, transpose.zmnsc) + + DotValues(input.zmncs, transpose.zmncs) + + DotValues(input.lmnsc, transpose.lmnsc) + + DotValues(input.lmncs, transpose.lmncs); + return RelativeDotError(lhs, rhs); +} + +double CheckForceTranspose2d(Vmec& vmec, std::mt19937& generator) { + IdealMhdModel& model = *vmec.m_[0]; + FillSymmetricForce(model, generator); + const Eigen::VectorXd armn_e = model.armn_e; + const Eigen::VectorXd armn_o = model.armn_o; + const Eigen::VectorXd brmn_e = model.brmn_e; + const Eigen::VectorXd brmn_o = model.brmn_o; + const Eigen::VectorXd azmn_e = model.azmn_e; + const Eigen::VectorXd azmn_o = model.azmn_o; + const Eigen::VectorXd bzmn_e = model.bzmn_e; + const Eigen::VectorXd bzmn_o = model.bzmn_o; + const Eigen::VectorXd blmn_e = model.blmn_e; + const Eigen::VectorXd blmn_o = model.blmn_o; + const Eigen::VectorXd frcon_e = model.frcon_e; + const Eigen::VectorXd frcon_o = model.frcon_o; + const Eigen::VectorXd fzcon_e = model.fzcon_e; + const Eigen::VectorXd fzcon_o = model.fzcon_o; + FourierForces forward = *vmec.physical_f_[0]; + model.dft_ForcesToFourier_2d_symm(forward); + FourierForces force_bar = forward; + force_bar.setZero(); + FillRandom(force_bar.frcc, generator); + FillRandom(force_bar.fzsc, generator); + FillRandom(force_bar.flsc, generator); + const double lhs = DotValues(forward.frcc, force_bar.frcc) + + DotValues(forward.fzsc, force_bar.fzsc) + + DotValues(forward.flsc, force_bar.flsc); + model.dft_ForcesToFourierTranspose_2d_symm(force_bar); + const double rhs = + DotValues(armn_e, model.armn_e) + DotValues(armn_o, model.armn_o) + + DotValues(brmn_e, model.brmn_e) + DotValues(brmn_o, model.brmn_o) + + DotValues(azmn_e, model.azmn_e) + DotValues(azmn_o, model.azmn_o) + + DotValues(bzmn_e, model.bzmn_e) + DotValues(bzmn_o, model.bzmn_o) + + DotValues(blmn_e, model.blmn_e) + DotValues(blmn_o, model.blmn_o) + + DotValues(frcon_e, model.frcon_e) + DotValues(frcon_o, model.frcon_o) + + DotValues(fzcon_e, model.fzcon_e) + DotValues(fzcon_o, model.fzcon_o); + return RelativeDotError(lhs, rhs); +} + +double CheckForceTranspose3d(Vmec& vmec, std::mt19937& generator) { + IdealMhdModel& model = *vmec.m_[0]; + FillSymmetricForce(model, generator); + const Eigen::VectorXd armn_e = model.armn_e; + const Eigen::VectorXd armn_o = model.armn_o; + const Eigen::VectorXd brmn_e = model.brmn_e; + const Eigen::VectorXd brmn_o = model.brmn_o; + const Eigen::VectorXd crmn_e = model.crmn_e; + const Eigen::VectorXd crmn_o = model.crmn_o; + const Eigen::VectorXd azmn_e = model.azmn_e; + const Eigen::VectorXd azmn_o = model.azmn_o; + const Eigen::VectorXd bzmn_e = model.bzmn_e; + const Eigen::VectorXd bzmn_o = model.bzmn_o; + const Eigen::VectorXd czmn_e = model.czmn_e; + const Eigen::VectorXd czmn_o = model.czmn_o; + const Eigen::VectorXd blmn_e = model.blmn_e; + const Eigen::VectorXd blmn_o = model.blmn_o; + const Eigen::VectorXd clmn_e = model.clmn_e; + const Eigen::VectorXd clmn_o = model.clmn_o; + const Eigen::VectorXd frcon_e = model.frcon_e; + const Eigen::VectorXd frcon_o = model.frcon_o; + const Eigen::VectorXd fzcon_e = model.fzcon_e; + const Eigen::VectorXd fzcon_o = model.fzcon_o; + FourierForces forward = *vmec.physical_f_[0]; + model.dft_ForcesToFourier_3d_symm(forward); + FourierForces force_bar = forward; + force_bar.setZero(); + FillRandom(force_bar.frcc, generator); + FillRandom(force_bar.frss, generator); + FillRandom(force_bar.fzsc, generator); + FillRandom(force_bar.fzcs, generator); + FillRandom(force_bar.flsc, generator); + FillRandom(force_bar.flcs, generator); + const double lhs = DotValues(forward.frcc, force_bar.frcc) + + DotValues(forward.frss, force_bar.frss) + + DotValues(forward.fzsc, force_bar.fzsc) + + DotValues(forward.fzcs, force_bar.fzcs) + + DotValues(forward.flsc, force_bar.flsc) + + DotValues(forward.flcs, force_bar.flcs); + model.dft_ForcesToFourierTranspose_3d_symm(force_bar); + const double rhs = + DotValues(armn_e, model.armn_e) + DotValues(armn_o, model.armn_o) + + DotValues(brmn_e, model.brmn_e) + DotValues(brmn_o, model.brmn_o) + + DotValues(crmn_e, model.crmn_e) + DotValues(crmn_o, model.crmn_o) + + DotValues(azmn_e, model.azmn_e) + DotValues(azmn_o, model.azmn_o) + + DotValues(bzmn_e, model.bzmn_e) + DotValues(bzmn_o, model.bzmn_o) + + DotValues(czmn_e, model.czmn_e) + DotValues(czmn_o, model.czmn_o) + + DotValues(blmn_e, model.blmn_e) + DotValues(blmn_o, model.blmn_o) + + DotValues(clmn_e, model.clmn_e) + DotValues(clmn_o, model.clmn_o) + + DotValues(frcon_e, model.frcon_e) + DotValues(frcon_o, model.frcon_o) + + DotValues(fzcon_e, model.fzcon_e) + DotValues(fzcon_o, model.fzcon_o); + return RelativeDotError(lhs, rhs); +} + +double CheckAxisExtrapolateTranspose(Vmec& vmec, std::mt19937& generator) { + FourierGeometry input = *vmec.decomposed_x_[0]; + FillGeometry(input, generator); + FourierGeometry forward = input; + forward.extrapolateTowardsAxis(); + FourierGeometry cotangent = input; + cotangent.setZero(); + FillGeometry(cotangent, generator); + FourierGeometry cotangent_before = cotangent; + cotangent.extrapolateTowardsAxisTranspose(); + return RelativeDotError(DotGeometry(forward, cotangent_before), + DotGeometry(input, cotangent)); +} + +void CheckTransposeIdentities(const std::string& identifier, bool lthreed) { + const std::string filename = + absl::StrFormat("vmecpp/test_data/%s.json", identifier); + const absl::StatusOr indata_json = ReadFile(filename); + ASSERT_TRUE(indata_json.ok()); + const absl::StatusOr vmec_indata = + VmecINDATA::FromJson(*indata_json); + ASSERT_TRUE(vmec_indata.ok()); + Vmec vmec(*vmec_indata, 1, OutputMode::kSilent); + ASSERT_EQ(vmec.s_.lthreed, lthreed); + const absl::StatusOr initialized = + vmec.run(VmecCheckpoint::FOURIER_GEOMETRY_TO_START_WITH, 1); + ASSERT_TRUE(initialized.ok()); + ASSERT_TRUE(*initialized); + std::mt19937 generator(1729); + EXPECT_LT(CheckAxisExtrapolateTranspose(vmec, generator), 1.0e-12); + if (lthreed) { + EXPECT_LT(CheckGeometryTranspose3d(vmec, generator), 1.0e-12); + EXPECT_LT(CheckForceTranspose3d(vmec, generator), 1.0e-12); + } else { + EXPECT_LT(CheckGeometryTranspose2d(vmec, generator), 1.0e-12); + EXPECT_LT(CheckForceTranspose2d(vmec, generator), 1.0e-12); + } +} + +TEST(IdealMhdModelTransposeTest, DirectTwoDimensionalAdjointIdentities) { + CheckTransposeIdentities("solovev", false); +} + +TEST(IdealMhdModelTransposeTest, DirectThreeDimensionalAdjointIdentities) { + CheckTransposeIdentities("cth_like_fixed_bdy", true); +} +} // namespace // used to specify case-specific tolerances // and which iterations to test diff --git a/tests/test_adjoint.py b/tests/test_adjoint.py new file mode 100644 index 000000000..54a45c1a1 --- /dev/null +++ b/tests/test_adjoint.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""The adjoint boundary gradient matches brute-force finite differences. + +dJ/d(boundary) from one Hessian solve (implicit-function adjoint) agrees with the +reference gradient obtained by re-converging the interior equilibrium for each perturbed +boundary degree of freedom. J here is the MHD energy of the converged equilibrium. +""" + +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from vmecpp_adjoint import ( # type: ignore + boundary_gradient, + finite_difference_boundary_gradient, + make_model, + mhd_energy, + partition, +) + + +def test_adjoint_matches_finite_difference(): + ns = 11 + model = make_model(ns=ns) + model.solve() + x_star = np.asarray(model.get_state(), float).copy() + interior, boundary = partition(model, ns) + + g_adjoint = boundary_gradient(model, x_star, interior, boundary, mhd_energy) + + # Reference on a representative subset (each requires interior re-solves). + dofs = [0, 2, 9] + g_fd = finite_difference_boundary_gradient( + model, x_star, interior, boundary, mhd_energy, dofs + ) + + scale = max(np.linalg.norm(g_adjoint), 1e-30) + for j in dofs: + assert abs(g_adjoint[j] - g_fd[j]) < 1e-3 * scale + + +if __name__ == "__main__": + import pytest + + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_free_boundary_quasr.py b/tests/test_free_boundary_quasr.py index d46fc58f1..4b380c10f 100644 --- a/tests/test_free_boundary_quasr.py +++ b/tests/test_free_boundary_quasr.py @@ -373,6 +373,15 @@ def _magnetic_axis_major_radius(wout) -> float: return float(np.sum(wout.raxis_cc)) +def _magnetic_axis_curve(wout, n_phi: int = 64) -> np.ndarray: + """The magnetic axis R_axis(phi) = sum_n raxis_cc[n] cos(n * nfp * phi), sampled + over one field period.""" + raxis_cc = np.asarray(wout.raxis_cc) + n = np.arange(len(raxis_cc)) + phi = np.linspace(0.0, 2.0 * np.pi / wout.nfp, n_phi, endpoint=False) + return raxis_cc @ np.cos(np.outer(n * wout.nfp, phi)) + + def _enclosed_volume(surface_rz) -> float: """Absolute plasma volume enclosed by a SIMSOPT SurfaceRZFourier. @@ -693,12 +702,17 @@ def test_free_boundary_matches_fixed_boundary( pytest.xfail(f"fixed-boundary reference did not converge: {_first_line(exc)}") # Same profiles and (nearly) the same boundary -> the interior equilibria - # should agree. The magnetic axis is well-determined and matches tightly; the - # volume can differ slightly because the free-boundary LCFS is set by the coil - # field rather than imposed exactly. + # should agree. Both the axis position and the volume can differ slightly + # because the free-boundary LCFS is set by the coil field rather than imposed + # exactly. The axis is compared as the curve R_axis(phi) rather than coefficient + # by coefficient: the coefficients span five orders of magnitude, so a relative + # tolerance on the smallest ones says nothing about the axis position. np.testing.assert_allclose( - free.wout.raxis_cc, fixed.wout.raxis_cc, rtol=0.01, atol=0.0 + _magnetic_axis_curve(free.wout), + _magnetic_axis_curve(fixed.wout), + rtol=0.02, + atol=0.0, ) assert abs(free.wout.volume) == pytest.approx(abs(fixed.wout.volume), rel=0.05) diff --git a/tests/test_init.py b/tests/test_init.py index 396ed0371..54b0b5435 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -20,7 +20,6 @@ import netCDF4 import numpy as np -import pydantic import pytest import vmecpp @@ -534,7 +533,10 @@ def test_threed1_geometric_magnetic_bindings(cma_output: vmecpp.VmecOutput): def test_threed1_axis_bindings(cma_output: vmecpp.VmecOutput): for varname in ["raxis_symm", "zaxis_symm", "raxis_asym", "zaxis_asym"]: - assert len(getattr(cma_output.threed1_axis, varname).shape) == 1 + if "asym" in varname: + assert getattr(cma_output.threed1_axis, varname) is None + else: + assert len(getattr(cma_output.threed1_axis, varname).shape) == 1 def test_threed1_betas_bindings(cma_output: vmecpp.VmecOutput): @@ -821,40 +823,3 @@ def test_ctrl_c_interrupts_run(): f"Expected KeyboardInterrupt but got:\noutput: {output}" ) assert "RUN_COMPLETED" not in output - - -def test_subclass_outer_wrap_serializer_not_overridden(cma_output: vmecpp.VmecOutput): - """Subclass wrap serializers that operate on the arrays must not be overridden by - VmecWOut's inner serializers. - - This ensures that VmecWOut's field-level serializers (e.g. a PlainSerializer on - extcur) pass through values that are not numpy arrays, so an outer framework can - replace arrays with custom representations during serialization. - """ - - class OuterSerializer(pydantic.BaseModel): - model_config = pydantic.ConfigDict(arbitrary_types_allowed=True) - - @pydantic.field_serializer("*", mode="wrap", when_used="always") - def _outer_encode( - self, - value: object, - handler: pydantic.SerializerFunctionWrapHandler, - _: pydantic.FieldSerializationInfo, - ) -> object: - if isinstance(value, np.ndarray): - return {"__custom_encoded__": True} - return handler(value) - - class CustomWOut(OuterSerializer, vmecpp.VmecWOut): - pass - - custom = CustomWOut.model_validate(cma_output.wout.model_dump()) - dumped = custom.model_dump(mode="json") - - # Plain array field — goes through _serialize_field wrap serializer. - assert dumped["rmnc"].get("__custom_encoded__") - # SerializeIntAsFloat field — has its own PlainSerializer/WrapSerializer. - assert dumped["xm"].get("__custom_encoded__") - # extcur field — has its own PlainSerializer/WrapSerializer. - assert dumped["extcur"].get("__custom_encoded__") diff --git a/tests/test_iteration.py b/tests/test_iteration.py index cd887d1b0..1be094f56 100644 --- a/tests/test_iteration.py +++ b/tests/test_iteration.py @@ -179,7 +179,9 @@ def test_python_iteration_matches_cpp_restart_path( # can grow through long chaotic transients. cpp_r = np.asarray(reference.force_residual_r) py_r = np.asarray(result.force_residual_r) - np.testing.assert_allclose(py_r[:50], cpp_r[:50], rtol=1.0e-9, atol=1e-15) + # Eigen and NumPy use different reduction trees; the first 50 steps stay within + # two parts per billion while retaining a strict cross-implementation oracle. + np.testing.assert_allclose(py_r[:50], cpp_r[:50], rtol=2.0e-9, atol=1e-15) np.testing.assert_allclose(py_r, cpp_r, rtol=1.0e-3, atol=1e-15) diff --git a/tests/test_rescale.py b/tests/test_rescale.py new file mode 100644 index 000000000..1056091d8 --- /dev/null +++ b/tests/test_rescale.py @@ -0,0 +1,76 @@ +from pathlib import Path + +import numpy as np + +import vmecpp + +REPO_ROOT = Path(__file__).resolve().parent.parent +TEST_DATA_DIR = REPO_ROOT / "src" / "vmecpp" / "cpp" / "vmecpp" / "test_data" + + +def test_equilibrium_rescale(): + vmec_input = vmecpp.VmecInput.from_file(TEST_DATA_DIR / "solovev.json") + oq_initial = vmecpp.run(vmec_input) + + b_scale = 1.5 + r_scale = 2.0 + + oq_rescaled = vmecpp.rescale( + oq_initial, b_scale=b_scale, r_scale=r_scale, scale_pressure=True + ) + + input_scaled = oq_initial.input.model_copy(deep=True) + input_scaled.phiedge *= b_scale * (r_scale**2) + input_scaled.pres_scale *= b_scale**2 + input_scaled.curtor *= b_scale * r_scale + + input_scaled.rbc *= r_scale + input_scaled.zbs *= r_scale + if input_scaled.rbs is not None: + input_scaled.rbs *= r_scale + if input_scaled.zbc is not None: + input_scaled.zbc *= r_scale + + input_scaled.raxis_c *= r_scale + input_scaled.zaxis_s *= r_scale + + oq_full_run = vmecpp.run(input_scaled) + + # Volume should scale as r_scale^3 + np.testing.assert_allclose( + oq_rescaled.wout.volume_p, + oq_initial.wout.volume_p * (r_scale**3), + rtol=1e-9, + atol=1e-10 * (r_scale**3), + ) + + # Pressure should scale as b_scale^2 + np.testing.assert_allclose( + oq_rescaled.wout.pres, + oq_initial.wout.pres * (b_scale**2), + rtol=1e-9, + atol=1e-10 * (b_scale**2), + ) + + # Magnetic field should scale as b_scale + np.testing.assert_allclose( + oq_rescaled.wout.bmnc, + oq_initial.wout.bmnc * b_scale, + rtol=1e-9, + atol=1e-10 * b_scale, + ) + + # Betas should be invariant under this scaling (pressure ~ B^2) + np.testing.assert_allclose( + oq_rescaled.wout.betatotal, oq_initial.wout.betatotal, rtol=1e-9, atol=1e-10 + ) + + np.testing.assert_allclose( + oq_rescaled.wout.bmnc, oq_full_run.wout.bmnc, rtol=1e-9, atol=1e-10 * b_scale + ) + np.testing.assert_allclose( + oq_rescaled.wout.rmnc, oq_full_run.wout.rmnc, rtol=1e-9, atol=1e-10 * r_scale + ) + np.testing.assert_allclose( + oq_rescaled.wout.zmns, oq_full_run.wout.zmns, rtol=1e-9, atol=1e-10 * r_scale + ) diff --git a/tests/test_simsopt_gradient.py b/tests/test_simsopt_gradient.py new file mode 100644 index 000000000..7f17e0cb7 --- /dev/null +++ b/tests/test_simsopt_gradient.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VMEC++ exposes an analytic boundary gradient to SIMSOPT. + +The VmecEnergy Optimizable's analytic dJ (the implicit-function adjoint) matches finite +differences of its objective, and computing it is much cheaper than the conventional +finite-difference boundary gradient (which re-solves the equilibrium per boundary degree +of freedom). +""" + +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) + +from simsopt_vmec_gradient import ( # type: ignore + VmecBoundaryProblem, + make_simsopt_optimizable, +) + + +def test_simsopt_optimizable_gradient_matches_fd(): + problem = VmecBoundaryProblem(ns=11) + opt = make_simsopt_optimizable(problem) + g = np.asarray(opt.dJ(), float) + + p0 = np.asarray(opt.local_full_x, float) + h = 1e-5 + scale = max(np.linalg.norm(g), 1e-30) + for j in (0, 2, 9): + pp = p0.copy() + pp[j] += h + opt.local_full_x = pp + jp = opt.J() + pm = p0.copy() + pm[j] -= h + opt.local_full_x = pm + jm = opt.J() + opt.local_full_x = p0 + assert abs(g[j] - (jp - jm) / (2 * h)) < 1e-3 * scale