diff --git a/pixi.lock b/pixi.lock index f775eb463f56..2b4a35c55fd5 100644 --- a/pixi.lock +++ b/pixi.lock @@ -153,6 +153,7 @@ environments: linux-64: - conda: https://prefix.dev/conda-forge/linux-64/_openmp_mutex-4.5-6_kmp_llvm.conda - conda: https://prefix.dev/conda-forge/noarch/array-api-strict-2.4.1-pyhe01879c_0.conda + - conda: https://prefix.dev/conda-forge/noarch/asttokens-3.0.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/attrs-25.4.0-pyh71513ae_0.conda - conda: https://prefix.dev/conda-forge/linux-64/blas-devel-3.9.0-38_hcf00494_mkl.conda - conda: https://prefix.dev/conda-forge/linux-64/brotli-python-1.2.0-py313h09d1b84_0.conda @@ -167,8 +168,10 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/cpython-3.13.9-py313hd8ed1ab_101.conda - conda: https://prefix.dev/conda-forge/noarch/dask-core-2025.11.0-pyhcf101f3_0.conda + - conda: https://prefix.dev/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/exceptiongroup-1.3.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/execnet-2.1.1-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/filelock-3.20.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/fsspec-2025.10.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/linux-64/gmp-6.3.0-hac33072_2.conda @@ -181,8 +184,11 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/idna-3.11-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.7.0-pyhe01879c_1.conda - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.3.0-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ipython-9.7.0-pyh53cf698_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ipython_pygments_lexers-1.1.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/jax-0.7.2-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/linux-64/jaxlib-0.7.2-cpu_py313h4c6af5e_2.conda + - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/linux-64/ld_impl_linux-64-2.44-h1aa0949_5.conda - conda: https://prefix.dev/conda-forge/linux-64/libabseil-20250512.1-cxx17_hba17884_0.conda @@ -216,6 +222,7 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/locket-1.0.0-pyhd8ed1ab_0.tar.bz2 - conda: https://prefix.dev/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_0.conda - conda: https://prefix.dev/conda-forge/noarch/marray-python-0.0.12-pyh332efcf_0.conda + - conda: https://prefix.dev/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/linux-64/mkl-2025.3.0-h0e700b2_462.conda - conda: https://prefix.dev/conda-forge/linux-64/mkl-devel-2025.3.0-ha770c72_462.conda - conda: https://prefix.dev/conda-forge/linux-64/mkl-include-2025.3.0-hf2ce2f3_462.conda @@ -230,10 +237,15 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/opt_einsum-3.4.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/linux-64/optree-0.17.0-py313h7037e92_2.conda - conda: https://prefix.dev/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda + - conda: https://prefix.dev/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/partd-1.4.2-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/platformdirs-4.5.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/pluggy-1.6.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/pooch-1.8.2-pyhd8ed1ab_3.conda + - conda: https://prefix.dev/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/pybind11-2.13.6-pyhc790b64_3.conda - conda: https://prefix.dev/conda-forge/noarch/pybind11-abi-4-hd8ed1ab_3.tar.bz2 - conda: https://prefix.dev/conda-forge/noarch/pybind11-global-2.13.6-pyh217bc35_3.conda @@ -256,22 +268,26 @@ environments: - conda: https://prefix.dev/conda-forge/linux-64/sleef-3.9.0-ha0421bc_0.conda - conda: https://prefix.dev/conda-forge/noarch/sortedcontainers-2.4.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/spin-0.15-pyh8f84b5b_0.conda + - conda: https://prefix.dev/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/sympy-1.14.0-pyh2585a3b_105.conda - conda: https://prefix.dev/conda-forge/linux-64/tbb-2022.3.0-h8d10470_1.conda - conda: https://prefix.dev/conda-forge/noarch/threadpoolctl-3.6.0-pyhecae5ae_0.conda - conda: https://prefix.dev/conda-forge/linux-64/tk-8.6.13-noxft_hd72426e_102.conda - conda: https://prefix.dev/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/toolz-1.1.0-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda - conda: https://prefix.dev/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/tzdata-2025b-h78e105d_0.conda - conda: https://prefix.dev/conda-forge/noarch/urllib3-2.5.0-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/linux-64/yaml-0.2.5-h280c20c_3.conda - conda: https://prefix.dev/conda-forge/noarch/zipp-3.23.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/linux-64/zstandard-0.25.0-py313h54dd161_1.conda - conda: https://prefix.dev/conda-forge/linux-64/zstd-1.5.7-hb8e6e7a_2.conda osx-arm64: - conda: https://prefix.dev/conda-forge/noarch/array-api-strict-2.4.1-pyhe01879c_0.conda + - conda: https://prefix.dev/conda-forge/noarch/asttokens-3.0.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/attrs-25.4.0-pyh71513ae_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/blas-devel-3.9.0-38_h11c0a38_openblas.conda - conda: https://prefix.dev/conda-forge/osx-arm64/brotli-python-1.2.0-py313h79bbab8_0.conda @@ -286,8 +302,10 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/cpython-3.13.9-py313hd8ed1ab_101.conda - conda: https://prefix.dev/conda-forge/noarch/dask-core-2025.11.0-pyhcf101f3_0.conda + - conda: https://prefix.dev/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/exceptiongroup-1.3.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/execnet-2.1.1-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/filelock-3.20.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/fsspec-2025.10.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/gmp-6.3.0-h7bae524_2.conda @@ -300,8 +318,11 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/idna-3.11-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.7.0-pyhe01879c_1.conda - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.3.0-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ipython-9.7.0-pyh53cf698_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ipython_pygments_lexers-1.1.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/jax-0.7.2-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/jaxlib-0.7.2-cpu_py313hf0aba26_2.conda + - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/libabseil-20250512.1-cxx17_hd41c47c_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/libblas-3.9.0-38_h51639a9_openblas.conda @@ -327,6 +348,7 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/locket-1.0.0-pyhd8ed1ab_0.tar.bz2 - conda: https://prefix.dev/conda-forge/osx-arm64/markupsafe-3.0.3-py313h7d74516_0.conda - conda: https://prefix.dev/conda-forge/noarch/marray-python-0.0.12-pyh332efcf_0.conda + - conda: https://prefix.dev/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/ml_dtypes-0.5.1-py313hd1f53c0_1.conda - conda: https://prefix.dev/conda-forge/osx-arm64/mpc-1.3.1-h8f1351a_1.conda - conda: https://prefix.dev/conda-forge/osx-arm64/mpfr-4.2.1-hb693164_3.conda @@ -340,10 +362,15 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/opt_einsum-3.4.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/osx-arm64/optree-0.17.0-py313ha61f8ec_2.conda - conda: https://prefix.dev/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda + - conda: https://prefix.dev/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/partd-1.4.2-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/pexpect-4.9.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/platformdirs-4.5.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/pluggy-1.6.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/pooch-1.8.2-pyhd8ed1ab_3.conda + - conda: https://prefix.dev/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ptyprocess-0.7.0-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/pybind11-3.0.1-pyh7a1b43c_0.conda - conda: https://prefix.dev/conda-forge/noarch/pybind11-global-3.0.1-pyhc7ab6ef_0.conda - conda: https://prefix.dev/conda-forge/noarch/pycparser-2.22-pyh29332c3_1.conda @@ -365,15 +392,18 @@ environments: - conda: https://prefix.dev/conda-forge/osx-arm64/sleef-3.9.0-hb028509_0.conda - conda: https://prefix.dev/conda-forge/noarch/sortedcontainers-2.4.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/spin-0.15-pyh8f84b5b_0.conda + - conda: https://prefix.dev/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/sympy-1.14.0-pyh2585a3b_105.conda - conda: https://prefix.dev/conda-forge/noarch/threadpoolctl-3.6.0-pyhecae5ae_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_2.conda - conda: https://prefix.dev/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/toolz-1.1.0-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda - conda: https://prefix.dev/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/tzdata-2025b-h78e105d_0.conda - conda: https://prefix.dev/conda-forge/noarch/urllib3-2.5.0-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/yaml-0.2.5-h925e9cb_3.conda - conda: https://prefix.dev/conda-forge/noarch/zipp-3.23.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/osx-arm64/zstandard-0.25.0-py313h9734d34_1.conda @@ -381,6 +411,7 @@ environments: win-64: - conda: https://prefix.dev/conda-forge/win-64/_openmp_mutex-4.5-2_gnu.conda - conda: https://prefix.dev/conda-forge/noarch/array-api-strict-2.4.1-pyhe01879c_0.conda + - conda: https://prefix.dev/conda-forge/noarch/asttokens-3.0.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/attrs-25.4.0-pyh71513ae_0.conda - conda: https://prefix.dev/conda-forge/win-64/blas-devel-3.9.0-38_h85df5b5_mkl.conda - conda: https://prefix.dev/conda-forge/win-64/brotli-python-1.2.0-py312h9d5906e_0.conda @@ -393,8 +424,10 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/cloudpickle-3.1.2-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/dask-core-2025.11.0-pyhcf101f3_0.conda + - conda: https://prefix.dev/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/exceptiongroup-1.3.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/execnet-2.1.1-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/executing-2.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/filelock-3.20.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/fsspec-2025.10.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/win-64/gmp-6.3.0-hfeafd45_2.conda @@ -406,6 +439,9 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/idna-3.11-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/importlib-metadata-8.7.0-pyhe01879c_1.conda - conda: https://prefix.dev/conda-forge/noarch/iniconfig-2.3.0-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ipython-9.7.0-pyhe2676ad_0.conda + - conda: https://prefix.dev/conda-forge/noarch/ipython_pygments_lexers-1.1.1-pyhd8ed1ab_0.conda + - conda: https://prefix.dev/conda-forge/noarch/jedi-0.19.2-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/jinja2-3.1.6-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/win-64/libabseil-20250512.1-cxx17_habfad5f_0.conda - conda: https://prefix.dev/conda-forge/win-64/libblas-3.9.0-38_hf2e6a31_mkl.conda @@ -431,6 +467,7 @@ environments: - conda: https://prefix.dev/conda-forge/noarch/locket-1.0.0-pyhd8ed1ab_0.tar.bz2 - conda: https://prefix.dev/conda-forge/win-64/markupsafe-3.0.3-py312h05f76fc_0.conda - conda: https://prefix.dev/conda-forge/noarch/marray-python-0.0.12-pyh332efcf_0.conda + - conda: https://prefix.dev/conda-forge/noarch/matplotlib-inline-0.2.1-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/win-64/mkl-2025.3.0-hac47afa_454.conda - conda: https://prefix.dev/conda-forge/win-64/mkl-devel-2025.3.0-h57928b3_454.conda - conda: https://prefix.dev/conda-forge/win-64/mkl-include-2025.3.0-h57928b3_454.conda @@ -442,10 +479,13 @@ environments: - conda: https://prefix.dev/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda - conda: https://prefix.dev/conda-forge/win-64/optree-0.17.0-py312hf90b1b7_2.conda - conda: https://prefix.dev/conda-forge/noarch/packaging-25.0-pyh29332c3_1.conda + - conda: https://prefix.dev/conda-forge/noarch/parso-0.8.5-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/partd-1.4.2-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/platformdirs-4.5.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/pluggy-1.6.0-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/pooch-1.8.2-pyhd8ed1ab_3.conda + - conda: https://prefix.dev/conda-forge/noarch/prompt-toolkit-3.0.52-pyha770c72_0.conda + - conda: https://prefix.dev/conda-forge/noarch/pure_eval-0.2.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/pybind11-2.13.6-pyhc790b64_3.conda - conda: https://prefix.dev/conda-forge/noarch/pybind11-abi-4-hd8ed1ab_3.tar.bz2 - conda: https://prefix.dev/conda-forge/noarch/pybind11-global-2.13.6-pyh6a1d191_3.conda @@ -465,12 +505,14 @@ environments: - conda: https://prefix.dev/conda-forge/win-64/sleef-3.9.0-h67fd636_0.conda - conda: https://prefix.dev/conda-forge/noarch/sortedcontainers-2.4.0-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/spin-0.15-pyha7b4d00_0.conda + - conda: https://prefix.dev/conda-forge/noarch/stack_data-0.6.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/sympy-1.14.0-pyh04b8f61_5.conda - conda: https://prefix.dev/conda-forge/win-64/tbb-2022.3.0-hd094cb3_1.conda - conda: https://prefix.dev/conda-forge/noarch/threadpoolctl-3.6.0-pyhecae5ae_0.conda - conda: https://prefix.dev/conda-forge/win-64/tk-8.6.13-h2c6b04d_2.conda - conda: https://prefix.dev/conda-forge/noarch/tomli-2.3.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/toolz-1.1.0-pyhd8ed1ab_1.conda + - conda: https://prefix.dev/conda-forge/noarch/traitlets-5.14.3-pyhd8ed1ab_1.conda - conda: https://prefix.dev/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda - conda: https://prefix.dev/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda - conda: https://prefix.dev/conda-forge/noarch/tzdata-2025b-h78e105d_0.conda @@ -479,6 +521,7 @@ environments: - conda: https://prefix.dev/conda-forge/win-64/vc-14.3-h2b53caa_32.conda - conda: https://prefix.dev/conda-forge/win-64/vc14_runtime-14.44.35208-h818238b_32.conda - conda: https://prefix.dev/conda-forge/win-64/vcomp14-14.44.35208-h818238b_32.conda + - conda: https://prefix.dev/conda-forge/noarch/wcwidth-0.2.14-pyhd8ed1ab_0.conda - conda: https://prefix.dev/conda-forge/noarch/win_inet_pton-1.1.0-pyh7428d3b_8.conda - conda: https://prefix.dev/conda-forge/win-64/yaml-0.2.5-h6a83c73_3.conda - conda: https://prefix.dev/conda-forge/noarch/zipp-3.23.0-pyhd8ed1ab_0.conda diff --git a/pixi.toml b/pixi.toml index b0a0282f31e0..4ce3a29f805c 100644 --- a/pixi.toml +++ b/pixi.toml @@ -26,7 +26,7 @@ solve-group = "default" [environments.ipython] # tasks: ipython -features = ["run-deps", "test-deps", "ipython"] +features = ["run-deps", "test-deps", "ipython-dep", "ipython-task"] solve-group = "default" [environments.build-debug] @@ -77,8 +77,8 @@ features = ["run-deps", "test-deps", "mkl", "torch-cpu", "torch-cpu-tasks"] solve-group = "array-api-cpu" [environments.array-api-cpu] -# tasks: test-cpu -features = ["run-deps", "test-deps", "test-cpu", "mkl", "array_api_strict", "dask", "jax-cpu", "marray", "torch-cpu"] +# tasks: test-cpu, ipython-cpu +features = ["run-deps", "test-deps", "test-cpu", "mkl", "array_api_strict", "dask", "jax-cpu", "marray", "torch-cpu", "ipython-dep", "ipython-cpu-task"] solve-group = "array-api-cpu" [environments.build-cuda] @@ -221,14 +221,20 @@ description = "Build the documentation" ### IPython ### -[feature.ipython.dependencies] +[feature.ipython-dep.dependencies] ipython = "*" -[feature.ipython.tasks.ipython] +[feature.ipython-task.tasks.ipython] cmd = "spin ipython --no-build" depends-on = "build" description = "Launch IPython" +[feature.ipython-cpu-task.tasks.ipython-cpu] +cmd = "spin ipython --build-dir=build-cpu --no-build" +depends-on = "build-cpu" +env.SCIPY_ARRAY_API = "1" +description = "Launch IPython" + ### Debugging ### diff --git a/scipy/_lib/_array_api.py b/scipy/_lib/_array_api.py index f133ee01818f..99312896af92 100644 --- a/scipy/_lib/_array_api.py +++ b/scipy/_lib/_array_api.py @@ -1028,3 +1028,7 @@ def xp_device_type(a: Array) -> Literal["cpu", "cuda", None]: return xp_device_type(a._meta) # array-api-strict is a stand-in for unknown libraries; don't special-case it return None + + +def xp_isscalar(x): + return np.isscalar(x) or (is_array_api_obj(x) and x.ndim == 0) diff --git a/scipy/sparse/_sputils.py b/scipy/sparse/_sputils.py index 326c68bab785..a1b9cb0a637e 100644 --- a/scipy/sparse/_sputils.py +++ b/scipy/sparse/_sputils.py @@ -369,14 +369,14 @@ def isintlike(x) -> bool: return True -def isshape(x, nonneg=False, *, allow_nd=(2,)) -> bool: +def isshape(x, nonneg=False, *, allow_nd=(2,), check_nd=True) -> bool: """Is x a valid tuple of dimensions? If nonneg, also checks that the dimensions are non-negative. Shapes of length in the tuple allow_nd are allowed. """ ndim = len(x) - if ndim not in allow_nd: + if check_nd and ndim not in allow_nd: return False for d in x: diff --git a/scipy/sparse/linalg/_eigen/_svds.py b/scipy/sparse/linalg/_eigen/_svds.py index f0591f8fe252..36dde26ab77f 100644 --- a/scipy/sparse/linalg/_eigen/_svds.py +++ b/scipy/sparse/linalg/_eigen/_svds.py @@ -34,6 +34,8 @@ def _iv(A, k, ncv, tol, which, v0, maxiter, if math.prod(A.shape) == 0: message = "`A` must not be empty." raise ValueError(message) + if len(A.shape) != 2: + raise ValueError("Only 2-D input is supported for `A` (a single matrix)") # input validation/standardization for `k` kmax = min(A.shape) if solver == 'propack' else min(A.shape) - 1 diff --git a/scipy/sparse/linalg/_eigen/tests/test_svds.py b/scipy/sparse/linalg/_eigen/tests/test_svds.py index c5755b747901..e75cc951273a 100644 --- a/scipy/sparse/linalg/_eigen/tests/test_svds.py +++ b/scipy/sparse/linalg/_eigen/tests/test_svds.py @@ -133,7 +133,7 @@ class SVDSCommonTests: _A_empty_msg = "`A` must not be empty." _A_dtype_msg = "`A` must be of numeric data type" _A_type_msg = "type not understood" - _A_ndim_msg = "array must have ndim <= 2" + _A_ndim_msg = "Only 2-D input" _A_validation_inputs = [ (np.asarray([[]]), ValueError, _A_empty_msg), (np.array([['a', 'b'], ['c', 'd']], dtype='object'), ValueError, _A_dtype_msg), diff --git a/scipy/sparse/linalg/_interface.py b/scipy/sparse/linalg/_interface.py index 5f9cfdff5564..1f2ef54a093e 100644 --- a/scipy/sparse/linalg/_interface.py +++ b/scipy/sparse/linalg/_interface.py @@ -47,8 +47,11 @@ import numpy as np +import scipy.sparse from scipy.sparse import issparse from scipy.sparse._sputils import isshape, isintlike, asmatrix, is_pydata_spmatrix +from scipy._lib._array_api import array_namespace, _asarray, is_lazy_array, xp_copy, xp_isscalar +from scipy._lib import array_api_extra as xpx __all__ = ['LinearOperator', 'aslinearoperator'] @@ -83,7 +86,7 @@ class LinearOperator: Parameters ---------- shape : tuple - Matrix dimensions ``(M, N)``. + Matrix dimensions ``(..., M, N)``. matvec : callable f(v) Returns returns ``A @ v``. rmatvec : callable f(v) @@ -148,7 +151,6 @@ class LinearOperator: """ - ndim = 2 # Necessary for right matmul with numpy arrays. __array_ufunc__ = None @@ -170,21 +172,26 @@ def __new__(cls, *args, **kwargs): return obj - def __init__(self, dtype, shape): + def __init__(self, dtype, shape, xp=None): """Initialize this LinearOperator. To be called by subclasses. ``dtype`` may be None; ``shape`` should be convertible to a length-2 tuple. """ + xp = np if xp is None else xp if dtype is not None: - dtype = np.dtype(dtype) + dtype = xp.empty(0, dtype=dtype).dtype shape = tuple(shape) - if not isshape(shape): - raise ValueError(f"invalid shape {shape!r} (must be 2-d)") + if len(shape) < 2: + raise ValueError(f"invalid shape {shape!r} (must be at least 2-d)") + if not is_lazy_array(xp.empty(0)) and not isshape(shape, check_nd=False): + raise ValueError(f"invalid shape {shape!r}") self.dtype = dtype self.shape = shape + self.ndim = len(shape) + self._xp = xp def _init_dtype(self): """Determine the dtype by executing `matvec` on an `int8` test vector. @@ -198,12 +205,13 @@ def _init_dtype(self): Called from subclasses at the end of the __init__ routine. """ if self.dtype is None: - v = np.zeros(self.shape[-1], dtype=np.int8) + batch_shape = self.shape[:-2] + N = self.shape[-1] + v = self._xp.zeros((*batch_shape, N), dtype=self._xp.int8) try: - matvec_v = np.asarray(self.matvec(v)) + matvec_v = self._xp.asarray(self.matvec(v)) except OverflowError: - # Python large `int` promoted to `np.int64`or `np.int32` - self.dtype = np.dtype(int) + self.dtype = xpx.default_dtype("integral") else: self.dtype = matvec_v.dtype @@ -214,19 +222,24 @@ def _matmat(self, X): define matrix multiplication (though in a very suboptimal way). """ - return np.hstack([self.matvec(col.reshape(-1,1)) for col in X.T]) + # X.mT here? + return self._xp.concat([self.matvec(col.reshape(-1, 1)) for col in X.T], axis=-1) def _matvec(self, x): """Default matrix-vector multiplication handler. - If self is a linear operator of shape (M, N), then this method will - be called on a shape (N,) or (N, 1) ndarray, and should return a - shape (M,) or (M, 1) ndarray. + If self is a linear operator of shape (..., M, N), then this method will + be called on a shape (..., N) or (..., N, 1) ndarray, and should return a + shape (..., M) or (..., M, 1) ndarray. This default implementation falls back on _matmat, so defining that will define matrix-vector multiplication as well. """ - return self.matmat(x.reshape(-1, 1)) + N = self.shape[-1] + if x.shape[-1] == N: + return self.matmat(self._xp.reshape(x, (*x.shape, 1))) + return self.matmat(x) + def matvec(self, x): """Matrix-vector multiplication. @@ -237,13 +250,14 @@ def matvec(self, x): Parameters ---------- x : {matrix, ndarray} - An array with shape (N,) or (N,1). + An array with shape (..., N) representing a row vector (or stack of row vectors), + or an array with shape (..., N, 1) representing a column vector (or stack of column vectors). Returns ------- y : {matrix, ndarray} - A matrix or ndarray with shape (M,) or (M,1) depending - on the type and shape of the x argument. + A matrix or ndarray with shape (..., M) or (..., M, 1) depending + on the type and shape of `x`. Notes ----- @@ -251,27 +265,33 @@ def matvec(self, x): _matvec method to ensure that y has the correct shape and type. """ + xp = self._xp + + x = _asarray(x, subok=True, xp=xp) - x = np.asanyarray(x) - - M,N = self.shape + *self_broadcast_dims, M, N = self.shape - if x.shape != (N,) and x.shape != (N,1): - raise ValueError('dimension mismatch') + x_broadcast_dims: tuple[int, ...] = () + row_vector: bool = False + if x.ndim >= 1 and (row_vector := x.shape[-1] == N): + x_broadcast_dims = x.shape[:-1] + if column_vector := x.shape[-2:] == (N, 1): + x_broadcast_dims = x.shape[:-2] + if not (row_vector or column_vector): + raise ValueError(f'Dimension mismatch: `x` must have a shape ending in `({N},)` or `({N}, 1)`. Given shape: {x.shape}') y = self._matvec(x) if isinstance(x, np.matrix): y = asmatrix(y) else: - y = np.asarray(y) + y = xp.asarray(y) - if x.ndim == 1: - y = y.reshape(M) - elif x.ndim == 2: - y = y.reshape(M,1) - else: - raise ValueError('invalid shape returned by user-defined matvec()') + broadcasted_dims = xpx.broadcast_shapes(self_broadcast_dims, x_broadcast_dims) + if row_vector: + y = xp.reshape(y, (*broadcasted_dims, M)) + elif column_vector: + y = y.reshape(*broadcasted_dims, M, 1) return y @@ -284,13 +304,13 @@ def rmatvec(self, x): Parameters ---------- x : {matrix, ndarray} - An array with shape (M,) or (M,1). + An array with shape (..., M) or (..., M, 1). Returns ------- y : {matrix, ndarray} - A matrix or ndarray with shape (N,) or (N,1) depending - on the type and shape of the x argument. + A matrix or ndarray with shape (..., N) or (..., N, 1) depending + on the type and shape of `x`. Notes ----- @@ -298,27 +318,32 @@ def rmatvec(self, x): _rmatvec method to ensure that y has the correct shape and type. """ + xp = self._xp + x = _asarray(x, subok=True, xp=xp) - x = np.asanyarray(x) + *self_broadcast_dims, M, N = self.shape - M,N = self.shape - - if x.shape != (M,) and x.shape != (M,1): - raise ValueError('dimension mismatch') + x_broadcast_dims: tuple[int, ...] = () + row_vector: bool = False + if x.ndim >= 1 and (row_vector := x.shape[-1] == M): + x_broadcast_dims = x.shape[:-1] + if column_vector := x.shape[-2:] == (M, 1): + x_broadcast_dims = x.shape[:-2] + if not (row_vector or column_vector): + raise ValueError(f'Dimension mismatch: `x` must have a shape ending in `({M},)` or `({M}, 1)`. Given shape: {x.shape}') y = self._rmatvec(x) if isinstance(x, np.matrix): y = asmatrix(y) else: - y = np.asarray(y) + y = xp.asarray(y) - if x.ndim == 1: - y = y.reshape(N) - elif x.ndim == 2: - y = y.reshape(N,1) - else: - raise ValueError('invalid shape returned by user-defined rmatvec()') + broadcasted_dims = xpx.broadcast_shapes(self_broadcast_dims, x_broadcast_dims) + if row_vector: + y = xp.reshape(y, (*broadcasted_dims, N)) + elif column_vector: + y = xp.reshape(y, (*broadcasted_dims, N, 1)) return y @@ -329,7 +354,11 @@ def _rmatvec(self, x): if (hasattr(self, "_rmatmat") and type(self)._rmatmat != LinearOperator._rmatmat): # Try to use _rmatmat as a fallback - return self._rmatmat(x.reshape(-1, 1)).reshape(-1) + xp = self._xp + if x.shape[-1] != 1: + return xp.reshape(self._rmatmat(xp.reshape(x, (*x.shape, 1))), *x.shape) + else: + return self._rmatmat(x) raise NotImplementedError else: return self.H.matvec(x) @@ -343,13 +372,13 @@ def matmat(self, X): Parameters ---------- X : {matrix, ndarray} - An array with shape (N,K). + An array with shape (..., N, K). Returns ------- Y : {matrix, ndarray} - A matrix or ndarray with shape (M,K) depending on - the type of the X argument. + A matrix or ndarray with shape (..., M, K) depending on + the type of `X`. Notes ----- @@ -358,12 +387,12 @@ def matmat(self, X): """ if not (issparse(X) or is_pydata_spmatrix(X)): - X = np.asanyarray(X) + X = _asarray(X, subok=True, xp=self._xp) - if X.ndim != 2: - raise ValueError(f'expected 2-d ndarray or matrix, not {X.ndim}-d') + if X.ndim < 2: + raise ValueError(f'expected at least 2-d ndarray or matrix, not {X.ndim}-d') - if X.shape[0] != self.shape[1]: + if X.shape[-2] != self.shape[-1]: raise ValueError(f'dimension mismatch: {self.shape}, {X.shape}') try: @@ -404,12 +433,12 @@ def rmatmat(self, X): """ if not (issparse(X) or is_pydata_spmatrix(X)): - X = np.asanyarray(X) + X = _asarray(X, subok=True, xp=self._xp) - if X.ndim != 2: - raise ValueError(f'expected 2-d ndarray or matrix, not {X.ndim}-d') + if X.ndim < 2: + raise ValueError(f'expected at least 2-d ndarray or matrix, not {X.ndim}-d') - if X.shape[0] != self.shape[0]: + if X.shape[-2] != self.shape[-2]: raise ValueError(f'dimension mismatch: {self.shape}, {X.shape}') try: @@ -428,22 +457,23 @@ def rmatmat(self, X): def _rmatmat(self, X): """Default implementation of _rmatmat defers to rmatvec or adjoint.""" + # X.mT here? if type(self)._adjoint == LinearOperator._adjoint: - return np.hstack([self.rmatvec(col.reshape(-1, 1)) for col in X.T]) + return self._xp.concat([self.rmatvec(col.reshape(-1, 1)) for col in X.T], axis=-1) else: return self.H.matmat(X) def __call__(self, x): - return self@x + return self @ x def __mul__(self, x): return self.dot(x) def __truediv__(self, other): - if not np.isscalar(other): + if not xp_isscalar(other): raise ValueError("Can only divide a linear operator by a scalar.") - return _ScaledLinearOperator(self, 1.0/other) + return _ScaledLinearOperator(self, 1.0/other, xp=self._xp) def dot(self, x): """Matrix-matrix or matrix-vector multiplication. @@ -461,36 +491,62 @@ def dot(self, x): """ if isinstance(x, LinearOperator): - return _ProductLinearOperator(self, x) - elif np.isscalar(x): - return _ScaledLinearOperator(self, x) + if (xp_x := getattr(x, "_xp", np)) != self._xp: + msg = ( + f"Mismatched array namespaces." + f"Namespace for self is {self._xp}, namespace for x is {xp_x}" + ) + raise TypeError(msg) + return _ProductLinearOperator(self, x, self._xp) + elif xp_isscalar(x): + if (xp_x := array_namespace(x, self._xp.empty(0))) != self._xp: + msg = ( + f"Mismatched array namespaces." + f"Namespace for self is {self._xp}, namespace for x is {xp_x}" + ) + raise TypeError(msg) + return _ScaledLinearOperator(self, x, self._xp) else: if not issparse(x) and not is_pydata_spmatrix(x): - # Sparse matrices shouldn't be converted to numpy arrays. - x = np.asarray(x) - - if x.ndim == 1 or x.ndim == 2 and x.shape[1] == 1: + x = self._xp.asarray(x) + + N = self.shape[-1] + + column_vector = x.shape[-2:] == (N, 1) # maintain column vector backwards-compatibility in 2-D case + matrix = x.ndim >= 2 and x.shape[-2] == N # maintain matmat backwards-compatibility in 2-D case + row_vector = x.shape[-1] == N # otherwise treat as a row-vector + + if not (row_vector or column_vector or matrix): + raise ValueError(f'Dimension mismatch: `x` must have a shape ending in `({N},)` or `({N}, 1)` or `({N}, K)` for some integer K. Given shape: {x.shape}') + + if column_vector: return self.matvec(x) - elif x.ndim == 2: + elif matrix: return self.matmat(x) - else: - raise ValueError(f'expected 1-d or 2-d array or matrix, got {x!r}') + elif row_vector: + return self.matvec(x) def __matmul__(self, other): - if np.isscalar(other): + if xp_isscalar(other): raise ValueError("Scalar operands are not allowed, " "use '*' instead") return self.__mul__(other) def __rmatmul__(self, other): - if np.isscalar(other): + if xp_isscalar(other): raise ValueError("Scalar operands are not allowed, " "use '*' instead") return self.__rmul__(other) def __rmul__(self, x): - if np.isscalar(x): - return _ScaledLinearOperator(self, x) + if xp_isscalar(x): + if (xp_x := array_namespace(x, self._xp.empty(0))) != self._xp: + msg = ( + f"Mismatched array namespaces." + f"Namespace for self is {self._xp}, namespace for x is {xp_x}" + ) + raise TypeError(msg) + return _ScaledLinearOperator(self, x, self._xp) else: return self._rdot(x) @@ -513,49 +569,83 @@ def _rdot(self, x): This is copied from dot to implement right multiplication. """ if isinstance(x, LinearOperator): - return _ProductLinearOperator(x, self) - elif np.isscalar(x): - return _ScaledLinearOperator(self, x) + if (xp_x := getattr(x, "_xp", np)) != self._xp: + msg = ( + f"Mismatched array namespaces." + f"Namespace for self is {self._xp}, namespace for x is {xp_x}" + ) + raise TypeError(msg) + return _ProductLinearOperator(x, self, self._xp) + elif xp_isscalar(x): + if (xp_x := array_namespace(x, self._xp.empty(0))) != self._xp: + msg = ( + f"Mismatched array namespaces." + f"Namespace for self is {self._xp}, namespace for x is {xp_x}" + ) + raise TypeError(msg) + return _ScaledLinearOperator(self, x, self._xp) else: if not issparse(x) and not is_pydata_spmatrix(x): - # Sparse matrices shouldn't be converted to numpy arrays. - x = np.asarray(x) - + x = self._xp.asarray(x) + + M = self.shape[-2] + + column_vector = x.shape[-2:] == (1, M) # maintain column vector backwards-compatibility in 2-D case + matrix = x.shape[-1] == M and x.ndim == 2 # maintain matmat backwards-compatibility in 2-D case + row_vector = x.shape[-1] == M # otherwise treat as a row-vector + # XXX: for `x.ndim > 2`, the equivalent `np.dot(a, b)` implements a sum product over the last axis of `a` and the second-to-last axis of `b` + # see https://numpy.org/doc/stable/reference/generated/numpy.dot.html + + if not (row_vector or column_vector or matrix): + raise ValueError(f'Dimension mismatch: `x` must have a shape ending in `({M},)` or `(1, {M})` or `(K, {M})` for some integer K. Given shape: {x.shape}') + # We use transpose instead of rmatvec/rmatmat to avoid # unnecessary complex conjugation if possible. - if x.ndim == 1 or x.ndim == 2 and x.shape[0] == 1: + if column_vector: return self.T.matvec(x.T).T - elif x.ndim == 2: + elif matrix: return self.T.matmat(x.T).T - else: - raise ValueError(f'expected 1-d or 2-d array or matrix, got {x!r}') + elif row_vector: + return self.T.matvec(x.T).T def __pow__(self, p): - if np.isscalar(p): - return _PowerLinearOperator(self, p) + if xp_isscalar(p): + if (xp_p := array_namespace(p, self._xp.empty(0))) != self._xp: + msg = ( + f"Mismatched array namespaces." + f"Namespace for self is {self._xp}, namespace for p is {xp_p}" + ) + raise TypeError(msg) + return _PowerLinearOperator(self, p, self._xp) else: return NotImplemented def __add__(self, x): if isinstance(x, LinearOperator): - return _SumLinearOperator(self, x) + if (xp_x := getattr(x, "_xp", np)) != self._xp: + msg = ( + f"Mismatched array namespaces." + f"Namespace for self is {self._xp}, namespace for x is {xp_x}" + ) + raise TypeError(msg) + return _SumLinearOperator(self, x, xp=self._xp) else: return NotImplemented def __neg__(self): - return _ScaledLinearOperator(self, -1) + return _ScaledLinearOperator(self, -1, xp=self._xp) def __sub__(self, x): return self.__add__(-x) def __repr__(self): - M,N = self.shape if self.dtype is None: dt = 'unspecified dtype' else: dt = 'dtype=' + str(self.dtype) - return f'<{M}x{N} {self.__class__.__name__} with {dt}>' + shape = 'x'.join(str(dim) for dim in self.shape) + return f'<{shape} {self.__class__.__name__} with {dt}>' def adjoint(self): """Hermitian adjoint. @@ -587,19 +677,19 @@ def transpose(self): def _adjoint(self): """Default implementation of _adjoint; defers to rmatvec.""" - return _AdjointLinearOperator(self) + return _AdjointLinearOperator(self, self._xp) def _transpose(self): """ Default implementation of _transpose; defers to rmatvec + conj""" - return _TransposedLinearOperator(self) + return _TransposedLinearOperator(self, self._xp) class _CustomLinearOperator(LinearOperator): """Linear operator defined in terms of user-specified operations.""" def __init__(self, shape, matvec, rmatvec=None, matmat=None, - dtype=None, rmatmat=None): - super().__init__(dtype, shape) + dtype=None, rmatmat=None, xp=None): + super().__init__(dtype, shape, xp) self.args = () @@ -632,20 +722,21 @@ def _rmatmat(self, X): return super()._rmatmat(X) def _adjoint(self): - return _CustomLinearOperator(shape=(self.shape[1], self.shape[0]), + return _CustomLinearOperator(shape=(*self.shape[:-2], self.shape[-1], self.shape[-2]), matvec=self.__rmatvec_impl, rmatvec=self.__matvec_impl, matmat=self.__rmatmat_impl, rmatmat=self.__matmat_impl, - dtype=self.dtype) + dtype=self.dtype, + xp=self._xp) class _AdjointLinearOperator(LinearOperator): """Adjoint of arbitrary Linear Operator""" - def __init__(self, A): - shape = (A.shape[1], A.shape[0]) - super().__init__(dtype=A.dtype, shape=shape) + def __init__(self, A, xp=None): + shape = (*A.shape[:-2], A.shape[-1], A.shape[-2]) + super().__init__(A.dtype, shape, xp) self.A = A self.args = (A,) @@ -664,44 +755,47 @@ def _rmatmat(self, x): class _TransposedLinearOperator(LinearOperator): """Transposition of arbitrary Linear Operator""" - def __init__(self, A): - shape = (A.shape[1], A.shape[0]) - super().__init__(dtype=A.dtype, shape=shape) + def __init__(self, A, xp=None): + shape = (*A.shape[:-2], A.shape[-1], A.shape[-2]) + super().__init__(A.dtype, shape, xp) self.A = A self.args = (A,) def _matvec(self, x): - # NB. np.conj works also on sparse matrices - return np.conj(self.A._rmatvec(np.conj(x))) + return self._xp.conj(self.A._rmatvec(self._xp.conj(x))) def _rmatvec(self, x): - return np.conj(self.A._matvec(np.conj(x))) + return self._xp.conj(self.A._matvec(self._xp.conj(x))) def _matmat(self, x): - # NB. np.conj works also on sparse matrices - return np.conj(self.A._rmatmat(np.conj(x))) + return self._xp.conj(self.A._rmatmat(self._xp.conj(x))) def _rmatmat(self, x): - return np.conj(self.A._matmat(np.conj(x))) + return self._xp.conj(self.A._matmat(self._xp.conj(x))) -def _get_dtype(operators, dtypes=None): + +def _get_dtype(operators, dtypes=None, xp=None): + xp = np if xp is None else xp if dtypes is None: dtypes = [] for obj in operators: if obj is not None and hasattr(obj, 'dtype'): dtypes.append(obj.dtype) - return np.result_type(*dtypes) + return xp.result_type(*dtypes) class _SumLinearOperator(LinearOperator): - def __init__(self, A, B): + def __init__(self, A, B, xp=None): if not isinstance(A, LinearOperator) or \ not isinstance(B, LinearOperator): raise ValueError('both operands have to be a LinearOperator') - if A.shape != B.shape: + *A_broadcast_dims, A_M, A_N = A.shape + *B_broadcast_dims, B_M, B_N = B.shape + if (A_M, A_N) != (B_M, B_N): raise ValueError(f'cannot add {A} and {B}: shape mismatch') + broadcasted_dims = xp.broadcast_shapes(A_broadcast_dims, B_broadcast_dims) self.args = (A, B) - super().__init__(_get_dtype([A, B]), A.shape) + super().__init__(_get_dtype([A, B]), (*broadcasted_dims, A_M, A_N), xp) def _matvec(self, x): return self.args[0].matvec(x) + self.args[1].matvec(x) @@ -721,14 +815,16 @@ def _adjoint(self): class _ProductLinearOperator(LinearOperator): - def __init__(self, A, B): + def __init__(self, A, B, xp=None): if not isinstance(A, LinearOperator) or \ not isinstance(B, LinearOperator): raise ValueError('both operands have to be a LinearOperator') - if A.shape[1] != B.shape[0]: + *A_broadcast_dims, A_M, A_N = A.shape + *B_broadcast_dims, B_M, B_N = B.shape + if A_N != B_M: raise ValueError(f'cannot multiply {A} and {B}: shape mismatch') - super().__init__(_get_dtype([A, B]), - (A.shape[0], B.shape[1])) + broadcasted_dims = np.broadcast_shapes(A_broadcast_dims, B_broadcast_dims) + super().__init__(_get_dtype([A, B]), (*broadcasted_dims, A_M, B_N), xp) self.args = (A, B) def _matvec(self, x): @@ -749,7 +845,7 @@ def _adjoint(self): class _ScaledLinearOperator(LinearOperator): - def __init__(self, A, alpha): + def __init__(self, A, alpha, xp=None): if not isinstance(A, LinearOperator): raise ValueError('LinearOperator expected as A') if not np.isscalar(alpha): @@ -761,7 +857,7 @@ def __init__(self, A, alpha): alpha = alpha * alpha_original dtype = _get_dtype([A], [type(alpha)]) - super().__init__(dtype, A.shape) + super().__init__(dtype, A.shape, xp) self.args = (A, alpha) # Note: args[1] is alpha (a scalar), so use `*` below, not `@` @@ -769,33 +865,33 @@ def _matvec(self, x): return self.args[1] * self.args[0].matvec(x) def _rmatvec(self, x): - return np.conj(self.args[1]) * self.args[0].rmatvec(x) + return self._xp.conj(self.args[1]) * self.args[0].rmatvec(x) def _rmatmat(self, x): - return np.conj(self.args[1]) * self.args[0].rmatmat(x) + return self._xp.conj(self.args[1]) * self.args[0].rmatmat(x) def _matmat(self, x): return self.args[1] * self.args[0].matmat(x) def _adjoint(self): A, alpha = self.args - return A.H * np.conj(alpha) + return A.H * self._xp.conj(alpha) class _PowerLinearOperator(LinearOperator): - def __init__(self, A, p): + def __init__(self, A, p, xp=None): if not isinstance(A, LinearOperator): raise ValueError('LinearOperator expected as A') - if A.shape[0] != A.shape[1]: - raise ValueError(f'square LinearOperator expected, got {A!r}') + if A.shape[-2] != A.shape[-1]: + raise ValueError(f'square core-dimensions of LinearOperator expected, got {A!r}') if not isintlike(p) or p < 0: raise ValueError('non-negative integer expected as p') - super().__init__(_get_dtype([A]), A.shape) + super().__init__(_get_dtype([A]), A.shape, xp) self.args = (A, p) def _power(self, fun, x): - res = np.array(x, copy=True) + res = xp_copy(x) for i in range(self.args[1]): res = fun(res) return res @@ -818,38 +914,39 @@ def _adjoint(self): class MatrixLinearOperator(LinearOperator): - def __init__(self, A): - super().__init__(A.dtype, A.shape) + def __init__(self, A, xp=None): + super().__init__(A.dtype, A.shape, xp) self.A = A self.__adj = None self.args = (A,) def _matmat(self, X): - return self.A.dot(X) + return self.A @ X def _adjoint(self): if self.__adj is None: - self.__adj = _AdjointMatrixOperator(self.A) + self.__adj = _AdjointMatrixOperator(self.A, self._xp) return self.__adj class _AdjointMatrixOperator(MatrixLinearOperator): - def __init__(self, adjoint_array): - self.A = adjoint_array.T.conj() + def __init__(self, adjoint_array, xp=None): + xp = np if xp is None else xp + self.A = xp.conj(adjoint_array.T) self.args = (adjoint_array,) - self.shape = adjoint_array.shape[1], adjoint_array.shape[0] + self.shape = *adjoint_array.shape[:-2], adjoint_array.shape[-1], adjoint_array.shape[-2] @property def dtype(self): return self.args[0].dtype def _adjoint(self): - return MatrixLinearOperator(self.args[0]) + return MatrixLinearOperator(self.args[0], self._xp) class IdentityOperator(LinearOperator): - def __init__(self, shape, dtype=None): - super().__init__(dtype, shape) + def __init__(self, shape, dtype=None, xp=None): + super().__init__(dtype, shape, xp) def _matvec(self, x): return x @@ -893,19 +990,31 @@ def aslinearoperator(A): >>> aslinearoperator(M) <2x3 MatrixLinearOperator with dtype=int32> """ + A, _ = _xp_aslinearoperator(A) + return A + +def _xp_aslinearoperator(A): + """ + Return `A` as a linear operator, + as well as a compatible array namespace `xp` for `A`. + Fallback to NumPy for unknown types. + """ if isinstance(A, LinearOperator): - return A + return A, getattr(A, "_xp", np) - elif isinstance(A, np.ndarray) or isinstance(A, np.matrix): - if A.ndim > 2: - raise ValueError('array must have ndim <= 2') + elif issparse(A): + return MatrixLinearOperator(A), scipy.sparse + + elif isinstance(A, np.matrix): A = np.atleast_2d(np.asarray(A)) - return MatrixLinearOperator(A) - - elif issparse(A) or is_pydata_spmatrix(A): - return MatrixLinearOperator(A) - - else: + return MatrixLinearOperator(A), np + + try: + xp = array_namespace(A) + A = xpx.atleast_nd(A, ndim=2, xp=xp) + return MatrixLinearOperator(A, xp=xp), xp + + except: if hasattr(A, 'shape') and hasattr(A, 'matvec'): rmatvec = None rmatmat = None @@ -917,8 +1026,9 @@ def aslinearoperator(A): rmatmat = A.rmatmat if hasattr(A, 'dtype'): dtype = A.dtype + xp = array_namespace(A) or np return LinearOperator(A.shape, A.matvec, rmatvec=rmatvec, - rmatmat=rmatmat, dtype=dtype) + rmatmat=rmatmat, dtype=dtype, xp=xp), xp else: raise TypeError('type not understood') diff --git a/scipy/sparse/linalg/_isolve/iterative.py b/scipy/sparse/linalg/_isolve/iterative.py index 637f2d021f62..c8448836bcd7 100644 --- a/scipy/sparse/linalg/_isolve/iterative.py +++ b/scipy/sparse/linalg/_isolve/iterative.py @@ -1,13 +1,18 @@ import warnings +import functools + import numpy as np from scipy.sparse.linalg._interface import LinearOperator from .utils import make_system from scipy.linalg import get_lapack_funcs +from scipy._lib import array_api_extra as xpx +from scipy._lib._array_api import xp_copy, xp_vector_norm + __all__ = ['bicg', 'bicgstab', 'cg', 'cgs', 'gmres', 'qmr'] -def _get_atol_rtol(name, b_norm, atol=0., rtol=1e-5): +def _get_atol_rtol(name, b_norm, atol=0., rtol=1e-5, xp=np): """ A helper function to handle tolerance normalization """ @@ -16,7 +21,7 @@ def _get_atol_rtol(name, b_norm, atol=0., rtol=1e-5): "if set, `atol` must be a real, non-negative number.") raise ValueError(msg) - atol = max(float(atol), float(rtol) * float(b_norm)) + atol = xp.max(xp.stack((xp.asarray(float(atol)), float(rtol) * xp.min(b_norm)))) return atol, rtol @@ -376,44 +381,62 @@ def cg(A, b, x0=None, *, rtol=1e-5, atol=0., maxiter=None, M=None, callback=None >>> np.allclose(A.dot(x), b) True """ - A, M, x, b = make_system(A, M, x0, b) - bnrm2 = np.linalg.norm(b) + A, M, x, b, xp = make_system(A, M, x0, b) + bnrm2 = xp_vector_norm(b, axis=-1) - atol, _ = _get_atol_rtol('cg', bnrm2, atol, rtol) + atol, _ = _get_atol_rtol('cg', bnrm2, atol, rtol, xp=xp) - if bnrm2 == 0: + if not xp.any(bnrm2): return b, 0 - n = len(b) - if maxiter is None: - maxiter = n*10 + maxiter = b.shape[-1] * 10 - dotprod = np.vdot if np.iscomplexobj(x) else np.dot + dotprod = np.vdot if xp.isdtype(x.dtype, "complex floating") else functools.partial(xp.vecdot, axis=-1) matvec = A.matvec psolve = M.matvec - r = b - matvec(x) if x.any() else b.copy() + r = b - matvec(x) if xp.any(x) else xp_copy(b) # Dummy value to initialize var, silences warnings rho_prev, p = None, None for iteration in range(maxiter): - if np.linalg.norm(r) < atol: # Are we done? + converged = xp_vector_norm(r, axis=-1) < atol + if xp.all(converged): return x, 0 z = psolve(r) rho_cur = dotprod(r, z) + if iteration > 0: - beta = rho_cur / rho_prev + beta = xpx.apply_where( + ~converged, + (rho_cur, rho_prev), + lambda cur, prev: cur / prev, + fill_value=0.0, + xp=xp + ) + beta = xp.expand_dims(beta, axis=-1) + p *= beta p += z else: # First spin - p = np.empty_like(r) - p[:] = z[:] + p = xp.empty_like(r) + p = xpx.at(p)[:, ...].set(z[:, ...]) q = matvec(p) - alpha = rho_cur / dotprod(p, q) + c = dotprod(p, q) + + alpha = xpx.apply_where( + ~converged , + (rho_cur, c), + lambda rc, c: rc / c, + fill_value=0.0, + xp=xp + ) + alpha = xp.expand_dims(alpha, axis=-1) + x += alpha*p r -= alpha*q rho_prev = rho_cur diff --git a/scipy/sparse/linalg/_isolve/utils.py b/scipy/sparse/linalg/_isolve/utils.py index 28925f48014b..bb71b5346915 100644 --- a/scipy/sparse/linalg/_isolve/utils.py +++ b/scipy/sparse/linalg/_isolve/utils.py @@ -3,31 +3,16 @@ __all__ = [] -from numpy import asanyarray, asarray, array, zeros - -from scipy.sparse.linalg._interface import aslinearoperator, LinearOperator, \ - IdentityOperator - -_coerce_rules = {('f','f'):'f', ('f','d'):'d', ('f','F'):'F', - ('f','D'):'D', ('d','f'):'d', ('d','d'):'d', - ('d','F'):'D', ('d','D'):'D', ('F','f'):'F', - ('F','d'):'D', ('F','F'):'F', ('F','D'):'D', - ('D','f'):'D', ('D','d'):'D', ('D','F'):'D', - ('D','D'):'D'} - - -def coerce(x,y): - if x not in 'fdFD': - x = 'd' - if y not in 'fdFD': - y = 'd' - return _coerce_rules[x,y] +import numpy as np +from scipy.sparse.linalg._interface import ( + _xp_aslinearoperator, LinearOperator, IdentityOperator +) +from scipy._lib._array_api import array_namespace, is_lazy_array, xp_copy, xp_ravel, xp_result_type, _asarray def id(x): return x - def make_system(A, M, x0, b): """Make a linear system Ax=b @@ -35,7 +20,7 @@ def make_system(A, M, x0, b): ---------- A : LinearOperator sparse or dense matrix (or any valid input to aslinearoperator) - M : {LinearOperator, Nones} + M : {LinearOperator, None} preconditioner sparse or dense matrix (or any valid input to aslinearoperator) x0 : {array_like, str, None} @@ -47,7 +32,7 @@ def make_system(A, M, x0, b): Returns ------- - (A, M, x, b) + (A, M, x, b, xp) A : LinearOperator matrix of the linear system M : LinearOperator @@ -56,33 +41,44 @@ def make_system(A, M, x0, b): initial guess b : rank 1 ndarray right hand side + xp : compatible array namespace """ A_ = A - A = aslinearoperator(A) - - if A.shape[0] != A.shape[1]: - raise ValueError(f'expected square matrix, but got shape={(A.shape,)}') - - N = A.shape[0] - - b = asanyarray(b) - - if not (b.shape == (N,1) or b.shape == (N,)): + A, xp = _xp_aslinearoperator(A) + lazy = is_lazy_array(xp.empty(0)) + + N = A.shape[-2] + if not lazy and N != A.shape[-1]: + raise ValueError(f'expected square matrix or stack of square matrices, but got shape={(A.shape,)}') + + xp_b = array_namespace(b) + if xp_b != xp: + msg = f"Mismatched array namespaces. Namespace for A is {xp}, namespace for b is {xp_b}" + raise TypeError(msg) + + b = _asarray(b, subok=True, xp=xp) + + column_vector = not lazy and b.ndim == 2 and b.shape[-2:] == (N, 1) # maintain column vector backwards-compatibility in 2-D case + row_vector = b.shape[-1] == N # otherwise treat as a row-vector + + if not lazy and not (column_vector or row_vector): raise ValueError(f'shapes of A {A.shape} and b {b.shape} are ' 'incompatible') - if b.dtype.char not in 'fdFD': - b = b.astype('d') # upcast non-FP types to double + if not xp.isdtype(b.dtype, ("real floating", "complex floating")): + b = xp.astype(b, xp.float64) # upcast non-FP types to float64 - if hasattr(A,'dtype'): - xtype = A.dtype.char + if hasattr(A, 'dtype'): + x_dtype = A.dtype else: - xtype = A.matvec(b).dtype.char - xtype = coerce(xtype, b.dtype.char) + x_dtype = A.matvec(b).dtype + # XXX: does this match the previous coercion? + x_dtype = xp_result_type(x_dtype, b.dtype, force_floating=True, xp=xp) - b = asarray(b,dtype=xtype) # make b the same type as x - b = b.ravel() + b = xp.astype(b, x_dtype) # make b the same type as x + if column_vector: + b = xp_ravel(b) # process preconditioner if M is None: @@ -95,27 +91,36 @@ def make_system(A, M, x0, b): else: rpsolve = id if psolve is id and rpsolve is id: - M = IdentityOperator(shape=A.shape, dtype=A.dtype) + M = IdentityOperator(shape=A.shape, dtype=A.dtype, xp=xp) else: M = LinearOperator(A.shape, matvec=psolve, rmatvec=rpsolve, - dtype=A.dtype) + dtype=A.dtype, xp=xp) else: - M = aslinearoperator(M) + M, xp_M = _xp_aslinearoperator(M) + if xp_M != xp: + msg = f"Mismatched array namespaces. Namespace for A is {xp}, namespace for M is {xp_M}" + raise TypeError(msg) if A.shape != M.shape: raise ValueError('matrix and preconditioner have different shapes') # set initial guess if x0 is None: - x = zeros(N, dtype=xtype) + x = xp.zeros((*M.shape[:-2], N), dtype=x_dtype) + # XXX: proper error handling for `x0` of type `str` but not equal to `'Mb'`? elif isinstance(x0, str): if x0 == 'Mb': # use nonzero initial guess ``M @ b`` - bCopy = b.copy() + bCopy = xp_copy(b) x = M.matvec(bCopy) else: - x = array(x0, dtype=xtype) - if not (x.shape == (N, 1) or x.shape == (N,)): + x = xp.asarray(x0, dtype=x_dtype) + + column_vector = x.ndim == 2 and x.shape[-2:] == (N, 1) # maintain column vector backwards-compatibility in 2-D case + row_vector = x.shape[-1] == N # otherwise treat as a row-vector + + if not (row_vector or column_vector): raise ValueError(f'shapes of A {A.shape} and ' f'x0 {x.shape} are incompatible') - x = x.ravel() + if column_vector: + x = xp_ravel(x) - return A, M, x, b + return A, M, x, b, xp