init
This commit is contained in:
@@ -0,0 +1,161 @@
|
||||
# This file is generated by SciPy's build process
|
||||
# It contains system_info results at the time of building this package.
|
||||
from enum import Enum
|
||||
|
||||
__all__ = ["show"]
|
||||
_built_with_meson = True
|
||||
|
||||
|
||||
class DisplayModes(Enum):
|
||||
stdout = "stdout"
|
||||
dicts = "dicts"
|
||||
|
||||
|
||||
def _cleanup(d):
|
||||
"""
|
||||
Removes empty values in a `dict` recursively
|
||||
This ensures we remove values that Meson could not provide to CONFIG
|
||||
"""
|
||||
if isinstance(d, dict):
|
||||
return { k: _cleanup(v) for k, v in d.items() if v != '' and _cleanup(v) != '' }
|
||||
else:
|
||||
return d
|
||||
|
||||
|
||||
CONFIG = _cleanup(
|
||||
{
|
||||
"Compilers": {
|
||||
"c": {
|
||||
"name": "gcc",
|
||||
"linker": r"ld.bfd",
|
||||
"version": "10.2.1",
|
||||
"commands": r"cc",
|
||||
"args": r"",
|
||||
"linker args": r"",
|
||||
},
|
||||
"cython": {
|
||||
"name": r"cython",
|
||||
"linker": r"cython",
|
||||
"version": r"3.1.6",
|
||||
"commands": r"cython",
|
||||
"args": r"",
|
||||
"linker args": r"",
|
||||
},
|
||||
"c++": {
|
||||
"name": "gcc",
|
||||
"linker": r"ld.bfd",
|
||||
"version": "10.2.1",
|
||||
"commands": r"c++",
|
||||
"args": r"",
|
||||
"linker args": r"",
|
||||
},
|
||||
"fortran": {
|
||||
"name": "gcc",
|
||||
"linker": r"ld.bfd",
|
||||
"version": "10.2.1",
|
||||
"commands": r"gfortran",
|
||||
"args": r"",
|
||||
"linker args": r"",
|
||||
},
|
||||
"pythran": {
|
||||
"version": r"0.18.0",
|
||||
"include directory": r"../../tmp/build-env-b9dts1ag/lib/python3.11/site-packages/pythran"
|
||||
},
|
||||
},
|
||||
"Machine Information": {
|
||||
"host": {
|
||||
"cpu": r"x86_64",
|
||||
"family": r"x86_64",
|
||||
"endian": r"little",
|
||||
"system": r"linux",
|
||||
},
|
||||
"build": {
|
||||
"cpu": r"x86_64",
|
||||
"family": r"x86_64",
|
||||
"endian": r"little",
|
||||
"system": r"linux",
|
||||
},
|
||||
"cross-compiled": bool("False".lower().replace('false', '')),
|
||||
},
|
||||
"Build Dependencies": {
|
||||
"blas": {
|
||||
"name": "scipy-openblas",
|
||||
"found": bool("True".lower().replace('false', '')),
|
||||
"version": "0.3.29.dev",
|
||||
"detection method": "pkgconfig",
|
||||
"include directory": r"/opt/_internal/cpython-3.11.13/lib/python3.11/site-packages/scipy_openblas32/include",
|
||||
"lib directory": r"/opt/_internal/cpython-3.11.13/lib/python3.11/site-packages/scipy_openblas32/lib",
|
||||
"openblas configuration": r"OpenBLAS 0.3.29.dev DYNAMIC_ARCH NO_AFFINITY SkylakeX MAX_THREADS=64",
|
||||
"pc file directory": r"/project",
|
||||
},
|
||||
"lapack": {
|
||||
"name": "scipy-openblas",
|
||||
"found": bool("True".lower().replace('false', '')),
|
||||
"version": "0.3.29.dev",
|
||||
"detection method": "pkgconfig",
|
||||
"include directory": r"/opt/_internal/cpython-3.11.13/lib/python3.11/site-packages/scipy_openblas32/include",
|
||||
"lib directory": r"/opt/_internal/cpython-3.11.13/lib/python3.11/site-packages/scipy_openblas32/lib",
|
||||
"openblas configuration": r"OpenBLAS 0.3.29.dev DYNAMIC_ARCH NO_AFFINITY SkylakeX MAX_THREADS=64",
|
||||
"pc file directory": r"/project",
|
||||
},
|
||||
"pybind11": {
|
||||
"name": "pybind11",
|
||||
"version": "3.0.1",
|
||||
"detection method": "config-tool",
|
||||
"include directory": r"unknown",
|
||||
},
|
||||
},
|
||||
"Python Information": {
|
||||
"path": r"/tmp/build-env-b9dts1ag/bin/python",
|
||||
"version": "3.11",
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _check_pyyaml():
|
||||
import yaml
|
||||
|
||||
return yaml
|
||||
|
||||
|
||||
def show(mode=DisplayModes.stdout.value):
|
||||
"""
|
||||
Show libraries and system information on which SciPy was built
|
||||
and is being used
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mode : {`'stdout'`, `'dicts'`}, optional.
|
||||
Indicates how to display the config information.
|
||||
`'stdout'` prints to console, `'dicts'` returns a dictionary
|
||||
of the configuration.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : {`dict`, `None`}
|
||||
If mode is `'dicts'`, a dict is returned, else None
|
||||
|
||||
Notes
|
||||
-----
|
||||
1. The `'stdout'` mode will give more readable
|
||||
output if ``pyyaml`` is installed
|
||||
|
||||
"""
|
||||
if mode == DisplayModes.stdout.value:
|
||||
try: # Non-standard library, check import
|
||||
yaml = _check_pyyaml()
|
||||
|
||||
print(yaml.dump(CONFIG))
|
||||
except ModuleNotFoundError:
|
||||
import warnings
|
||||
import json
|
||||
|
||||
warnings.warn("Install `pyyaml` for better output", stacklevel=1)
|
||||
print(json.dumps(CONFIG, indent=2))
|
||||
elif mode == DisplayModes.dicts.value:
|
||||
return CONFIG
|
||||
else:
|
||||
raise AttributeError(
|
||||
f"Invalid `mode`, use one of: {', '.join([e.value for e in DisplayModes])}"
|
||||
)
|
||||
@@ -0,0 +1,138 @@
|
||||
"""
|
||||
SciPy: A scientific computing package for Python
|
||||
================================================
|
||||
|
||||
Documentation is available in the docstrings and
|
||||
online at https://docs.scipy.org/doc/scipy/
|
||||
|
||||
Subpackages
|
||||
-----------
|
||||
::
|
||||
|
||||
cluster --- Vector Quantization / Kmeans
|
||||
constants --- Physical and mathematical constants and units
|
||||
datasets --- Dataset methods
|
||||
differentiate --- Finite difference differentiation tools
|
||||
fft --- Discrete Fourier transforms
|
||||
fftpack --- Legacy discrete Fourier transforms
|
||||
integrate --- Integration routines
|
||||
interpolate --- Interpolation Tools
|
||||
io --- Data input and output
|
||||
linalg --- Linear algebra routines
|
||||
ndimage --- N-D image package
|
||||
odr --- Orthogonal Distance Regression
|
||||
optimize --- Optimization Tools
|
||||
signal --- Signal Processing Tools
|
||||
sparse --- Sparse Matrices
|
||||
spatial --- Spatial data structures and algorithms
|
||||
special --- Special functions
|
||||
stats --- Statistical Functions
|
||||
|
||||
Public API in the main SciPy namespace
|
||||
--------------------------------------
|
||||
::
|
||||
|
||||
__version__ --- SciPy version string
|
||||
LowLevelCallable --- Low-level callback function
|
||||
show_config --- Show scipy build configuration
|
||||
test --- Run scipy unittests
|
||||
|
||||
"""
|
||||
|
||||
import importlib as _importlib
|
||||
|
||||
from numpy import __version__ as __numpy_version__
|
||||
|
||||
|
||||
try:
|
||||
from scipy.__config__ import show as show_config
|
||||
except ImportError as e:
|
||||
msg = """Error importing SciPy: you cannot import SciPy while
|
||||
being in scipy source directory; please exit the SciPy source
|
||||
tree first and relaunch your Python interpreter."""
|
||||
raise ImportError(msg) from e
|
||||
|
||||
|
||||
from scipy.version import version as __version__
|
||||
|
||||
|
||||
# Allow distributors to run custom init code
|
||||
from . import _distributor_init
|
||||
del _distributor_init
|
||||
|
||||
|
||||
from scipy._lib import _pep440
|
||||
# In maintenance branch, change to np_maxversion N+3 if numpy is at N
|
||||
np_minversion = '1.25.2'
|
||||
np_maxversion = '2.6.0'
|
||||
if (_pep440.parse(__numpy_version__) < _pep440.Version(np_minversion) or
|
||||
_pep440.parse(__numpy_version__) >= _pep440.Version(np_maxversion)):
|
||||
import warnings
|
||||
warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
|
||||
f" is required for this version of SciPy (detected "
|
||||
f"version {__numpy_version__})",
|
||||
UserWarning, stacklevel=2)
|
||||
del _pep440
|
||||
|
||||
|
||||
# This is the first import of an extension module within SciPy. If there's
|
||||
# a general issue with the install, such that extension modules are missing
|
||||
# or cannot be imported, this is where we'll get a failure - so give an
|
||||
# informative error message.
|
||||
try:
|
||||
from scipy._lib._ccallback import LowLevelCallable
|
||||
except ImportError as e:
|
||||
msg = "The `scipy` install you are using seems to be broken, " + \
|
||||
"(extension modules cannot be imported), " + \
|
||||
"please try reinstalling."
|
||||
raise ImportError(msg) from e
|
||||
|
||||
|
||||
from scipy._lib._testutils import PytestTester
|
||||
test = PytestTester(__name__)
|
||||
del PytestTester
|
||||
|
||||
|
||||
submodules = [
|
||||
'cluster',
|
||||
'constants',
|
||||
'datasets',
|
||||
'differentiate',
|
||||
'fft',
|
||||
'fftpack',
|
||||
'integrate',
|
||||
'interpolate',
|
||||
'io',
|
||||
'linalg',
|
||||
'ndimage',
|
||||
'odr',
|
||||
'optimize',
|
||||
'signal',
|
||||
'sparse',
|
||||
'spatial',
|
||||
'special',
|
||||
'stats'
|
||||
]
|
||||
|
||||
__all__ = submodules + [
|
||||
'LowLevelCallable',
|
||||
'test',
|
||||
'show_config',
|
||||
'__version__',
|
||||
]
|
||||
|
||||
|
||||
def __dir__():
|
||||
return __all__
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
if name in submodules:
|
||||
return _importlib.import_module(f'scipy.{name}')
|
||||
else:
|
||||
try:
|
||||
return globals()[name]
|
||||
except KeyError:
|
||||
raise AttributeError(
|
||||
f"Module 'scipy' has no attribute '{name}'"
|
||||
)
|
||||
Binary file not shown.
@@ -0,0 +1,18 @@
|
||||
""" Distributor init file
|
||||
|
||||
Distributors: you can replace the contents of this file with your own custom
|
||||
code to support particular distributions of SciPy.
|
||||
|
||||
For example, this is a good place to put any checks for hardware requirements
|
||||
or BLAS/LAPACK library initialization.
|
||||
|
||||
The SciPy standard source distribution will not put code in this file beyond
|
||||
the try-except import of `_distributor_init_local` (which is not part of a
|
||||
standard source distribution), so you can safely replace this file with your
|
||||
own version.
|
||||
"""
|
||||
|
||||
try:
|
||||
from . import _distributor_init_local # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
Module containing private utility functions
|
||||
===========================================
|
||||
|
||||
The ``scipy._lib`` namespace is empty (for now). Tests for all
|
||||
utilities in submodules of ``_lib`` can be run with::
|
||||
|
||||
from scipy import _lib
|
||||
_lib.test()
|
||||
|
||||
"""
|
||||
from scipy._lib._testutils import PytestTester
|
||||
test = PytestTester(__name__)
|
||||
del PytestTester
|
||||
@@ -0,0 +1,931 @@
|
||||
"""Utility functions to use Python Array API compatible libraries.
|
||||
|
||||
For the context about the Array API see:
|
||||
https://data-apis.org/array-api/latest/purpose_and_scope.html
|
||||
|
||||
The SciPy use case of the Array API is described on the following page:
|
||||
https://data-apis.org/array-api/latest/use_cases.html#use-case-scipy
|
||||
"""
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import functools
|
||||
import os
|
||||
import textwrap
|
||||
|
||||
from collections.abc import Generator, Iterable, Iterator
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from types import ModuleType
|
||||
from typing import Any, Literal, TypeAlias
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
from scipy._lib import array_api_compat
|
||||
from scipy._lib.array_api_compat import (
|
||||
is_array_api_obj,
|
||||
is_lazy_array,
|
||||
size as xp_size,
|
||||
numpy as np_compat,
|
||||
device as xp_device,
|
||||
is_numpy_namespace as is_numpy,
|
||||
is_cupy_namespace as is_cupy,
|
||||
is_torch_namespace as is_torch,
|
||||
is_jax_namespace as is_jax,
|
||||
is_dask_namespace as is_dask,
|
||||
is_array_api_strict_namespace as is_array_api_strict
|
||||
)
|
||||
from scipy._lib._sparse import issparse
|
||||
from scipy._lib._docscrape import FunctionDoc
|
||||
|
||||
__all__ = [
|
||||
'_asarray', 'array_namespace', 'assert_almost_equal', 'assert_array_almost_equal',
|
||||
'default_xp', 'eager_warns', 'is_lazy_array', 'is_marray',
|
||||
'is_array_api_strict', 'is_complex', 'is_cupy', 'is_jax', 'is_numpy', 'is_torch',
|
||||
'SCIPY_ARRAY_API', 'SCIPY_DEVICE', 'scipy_namespace_for',
|
||||
'xp_assert_close', 'xp_assert_equal', 'xp_assert_less',
|
||||
'xp_copy', 'xp_device', 'xp_ravel', 'xp_size',
|
||||
'xp_unsupported_param_msg', 'xp_vector_norm', 'xp_capabilities',
|
||||
'xp_result_type', 'xp_promote'
|
||||
]
|
||||
|
||||
|
||||
# To enable array API and strict array-like input validation
|
||||
SCIPY_ARRAY_API: str | bool = os.environ.get("SCIPY_ARRAY_API", False)
|
||||
# To control the default device - for use in the test suite only
|
||||
SCIPY_DEVICE = os.environ.get("SCIPY_DEVICE", "cpu")
|
||||
|
||||
_GLOBAL_CONFIG = {
|
||||
"SCIPY_ARRAY_API": SCIPY_ARRAY_API,
|
||||
"SCIPY_DEVICE": SCIPY_DEVICE,
|
||||
}
|
||||
|
||||
|
||||
Array: TypeAlias = Any # To be changed to a Protocol later (see array-api#589)
|
||||
ArrayLike: TypeAlias = Array | npt.ArrayLike
|
||||
|
||||
|
||||
def _compliance_scipy(arrays: Iterable[ArrayLike]) -> Iterator[Array]:
|
||||
"""Raise exceptions on known-bad subclasses. Discard 0-dimensional ArrayLikes
|
||||
and convert 1+-dimensional ArrayLikes to numpy.
|
||||
|
||||
The following subclasses are not supported and raise and error:
|
||||
- `numpy.ma.MaskedArray`
|
||||
- `numpy.matrix`
|
||||
- NumPy arrays which do not have a boolean or numerical dtype
|
||||
- Any array-like which is neither array API compatible nor coercible by NumPy
|
||||
- Any array-like which is coerced by NumPy to an unsupported dtype
|
||||
"""
|
||||
for array in arrays:
|
||||
if array is None:
|
||||
continue
|
||||
|
||||
# this comes from `_util._asarray_validated`
|
||||
if issparse(array):
|
||||
msg = ('Sparse arrays/matrices are not supported by this function. '
|
||||
'Perhaps one of the `scipy.sparse.linalg` functions '
|
||||
'would work instead.')
|
||||
raise ValueError(msg)
|
||||
|
||||
if isinstance(array, np.ma.MaskedArray):
|
||||
raise TypeError("Inputs of type `numpy.ma.MaskedArray` are not supported.")
|
||||
|
||||
if isinstance(array, np.matrix):
|
||||
raise TypeError("Inputs of type `numpy.matrix` are not supported.")
|
||||
|
||||
if isinstance(array, np.ndarray | np.generic):
|
||||
dtype = array.dtype
|
||||
if not (np.issubdtype(dtype, np.number) or np.issubdtype(dtype, np.bool_)):
|
||||
raise TypeError(f"An argument has dtype `{dtype!r}`; "
|
||||
f"only boolean and numerical dtypes are supported.")
|
||||
|
||||
if is_array_api_obj(array):
|
||||
yield array
|
||||
else:
|
||||
try:
|
||||
array = np.asanyarray(array)
|
||||
except TypeError:
|
||||
raise TypeError("An argument is neither array API compatible nor "
|
||||
"coercible by NumPy.")
|
||||
dtype = array.dtype
|
||||
if not (np.issubdtype(dtype, np.number) or np.issubdtype(dtype, np.bool_)):
|
||||
message = (
|
||||
f"An argument was coerced to an unsupported dtype `{dtype!r}`; "
|
||||
f"only boolean and numerical dtypes are supported."
|
||||
)
|
||||
raise TypeError(message)
|
||||
# Ignore 0-dimensional arrays, coherently with array-api-compat.
|
||||
# Raise if there are 1+-dimensional array-likes mixed with non-numpy
|
||||
# Array API objects.
|
||||
if array.ndim:
|
||||
yield array
|
||||
|
||||
|
||||
def _check_finite(array: Array, xp: ModuleType) -> None:
|
||||
"""Check for NaNs or Infs."""
|
||||
if not xp.all(xp.isfinite(array)):
|
||||
msg = "array must not contain infs or NaNs"
|
||||
raise ValueError(msg)
|
||||
|
||||
|
||||
def array_namespace(*arrays: Array) -> ModuleType:
|
||||
"""Get the array API compatible namespace for the arrays xs.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
*arrays : sequence of array_like
|
||||
Arrays used to infer the common namespace.
|
||||
|
||||
Returns
|
||||
-------
|
||||
namespace : module
|
||||
Common namespace.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Thin wrapper around `array_api_compat.array_namespace`.
|
||||
|
||||
1. Check for the global switch: SCIPY_ARRAY_API. This can also be accessed
|
||||
dynamically through ``_GLOBAL_CONFIG['SCIPY_ARRAY_API']``.
|
||||
2. `_compliance_scipy` raise exceptions on known-bad subclasses. See
|
||||
its definition for more details.
|
||||
|
||||
When the global switch is False, it defaults to the `numpy` namespace.
|
||||
In that case, there is no compliance check. This is a convenience to
|
||||
ease the adoption. Otherwise, arrays must comply with the new rules.
|
||||
"""
|
||||
if not _GLOBAL_CONFIG["SCIPY_ARRAY_API"]:
|
||||
# here we could wrap the namespace if needed
|
||||
return np_compat
|
||||
|
||||
api_arrays = list(_compliance_scipy(arrays))
|
||||
# In case of a mix of array API compliant arrays and scalars, return
|
||||
# the array API namespace. If there are only ArrayLikes (e.g. lists),
|
||||
# return NumPy (wrapped by array-api-compat).
|
||||
if api_arrays:
|
||||
return array_api_compat.array_namespace(*api_arrays)
|
||||
return np_compat
|
||||
|
||||
|
||||
def _asarray(
|
||||
array: ArrayLike,
|
||||
dtype: Any = None,
|
||||
order: Literal['K', 'A', 'C', 'F'] | None = None,
|
||||
copy: bool | None = None,
|
||||
*,
|
||||
xp: ModuleType | None = None,
|
||||
check_finite: bool = False,
|
||||
subok: bool = False,
|
||||
) -> Array:
|
||||
"""SciPy-specific replacement for `np.asarray` with `order`, `check_finite`, and
|
||||
`subok`.
|
||||
|
||||
Memory layout parameter `order` is not exposed in the Array API standard.
|
||||
`order` is only enforced if the input array implementation
|
||||
is NumPy based, otherwise `order` is just silently ignored.
|
||||
|
||||
`check_finite` is also not a keyword in the array API standard; included
|
||||
here for convenience rather than that having to be a separate function
|
||||
call inside SciPy functions.
|
||||
|
||||
`subok` is included to allow this function to preserve the behaviour of
|
||||
`np.asanyarray` for NumPy based inputs.
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(array)
|
||||
if is_numpy(xp):
|
||||
# Use NumPy API to support order
|
||||
if copy is True:
|
||||
array = np.array(array, order=order, dtype=dtype, subok=subok)
|
||||
elif subok:
|
||||
array = np.asanyarray(array, order=order, dtype=dtype)
|
||||
else:
|
||||
array = np.asarray(array, order=order, dtype=dtype)
|
||||
else:
|
||||
try:
|
||||
array = xp.asarray(array, dtype=dtype, copy=copy)
|
||||
except TypeError:
|
||||
coerced_xp = array_namespace(xp.asarray(3))
|
||||
array = coerced_xp.asarray(array, dtype=dtype, copy=copy)
|
||||
|
||||
if check_finite:
|
||||
_check_finite(array, xp)
|
||||
|
||||
return array
|
||||
|
||||
|
||||
def xp_copy(x: Array, *, xp: ModuleType | None = None) -> Array:
|
||||
"""
|
||||
Copies an array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array
|
||||
|
||||
xp : array_namespace
|
||||
|
||||
Returns
|
||||
-------
|
||||
copy : array
|
||||
Copied array
|
||||
|
||||
Notes
|
||||
-----
|
||||
This copy function does not offer all the semantics of `np.copy`, i.e. the
|
||||
`subok` and `order` keywords are not used.
|
||||
"""
|
||||
# Note: for older NumPy versions, `np.asarray` did not support the `copy` kwarg,
|
||||
# so this uses our other helper `_asarray`.
|
||||
if xp is None:
|
||||
xp = array_namespace(x)
|
||||
|
||||
return _asarray(x, copy=True, xp=xp)
|
||||
|
||||
|
||||
_default_xp_ctxvar: ContextVar[ModuleType] = ContextVar("_default_xp")
|
||||
|
||||
@contextmanager
|
||||
def default_xp(xp: ModuleType) -> Generator[None, None, None]:
|
||||
"""In all ``xp_assert_*`` and ``assert_*`` function calls executed within this
|
||||
context manager, test by default that the array namespace is
|
||||
the provided across all arrays, unless one explicitly passes the ``xp=``
|
||||
parameter or ``check_namespace=False``.
|
||||
|
||||
Without this context manager, the default value for `xp` is the namespace
|
||||
for the desired array (the second parameter of the tests).
|
||||
"""
|
||||
token = _default_xp_ctxvar.set(xp)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_default_xp_ctxvar.reset(token)
|
||||
|
||||
|
||||
def eager_warns(x, warning_type, match=None):
|
||||
"""pytest.warns context manager, but only if x is not a lazy array."""
|
||||
import pytest
|
||||
# This attribute is interpreted by pytest-run-parallel, ensuring that tests that use
|
||||
# `eager_warns` aren't run in parallel (since pytest.warns isn't thread-safe).
|
||||
__thread_safe__ = False # noqa: F841
|
||||
if is_lazy_array(x):
|
||||
return contextlib.nullcontext()
|
||||
return pytest.warns(warning_type, match=match)
|
||||
|
||||
|
||||
def _strict_check(actual, desired, xp, *,
|
||||
check_namespace=True, check_dtype=True, check_shape=True,
|
||||
check_0d=True):
|
||||
__tracebackhide__ = True # Hide traceback for py.test
|
||||
|
||||
if xp is None:
|
||||
try:
|
||||
xp = _default_xp_ctxvar.get()
|
||||
except LookupError:
|
||||
xp = array_namespace(desired)
|
||||
|
||||
if check_namespace:
|
||||
_assert_matching_namespace(actual, desired, xp)
|
||||
|
||||
# only NumPy distinguishes between scalars and arrays; we do if check_0d=True.
|
||||
# do this first so we can then cast to array (and thus use the array API) below.
|
||||
if is_numpy(xp) and check_0d:
|
||||
_msg = ("Array-ness does not match:\n Actual: "
|
||||
f"{type(actual)}\n Desired: {type(desired)}")
|
||||
assert ((xp.isscalar(actual) and xp.isscalar(desired))
|
||||
or (not xp.isscalar(actual) and not xp.isscalar(desired))), _msg
|
||||
|
||||
actual = xp.asarray(actual)
|
||||
desired = xp.asarray(desired)
|
||||
|
||||
if check_dtype:
|
||||
_msg = f"dtypes do not match.\nActual: {actual.dtype}\nDesired: {desired.dtype}"
|
||||
assert actual.dtype == desired.dtype, _msg
|
||||
|
||||
if check_shape:
|
||||
if is_dask(xp):
|
||||
actual.compute_chunk_sizes()
|
||||
desired.compute_chunk_sizes()
|
||||
_msg = f"Shapes do not match.\nActual: {actual.shape}\nDesired: {desired.shape}"
|
||||
assert actual.shape == desired.shape, _msg
|
||||
|
||||
desired = xp.broadcast_to(desired, actual.shape)
|
||||
return actual, desired, xp
|
||||
|
||||
|
||||
def _assert_matching_namespace(actual, desired, xp):
|
||||
__tracebackhide__ = True # Hide traceback for py.test
|
||||
|
||||
desired_arr_space = array_namespace(desired)
|
||||
_msg = ("Namespace of desired array does not match expectations "
|
||||
"set by the `default_xp` context manager or by the `xp`"
|
||||
"pytest fixture.\n"
|
||||
f"Desired array's space: {desired_arr_space.__name__}\n"
|
||||
f"Expected namespace: {xp.__name__}")
|
||||
assert desired_arr_space == xp, _msg
|
||||
|
||||
actual_arr_space = array_namespace(actual)
|
||||
_msg = ("Namespace of actual and desired arrays do not match.\n"
|
||||
f"Actual: {actual_arr_space.__name__}\n"
|
||||
f"Desired: {xp.__name__}")
|
||||
assert actual_arr_space == xp, _msg
|
||||
|
||||
|
||||
def xp_assert_equal(actual, desired, *, check_namespace=True, check_dtype=True,
|
||||
check_shape=True, check_0d=True, err_msg='', xp=None):
|
||||
__tracebackhide__ = True # Hide traceback for py.test
|
||||
|
||||
actual, desired, xp = _strict_check(
|
||||
actual, desired, xp, check_namespace=check_namespace,
|
||||
check_dtype=check_dtype, check_shape=check_shape,
|
||||
check_0d=check_0d
|
||||
)
|
||||
|
||||
if is_cupy(xp):
|
||||
return xp.testing.assert_array_equal(actual, desired, err_msg=err_msg)
|
||||
elif is_torch(xp):
|
||||
# PyTorch recommends using `rtol=0, atol=0` like this
|
||||
# to test for exact equality
|
||||
err_msg = None if err_msg == '' else err_msg
|
||||
return xp.testing.assert_close(actual, desired, rtol=0, atol=0, equal_nan=True,
|
||||
check_dtype=False, msg=err_msg)
|
||||
# JAX uses `np.testing`
|
||||
return np.testing.assert_array_equal(actual, desired, err_msg=err_msg)
|
||||
|
||||
|
||||
def xp_assert_close(actual, desired, *, rtol=None, atol=0, check_namespace=True,
|
||||
check_dtype=True, check_shape=True, check_0d=True,
|
||||
err_msg='', xp=None):
|
||||
__tracebackhide__ = True # Hide traceback for py.test
|
||||
|
||||
actual, desired, xp = _strict_check(
|
||||
actual, desired, xp,
|
||||
check_namespace=check_namespace, check_dtype=check_dtype,
|
||||
check_shape=check_shape, check_0d=check_0d
|
||||
)
|
||||
|
||||
floating = xp.isdtype(actual.dtype, ('real floating', 'complex floating'))
|
||||
if rtol is None and floating:
|
||||
# multiplier of 4 is used as for `np.float64` this puts the default `rtol`
|
||||
# roughly half way between sqrt(eps) and the default for
|
||||
# `numpy.testing.assert_allclose`, 1e-7
|
||||
rtol = xp.finfo(actual.dtype).eps**0.5 * 4
|
||||
elif rtol is None:
|
||||
rtol = 1e-7
|
||||
|
||||
if is_cupy(xp):
|
||||
return xp.testing.assert_allclose(actual, desired, rtol=rtol,
|
||||
atol=atol, err_msg=err_msg)
|
||||
elif is_torch(xp):
|
||||
err_msg = None if err_msg == '' else err_msg
|
||||
return xp.testing.assert_close(actual, desired, rtol=rtol, atol=atol,
|
||||
equal_nan=True, check_dtype=False, msg=err_msg)
|
||||
# JAX uses `np.testing`
|
||||
return np.testing.assert_allclose(actual, desired, rtol=rtol,
|
||||
atol=atol, err_msg=err_msg)
|
||||
|
||||
|
||||
def xp_assert_less(actual, desired, *, check_namespace=True, check_dtype=True,
|
||||
check_shape=True, check_0d=True, err_msg='', verbose=True, xp=None):
|
||||
__tracebackhide__ = True # Hide traceback for py.test
|
||||
|
||||
actual, desired, xp = _strict_check(
|
||||
actual, desired, xp, check_namespace=check_namespace,
|
||||
check_dtype=check_dtype, check_shape=check_shape,
|
||||
check_0d=check_0d
|
||||
)
|
||||
|
||||
if is_cupy(xp):
|
||||
return xp.testing.assert_array_less(actual, desired,
|
||||
err_msg=err_msg, verbose=verbose)
|
||||
elif is_torch(xp):
|
||||
if actual.device.type != 'cpu':
|
||||
actual = actual.cpu()
|
||||
if desired.device.type != 'cpu':
|
||||
desired = desired.cpu()
|
||||
# JAX uses `np.testing`
|
||||
return np.testing.assert_array_less(actual, desired,
|
||||
err_msg=err_msg, verbose=verbose)
|
||||
|
||||
|
||||
def assert_array_almost_equal(actual, desired, decimal=6, *args, **kwds):
|
||||
"""Backwards compatible replacement. In new code, use xp_assert_close instead.
|
||||
"""
|
||||
rtol, atol = 0, 1.5*10**(-decimal)
|
||||
return xp_assert_close(actual, desired,
|
||||
atol=atol, rtol=rtol, check_dtype=False, check_shape=False,
|
||||
*args, **kwds)
|
||||
|
||||
|
||||
def assert_almost_equal(actual, desired, decimal=7, *args, **kwds):
|
||||
"""Backwards compatible replacement. In new code, use xp_assert_close instead.
|
||||
"""
|
||||
rtol, atol = 0, 1.5*10**(-decimal)
|
||||
return xp_assert_close(actual, desired,
|
||||
atol=atol, rtol=rtol, check_dtype=False, check_shape=False,
|
||||
*args, **kwds)
|
||||
|
||||
|
||||
def xp_unsupported_param_msg(param: Any) -> str:
|
||||
return f'Providing {param!r} is only supported for numpy arrays.'
|
||||
|
||||
|
||||
def is_complex(x: Array, xp: ModuleType) -> bool:
|
||||
return xp.isdtype(x.dtype, 'complex floating')
|
||||
|
||||
|
||||
def scipy_namespace_for(xp: ModuleType) -> ModuleType | None:
|
||||
"""Return the `scipy`-like namespace of a non-NumPy backend
|
||||
|
||||
That is, return the namespace corresponding with backend `xp` that contains
|
||||
`scipy` sub-namespaces like `linalg` and `special`. If no such namespace
|
||||
exists, return ``None``. Useful for dispatching.
|
||||
"""
|
||||
|
||||
if is_cupy(xp):
|
||||
import cupyx # type: ignore[import-not-found,import-untyped]
|
||||
return cupyx.scipy
|
||||
|
||||
if is_jax(xp):
|
||||
import jax # type: ignore[import-not-found]
|
||||
return jax.scipy
|
||||
|
||||
if is_torch(xp):
|
||||
return xp
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# maybe use `scipy.linalg` if/when array API support is added
|
||||
def xp_vector_norm(x: Array, /, *,
|
||||
axis: int | tuple[int] | None = None,
|
||||
keepdims: bool = False,
|
||||
ord: int | float = 2,
|
||||
xp: ModuleType | None = None) -> Array:
|
||||
xp = array_namespace(x) if xp is None else xp
|
||||
|
||||
if SCIPY_ARRAY_API:
|
||||
# check for optional `linalg` extension
|
||||
if hasattr(xp, 'linalg'):
|
||||
return xp.linalg.vector_norm(x, axis=axis, keepdims=keepdims, ord=ord)
|
||||
else:
|
||||
if ord != 2:
|
||||
raise ValueError(
|
||||
"only the Euclidean norm (`ord=2`) is currently supported in "
|
||||
"`xp_vector_norm` for backends not implementing the `linalg` "
|
||||
"extension."
|
||||
)
|
||||
# return (x @ x)**0.5
|
||||
# or to get the right behavior with nd, complex arrays
|
||||
return xp.sum(xp.conj(x) * x, axis=axis, keepdims=keepdims)**0.5
|
||||
else:
|
||||
# to maintain backwards compatibility
|
||||
return np.linalg.norm(x, ord=ord, axis=axis, keepdims=keepdims)
|
||||
|
||||
|
||||
def xp_ravel(x: Array, /, *, xp: ModuleType | None = None) -> Array:
|
||||
# Equivalent of np.ravel written in terms of array API
|
||||
# Even though it's one line, it comes up so often that it's worth having
|
||||
# this function for readability
|
||||
xp = array_namespace(x) if xp is None else xp
|
||||
return xp.reshape(x, (-1,))
|
||||
|
||||
|
||||
def xp_swapaxes(a, axis1, axis2, xp=None):
|
||||
# Equivalent of np.swapaxes written in terms of array API
|
||||
xp = array_namespace(a) if xp is None else xp
|
||||
axes = list(range(a.ndim))
|
||||
axes[axis1], axes[axis2] = axes[axis2], axes[axis1]
|
||||
a = xp.permute_dims(a, axes)
|
||||
return a
|
||||
|
||||
|
||||
# utility to find common dtype with option to force floating
|
||||
def xp_result_type(*args, force_floating=False, xp):
|
||||
"""
|
||||
Returns the dtype that results from applying type promotion rules
|
||||
(see Array API Standard Type Promotion Rules) to the arguments. Augments
|
||||
standard `result_type` in a few ways:
|
||||
|
||||
- There is a `force_floating` argument that ensures that the result type
|
||||
is floating point, even when all args are integer.
|
||||
- When a TypeError is raised (e.g. due to an unsupported promotion)
|
||||
and `force_floating=True`, we define a custom rule: use the result type
|
||||
of the default float and any other floats passed. See
|
||||
https://github.com/scipy/scipy/pull/22695/files#r1997905891
|
||||
for rationale.
|
||||
- This function accepts array-like iterables, which are immediately converted
|
||||
to the namespace's arrays before result type calculation. Consequently, the
|
||||
result dtype may be different when an argument is `1.` vs `[1.]`.
|
||||
|
||||
Typically, this function will be called shortly after `array_namespace`
|
||||
on a subset of the arguments passed to `array_namespace`.
|
||||
"""
|
||||
args = [(_asarray(arg, subok=True, xp=xp) if np.iterable(arg) else arg)
|
||||
for arg in args]
|
||||
args_not_none = [arg for arg in args if arg is not None]
|
||||
if force_floating:
|
||||
args_not_none.append(1.0)
|
||||
|
||||
if is_numpy(xp) and xp.__version__ < '2.0':
|
||||
# Follow NEP 50 promotion rules anyway
|
||||
args_not_none = [arg.dtype if getattr(arg, 'size', 0) == 1 else arg
|
||||
for arg in args_not_none]
|
||||
return xp.result_type(*args_not_none)
|
||||
|
||||
try: # follow library's preferred promotion rules
|
||||
return xp.result_type(*args_not_none)
|
||||
except TypeError: # mixed type promotion isn't defined
|
||||
if not force_floating:
|
||||
raise
|
||||
# use `result_type` of default floating point type and any floats present
|
||||
# This can be revisited, but right now, the only backends that get here
|
||||
# are array-api-strict (which is not for production use) and PyTorch
|
||||
# (due to data-apis/array-api-compat#279).
|
||||
float_args = []
|
||||
for arg in args_not_none:
|
||||
arg_array = xp.asarray(arg) if np.isscalar(arg) else arg
|
||||
dtype = getattr(arg_array, 'dtype', arg)
|
||||
if xp.isdtype(dtype, ('real floating', 'complex floating')):
|
||||
float_args.append(arg)
|
||||
return xp.result_type(*float_args, xp_default_dtype(xp))
|
||||
|
||||
|
||||
def xp_promote(*args, broadcast=False, force_floating=False, xp):
|
||||
"""
|
||||
Promotes elements of *args to result dtype, ignoring `None`s.
|
||||
Includes options for forcing promotion to floating point and
|
||||
broadcasting the arrays, again ignoring `None`s.
|
||||
Type promotion rules follow `xp_result_type` instead of `xp.result_type`.
|
||||
|
||||
Typically, this function will be called shortly after `array_namespace`
|
||||
on a subset of the arguments passed to `array_namespace`.
|
||||
|
||||
This function accepts array-like iterables, which are immediately converted
|
||||
to the namespace's arrays before result type calculation. Consequently, the
|
||||
result dtype may be different when an argument is `1.` vs `[1.]`.
|
||||
|
||||
See Also
|
||||
--------
|
||||
xp_result_type
|
||||
"""
|
||||
args = [(_asarray(arg, subok=True, xp=xp) if np.iterable(arg) else arg)
|
||||
for arg in args] # solely to prevent double conversion of iterable to array
|
||||
|
||||
dtype = xp_result_type(*args, force_floating=force_floating, xp=xp)
|
||||
|
||||
args = [(_asarray(arg, dtype=dtype, subok=True, xp=xp) if arg is not None else arg)
|
||||
for arg in args]
|
||||
|
||||
if not broadcast:
|
||||
return args[0] if len(args)==1 else tuple(args)
|
||||
|
||||
args_not_none = [arg for arg in args if arg is not None]
|
||||
|
||||
# determine result shape
|
||||
shapes = {arg.shape for arg in args_not_none}
|
||||
try:
|
||||
shape = (np.broadcast_shapes(*shapes) if len(shapes) != 1
|
||||
else args_not_none[0].shape)
|
||||
except ValueError as e:
|
||||
message = "Array shapes are incompatible for broadcasting."
|
||||
raise ValueError(message) from e
|
||||
|
||||
out = []
|
||||
for arg in args:
|
||||
if arg is None:
|
||||
out.append(arg)
|
||||
continue
|
||||
|
||||
# broadcast only if needed
|
||||
# Even if two arguments need broadcasting, this is faster than
|
||||
# `broadcast_arrays`, especially since we've already determined `shape`
|
||||
if arg.shape != shape:
|
||||
kwargs = {'subok': True} if is_numpy(xp) else {}
|
||||
arg = xp.broadcast_to(arg, shape, **kwargs)
|
||||
|
||||
# This is much faster than xp.astype(arg, dtype, copy=False)
|
||||
if arg.dtype != dtype:
|
||||
arg = xp.astype(arg, dtype)
|
||||
|
||||
out.append(arg)
|
||||
|
||||
return out[0] if len(out)==1 else tuple(out)
|
||||
|
||||
|
||||
def xp_float_to_complex(arr: Array, xp: ModuleType | None = None) -> Array:
|
||||
xp = array_namespace(arr) if xp is None else xp
|
||||
arr_dtype = arr.dtype
|
||||
# The standard float dtypes are float32 and float64.
|
||||
# Convert float32 to complex64,
|
||||
# and float64 (and non-standard real dtypes) to complex128
|
||||
if xp.isdtype(arr_dtype, xp.float32):
|
||||
arr = xp.astype(arr, xp.complex64)
|
||||
elif xp.isdtype(arr_dtype, 'real floating'):
|
||||
arr = xp.astype(arr, xp.complex128)
|
||||
|
||||
return arr
|
||||
|
||||
|
||||
def xp_default_dtype(xp):
|
||||
"""Query the namespace-dependent default floating-point dtype.
|
||||
"""
|
||||
if is_torch(xp):
|
||||
# historically, we allow pytorch to keep its default of float32
|
||||
return xp.get_default_dtype()
|
||||
else:
|
||||
# we default to float64
|
||||
return xp.float64
|
||||
|
||||
|
||||
def xp_result_device(*args):
|
||||
"""Return the device of an array in `args`, for the purpose of
|
||||
input-output device propagation.
|
||||
If there are multiple devices, return an arbitrary one.
|
||||
If there are no arrays, return None (this typically happens only on NumPy).
|
||||
"""
|
||||
for arg in args:
|
||||
# Do not do a duck-type test for the .device attribute, as many backends today
|
||||
# don't have it yet. See workarouunds in array_api_compat.device().
|
||||
if is_array_api_obj(arg):
|
||||
return xp_device(arg)
|
||||
return None
|
||||
|
||||
|
||||
def is_marray(xp):
|
||||
"""Returns True if `xp` is an MArray namespace; False otherwise."""
|
||||
return "marray" in xp.__name__
|
||||
|
||||
|
||||
@dataclasses.dataclass(repr=False)
|
||||
class _XPSphinxCapability:
|
||||
cpu: bool | None # None if not applicable
|
||||
gpu: bool | None
|
||||
warnings: list[str] = dataclasses.field(default_factory=list)
|
||||
|
||||
def _render(self, value):
|
||||
if value is None:
|
||||
return "n/a"
|
||||
if not value:
|
||||
return "⛔"
|
||||
if self.warnings:
|
||||
res = "⚠️ " + '; '.join(self.warnings)
|
||||
assert len(res) <= 20, "Warnings too long"
|
||||
return res
|
||||
return "✅"
|
||||
|
||||
def __str__(self):
|
||||
cpu = self._render(self.cpu)
|
||||
gpu = self._render(self.gpu)
|
||||
return f"{cpu:20} {gpu:20}"
|
||||
|
||||
|
||||
def _make_sphinx_capabilities(
|
||||
# lists of tuples [(module name, reason), ...]
|
||||
skip_backends=(), xfail_backends=(),
|
||||
# @pytest.mark.skip/xfail_xp_backends kwargs
|
||||
cpu_only=False, np_only=False, exceptions=(),
|
||||
# xpx.lazy_xp_backends kwargs
|
||||
allow_dask_compute=False, jax_jit=True,
|
||||
# list of tuples [(module name, reason), ...]
|
||||
warnings = (),
|
||||
# unused in documentation
|
||||
reason=None,
|
||||
):
|
||||
exceptions = set(exceptions)
|
||||
|
||||
# Default capabilities
|
||||
capabilities = {
|
||||
"numpy": _XPSphinxCapability(cpu=True, gpu=None),
|
||||
"array_api_strict": _XPSphinxCapability(cpu=True, gpu=None),
|
||||
"cupy": _XPSphinxCapability(cpu=None, gpu=True),
|
||||
"torch": _XPSphinxCapability(cpu=True, gpu=True),
|
||||
"jax.numpy": _XPSphinxCapability(cpu=True, gpu=True,
|
||||
warnings=[] if jax_jit else ["no JIT"]),
|
||||
# Note: Dask+CuPy is currently untested and unsupported
|
||||
"dask.array": _XPSphinxCapability(cpu=True, gpu=None,
|
||||
warnings=["computes graph"] if allow_dask_compute else []),
|
||||
}
|
||||
|
||||
# documentation doesn't display the reason
|
||||
for module, _ in list(skip_backends) + list(xfail_backends):
|
||||
backend = capabilities[module]
|
||||
if backend.cpu is not None:
|
||||
backend.cpu = False
|
||||
if backend.gpu is not None:
|
||||
backend.gpu = False
|
||||
|
||||
for module, backend in capabilities.items():
|
||||
if np_only and module not in exceptions | {"numpy"}:
|
||||
if backend.cpu is not None:
|
||||
backend.cpu = False
|
||||
if backend.gpu is not None:
|
||||
backend.gpu = False
|
||||
elif cpu_only and module not in exceptions and backend.gpu is not None:
|
||||
backend.gpu = False
|
||||
|
||||
for module, warning in warnings:
|
||||
backend = capabilities[module]
|
||||
backend.warnings.append(warning)
|
||||
|
||||
return capabilities
|
||||
|
||||
|
||||
def _make_capabilities_note(fun_name, capabilities):
|
||||
# Note: deliberately not documenting array-api-strict
|
||||
note = f"""
|
||||
`{fun_name}` has experimental support for Python Array API Standard compatible
|
||||
backends in addition to NumPy. Please consider testing these features
|
||||
by setting an environment variable ``SCIPY_ARRAY_API=1`` and providing
|
||||
CuPy, PyTorch, JAX, or Dask arrays as array arguments. The following
|
||||
combinations of backend and device (or other capability) are supported.
|
||||
|
||||
==================== ==================== ====================
|
||||
Library CPU GPU
|
||||
==================== ==================== ====================
|
||||
NumPy {capabilities['numpy'] }
|
||||
CuPy {capabilities['cupy'] }
|
||||
PyTorch {capabilities['torch'] }
|
||||
JAX {capabilities['jax.numpy'] }
|
||||
Dask {capabilities['dask.array'] }
|
||||
==================== ==================== ====================
|
||||
|
||||
See :ref:`dev-arrayapi` for more information.
|
||||
"""
|
||||
return textwrap.dedent(note)
|
||||
|
||||
|
||||
def xp_capabilities(
|
||||
*,
|
||||
# Alternative capabilities table.
|
||||
# Used only for testing this decorator.
|
||||
capabilities_table=None,
|
||||
# Generate pytest.mark.skip/xfail_xp_backends.
|
||||
# See documentation in conftest.py.
|
||||
# lists of tuples [(module name, reason), ...]
|
||||
skip_backends=(), xfail_backends=(),
|
||||
cpu_only=False, np_only=False, reason=None, exceptions=(),
|
||||
# lists of tuples [(module name, reason), ...]
|
||||
warnings=(),
|
||||
# xpx.testing.lazy_xp_function kwargs.
|
||||
# Refer to array-api-extra documentation.
|
||||
allow_dask_compute=False, jax_jit=True,
|
||||
):
|
||||
"""Decorator for a function that states its support among various
|
||||
Array API compatible backends.
|
||||
|
||||
This decorator has two effects:
|
||||
1. It allows tagging tests with ``@make_xp_test_case`` or
|
||||
``make_xp_pytest_param`` (see below) to automatically generate
|
||||
SKIP/XFAIL markers and perform additional backend-specific
|
||||
testing, such as extra validation for Dask and JAX;
|
||||
2. It automatically adds a note to the function's docstring, containing
|
||||
a table matching what has been tested.
|
||||
|
||||
See Also
|
||||
--------
|
||||
make_xp_test_case
|
||||
make_xp_pytest_param
|
||||
array_api_extra.testing.lazy_xp_function
|
||||
"""
|
||||
capabilities_table = (xp_capabilities_table if capabilities_table is None
|
||||
else capabilities_table)
|
||||
|
||||
capabilities = dict(
|
||||
skip_backends=skip_backends,
|
||||
xfail_backends=xfail_backends,
|
||||
cpu_only=cpu_only,
|
||||
np_only=np_only,
|
||||
reason=reason,
|
||||
exceptions=exceptions,
|
||||
allow_dask_compute=allow_dask_compute,
|
||||
jax_jit=jax_jit,
|
||||
warnings=warnings,
|
||||
)
|
||||
sphinx_capabilities = _make_sphinx_capabilities(**capabilities)
|
||||
|
||||
def decorator(f):
|
||||
# Don't use a wrapper, as in some cases @xp_capabilities is
|
||||
# applied to a ufunc
|
||||
capabilities_table[f] = capabilities
|
||||
note = _make_capabilities_note(f.__name__, sphinx_capabilities)
|
||||
doc = FunctionDoc(f)
|
||||
doc['Notes'].append(note)
|
||||
doc = str(doc).split("\n", 1)[1] # remove signature
|
||||
try:
|
||||
f.__doc__ = doc
|
||||
except AttributeError:
|
||||
# Can't update __doc__ on ufuncs if SciPy
|
||||
# was compiled against NumPy < 2.2.
|
||||
pass
|
||||
|
||||
return f
|
||||
return decorator
|
||||
|
||||
|
||||
def _make_xp_pytest_marks(*funcs, capabilities_table=None):
|
||||
capabilities_table = (xp_capabilities_table if capabilities_table is None
|
||||
else capabilities_table)
|
||||
import pytest
|
||||
from scipy._lib.array_api_extra.testing import lazy_xp_function
|
||||
|
||||
marks = []
|
||||
for func in funcs:
|
||||
capabilities = capabilities_table[func]
|
||||
exceptions = capabilities['exceptions']
|
||||
reason = capabilities['reason']
|
||||
|
||||
if capabilities['cpu_only']:
|
||||
marks.append(pytest.mark.skip_xp_backends(
|
||||
cpu_only=True, exceptions=exceptions, reason=reason))
|
||||
if capabilities['np_only']:
|
||||
marks.append(pytest.mark.skip_xp_backends(
|
||||
np_only=True, exceptions=exceptions, reason=reason))
|
||||
|
||||
for mod_name, reason in capabilities['skip_backends']:
|
||||
marks.append(pytest.mark.skip_xp_backends(mod_name, reason=reason))
|
||||
for mod_name, reason in capabilities['xfail_backends']:
|
||||
marks.append(pytest.mark.xfail_xp_backends(mod_name, reason=reason))
|
||||
|
||||
lazy_kwargs = {k: capabilities[k]
|
||||
for k in ('allow_dask_compute', 'jax_jit')}
|
||||
lazy_xp_function(func, **lazy_kwargs)
|
||||
|
||||
return marks
|
||||
|
||||
|
||||
def make_xp_test_case(*funcs, capabilities_table=None):
|
||||
capabilities_table = (xp_capabilities_table if capabilities_table is None
|
||||
else capabilities_table)
|
||||
"""Generate pytest decorator for a test function that tests functionality
|
||||
of one or more Array API compatible functions.
|
||||
|
||||
Read the parameters of the ``@xp_capabilities`` decorator applied to the
|
||||
listed functions and:
|
||||
|
||||
- Generate the ``@pytest.mark.skip_xp_backends`` and
|
||||
``@pytest.mark.xfail_xp_backends`` decorators
|
||||
for the decorated test function
|
||||
- Tag the function with `xpx.testing.lazy_xp_function`
|
||||
|
||||
See Also
|
||||
--------
|
||||
xp_capabilities
|
||||
make_xp_pytest_param
|
||||
array_api_extra.testing.lazy_xp_function
|
||||
"""
|
||||
marks = _make_xp_pytest_marks(*funcs, capabilities_table=capabilities_table)
|
||||
return lambda func: functools.reduce(lambda f, g: g(f), marks, func)
|
||||
|
||||
|
||||
def make_xp_pytest_param(func, *args, capabilities_table=None):
|
||||
"""Variant of ``make_xp_test_case`` that returns a pytest.param for a function,
|
||||
with all necessary skip_xp_backends and xfail_xp_backends marks applied::
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"func", [make_xp_pytest_param(f1), make_xp_pytest_param(f2)]
|
||||
)
|
||||
def test(func, xp):
|
||||
...
|
||||
|
||||
The above is equivalent to::
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"func", [
|
||||
pytest.param(f1, marks=[
|
||||
pytest.mark.skip_xp_backends(...),
|
||||
pytest.mark.xfail_xp_backends(...), ...]),
|
||||
pytest.param(f2, marks=[
|
||||
pytest.mark.skip_xp_backends(...),
|
||||
pytest.mark.xfail_xp_backends(...), ...]),
|
||||
)
|
||||
def test(func, xp):
|
||||
...
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : Callable
|
||||
Function to be tested. It must be decorated with ``@xp_capabilities``.
|
||||
*args : Any, optional
|
||||
Extra pytest parameters for the use case, e.g.::
|
||||
|
||||
@pytest.mark.parametrize("func,verb", [
|
||||
make_xp_pytest_param(f1, "hello"),
|
||||
make_xp_pytest_param(f2, "world")])
|
||||
def test(func, verb, xp):
|
||||
# iterates on (func=f1, verb="hello")
|
||||
# and (func=f2, verb="world")
|
||||
|
||||
See Also
|
||||
--------
|
||||
xp_capabilities
|
||||
make_xp_test_case
|
||||
array_api_extra.testing.lazy_xp_function
|
||||
"""
|
||||
import pytest
|
||||
|
||||
marks = _make_xp_pytest_marks(func, capabilities_table=capabilities_table)
|
||||
return pytest.param(func, *args, marks=marks, id=func.__name__)
|
||||
|
||||
|
||||
# Is it OK to have a dictionary that is mutated (once upon import) in many places?
|
||||
xp_capabilities_table = {} # type: ignore[var-annotated]
|
||||
@@ -0,0 +1,9 @@
|
||||
# DO NOT RENAME THIS FILE
|
||||
# This is a hook for array_api_extra/src/array_api_extra/_lib/_compat.py
|
||||
# to override functions of array_api_compat.
|
||||
|
||||
from .array_api_compat import * # noqa: F403
|
||||
from ._array_api import array_namespace as scipy_array_namespace
|
||||
|
||||
# overrides array_api_compat.array_namespace inside array-api-extra
|
||||
array_namespace = scipy_array_namespace # type: ignore[assignment]
|
||||
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Extra testing functions that forbid 0d-input, see #21044
|
||||
|
||||
While the xp_assert_* functions generally aim to follow the conventions of the
|
||||
underlying `xp` library, NumPy in particular is inconsistent in its handling
|
||||
of scalars vs. 0d-arrays, see https://github.com/numpy/numpy/issues/24897.
|
||||
|
||||
For example, this means that the following operations (as of v2.0.1) currently
|
||||
return scalars, even though a 0d-array would often be more appropriate:
|
||||
|
||||
import numpy as np
|
||||
np.array(0) * 2 # scalar, not 0d array
|
||||
- np.array(0) # scalar, not 0d-array
|
||||
np.sin(np.array(0)) # scalar, not 0d array
|
||||
np.mean([1, 2, 3]) # scalar, not 0d array
|
||||
|
||||
Libraries like CuPy tend to return a 0d-array in scenarios like those above,
|
||||
and even `xp.asarray(0)[()]` remains a 0d-array there. To deal with the reality
|
||||
of the inconsistencies present in NumPy, as well as 20+ years of code on top,
|
||||
the `xp_assert_*` functions here enforce consistency in the only way that
|
||||
doesn't go against the tide, i.e. by forbidding 0d-arrays as the return type.
|
||||
|
||||
However, when scalars are not generally the expected NumPy return type,
|
||||
it remains preferable to use the assert functions from
|
||||
the `scipy._lib._array_api` module, which have less surprising behaviour.
|
||||
"""
|
||||
from scipy._lib._array_api import array_namespace, is_numpy
|
||||
from scipy._lib._array_api import (xp_assert_close as xp_assert_close_base,
|
||||
xp_assert_equal as xp_assert_equal_base,
|
||||
xp_assert_less as xp_assert_less_base)
|
||||
|
||||
__all__: list[str] = []
|
||||
|
||||
|
||||
def _check_scalar(actual, desired, *, xp=None, **kwargs):
|
||||
__tracebackhide__ = True # Hide traceback for py.test
|
||||
|
||||
if xp is None:
|
||||
xp = array_namespace(actual)
|
||||
|
||||
# necessary to handle non-numpy scalars, e.g. bare `0.0` has no shape
|
||||
desired = xp.asarray(desired)
|
||||
|
||||
# Only NumPy distinguishes between scalars and arrays;
|
||||
# shape check in xp_assert_* is sufficient except for shape == ()
|
||||
if not (is_numpy(xp) and desired.shape == ()):
|
||||
return
|
||||
|
||||
_msg = ("Result is a NumPy 0d-array. Many SciPy functions intend to follow "
|
||||
"the convention of many NumPy functions, returning a scalar when a "
|
||||
"0d-array would be correct. The specialized `xp_assert_*` functions "
|
||||
"in the `scipy._lib._array_api_no_0d` module err on the side of "
|
||||
"caution and do not accept 0d-arrays by default. If the correct "
|
||||
"result may legitimately be a 0d-array, pass `check_0d=True`, "
|
||||
"or use the `xp_assert_*` functions from `scipy._lib._array_api`.")
|
||||
assert xp.isscalar(actual), _msg
|
||||
|
||||
|
||||
def xp_assert_equal(actual, desired, *, check_0d=False, **kwargs):
|
||||
# in contrast to xp_assert_equal_base, this defaults to check_0d=False,
|
||||
# but will do an extra check in that case, which forbids 0d-arrays for `actual`
|
||||
__tracebackhide__ = True # Hide traceback for py.test
|
||||
|
||||
# array-ness (check_0d == True) is taken care of by the *_base functions
|
||||
if not check_0d:
|
||||
_check_scalar(actual, desired, **kwargs)
|
||||
return xp_assert_equal_base(actual, desired, check_0d=check_0d, **kwargs)
|
||||
|
||||
|
||||
def xp_assert_close(actual, desired, *, check_0d=False, **kwargs):
|
||||
# as for xp_assert_equal
|
||||
__tracebackhide__ = True
|
||||
|
||||
if not check_0d:
|
||||
_check_scalar(actual, desired, **kwargs)
|
||||
return xp_assert_close_base(actual, desired, check_0d=check_0d, **kwargs)
|
||||
|
||||
|
||||
def xp_assert_less(actual, desired, *, check_0d=False, **kwargs):
|
||||
# as for xp_assert_equal
|
||||
__tracebackhide__ = True
|
||||
|
||||
if not check_0d:
|
||||
_check_scalar(actual, desired, **kwargs)
|
||||
return xp_assert_less_base(actual, desired, check_0d=check_0d, **kwargs)
|
||||
|
||||
|
||||
def assert_array_almost_equal(actual, desired, decimal=6, *args, **kwds):
|
||||
"""Backwards compatible replacement. In new code, use xp_assert_close instead.
|
||||
"""
|
||||
rtol, atol = 0, 1.5*10**(-decimal)
|
||||
return xp_assert_close(actual, desired,
|
||||
atol=atol, rtol=rtol, check_dtype=False, check_shape=False,
|
||||
*args, **kwds)
|
||||
|
||||
|
||||
def assert_almost_equal(actual, desired, decimal=7, *args, **kwds):
|
||||
"""Backwards compatible replacement. In new code, use xp_assert_close instead.
|
||||
"""
|
||||
rtol, atol = 0, 1.5*10**(-decimal)
|
||||
return xp_assert_close(actual, desired,
|
||||
atol=atol, rtol=rtol, check_dtype=False, check_shape=False,
|
||||
*args, **kwds)
|
||||
@@ -0,0 +1,229 @@
|
||||
import sys as _sys
|
||||
from keyword import iskeyword as _iskeyword
|
||||
|
||||
|
||||
def _validate_names(typename, field_names, extra_field_names):
|
||||
"""
|
||||
Ensure that all the given names are valid Python identifiers that
|
||||
do not start with '_'. Also check that there are no duplicates
|
||||
among field_names + extra_field_names.
|
||||
"""
|
||||
for name in [typename] + field_names + extra_field_names:
|
||||
if not isinstance(name, str):
|
||||
raise TypeError('typename and all field names must be strings')
|
||||
if not name.isidentifier():
|
||||
raise ValueError('typename and all field names must be valid '
|
||||
f'identifiers: {name!r}')
|
||||
if _iskeyword(name):
|
||||
raise ValueError('typename and all field names cannot be a '
|
||||
f'keyword: {name!r}')
|
||||
|
||||
seen = set()
|
||||
for name in field_names + extra_field_names:
|
||||
if name.startswith('_'):
|
||||
raise ValueError('Field names cannot start with an underscore: '
|
||||
f'{name!r}')
|
||||
if name in seen:
|
||||
raise ValueError(f'Duplicate field name: {name!r}')
|
||||
seen.add(name)
|
||||
|
||||
|
||||
# Note: This code is adapted from CPython:Lib/collections/__init__.py
|
||||
def _make_tuple_bunch(typename, field_names, extra_field_names=None,
|
||||
module=None):
|
||||
"""
|
||||
Create a namedtuple-like class with additional attributes.
|
||||
|
||||
This function creates a subclass of tuple that acts like a namedtuple
|
||||
and that has additional attributes.
|
||||
|
||||
The additional attributes are listed in `extra_field_names`. The
|
||||
values assigned to these attributes are not part of the tuple.
|
||||
|
||||
The reason this function exists is to allow functions in SciPy
|
||||
that currently return a tuple or a namedtuple to returned objects
|
||||
that have additional attributes, while maintaining backwards
|
||||
compatibility.
|
||||
|
||||
This should only be used to enhance *existing* functions in SciPy.
|
||||
New functions are free to create objects as return values without
|
||||
having to maintain backwards compatibility with an old tuple or
|
||||
namedtuple return value.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
typename : str
|
||||
The name of the type.
|
||||
field_names : list of str
|
||||
List of names of the values to be stored in the tuple. These names
|
||||
will also be attributes of instances, so the values in the tuple
|
||||
can be accessed by indexing or as attributes. At least one name
|
||||
is required. See the Notes for additional restrictions.
|
||||
extra_field_names : list of str, optional
|
||||
List of names of values that will be stored as attributes of the
|
||||
object. See the notes for additional restrictions.
|
||||
|
||||
Returns
|
||||
-------
|
||||
cls : type
|
||||
The new class.
|
||||
|
||||
Notes
|
||||
-----
|
||||
There are restrictions on the names that may be used in `field_names`
|
||||
and `extra_field_names`:
|
||||
|
||||
* The names must be unique--no duplicates allowed.
|
||||
* The names must be valid Python identifiers, and must not begin with
|
||||
an underscore.
|
||||
* The names must not be Python keywords (e.g. 'def', 'and', etc., are
|
||||
not allowed).
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> from scipy._lib._bunch import _make_tuple_bunch
|
||||
|
||||
Create a class that acts like a namedtuple with length 2 (with field
|
||||
names `x` and `y`) that will also have the attributes `w` and `beta`:
|
||||
|
||||
>>> Result = _make_tuple_bunch('Result', ['x', 'y'], ['w', 'beta'])
|
||||
|
||||
`Result` is the new class. We call it with keyword arguments to create
|
||||
a new instance with given values.
|
||||
|
||||
>>> result1 = Result(x=1, y=2, w=99, beta=0.5)
|
||||
>>> result1
|
||||
Result(x=1, y=2, w=99, beta=0.5)
|
||||
|
||||
`result1` acts like a tuple of length 2:
|
||||
|
||||
>>> len(result1)
|
||||
2
|
||||
>>> result1[:]
|
||||
(1, 2)
|
||||
|
||||
The values assigned when the instance was created are available as
|
||||
attributes:
|
||||
|
||||
>>> result1.y
|
||||
2
|
||||
>>> result1.beta
|
||||
0.5
|
||||
"""
|
||||
if len(field_names) == 0:
|
||||
raise ValueError('field_names must contain at least one name')
|
||||
|
||||
if extra_field_names is None:
|
||||
extra_field_names = []
|
||||
_validate_names(typename, field_names, extra_field_names)
|
||||
|
||||
typename = _sys.intern(str(typename))
|
||||
field_names = tuple(map(_sys.intern, field_names))
|
||||
extra_field_names = tuple(map(_sys.intern, extra_field_names))
|
||||
|
||||
all_names = field_names + extra_field_names
|
||||
arg_list = ', '.join(field_names)
|
||||
full_list = ', '.join(all_names)
|
||||
repr_fmt = ''.join(('(',
|
||||
', '.join(f'{name}=%({name})r' for name in all_names),
|
||||
')'))
|
||||
tuple_new = tuple.__new__
|
||||
_dict, _tuple, _zip = dict, tuple, zip
|
||||
|
||||
# Create all the named tuple methods to be added to the class namespace
|
||||
|
||||
s = f"""\
|
||||
def __new__(_cls, {arg_list}, **extra_fields):
|
||||
return _tuple_new(_cls, ({arg_list},))
|
||||
|
||||
def __init__(self, {arg_list}, **extra_fields):
|
||||
for key in self._extra_fields:
|
||||
if key not in extra_fields:
|
||||
raise TypeError("missing keyword argument '%s'" % (key,))
|
||||
for key, val in extra_fields.items():
|
||||
if key not in self._extra_fields:
|
||||
raise TypeError("unexpected keyword argument '%s'" % (key,))
|
||||
self.__dict__[key] = val
|
||||
|
||||
def __setattr__(self, key, val):
|
||||
if key in {repr(field_names)}:
|
||||
raise AttributeError("can't set attribute %r of class %r"
|
||||
% (key, self.__class__.__name__))
|
||||
else:
|
||||
self.__dict__[key] = val
|
||||
"""
|
||||
del arg_list
|
||||
namespace = {'_tuple_new': tuple_new,
|
||||
'__builtins__': dict(TypeError=TypeError,
|
||||
AttributeError=AttributeError),
|
||||
'__name__': f'namedtuple_{typename}'}
|
||||
exec(s, namespace)
|
||||
__new__ = namespace['__new__']
|
||||
__new__.__doc__ = f'Create new instance of {typename}({full_list})'
|
||||
__init__ = namespace['__init__']
|
||||
__init__.__doc__ = f'Instantiate instance of {typename}({full_list})'
|
||||
__setattr__ = namespace['__setattr__']
|
||||
|
||||
def __repr__(self):
|
||||
'Return a nicely formatted representation string'
|
||||
return self.__class__.__name__ + repr_fmt % self._asdict()
|
||||
|
||||
def _asdict(self):
|
||||
'Return a new dict which maps field names to their values.'
|
||||
out = _dict(_zip(self._fields, self))
|
||||
out.update(self.__dict__)
|
||||
return out
|
||||
|
||||
def __getnewargs_ex__(self):
|
||||
'Return self as a plain tuple. Used by copy and pickle.'
|
||||
return _tuple(self), self.__dict__
|
||||
|
||||
# Modify function metadata to help with introspection and debugging
|
||||
for method in (__new__, __repr__, _asdict, __getnewargs_ex__):
|
||||
method.__qualname__ = f'{typename}.{method.__name__}'
|
||||
|
||||
# Build-up the class namespace dictionary
|
||||
# and use type() to build the result class
|
||||
class_namespace = {
|
||||
'__doc__': f'{typename}({full_list})',
|
||||
'_fields': field_names,
|
||||
'__new__': __new__,
|
||||
'__init__': __init__,
|
||||
'__repr__': __repr__,
|
||||
'__setattr__': __setattr__,
|
||||
'_asdict': _asdict,
|
||||
'_extra_fields': extra_field_names,
|
||||
'__getnewargs_ex__': __getnewargs_ex__,
|
||||
# _field_defaults and _replace are added to get Polars to detect
|
||||
# a bunch object as a namedtuple. See gh-22450
|
||||
'_field_defaults': {},
|
||||
'_replace': None,
|
||||
}
|
||||
for index, name in enumerate(field_names):
|
||||
|
||||
def _get(self, index=index):
|
||||
return self[index]
|
||||
class_namespace[name] = property(_get)
|
||||
for name in extra_field_names:
|
||||
|
||||
def _get(self, name=name):
|
||||
return self.__dict__[name]
|
||||
class_namespace[name] = property(_get)
|
||||
|
||||
result = type(typename, (tuple,), class_namespace)
|
||||
|
||||
# For pickling to work, the __module__ variable needs to be set to the
|
||||
# frame where the named tuple is created. Bypass this step in environments
|
||||
# where sys._getframe is not defined (Jython for example) or sys._getframe
|
||||
# is not defined for arguments greater than 0 (IronPython), or where the
|
||||
# user has specified a particular module.
|
||||
if module is None:
|
||||
try:
|
||||
module = _sys._getframe(1).f_globals.get('__name__', '__main__')
|
||||
except (AttributeError, ValueError):
|
||||
pass
|
||||
if module is not None:
|
||||
result.__module__ = module
|
||||
__new__.__module__ = module
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,251 @@
|
||||
from . import _ccallback_c
|
||||
|
||||
import ctypes
|
||||
|
||||
PyCFuncPtr = ctypes.CFUNCTYPE(ctypes.c_void_p).__bases__[0]
|
||||
|
||||
ffi = None
|
||||
|
||||
class CData:
|
||||
pass
|
||||
|
||||
def _import_cffi():
|
||||
global ffi, CData
|
||||
|
||||
if ffi is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
import cffi
|
||||
ffi = cffi.FFI()
|
||||
CData = ffi.CData
|
||||
except ImportError:
|
||||
ffi = False
|
||||
|
||||
|
||||
class LowLevelCallable(tuple):
|
||||
"""
|
||||
Low-level callback function.
|
||||
|
||||
Some functions in SciPy take as arguments callback functions, which
|
||||
can either be python callables or low-level compiled functions. Using
|
||||
compiled callback functions can improve performance somewhat by
|
||||
avoiding wrapping data in Python objects.
|
||||
|
||||
Such low-level functions in SciPy are wrapped in `LowLevelCallable`
|
||||
objects, which can be constructed from function pointers obtained from
|
||||
ctypes, cffi, Cython, or contained in Python `PyCapsule` objects.
|
||||
|
||||
.. seealso::
|
||||
|
||||
Functions accepting low-level callables:
|
||||
|
||||
`scipy.integrate.quad`, `scipy.ndimage.generic_filter`,
|
||||
`scipy.ndimage.generic_filter1d`, `scipy.ndimage.geometric_transform`
|
||||
|
||||
Usage examples:
|
||||
|
||||
:ref:`ndimage-ccallbacks`, :ref:`quad-callbacks`
|
||||
|
||||
Parameters
|
||||
----------
|
||||
function : {PyCapsule, ctypes function pointer, cffi function pointer}
|
||||
Low-level callback function.
|
||||
user_data : {PyCapsule, ctypes void pointer, cffi void pointer}
|
||||
User data to pass on to the callback function.
|
||||
signature : str, optional
|
||||
Signature of the function. If omitted, determined from *function*,
|
||||
if possible.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
function
|
||||
Callback function given.
|
||||
user_data
|
||||
User data given.
|
||||
signature
|
||||
Signature of the function.
|
||||
|
||||
Methods
|
||||
-------
|
||||
from_cython
|
||||
Class method for constructing callables from Cython C-exported
|
||||
functions.
|
||||
|
||||
Notes
|
||||
-----
|
||||
The argument ``function`` can be one of:
|
||||
|
||||
- PyCapsule, whose name contains the C function signature
|
||||
- ctypes function pointer
|
||||
- cffi function pointer
|
||||
|
||||
The signature of the low-level callback must match one of those expected
|
||||
by the routine it is passed to.
|
||||
|
||||
If constructing low-level functions from a PyCapsule, the name of the
|
||||
capsule must be the corresponding signature, in the format::
|
||||
|
||||
return_type (arg1_type, arg2_type, ...)
|
||||
|
||||
For example::
|
||||
|
||||
"void (double)"
|
||||
"double (double, int *, void *)"
|
||||
|
||||
The context of a PyCapsule passed in as ``function`` is used as ``user_data``,
|
||||
if an explicit value for ``user_data`` was not given.
|
||||
|
||||
"""
|
||||
|
||||
# Make the class immutable
|
||||
__slots__ = ()
|
||||
|
||||
def __new__(cls, function, user_data=None, signature=None):
|
||||
# We need to hold a reference to the function & user data,
|
||||
# to prevent them going out of scope
|
||||
item = cls._parse_callback(function, user_data, signature)
|
||||
return tuple.__new__(cls, (item, function, user_data))
|
||||
|
||||
def __repr__(self):
|
||||
return f"LowLevelCallable({self.function!r}, {self.user_data!r})"
|
||||
|
||||
@property
|
||||
def function(self):
|
||||
return tuple.__getitem__(self, 1)
|
||||
|
||||
@property
|
||||
def user_data(self):
|
||||
return tuple.__getitem__(self, 2)
|
||||
|
||||
@property
|
||||
def signature(self):
|
||||
return _ccallback_c.get_capsule_signature(tuple.__getitem__(self, 0))
|
||||
|
||||
def __getitem__(self, idx):
|
||||
raise ValueError()
|
||||
|
||||
@classmethod
|
||||
def from_cython(cls, module, name, user_data=None, signature=None):
|
||||
"""
|
||||
Create a low-level callback function from an exported Cython function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
module : module
|
||||
Cython module where the exported function resides
|
||||
name : str
|
||||
Name of the exported function
|
||||
user_data : {PyCapsule, ctypes void pointer, cffi void pointer}, optional
|
||||
User data to pass on to the callback function.
|
||||
signature : str, optional
|
||||
Signature of the function. If omitted, determined from *function*.
|
||||
|
||||
"""
|
||||
try:
|
||||
function = module.__pyx_capi__[name]
|
||||
except AttributeError as e:
|
||||
message = "Given module is not a Cython module with __pyx_capi__ attribute"
|
||||
raise ValueError(message) from e
|
||||
except KeyError as e:
|
||||
message = f"No function {name!r} found in __pyx_capi__ of the module"
|
||||
raise ValueError(message) from e
|
||||
return cls(function, user_data, signature)
|
||||
|
||||
@classmethod
|
||||
def _parse_callback(cls, obj, user_data=None, signature=None):
|
||||
_import_cffi()
|
||||
|
||||
if isinstance(obj, LowLevelCallable):
|
||||
func = tuple.__getitem__(obj, 0)
|
||||
elif isinstance(obj, PyCFuncPtr):
|
||||
func, signature = _get_ctypes_func(obj, signature)
|
||||
elif isinstance(obj, CData):
|
||||
func, signature = _get_cffi_func(obj, signature)
|
||||
elif _ccallback_c.check_capsule(obj):
|
||||
func = obj
|
||||
else:
|
||||
raise ValueError("Given input is not a callable or a "
|
||||
"low-level callable (pycapsule/ctypes/cffi)")
|
||||
|
||||
if isinstance(user_data, ctypes.c_void_p):
|
||||
context = _get_ctypes_data(user_data)
|
||||
elif isinstance(user_data, CData):
|
||||
context = _get_cffi_data(user_data)
|
||||
elif user_data is None:
|
||||
context = 0
|
||||
elif _ccallback_c.check_capsule(user_data):
|
||||
context = user_data
|
||||
else:
|
||||
raise ValueError("Given user data is not a valid "
|
||||
"low-level void* pointer (pycapsule/ctypes/cffi)")
|
||||
|
||||
return _ccallback_c.get_raw_capsule(func, signature, context)
|
||||
|
||||
|
||||
#
|
||||
# ctypes helpers
|
||||
#
|
||||
|
||||
def _get_ctypes_func(func, signature=None):
|
||||
# Get function pointer
|
||||
func_ptr = ctypes.cast(func, ctypes.c_void_p).value
|
||||
|
||||
# Construct function signature
|
||||
if signature is None:
|
||||
signature = _typename_from_ctypes(func.restype) + " ("
|
||||
for j, arg in enumerate(func.argtypes):
|
||||
if j == 0:
|
||||
signature += _typename_from_ctypes(arg)
|
||||
else:
|
||||
signature += ", " + _typename_from_ctypes(arg)
|
||||
signature += ")"
|
||||
|
||||
return func_ptr, signature
|
||||
|
||||
|
||||
def _typename_from_ctypes(item):
|
||||
if item is None:
|
||||
return "void"
|
||||
elif item is ctypes.c_void_p:
|
||||
return "void *"
|
||||
|
||||
name = item.__name__
|
||||
|
||||
pointer_level = 0
|
||||
while name.startswith("LP_"):
|
||||
pointer_level += 1
|
||||
name = name[3:]
|
||||
|
||||
if name.startswith('c_'):
|
||||
name = name[2:]
|
||||
|
||||
if pointer_level > 0:
|
||||
name += " " + "*"*pointer_level
|
||||
|
||||
return name
|
||||
|
||||
|
||||
def _get_ctypes_data(data):
|
||||
# Get voidp pointer
|
||||
return ctypes.cast(data, ctypes.c_void_p).value
|
||||
|
||||
|
||||
#
|
||||
# CFFI helpers
|
||||
#
|
||||
|
||||
def _get_cffi_func(func, signature=None):
|
||||
# Get function pointer
|
||||
func_ptr = ffi.cast('uintptr_t', func)
|
||||
|
||||
# Get signature
|
||||
if signature is None:
|
||||
signature = ffi.getctype(ffi.typeof(func)).replace('(*)', ' ')
|
||||
|
||||
return func_ptr, signature
|
||||
|
||||
|
||||
def _get_cffi_data(data):
|
||||
# Get pointer
|
||||
return ffi.cast('uintptr_t', data)
|
||||
BIN
Binary file not shown.
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
Disjoint set data structure
|
||||
"""
|
||||
|
||||
|
||||
class DisjointSet:
|
||||
""" Disjoint set data structure for incremental connectivity queries.
|
||||
|
||||
.. versionadded:: 1.6.0
|
||||
|
||||
Attributes
|
||||
----------
|
||||
n_subsets : int
|
||||
The number of subsets.
|
||||
|
||||
Methods
|
||||
-------
|
||||
add
|
||||
merge
|
||||
connected
|
||||
subset
|
||||
subset_size
|
||||
subsets
|
||||
__getitem__
|
||||
|
||||
Notes
|
||||
-----
|
||||
This class implements the disjoint set [1]_, also known as the *union-find*
|
||||
or *merge-find* data structure. The *find* operation (implemented in
|
||||
`__getitem__`) implements the *path halving* variant. The *merge* method
|
||||
implements the *merge by size* variant.
|
||||
|
||||
References
|
||||
----------
|
||||
.. [1] https://en.wikipedia.org/wiki/Disjoint-set_data_structure
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> from scipy.cluster.hierarchy import DisjointSet
|
||||
|
||||
Initialize a disjoint set:
|
||||
|
||||
>>> disjoint_set = DisjointSet([1, 2, 3, 'a', 'b'])
|
||||
|
||||
Merge some subsets:
|
||||
|
||||
>>> disjoint_set.merge(1, 2)
|
||||
True
|
||||
>>> disjoint_set.merge(3, 'a')
|
||||
True
|
||||
>>> disjoint_set.merge('a', 'b')
|
||||
True
|
||||
>>> disjoint_set.merge('b', 'b')
|
||||
False
|
||||
|
||||
Find root elements:
|
||||
|
||||
>>> disjoint_set[2]
|
||||
1
|
||||
>>> disjoint_set['b']
|
||||
3
|
||||
|
||||
Test connectivity:
|
||||
|
||||
>>> disjoint_set.connected(1, 2)
|
||||
True
|
||||
>>> disjoint_set.connected(1, 'b')
|
||||
False
|
||||
|
||||
List elements in disjoint set:
|
||||
|
||||
>>> list(disjoint_set)
|
||||
[1, 2, 3, 'a', 'b']
|
||||
|
||||
Get the subset containing 'a':
|
||||
|
||||
>>> disjoint_set.subset('a')
|
||||
{'a', 3, 'b'}
|
||||
|
||||
Get the size of the subset containing 'a' (without actually instantiating
|
||||
the subset):
|
||||
|
||||
>>> disjoint_set.subset_size('a')
|
||||
3
|
||||
|
||||
Get all subsets in the disjoint set:
|
||||
|
||||
>>> disjoint_set.subsets()
|
||||
[{1, 2}, {'a', 3, 'b'}]
|
||||
"""
|
||||
def __init__(self, elements=None):
|
||||
self.n_subsets = 0
|
||||
self._sizes = {}
|
||||
self._parents = {}
|
||||
# _nbrs is a circular linked list which links connected elements.
|
||||
self._nbrs = {}
|
||||
# _indices tracks the element insertion order in `__iter__`.
|
||||
self._indices = {}
|
||||
if elements is not None:
|
||||
for x in elements:
|
||||
self.add(x)
|
||||
|
||||
def __iter__(self):
|
||||
"""Returns an iterator of the elements in the disjoint set.
|
||||
|
||||
Elements are ordered by insertion order.
|
||||
"""
|
||||
return iter(self._indices)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._indices)
|
||||
|
||||
def __contains__(self, x):
|
||||
return x in self._indices
|
||||
|
||||
def __getitem__(self, x):
|
||||
"""Find the root element of `x`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : hashable object
|
||||
Input element.
|
||||
|
||||
Returns
|
||||
-------
|
||||
root : hashable object
|
||||
Root element of `x`.
|
||||
"""
|
||||
if x not in self._indices:
|
||||
raise KeyError(x)
|
||||
|
||||
# find by "path halving"
|
||||
parents = self._parents
|
||||
while self._indices[x] != self._indices[parents[x]]:
|
||||
parents[x] = parents[parents[x]]
|
||||
x = parents[x]
|
||||
return x
|
||||
|
||||
def add(self, x):
|
||||
"""Add element `x` to disjoint set
|
||||
"""
|
||||
if x in self._indices:
|
||||
return
|
||||
|
||||
self._sizes[x] = 1
|
||||
self._parents[x] = x
|
||||
self._nbrs[x] = x
|
||||
self._indices[x] = len(self._indices)
|
||||
self.n_subsets += 1
|
||||
|
||||
def merge(self, x, y):
|
||||
"""Merge the subsets of `x` and `y`.
|
||||
|
||||
The smaller subset (the child) is merged into the larger subset (the
|
||||
parent). If the subsets are of equal size, the root element which was
|
||||
first inserted into the disjoint set is selected as the parent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x, y : hashable object
|
||||
Elements to merge.
|
||||
|
||||
Returns
|
||||
-------
|
||||
merged : bool
|
||||
True if `x` and `y` were in disjoint sets, False otherwise.
|
||||
"""
|
||||
xr = self[x]
|
||||
yr = self[y]
|
||||
if self._indices[xr] == self._indices[yr]:
|
||||
return False
|
||||
|
||||
sizes = self._sizes
|
||||
if (sizes[xr], self._indices[yr]) < (sizes[yr], self._indices[xr]):
|
||||
xr, yr = yr, xr
|
||||
self._parents[yr] = xr
|
||||
self._sizes[xr] += self._sizes[yr]
|
||||
self._nbrs[xr], self._nbrs[yr] = self._nbrs[yr], self._nbrs[xr]
|
||||
self.n_subsets -= 1
|
||||
return True
|
||||
|
||||
def connected(self, x, y):
|
||||
"""Test whether `x` and `y` are in the same subset.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x, y : hashable object
|
||||
Elements to test.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : bool
|
||||
True if `x` and `y` are in the same set, False otherwise.
|
||||
"""
|
||||
return self._indices[self[x]] == self._indices[self[y]]
|
||||
|
||||
def subset(self, x):
|
||||
"""Get the subset containing `x`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : hashable object
|
||||
Input element.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : set
|
||||
Subset containing `x`.
|
||||
"""
|
||||
if x not in self._indices:
|
||||
raise KeyError(x)
|
||||
|
||||
result = [x]
|
||||
nxt = self._nbrs[x]
|
||||
while self._indices[nxt] != self._indices[x]:
|
||||
result.append(nxt)
|
||||
nxt = self._nbrs[nxt]
|
||||
return set(result)
|
||||
|
||||
def subset_size(self, x):
|
||||
"""Get the size of the subset containing `x`.
|
||||
|
||||
Note that this method is faster than ``len(self.subset(x))`` because
|
||||
the size is directly read off an internal field, without the need to
|
||||
instantiate the full subset.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : hashable object
|
||||
Input element.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : int
|
||||
Size of the subset containing `x`.
|
||||
"""
|
||||
return self._sizes[self[x]]
|
||||
|
||||
def subsets(self):
|
||||
"""Get all the subsets in the disjoint set.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : list
|
||||
Subsets in the disjoint set.
|
||||
"""
|
||||
result = []
|
||||
visited = set()
|
||||
for x in self:
|
||||
if x not in visited:
|
||||
xset = self.subset(x)
|
||||
visited.update(xset)
|
||||
result.append(xset)
|
||||
return result
|
||||
@@ -0,0 +1,761 @@
|
||||
# copied from numpydoc/docscrape.py, commit 97a6026508e0dd5382865672e9563a72cc113bd2
|
||||
"""Extract reference documentation from the NumPy source tree."""
|
||||
|
||||
import copy
|
||||
import inspect
|
||||
import pydoc
|
||||
import re
|
||||
import sys
|
||||
import textwrap
|
||||
from collections import namedtuple
|
||||
from collections.abc import Callable, Mapping
|
||||
from functools import cached_property
|
||||
from warnings import warn
|
||||
|
||||
|
||||
def strip_blank_lines(l):
|
||||
"Remove leading and trailing blank lines from a list of lines"
|
||||
while l and not l[0].strip():
|
||||
del l[0]
|
||||
while l and not l[-1].strip():
|
||||
del l[-1]
|
||||
return l
|
||||
|
||||
|
||||
class Reader:
|
||||
"""A line-based string reader."""
|
||||
|
||||
def __init__(self, data):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
data : str
|
||||
String with lines separated by '\\n'.
|
||||
|
||||
"""
|
||||
if isinstance(data, list):
|
||||
self._str = data
|
||||
else:
|
||||
self._str = data.split("\n") # store string as list of lines
|
||||
|
||||
self.reset()
|
||||
|
||||
def __getitem__(self, n):
|
||||
return self._str[n]
|
||||
|
||||
def reset(self):
|
||||
self._l = 0 # current line nr
|
||||
|
||||
def read(self):
|
||||
if not self.eof():
|
||||
out = self[self._l]
|
||||
self._l += 1
|
||||
return out
|
||||
else:
|
||||
return ""
|
||||
|
||||
def seek_next_non_empty_line(self):
|
||||
for l in self[self._l :]:
|
||||
if l.strip():
|
||||
break
|
||||
else:
|
||||
self._l += 1
|
||||
|
||||
def eof(self):
|
||||
return self._l >= len(self._str)
|
||||
|
||||
def read_to_condition(self, condition_func):
|
||||
start = self._l
|
||||
for line in self[start:]:
|
||||
if condition_func(line):
|
||||
return self[start : self._l]
|
||||
self._l += 1
|
||||
if self.eof():
|
||||
return self[start : self._l + 1]
|
||||
return []
|
||||
|
||||
def read_to_next_empty_line(self):
|
||||
self.seek_next_non_empty_line()
|
||||
|
||||
def is_empty(line):
|
||||
return not line.strip()
|
||||
|
||||
return self.read_to_condition(is_empty)
|
||||
|
||||
def read_to_next_unindented_line(self):
|
||||
def is_unindented(line):
|
||||
return line.strip() and (len(line.lstrip()) == len(line))
|
||||
|
||||
return self.read_to_condition(is_unindented)
|
||||
|
||||
def peek(self, n=0):
|
||||
if self._l + n < len(self._str):
|
||||
return self[self._l + n]
|
||||
else:
|
||||
return ""
|
||||
|
||||
def is_empty(self):
|
||||
return not "".join(self._str).strip()
|
||||
|
||||
|
||||
class ParseError(Exception):
|
||||
def __str__(self):
|
||||
message = self.args[0]
|
||||
if hasattr(self, "docstring"):
|
||||
message = f"{message} in {self.docstring!r}"
|
||||
return message
|
||||
|
||||
|
||||
Parameter = namedtuple("Parameter", ["name", "type", "desc"])
|
||||
|
||||
|
||||
class NumpyDocString(Mapping):
|
||||
"""Parses a numpydoc string to an abstract representation
|
||||
|
||||
Instances define a mapping from section title to structured data.
|
||||
|
||||
"""
|
||||
|
||||
sections = {
|
||||
"Signature": "",
|
||||
"Summary": [""],
|
||||
"Extended Summary": [],
|
||||
"Parameters": [],
|
||||
"Attributes": [],
|
||||
"Methods": [],
|
||||
"Returns": [],
|
||||
"Yields": [],
|
||||
"Receives": [],
|
||||
"Other Parameters": [],
|
||||
"Raises": [],
|
||||
"Warns": [],
|
||||
"Warnings": [],
|
||||
"See Also": [],
|
||||
"Notes": [],
|
||||
"References": "",
|
||||
"Examples": "",
|
||||
"index": {},
|
||||
}
|
||||
|
||||
def __init__(self, docstring, config=None):
|
||||
orig_docstring = docstring
|
||||
docstring = textwrap.dedent(docstring).split("\n")
|
||||
|
||||
self._doc = Reader(docstring)
|
||||
self._parsed_data = copy.deepcopy(self.sections)
|
||||
|
||||
try:
|
||||
self._parse()
|
||||
except ParseError as e:
|
||||
e.docstring = orig_docstring
|
||||
raise
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._parsed_data[key]
|
||||
|
||||
def __setitem__(self, key, val):
|
||||
if key not in self._parsed_data:
|
||||
self._error_location(f"Unknown section {key}", error=False)
|
||||
else:
|
||||
self._parsed_data[key] = val
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._parsed_data)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._parsed_data)
|
||||
|
||||
def _is_at_section(self):
|
||||
self._doc.seek_next_non_empty_line()
|
||||
|
||||
if self._doc.eof():
|
||||
return False
|
||||
|
||||
l1 = self._doc.peek().strip() # e.g. Parameters
|
||||
|
||||
if l1.startswith(".. index::"):
|
||||
return True
|
||||
|
||||
l2 = self._doc.peek(1).strip() # ---------- or ==========
|
||||
if len(l2) >= 3 and (set(l2) in ({"-"}, {"="})) and len(l2) != len(l1):
|
||||
snip = "\n".join(self._doc._str[:2]) + "..."
|
||||
self._error_location(
|
||||
f"potentially wrong underline length... \n{l1} \n{l2} in \n{snip}",
|
||||
error=False,
|
||||
)
|
||||
return l2.startswith("-" * len(l1)) or l2.startswith("=" * len(l1))
|
||||
|
||||
def _strip(self, doc):
|
||||
i = 0
|
||||
j = 0
|
||||
for i, line in enumerate(doc):
|
||||
if line.strip():
|
||||
break
|
||||
|
||||
for j, line in enumerate(doc[::-1]):
|
||||
if line.strip():
|
||||
break
|
||||
|
||||
return doc[i : len(doc) - j]
|
||||
|
||||
def _read_to_next_section(self):
|
||||
section = self._doc.read_to_next_empty_line()
|
||||
|
||||
while not self._is_at_section() and not self._doc.eof():
|
||||
if not self._doc.peek(-1).strip(): # previous line was empty
|
||||
section += [""]
|
||||
|
||||
section += self._doc.read_to_next_empty_line()
|
||||
|
||||
return section
|
||||
|
||||
def _read_sections(self):
|
||||
while not self._doc.eof():
|
||||
data = self._read_to_next_section()
|
||||
name = data[0].strip()
|
||||
|
||||
if name.startswith(".."): # index section
|
||||
yield name, data[1:]
|
||||
elif len(data) < 2:
|
||||
yield StopIteration
|
||||
else:
|
||||
yield name, self._strip(data[2:])
|
||||
|
||||
def _parse_param_list(self, content, single_element_is_type=False):
|
||||
content = dedent_lines(content)
|
||||
r = Reader(content)
|
||||
params = []
|
||||
while not r.eof():
|
||||
header = r.read().strip()
|
||||
if " : " in header:
|
||||
arg_name, arg_type = header.split(" : ", maxsplit=1)
|
||||
else:
|
||||
# NOTE: param line with single element should never have a
|
||||
# a " :" before the description line, so this should probably
|
||||
# warn.
|
||||
if header.endswith(" :"):
|
||||
header = header[:-2]
|
||||
if single_element_is_type:
|
||||
arg_name, arg_type = "", header
|
||||
else:
|
||||
arg_name, arg_type = header, ""
|
||||
|
||||
desc = r.read_to_next_unindented_line()
|
||||
desc = dedent_lines(desc)
|
||||
desc = strip_blank_lines(desc)
|
||||
|
||||
params.append(Parameter(arg_name, arg_type, desc))
|
||||
|
||||
return params
|
||||
|
||||
# See also supports the following formats.
|
||||
#
|
||||
# <FUNCNAME>
|
||||
# <FUNCNAME> SPACE* COLON SPACE+ <DESC> SPACE*
|
||||
# <FUNCNAME> ( COMMA SPACE+ <FUNCNAME>)+ (COMMA | PERIOD)? SPACE*
|
||||
# <FUNCNAME> ( COMMA SPACE+ <FUNCNAME>)* SPACE* COLON SPACE+ <DESC> SPACE*
|
||||
|
||||
# <FUNCNAME> is one of
|
||||
# <PLAIN_FUNCNAME>
|
||||
# COLON <ROLE> COLON BACKTICK <PLAIN_FUNCNAME> BACKTICK
|
||||
# where
|
||||
# <PLAIN_FUNCNAME> is a legal function name, and
|
||||
# <ROLE> is any nonempty sequence of word characters.
|
||||
# Examples: func_f1 :meth:`func_h1` :obj:`~baz.obj_r` :class:`class_j`
|
||||
# <DESC> is a string describing the function.
|
||||
|
||||
_role = r":(?P<role>(py:)?\w+):"
|
||||
_funcbacktick = r"`(?P<name>(?:~\w+\.)?[a-zA-Z0-9_\.-]+)`"
|
||||
_funcplain = r"(?P<name2>[a-zA-Z0-9_\.-]+)"
|
||||
_funcname = r"(" + _role + _funcbacktick + r"|" + _funcplain + r")"
|
||||
_funcnamenext = _funcname.replace("role", "rolenext")
|
||||
_funcnamenext = _funcnamenext.replace("name", "namenext")
|
||||
_description = r"(?P<description>\s*:(\s+(?P<desc>\S+.*))?)?\s*$"
|
||||
_func_rgx = re.compile(r"^\s*" + _funcname + r"\s*")
|
||||
_line_rgx = re.compile(
|
||||
r"^\s*"
|
||||
+ r"(?P<allfuncs>"
|
||||
+ _funcname # group for all function names
|
||||
+ r"(?P<morefuncs>([,]\s+"
|
||||
+ _funcnamenext
|
||||
+ r")*)"
|
||||
+ r")"
|
||||
+ r"(?P<trailing>[,\.])?" # end of "allfuncs"
|
||||
+ _description # Some function lists have a trailing comma (or period) '\s*'
|
||||
)
|
||||
|
||||
# Empty <DESC> elements are replaced with '..'
|
||||
empty_description = ".."
|
||||
|
||||
def _parse_see_also(self, content):
|
||||
"""
|
||||
func_name : Descriptive text
|
||||
continued text
|
||||
another_func_name : Descriptive text
|
||||
func_name1, func_name2, :meth:`func_name`, func_name3
|
||||
|
||||
"""
|
||||
|
||||
content = dedent_lines(content)
|
||||
|
||||
items = []
|
||||
|
||||
def parse_item_name(text):
|
||||
"""Match ':role:`name`' or 'name'."""
|
||||
m = self._func_rgx.match(text)
|
||||
if not m:
|
||||
self._error_location(f"Error parsing See Also entry {line!r}")
|
||||
role = m.group("role")
|
||||
name = m.group("name") if role else m.group("name2")
|
||||
return name, role, m.end()
|
||||
|
||||
rest = []
|
||||
for line in content:
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
line_match = self._line_rgx.match(line)
|
||||
description = None
|
||||
if line_match:
|
||||
description = line_match.group("desc")
|
||||
if line_match.group("trailing") and description:
|
||||
self._error_location(
|
||||
"Unexpected comma or period after function list at index %d of "
|
||||
'line "%s"' % (line_match.end("trailing"), line),
|
||||
error=False,
|
||||
)
|
||||
if not description and line.startswith(" "):
|
||||
rest.append(line.strip())
|
||||
elif line_match:
|
||||
funcs = []
|
||||
text = line_match.group("allfuncs")
|
||||
while True:
|
||||
if not text.strip():
|
||||
break
|
||||
name, role, match_end = parse_item_name(text)
|
||||
funcs.append((name, role))
|
||||
text = text[match_end:].strip()
|
||||
if text and text[0] == ",":
|
||||
text = text[1:].strip()
|
||||
rest = list(filter(None, [description]))
|
||||
items.append((funcs, rest))
|
||||
else:
|
||||
self._error_location(f"Error parsing See Also entry {line!r}")
|
||||
return items
|
||||
|
||||
def _parse_index(self, section, content):
|
||||
"""
|
||||
.. index:: default
|
||||
:refguide: something, else, and more
|
||||
|
||||
"""
|
||||
|
||||
def strip_each_in(lst):
|
||||
return [s.strip() for s in lst]
|
||||
|
||||
out = {}
|
||||
section = section.split("::")
|
||||
if len(section) > 1:
|
||||
out["default"] = strip_each_in(section[1].split(","))[0]
|
||||
for line in content:
|
||||
line = line.split(":")
|
||||
if len(line) > 2:
|
||||
out[line[1]] = strip_each_in(line[2].split(","))
|
||||
return out
|
||||
|
||||
def _parse_summary(self):
|
||||
"""Grab signature (if given) and summary"""
|
||||
if self._is_at_section():
|
||||
return
|
||||
|
||||
# If several signatures present, take the last one
|
||||
while True:
|
||||
summary = self._doc.read_to_next_empty_line()
|
||||
summary_str = " ".join([s.strip() for s in summary]).strip()
|
||||
compiled = re.compile(r"^([\w., ]+=)?\s*[\w\.]+\(.*\)$")
|
||||
if compiled.match(summary_str):
|
||||
self["Signature"] = summary_str
|
||||
if not self._is_at_section():
|
||||
continue
|
||||
break
|
||||
|
||||
if summary is not None:
|
||||
self["Summary"] = summary
|
||||
|
||||
if not self._is_at_section():
|
||||
self["Extended Summary"] = self._read_to_next_section()
|
||||
|
||||
def _parse(self):
|
||||
self._doc.reset()
|
||||
self._parse_summary()
|
||||
|
||||
sections = list(self._read_sections())
|
||||
section_names = {section for section, content in sections}
|
||||
|
||||
has_yields = "Yields" in section_names
|
||||
# We could do more tests, but we are not. Arbitrarily.
|
||||
if not has_yields and "Receives" in section_names:
|
||||
msg = "Docstring contains a Receives section but not Yields."
|
||||
raise ValueError(msg)
|
||||
|
||||
for section, content in sections:
|
||||
if not section.startswith(".."):
|
||||
section = (s.capitalize() for s in section.split(" "))
|
||||
section = " ".join(section)
|
||||
if self.get(section):
|
||||
self._error_location(
|
||||
"The section %s appears twice in %s"
|
||||
% (section, "\n".join(self._doc._str))
|
||||
)
|
||||
|
||||
if section in ("Parameters", "Other Parameters", "Attributes", "Methods"):
|
||||
self[section] = self._parse_param_list(content)
|
||||
elif section in ("Returns", "Yields", "Raises", "Warns", "Receives"):
|
||||
self[section] = self._parse_param_list(
|
||||
content, single_element_is_type=True
|
||||
)
|
||||
elif section.startswith(".. index::"):
|
||||
self["index"] = self._parse_index(section, content)
|
||||
elif section == "See Also":
|
||||
self["See Also"] = self._parse_see_also(content)
|
||||
else:
|
||||
self[section] = content
|
||||
|
||||
@property
|
||||
def _obj(self):
|
||||
if hasattr(self, "_cls"):
|
||||
return self._cls
|
||||
elif hasattr(self, "_f"):
|
||||
return self._f
|
||||
return None
|
||||
|
||||
def _error_location(self, msg, error=True):
|
||||
if self._obj is not None:
|
||||
# we know where the docs came from:
|
||||
try:
|
||||
filename = inspect.getsourcefile(self._obj)
|
||||
except TypeError:
|
||||
filename = None
|
||||
# Make UserWarning more descriptive via object introspection.
|
||||
# Skip if introspection fails
|
||||
name = getattr(self._obj, "__name__", None)
|
||||
if name is None:
|
||||
name = getattr(getattr(self._obj, "__class__", None), "__name__", None)
|
||||
if name is not None:
|
||||
msg += f" in the docstring of {name}"
|
||||
msg += f" in {filename}." if filename else ""
|
||||
if error:
|
||||
raise ValueError(msg)
|
||||
else:
|
||||
warn(msg, stacklevel=3)
|
||||
|
||||
# string conversion routines
|
||||
|
||||
def _str_header(self, name, symbol="-"):
|
||||
return [name, len(name) * symbol]
|
||||
|
||||
def _str_indent(self, doc, indent=4):
|
||||
return [" " * indent + line for line in doc]
|
||||
|
||||
def _str_signature(self):
|
||||
if self["Signature"]:
|
||||
return [self["Signature"].replace("*", r"\*")] + [""]
|
||||
return [""]
|
||||
|
||||
def _str_summary(self):
|
||||
if self["Summary"]:
|
||||
return self["Summary"] + [""]
|
||||
return []
|
||||
|
||||
def _str_extended_summary(self):
|
||||
if self["Extended Summary"]:
|
||||
return self["Extended Summary"] + [""]
|
||||
return []
|
||||
|
||||
def _str_param_list(self, name):
|
||||
out = []
|
||||
if self[name]:
|
||||
out += self._str_header(name)
|
||||
for param in self[name]:
|
||||
parts = []
|
||||
if param.name:
|
||||
parts.append(param.name)
|
||||
if param.type:
|
||||
parts.append(param.type)
|
||||
out += [" : ".join(parts)]
|
||||
if param.desc and "".join(param.desc).strip():
|
||||
out += self._str_indent(param.desc)
|
||||
out += [""]
|
||||
return out
|
||||
|
||||
def _str_section(self, name):
|
||||
out = []
|
||||
if self[name]:
|
||||
out += self._str_header(name)
|
||||
out += self[name]
|
||||
out += [""]
|
||||
return out
|
||||
|
||||
def _str_see_also(self, func_role):
|
||||
if not self["See Also"]:
|
||||
return []
|
||||
out = []
|
||||
out += self._str_header("See Also")
|
||||
out += [""]
|
||||
last_had_desc = True
|
||||
for funcs, desc in self["See Also"]:
|
||||
assert isinstance(funcs, list)
|
||||
links = []
|
||||
for func, role in funcs:
|
||||
if role:
|
||||
link = f":{role}:`{func}`"
|
||||
elif func_role:
|
||||
link = f":{func_role}:`{func}`"
|
||||
else:
|
||||
link = f"`{func}`_"
|
||||
links.append(link)
|
||||
link = ", ".join(links)
|
||||
out += [link]
|
||||
if desc:
|
||||
out += self._str_indent([" ".join(desc)])
|
||||
last_had_desc = True
|
||||
else:
|
||||
last_had_desc = False
|
||||
out += self._str_indent([self.empty_description])
|
||||
|
||||
if last_had_desc:
|
||||
out += [""]
|
||||
out += [""]
|
||||
return out
|
||||
|
||||
def _str_index(self):
|
||||
idx = self["index"]
|
||||
out = []
|
||||
output_index = False
|
||||
default_index = idx.get("default", "")
|
||||
if default_index:
|
||||
output_index = True
|
||||
out += [f".. index:: {default_index}"]
|
||||
for section, references in idx.items():
|
||||
if section == "default":
|
||||
continue
|
||||
output_index = True
|
||||
out += [f" :{section}: {', '.join(references)}"]
|
||||
if output_index:
|
||||
return out
|
||||
return ""
|
||||
|
||||
def __str__(self, func_role=""):
|
||||
out = []
|
||||
out += self._str_signature()
|
||||
out += self._str_summary()
|
||||
out += self._str_extended_summary()
|
||||
out += self._str_param_list("Parameters")
|
||||
for param_list in ("Attributes", "Methods"):
|
||||
out += self._str_param_list(param_list)
|
||||
for param_list in (
|
||||
"Returns",
|
||||
"Yields",
|
||||
"Receives",
|
||||
"Other Parameters",
|
||||
"Raises",
|
||||
"Warns",
|
||||
):
|
||||
out += self._str_param_list(param_list)
|
||||
out += self._str_section("Warnings")
|
||||
out += self._str_see_also(func_role)
|
||||
for s in ("Notes", "References", "Examples"):
|
||||
out += self._str_section(s)
|
||||
out += self._str_index()
|
||||
return "\n".join(out)
|
||||
|
||||
|
||||
def dedent_lines(lines):
|
||||
"""Deindent a list of lines maximally"""
|
||||
return textwrap.dedent("\n".join(lines)).split("\n")
|
||||
|
||||
|
||||
class FunctionDoc(NumpyDocString):
|
||||
def __init__(self, func, role="func", doc=None, config=None):
|
||||
self._f = func
|
||||
self._role = role # e.g. "func" or "meth"
|
||||
|
||||
if doc is None:
|
||||
if func is None:
|
||||
raise ValueError("No function or docstring given")
|
||||
doc = inspect.getdoc(func) or ""
|
||||
if config is None:
|
||||
config = {}
|
||||
NumpyDocString.__init__(self, doc, config)
|
||||
|
||||
def get_func(self):
|
||||
func_name = getattr(self._f, "__name__", self.__class__.__name__)
|
||||
if inspect.isclass(self._f):
|
||||
func = getattr(self._f, "__call__", self._f.__init__)
|
||||
else:
|
||||
func = self._f
|
||||
return func, func_name
|
||||
|
||||
def __str__(self):
|
||||
out = ""
|
||||
|
||||
func, func_name = self.get_func()
|
||||
|
||||
roles = {"func": "function", "meth": "method"}
|
||||
|
||||
if self._role:
|
||||
if self._role not in roles:
|
||||
print(f"Warning: invalid role {self._role}")
|
||||
out += f".. {roles.get(self._role, '')}:: {func_name}\n \n\n"
|
||||
|
||||
out += super().__str__(func_role=self._role)
|
||||
return out
|
||||
|
||||
|
||||
class ObjDoc(NumpyDocString):
|
||||
def __init__(self, obj, doc=None, config=None):
|
||||
self._f = obj
|
||||
if config is None:
|
||||
config = {}
|
||||
NumpyDocString.__init__(self, doc, config=config)
|
||||
|
||||
|
||||
class ClassDoc(NumpyDocString):
|
||||
extra_public_methods = ["__call__"]
|
||||
|
||||
def __init__(self, cls, doc=None, modulename="", func_doc=FunctionDoc, config=None):
|
||||
if not inspect.isclass(cls) and cls is not None:
|
||||
raise ValueError(f"Expected a class or None, but got {cls!r}")
|
||||
self._cls = cls
|
||||
|
||||
if "sphinx" in sys.modules:
|
||||
from sphinx.ext.autodoc import ALL
|
||||
else:
|
||||
ALL = object()
|
||||
|
||||
if config is None:
|
||||
config = {}
|
||||
self.show_inherited_members = config.get("show_inherited_class_members", True)
|
||||
|
||||
if modulename and not modulename.endswith("."):
|
||||
modulename += "."
|
||||
self._mod = modulename
|
||||
|
||||
if doc is None:
|
||||
if cls is None:
|
||||
raise ValueError("No class or documentation string given")
|
||||
doc = pydoc.getdoc(cls)
|
||||
|
||||
NumpyDocString.__init__(self, doc)
|
||||
|
||||
_members = config.get("members", [])
|
||||
if _members is ALL:
|
||||
_members = None
|
||||
_exclude = config.get("exclude-members", [])
|
||||
|
||||
if config.get("show_class_members", True) and _exclude is not ALL:
|
||||
|
||||
def splitlines_x(s):
|
||||
if not s:
|
||||
return []
|
||||
else:
|
||||
return s.splitlines()
|
||||
|
||||
for field, items in [
|
||||
("Methods", self.methods),
|
||||
("Attributes", self.properties),
|
||||
]:
|
||||
if not self[field]:
|
||||
doc_list = []
|
||||
for name in sorted(items):
|
||||
if name in _exclude or (_members and name not in _members):
|
||||
continue
|
||||
try:
|
||||
doc_item = pydoc.getdoc(getattr(self._cls, name))
|
||||
doc_list.append(Parameter(name, "", splitlines_x(doc_item)))
|
||||
except AttributeError:
|
||||
pass # method doesn't exist
|
||||
self[field] = doc_list
|
||||
|
||||
@property
|
||||
def methods(self):
|
||||
if self._cls is None:
|
||||
return []
|
||||
return [
|
||||
name
|
||||
for name, func in inspect.getmembers(self._cls)
|
||||
if (
|
||||
(not name.startswith("_") or name in self.extra_public_methods)
|
||||
and isinstance(func, Callable)
|
||||
and self._is_show_member(name)
|
||||
)
|
||||
]
|
||||
|
||||
@property
|
||||
def properties(self):
|
||||
if self._cls is None:
|
||||
return []
|
||||
return [
|
||||
name
|
||||
for name, func in inspect.getmembers(self._cls)
|
||||
if (
|
||||
not name.startswith("_")
|
||||
and not self._should_skip_member(name, self._cls)
|
||||
and (
|
||||
func is None
|
||||
or isinstance(func, property | cached_property)
|
||||
or inspect.isdatadescriptor(func)
|
||||
)
|
||||
and self._is_show_member(name)
|
||||
)
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _should_skip_member(name, klass):
|
||||
return (
|
||||
# Namedtuples should skip everything in their ._fields as the
|
||||
# docstrings for each of the members is: "Alias for field number X"
|
||||
issubclass(klass, tuple)
|
||||
and hasattr(klass, "_asdict")
|
||||
and hasattr(klass, "_fields")
|
||||
and name in klass._fields
|
||||
)
|
||||
|
||||
def _is_show_member(self, name):
|
||||
return (
|
||||
# show all class members
|
||||
self.show_inherited_members
|
||||
# or class member is not inherited
|
||||
or name in self._cls.__dict__
|
||||
)
|
||||
|
||||
|
||||
def get_doc_object(
|
||||
obj,
|
||||
what=None,
|
||||
doc=None,
|
||||
config=None,
|
||||
class_doc=ClassDoc,
|
||||
func_doc=FunctionDoc,
|
||||
obj_doc=ObjDoc,
|
||||
):
|
||||
if what is None:
|
||||
if inspect.isclass(obj):
|
||||
what = "class"
|
||||
elif inspect.ismodule(obj):
|
||||
what = "module"
|
||||
elif isinstance(obj, Callable):
|
||||
what = "function"
|
||||
else:
|
||||
what = "object"
|
||||
if config is None:
|
||||
config = {}
|
||||
|
||||
if what == "class":
|
||||
return class_doc(obj, func_doc=func_doc, doc=doc, config=config)
|
||||
elif what in ("function", "method"):
|
||||
return func_doc(obj, doc=doc, config=config)
|
||||
else:
|
||||
if doc is None:
|
||||
doc = pydoc.getdoc(obj)
|
||||
return obj_doc(obj, doc, config=config)
|
||||
@@ -0,0 +1,346 @@
|
||||
# `_elementwise_iterative_method.py` includes tools for writing functions that
|
||||
# - are vectorized to work elementwise on arrays,
|
||||
# - implement non-trivial, iterative algorithms with a callback interface, and
|
||||
# - return rich objects with iteration count, termination status, etc.
|
||||
#
|
||||
# Examples include:
|
||||
# `scipy.optimize._chandrupatla._chandrupatla for scalar rootfinding,
|
||||
# `scipy.optimize._chandrupatla._chandrupatla_minimize for scalar minimization,
|
||||
# `scipy.optimize._differentiate._differentiate for numerical differentiation,
|
||||
# `scipy.optimize._bracket._bracket_root for finding rootfinding brackets,
|
||||
# `scipy.optimize._bracket._bracket_minimize for finding minimization brackets,
|
||||
# `scipy.integrate._tanhsinh._tanhsinh` for numerical quadrature,
|
||||
# `scipy.differentiate.derivative` for finite difference based differentiation.
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
from ._util import _RichResult, _call_callback_maybe_halt
|
||||
from ._array_api import array_namespace, xp_size, xp_result_type
|
||||
import scipy._lib.array_api_extra as xpx
|
||||
|
||||
_ESIGNERR = -1
|
||||
_ECONVERR = -2
|
||||
_EVALUEERR = -3
|
||||
_ECALLBACK = -4
|
||||
_EINPUTERR = -5
|
||||
_ECONVERGED = 0
|
||||
_EINPROGRESS = 1
|
||||
|
||||
def _initialize(func, xs, args, complex_ok=False, preserve_shape=None, xp=None):
|
||||
"""Initialize abscissa, function, and args arrays for elementwise function
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : callable
|
||||
An elementwise function with signature
|
||||
|
||||
func(x: ndarray, *args) -> ndarray
|
||||
|
||||
where each element of ``x`` is a finite real and ``args`` is a tuple,
|
||||
which may contain an arbitrary number of arrays that are broadcastable
|
||||
with ``x``.
|
||||
xs : tuple of arrays
|
||||
Finite real abscissa arrays. Must be broadcastable.
|
||||
args : tuple, optional
|
||||
Additional positional arguments to be passed to `func`.
|
||||
preserve_shape : bool, default:False
|
||||
When ``preserve_shape=False`` (default), `func` may be passed
|
||||
arguments of any shape; `_scalar_optimization_loop` is permitted
|
||||
to reshape and compress arguments at will. When
|
||||
``preserve_shape=False``, arguments passed to `func` must have shape
|
||||
`shape` or ``shape + (n,)``, where ``n`` is any integer.
|
||||
xp : namespace
|
||||
Namespace of array arguments in `xs`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
xs, fs, args : tuple of arrays
|
||||
Broadcasted, writeable, 1D abscissa and function value arrays (or
|
||||
NumPy floats, if appropriate). The dtypes of the `xs` and `fs` are
|
||||
`xfat`; the dtype of the `args` are unchanged.
|
||||
shape : tuple of ints
|
||||
Original shape of broadcasted arrays.
|
||||
xfat : NumPy dtype
|
||||
Result dtype of abscissae, function values, and args determined using
|
||||
`np.result_type`, except integer types are promoted to `np.float64`.
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If the result dtype is not that of a real scalar
|
||||
|
||||
Notes
|
||||
-----
|
||||
Useful for initializing the input of SciPy functions that accept
|
||||
an elementwise callable, abscissae, and arguments; e.g.
|
||||
`scipy.optimize._chandrupatla`.
|
||||
"""
|
||||
nx = len(xs)
|
||||
xp = array_namespace(*xs) if xp is None else xp
|
||||
|
||||
# Try to preserve `dtype`, but we need to ensure that the arguments are at
|
||||
# least floats before passing them into the function; integers can overflow
|
||||
# and cause failure.
|
||||
# There might be benefit to combining the `xs` into a single array and
|
||||
# calling `func` once on the combined array. For now, keep them separate.
|
||||
xat = xp_result_type(*xs, force_floating=True, xp=xp)
|
||||
xas = xp.broadcast_arrays(*xs, *args) # broadcast and rename
|
||||
xs, args = xas[:nx], xas[nx:]
|
||||
xs = [xp.asarray(x, dtype=xat) for x in xs] # use copy=False when implemented
|
||||
fs = [xp.asarray(func(x, *args)) for x in xs]
|
||||
shape = xs[0].shape
|
||||
fshape = fs[0].shape
|
||||
|
||||
if preserve_shape:
|
||||
# bind original shape/func now to avoid late-binding gotcha
|
||||
def func(x, *args, shape=shape, func=func, **kwargs):
|
||||
i = (0,)*(len(fshape) - len(shape))
|
||||
return func(x[i], *args, **kwargs)
|
||||
shape = np.broadcast_shapes(fshape, shape) # just shapes; use of NumPy OK
|
||||
xs = [xp.broadcast_to(x, shape) for x in xs]
|
||||
args = [xp.broadcast_to(arg, shape) for arg in args]
|
||||
|
||||
message = ("The shape of the array returned by `func` must be the same as "
|
||||
"the broadcasted shape of `x` and all other `args`.")
|
||||
if preserve_shape is not None: # only in tanhsinh for now
|
||||
message = f"When `preserve_shape=False`, {message.lower()}"
|
||||
shapes_equal = [f.shape == shape for f in fs]
|
||||
if not all(shapes_equal): # use Python all to reduce overhead
|
||||
raise ValueError(message)
|
||||
|
||||
# These algorithms tend to mix the dtypes of the abscissae and function
|
||||
# values, so figure out what the result will be and convert them all to
|
||||
# that type from the outset.
|
||||
xfat = xp.result_type(*([f.dtype for f in fs] + [xat]))
|
||||
if not complex_ok and not xp.isdtype(xfat, "real floating"):
|
||||
raise ValueError("Abscissae and function output must be real numbers.")
|
||||
xs = [xp.asarray(x, dtype=xfat, copy=True) for x in xs]
|
||||
fs = [xp.asarray(f, dtype=xfat, copy=True) for f in fs]
|
||||
|
||||
# To ensure that we can do indexing, we'll work with at least 1d arrays,
|
||||
# but remember the appropriate shape of the output.
|
||||
xs = [xp.reshape(x, (-1,)) for x in xs]
|
||||
fs = [xp.reshape(f, (-1,)) for f in fs]
|
||||
args = [xp.reshape(xp.asarray(arg, copy=True), (-1,)) for arg in args]
|
||||
return func, xs, fs, args, shape, xfat, xp
|
||||
|
||||
|
||||
def _loop(work, callback, shape, maxiter, func, args, dtype, pre_func_eval,
|
||||
post_func_eval, check_termination, post_termination_check,
|
||||
customize_result, res_work_pairs, xp, preserve_shape=False):
|
||||
"""Main loop of a vectorized scalar optimization algorithm
|
||||
|
||||
Parameters
|
||||
----------
|
||||
work : _RichResult
|
||||
All variables that need to be retained between iterations. Must
|
||||
contain attributes `nit`, `nfev`, and `success`. All arrays are
|
||||
subject to being "compressed" if `preserve_shape is False`; nest
|
||||
arrays that should not be compressed inside another object (e.g.
|
||||
`dict` or `_RichResult`).
|
||||
callback : callable
|
||||
User-specified callback function
|
||||
shape : tuple of ints
|
||||
The shape of all output arrays
|
||||
maxiter :
|
||||
Maximum number of iterations of the algorithm
|
||||
func : callable
|
||||
The user-specified callable that is being optimized or solved
|
||||
args : tuple
|
||||
Additional positional arguments to be passed to `func`.
|
||||
dtype : NumPy dtype
|
||||
The common dtype of all abscissae and function values
|
||||
pre_func_eval : callable
|
||||
A function that accepts `work` and returns `x`, the active elements
|
||||
of `x` at which `func` will be evaluated. May modify attributes
|
||||
of `work` with any algorithmic steps that need to happen
|
||||
at the beginning of an iteration, before `func` is evaluated,
|
||||
post_func_eval : callable
|
||||
A function that accepts `x`, `func(x)`, and `work`. May modify
|
||||
attributes of `work` with any algorithmic steps that need to happen
|
||||
in the middle of an iteration, after `func` is evaluated but before
|
||||
the termination check.
|
||||
check_termination : callable
|
||||
A function that accepts `work` and returns `stop`, a boolean array
|
||||
indicating which of the active elements have met a termination
|
||||
condition.
|
||||
post_termination_check : callable
|
||||
A function that accepts `work`. May modify `work` with any algorithmic
|
||||
steps that need to happen after the termination check and before the
|
||||
end of the iteration.
|
||||
customize_result : callable
|
||||
A function that accepts `res` and `shape` and returns `shape`. May
|
||||
modify `res` (in-place) according to preferences (e.g. rearrange
|
||||
elements between attributes) and modify `shape` if needed.
|
||||
res_work_pairs : list of (str, str)
|
||||
Identifies correspondence between attributes of `res` and attributes
|
||||
of `work`; i.e., attributes of active elements of `work` will be
|
||||
copied to the appropriate indices of `res` when appropriate. The order
|
||||
determines the order in which _RichResult attributes will be
|
||||
pretty-printed.
|
||||
preserve_shape : bool, default: False
|
||||
Whether to compress the attributes of `work` (to avoid unnecessary
|
||||
computation on elements that have already converged).
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : _RichResult
|
||||
The final result object
|
||||
|
||||
Notes
|
||||
-----
|
||||
Besides providing structure, this framework provides several important
|
||||
services for a vectorized optimization algorithm.
|
||||
|
||||
- It handles common tasks involving iteration count, function evaluation
|
||||
count, a user-specified callback, and associated termination conditions.
|
||||
- It compresses the attributes of `work` to eliminate unnecessary
|
||||
computation on elements that have already converged.
|
||||
|
||||
"""
|
||||
if xp is None:
|
||||
raise NotImplementedError("Must provide xp.")
|
||||
|
||||
cb_terminate = False
|
||||
|
||||
# Initialize the result object and active element index array
|
||||
n_elements = math.prod(shape)
|
||||
active = xp.arange(n_elements) # in-progress element indices
|
||||
res_dict = {i: xp.zeros(n_elements, dtype=dtype) for i, j in res_work_pairs}
|
||||
res_dict['success'] = xp.zeros(n_elements, dtype=xp.bool)
|
||||
res_dict['status'] = xp.full(n_elements, xp.asarray(_EINPROGRESS), dtype=xp.int32)
|
||||
res_dict['nit'] = xp.zeros(n_elements, dtype=xp.int32)
|
||||
res_dict['nfev'] = xp.zeros(n_elements, dtype=xp.int32)
|
||||
res = _RichResult(res_dict)
|
||||
work.args = args
|
||||
|
||||
active = _check_termination(work, res, res_work_pairs, active,
|
||||
check_termination, preserve_shape, xp)
|
||||
|
||||
if callback is not None:
|
||||
temp = _prepare_result(work, res, res_work_pairs, active, shape,
|
||||
customize_result, preserve_shape, xp)
|
||||
if _call_callback_maybe_halt(callback, temp):
|
||||
cb_terminate = True
|
||||
|
||||
while work.nit < maxiter and xp_size(active) and not cb_terminate and n_elements:
|
||||
x = pre_func_eval(work)
|
||||
|
||||
if work.args and work.args[0].ndim != x.ndim:
|
||||
# `x` always starts as 1D. If the SciPy function that uses
|
||||
# _loop added dimensions to `x`, we need to
|
||||
# add them to the elements of `args`.
|
||||
args = []
|
||||
for arg in work.args:
|
||||
n_new_dims = x.ndim - arg.ndim
|
||||
new_shape = arg.shape + (1,)*n_new_dims
|
||||
args.append(xp.reshape(arg, new_shape))
|
||||
work.args = args
|
||||
|
||||
x_shape = x.shape
|
||||
if preserve_shape:
|
||||
x = xp.reshape(x, (shape + (-1,)))
|
||||
f = func(x, *work.args)
|
||||
f = xp.asarray(f, dtype=dtype)
|
||||
if preserve_shape:
|
||||
x = xp.reshape(x, x_shape)
|
||||
f = xp.reshape(f, x_shape)
|
||||
work.nfev += 1 if x.ndim == 1 else x.shape[-1]
|
||||
|
||||
post_func_eval(x, f, work)
|
||||
|
||||
work.nit += 1
|
||||
active = _check_termination(work, res, res_work_pairs, active,
|
||||
check_termination, preserve_shape, xp)
|
||||
|
||||
if callback is not None:
|
||||
temp = _prepare_result(work, res, res_work_pairs, active, shape,
|
||||
customize_result, preserve_shape, xp)
|
||||
if _call_callback_maybe_halt(callback, temp):
|
||||
cb_terminate = True
|
||||
break
|
||||
if xp_size(active) == 0:
|
||||
break
|
||||
|
||||
post_termination_check(work)
|
||||
|
||||
work.status = xpx.at(work.status)[:].set(_ECALLBACK if cb_terminate else _ECONVERR)
|
||||
return _prepare_result(work, res, res_work_pairs, active, shape,
|
||||
customize_result, preserve_shape, xp)
|
||||
|
||||
|
||||
def _check_termination(work, res, res_work_pairs, active, check_termination,
|
||||
preserve_shape, xp):
|
||||
# Checks termination conditions, updates elements of `res` with
|
||||
# corresponding elements of `work`, and compresses `work`.
|
||||
|
||||
stop = check_termination(work)
|
||||
|
||||
if xp.any(stop):
|
||||
# update the active elements of the result object with the active
|
||||
# elements for which a termination condition has been met
|
||||
_update_active(work, res, res_work_pairs, active, stop, preserve_shape, xp)
|
||||
|
||||
if preserve_shape:
|
||||
stop = stop[active]
|
||||
|
||||
proceed = ~stop
|
||||
active = active[proceed]
|
||||
|
||||
if not preserve_shape:
|
||||
# compress the arrays to avoid unnecessary computation
|
||||
for key, val in work.items():
|
||||
# `continued_fraction` hacks `n`; improve if this becomes a problem
|
||||
if key in {'args', 'n'}:
|
||||
continue
|
||||
work[key] = val[proceed] if getattr(val, 'ndim', 0) > 0 else val
|
||||
work.args = [arg[proceed] for arg in work.args]
|
||||
|
||||
return active
|
||||
|
||||
|
||||
def _update_active(work, res, res_work_pairs, active, mask, preserve_shape, xp):
|
||||
# Update `active` indices of the arrays in result object `res` with the
|
||||
# contents of the scalars and arrays in `update_dict`. When provided,
|
||||
# `mask` is a boolean array applied both to the arrays in `update_dict`
|
||||
# that are to be used and to the arrays in `res` that are to be updated.
|
||||
update_dict = {key1: work[key2] for key1, key2 in res_work_pairs}
|
||||
update_dict['success'] = work.status == 0
|
||||
|
||||
if mask is not None:
|
||||
if preserve_shape:
|
||||
active_mask = xp.zeros_like(mask)
|
||||
active_mask = xpx.at(active_mask)[active].set(True)
|
||||
active_mask = active_mask & mask
|
||||
for key, val in update_dict.items():
|
||||
val = val[active_mask] if getattr(val, 'ndim', 0) > 0 else val
|
||||
res[key] = xpx.at(res[key])[active_mask].set(val)
|
||||
else:
|
||||
active_mask = active[mask]
|
||||
for key, val in update_dict.items():
|
||||
val = val[mask] if getattr(val, 'ndim', 0) > 0 else val
|
||||
res[key] = xpx.at(res[key])[active_mask].set(val)
|
||||
else:
|
||||
for key, val in update_dict.items():
|
||||
if preserve_shape and getattr(val, 'ndim', 0) > 0:
|
||||
val = val[active]
|
||||
res[key] = xpx.at(res[key])[active].set(val)
|
||||
|
||||
|
||||
def _prepare_result(work, res, res_work_pairs, active, shape, customize_result,
|
||||
preserve_shape, xp):
|
||||
# Prepare the result object `res` by creating a copy, copying the latest
|
||||
# data from work, running the provided result customization function,
|
||||
# and reshaping the data to the original shapes.
|
||||
res = res.copy()
|
||||
_update_active(work, res, res_work_pairs, active, None, preserve_shape, xp)
|
||||
|
||||
shape = customize_result(res, shape)
|
||||
|
||||
for key, val in res.items():
|
||||
# this looks like it won't work for xp != np if val is not numeric
|
||||
temp = xp.reshape(val, shape)
|
||||
res[key] = temp[()] if temp.ndim == 0 else temp
|
||||
|
||||
res['_order_keys'] = ['success'] + [i for i, j in res_work_pairs]
|
||||
return _RichResult(**res)
|
||||
Binary file not shown.
@@ -0,0 +1,105 @@
|
||||
"""
|
||||
Module for testing automatic garbage collection of objects
|
||||
|
||||
.. autosummary::
|
||||
:toctree: generated/
|
||||
|
||||
set_gc_state - enable or disable garbage collection
|
||||
gc_state - context manager for given state of garbage collector
|
||||
assert_deallocated - context manager to check for circular references on object
|
||||
|
||||
"""
|
||||
import weakref
|
||||
import gc
|
||||
|
||||
from contextlib import contextmanager
|
||||
from platform import python_implementation
|
||||
|
||||
__all__ = ['set_gc_state', 'gc_state', 'assert_deallocated']
|
||||
|
||||
|
||||
IS_PYPY = python_implementation() == 'PyPy'
|
||||
|
||||
|
||||
class ReferenceError(AssertionError):
|
||||
pass
|
||||
|
||||
|
||||
def set_gc_state(state):
|
||||
""" Set status of garbage collector """
|
||||
if gc.isenabled() == state:
|
||||
return
|
||||
if state:
|
||||
gc.enable()
|
||||
else:
|
||||
gc.disable()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def gc_state(state):
|
||||
""" Context manager to set state of garbage collector to `state`
|
||||
|
||||
Parameters
|
||||
----------
|
||||
state : bool
|
||||
True for gc enabled, False for disabled
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> with gc_state(False):
|
||||
... assert not gc.isenabled()
|
||||
>>> with gc_state(True):
|
||||
... assert gc.isenabled()
|
||||
"""
|
||||
orig_state = gc.isenabled()
|
||||
set_gc_state(state)
|
||||
yield
|
||||
set_gc_state(orig_state)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def assert_deallocated(func, *args, **kwargs):
|
||||
"""Context manager to check that object is deallocated
|
||||
|
||||
This is useful for checking that an object can be freed directly by
|
||||
reference counting, without requiring gc to break reference cycles.
|
||||
GC is disabled inside the context manager.
|
||||
|
||||
This check is not available on PyPy.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : callable
|
||||
Callable to create object to check
|
||||
\\*args : sequence
|
||||
positional arguments to `func` in order to create object to check
|
||||
\\*\\*kwargs : dict
|
||||
keyword arguments to `func` in order to create object to check
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> class C: pass
|
||||
>>> with assert_deallocated(C) as c:
|
||||
... # do something
|
||||
... del c
|
||||
|
||||
>>> class C:
|
||||
... def __init__(self):
|
||||
... self._circular = self # Make circular reference
|
||||
>>> with assert_deallocated(C) as c: #doctest: +IGNORE_EXCEPTION_DETAIL
|
||||
... # do something
|
||||
... del c
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
ReferenceError: Remaining reference(s) to object
|
||||
"""
|
||||
if IS_PYPY:
|
||||
raise RuntimeError("assert_deallocated is unavailable on PyPy")
|
||||
|
||||
with gc_state(False):
|
||||
obj = func(*args, **kwargs)
|
||||
ref = weakref.ref(obj)
|
||||
yield obj
|
||||
del obj
|
||||
if ref() is not None:
|
||||
raise ReferenceError("Remaining reference(s) to object")
|
||||
@@ -0,0 +1,487 @@
|
||||
"""Utility to compare pep440 compatible version strings.
|
||||
|
||||
The LooseVersion and StrictVersion classes that distutils provides don't
|
||||
work; they don't recognize anything like alpha/beta/rc/dev versions.
|
||||
"""
|
||||
|
||||
# Copyright (c) Donald Stufft and individual contributors.
|
||||
# All rights reserved.
|
||||
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
|
||||
# 1. Redistributions of source code must retain the above copyright notice,
|
||||
# this list of conditions and the following disclaimer.
|
||||
|
||||
# 2. Redistributions in binary form must reproduce the above copyright
|
||||
# notice, this list of conditions and the following disclaimer in the
|
||||
# documentation and/or other materials provided with the distribution.
|
||||
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
||||
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
|
||||
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
||||
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
||||
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
||||
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
||||
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
||||
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
# POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
import collections
|
||||
import itertools
|
||||
import re
|
||||
|
||||
|
||||
__all__ = [
|
||||
"parse", "Version", "LegacyVersion", "InvalidVersion", "VERSION_PATTERN",
|
||||
]
|
||||
|
||||
|
||||
# BEGIN packaging/_structures.py
|
||||
|
||||
|
||||
class Infinity:
|
||||
def __repr__(self):
|
||||
return "Infinity"
|
||||
|
||||
def __hash__(self):
|
||||
return hash(repr(self))
|
||||
|
||||
def __lt__(self, other):
|
||||
return False
|
||||
|
||||
def __le__(self, other):
|
||||
return False
|
||||
|
||||
def __eq__(self, other):
|
||||
return isinstance(other, self.__class__)
|
||||
|
||||
def __ne__(self, other):
|
||||
return not isinstance(other, self.__class__)
|
||||
|
||||
def __gt__(self, other):
|
||||
return True
|
||||
|
||||
def __ge__(self, other):
|
||||
return True
|
||||
|
||||
def __neg__(self):
|
||||
return NegativeInfinity
|
||||
|
||||
|
||||
Infinity = Infinity()
|
||||
|
||||
|
||||
class NegativeInfinity:
|
||||
def __repr__(self):
|
||||
return "-Infinity"
|
||||
|
||||
def __hash__(self):
|
||||
return hash(repr(self))
|
||||
|
||||
def __lt__(self, other):
|
||||
return True
|
||||
|
||||
def __le__(self, other):
|
||||
return True
|
||||
|
||||
def __eq__(self, other):
|
||||
return isinstance(other, self.__class__)
|
||||
|
||||
def __ne__(self, other):
|
||||
return not isinstance(other, self.__class__)
|
||||
|
||||
def __gt__(self, other):
|
||||
return False
|
||||
|
||||
def __ge__(self, other):
|
||||
return False
|
||||
|
||||
def __neg__(self):
|
||||
return Infinity
|
||||
|
||||
|
||||
# BEGIN packaging/version.py
|
||||
|
||||
|
||||
NegativeInfinity = NegativeInfinity()
|
||||
|
||||
_Version = collections.namedtuple(
|
||||
"_Version",
|
||||
["epoch", "release", "dev", "pre", "post", "local"],
|
||||
)
|
||||
|
||||
|
||||
def parse(version):
|
||||
"""
|
||||
Parse the given version string and return either a :class:`Version` object
|
||||
or a :class:`LegacyVersion` object depending on if the given version is
|
||||
a valid PEP 440 version or a legacy version.
|
||||
"""
|
||||
try:
|
||||
return Version(version)
|
||||
except InvalidVersion:
|
||||
return LegacyVersion(version)
|
||||
|
||||
|
||||
class InvalidVersion(ValueError):
|
||||
"""
|
||||
An invalid version was found, users should refer to PEP 440.
|
||||
"""
|
||||
|
||||
|
||||
class _BaseVersion:
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self._key)
|
||||
|
||||
def __lt__(self, other):
|
||||
return self._compare(other, lambda s, o: s < o)
|
||||
|
||||
def __le__(self, other):
|
||||
return self._compare(other, lambda s, o: s <= o)
|
||||
|
||||
def __eq__(self, other):
|
||||
return self._compare(other, lambda s, o: s == o)
|
||||
|
||||
def __ge__(self, other):
|
||||
return self._compare(other, lambda s, o: s >= o)
|
||||
|
||||
def __gt__(self, other):
|
||||
return self._compare(other, lambda s, o: s > o)
|
||||
|
||||
def __ne__(self, other):
|
||||
return self._compare(other, lambda s, o: s != o)
|
||||
|
||||
def _compare(self, other, method):
|
||||
if not isinstance(other, _BaseVersion):
|
||||
return NotImplemented
|
||||
|
||||
return method(self._key, other._key)
|
||||
|
||||
|
||||
class LegacyVersion(_BaseVersion):
|
||||
|
||||
def __init__(self, version):
|
||||
self._version = str(version)
|
||||
self._key = _legacy_cmpkey(self._version)
|
||||
|
||||
def __str__(self):
|
||||
return self._version
|
||||
|
||||
def __repr__(self):
|
||||
return f"<LegacyVersion({repr(str(self))})>"
|
||||
|
||||
@property
|
||||
def public(self):
|
||||
return self._version
|
||||
|
||||
@property
|
||||
def base_version(self):
|
||||
return self._version
|
||||
|
||||
@property
|
||||
def local(self):
|
||||
return None
|
||||
|
||||
@property
|
||||
def is_prerelease(self):
|
||||
return False
|
||||
|
||||
@property
|
||||
def is_postrelease(self):
|
||||
return False
|
||||
|
||||
|
||||
_legacy_version_component_re = re.compile(
|
||||
r"(\d+ | [a-z]+ | \.| -)", re.VERBOSE,
|
||||
)
|
||||
|
||||
_legacy_version_replacement_map = {
|
||||
"pre": "c", "preview": "c", "-": "final-", "rc": "c", "dev": "@",
|
||||
}
|
||||
|
||||
|
||||
def _parse_version_parts(s):
|
||||
for part in _legacy_version_component_re.split(s):
|
||||
part = _legacy_version_replacement_map.get(part, part)
|
||||
|
||||
if not part or part == ".":
|
||||
continue
|
||||
|
||||
if part[:1] in "0123456789":
|
||||
# pad for numeric comparison
|
||||
yield part.zfill(8)
|
||||
else:
|
||||
yield "*" + part
|
||||
|
||||
# ensure that alpha/beta/candidate are before final
|
||||
yield "*final"
|
||||
|
||||
|
||||
def _legacy_cmpkey(version):
|
||||
# We hardcode an epoch of -1 here. A PEP 440 version can only have an epoch
|
||||
# greater than or equal to 0. This will effectively put the LegacyVersion,
|
||||
# which uses the defacto standard originally implemented by setuptools,
|
||||
# as before all PEP 440 versions.
|
||||
epoch = -1
|
||||
|
||||
# This scheme is taken from pkg_resources.parse_version setuptools prior to
|
||||
# its adoption of the packaging library.
|
||||
parts = []
|
||||
for part in _parse_version_parts(version.lower()):
|
||||
if part.startswith("*"):
|
||||
# remove "-" before a prerelease tag
|
||||
if part < "*final":
|
||||
while parts and parts[-1] == "*final-":
|
||||
parts.pop()
|
||||
|
||||
# remove trailing zeros from each series of numeric parts
|
||||
while parts and parts[-1] == "00000000":
|
||||
parts.pop()
|
||||
|
||||
parts.append(part)
|
||||
parts = tuple(parts)
|
||||
|
||||
return epoch, parts
|
||||
|
||||
|
||||
# Deliberately not anchored to the start and end of the string, to make it
|
||||
# easier for 3rd party code to reuse
|
||||
VERSION_PATTERN = r"""
|
||||
v?
|
||||
(?:
|
||||
(?:(?P<epoch>[0-9]+)!)? # epoch
|
||||
(?P<release>[0-9]+(?:\.[0-9]+)*) # release segment
|
||||
(?P<pre> # pre-release
|
||||
[-_\.]?
|
||||
(?P<pre_l>(a|b|c|rc|alpha|beta|pre|preview))
|
||||
[-_\.]?
|
||||
(?P<pre_n>[0-9]+)?
|
||||
)?
|
||||
(?P<post> # post release
|
||||
(?:-(?P<post_n1>[0-9]+))
|
||||
|
|
||||
(?:
|
||||
[-_\.]?
|
||||
(?P<post_l>post|rev|r)
|
||||
[-_\.]?
|
||||
(?P<post_n2>[0-9]+)?
|
||||
)
|
||||
)?
|
||||
(?P<dev> # dev release
|
||||
[-_\.]?
|
||||
(?P<dev_l>dev)
|
||||
[-_\.]?
|
||||
(?P<dev_n>[0-9]+)?
|
||||
)?
|
||||
)
|
||||
(?:\+(?P<local>[a-z0-9]+(?:[-_\.][a-z0-9]+)*))? # local version
|
||||
"""
|
||||
|
||||
|
||||
class Version(_BaseVersion):
|
||||
|
||||
_regex = re.compile(
|
||||
r"^\s*" + VERSION_PATTERN + r"\s*$",
|
||||
re.VERBOSE | re.IGNORECASE,
|
||||
)
|
||||
|
||||
def __init__(self, version):
|
||||
# Validate the version and parse it into pieces
|
||||
match = self._regex.search(version)
|
||||
if not match:
|
||||
raise InvalidVersion(f"Invalid version: '{version}'")
|
||||
|
||||
# Store the parsed out pieces of the version
|
||||
self._version = _Version(
|
||||
epoch=int(match.group("epoch")) if match.group("epoch") else 0,
|
||||
release=tuple(int(i) for i in match.group("release").split(".")),
|
||||
pre=_parse_letter_version(
|
||||
match.group("pre_l"),
|
||||
match.group("pre_n"),
|
||||
),
|
||||
post=_parse_letter_version(
|
||||
match.group("post_l"),
|
||||
match.group("post_n1") or match.group("post_n2"),
|
||||
),
|
||||
dev=_parse_letter_version(
|
||||
match.group("dev_l"),
|
||||
match.group("dev_n"),
|
||||
),
|
||||
local=_parse_local_version(match.group("local")),
|
||||
)
|
||||
|
||||
# Generate a key which will be used for sorting
|
||||
self._key = _cmpkey(
|
||||
self._version.epoch,
|
||||
self._version.release,
|
||||
self._version.pre,
|
||||
self._version.post,
|
||||
self._version.dev,
|
||||
self._version.local,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return f"<Version({repr(str(self))})>"
|
||||
|
||||
def __str__(self):
|
||||
parts = []
|
||||
|
||||
# Epoch
|
||||
if self._version.epoch != 0:
|
||||
parts.append(f"{self._version.epoch}!")
|
||||
|
||||
# Release segment
|
||||
parts.append(".".join(str(x) for x in self._version.release))
|
||||
|
||||
# Pre-release
|
||||
if self._version.pre is not None:
|
||||
parts.append("".join(str(x) for x in self._version.pre))
|
||||
|
||||
# Post-release
|
||||
if self._version.post is not None:
|
||||
parts.append(f".post{self._version.post[1]}")
|
||||
|
||||
# Development release
|
||||
if self._version.dev is not None:
|
||||
parts.append(f".dev{self._version.dev[1]}")
|
||||
|
||||
# Local version segment
|
||||
if self._version.local is not None:
|
||||
parts.append(
|
||||
"+{}".format(".".join(str(x) for x in self._version.local))
|
||||
)
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
@property
|
||||
def public(self):
|
||||
return str(self).split("+", 1)[0]
|
||||
|
||||
@property
|
||||
def base_version(self):
|
||||
parts = []
|
||||
|
||||
# Epoch
|
||||
if self._version.epoch != 0:
|
||||
parts.append(f"{self._version.epoch}!")
|
||||
|
||||
# Release segment
|
||||
parts.append(".".join(str(x) for x in self._version.release))
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
@property
|
||||
def local(self):
|
||||
version_string = str(self)
|
||||
if "+" in version_string:
|
||||
return version_string.split("+", 1)[1]
|
||||
|
||||
@property
|
||||
def is_prerelease(self):
|
||||
return bool(self._version.dev or self._version.pre)
|
||||
|
||||
@property
|
||||
def is_postrelease(self):
|
||||
return bool(self._version.post)
|
||||
|
||||
|
||||
def _parse_letter_version(letter, number):
|
||||
if letter:
|
||||
# We assume there is an implicit 0 in a pre-release if there is
|
||||
# no numeral associated with it.
|
||||
if number is None:
|
||||
number = 0
|
||||
|
||||
# We normalize any letters to their lower-case form
|
||||
letter = letter.lower()
|
||||
|
||||
# We consider some words to be alternate spellings of other words and
|
||||
# in those cases we want to normalize the spellings to our preferred
|
||||
# spelling.
|
||||
if letter == "alpha":
|
||||
letter = "a"
|
||||
elif letter == "beta":
|
||||
letter = "b"
|
||||
elif letter in ["c", "pre", "preview"]:
|
||||
letter = "rc"
|
||||
elif letter in ["rev", "r"]:
|
||||
letter = "post"
|
||||
|
||||
return letter, int(number)
|
||||
if not letter and number:
|
||||
# We assume that if we are given a number but not given a letter,
|
||||
# then this is using the implicit post release syntax (e.g., 1.0-1)
|
||||
letter = "post"
|
||||
|
||||
return letter, int(number)
|
||||
|
||||
|
||||
_local_version_seperators = re.compile(r"[\._-]")
|
||||
|
||||
|
||||
def _parse_local_version(local):
|
||||
"""
|
||||
Takes a string like abc.1.twelve and turns it into ("abc", 1, "twelve").
|
||||
"""
|
||||
if local is not None:
|
||||
return tuple(
|
||||
part.lower() if not part.isdigit() else int(part)
|
||||
for part in _local_version_seperators.split(local)
|
||||
)
|
||||
|
||||
|
||||
def _cmpkey(epoch, release, pre, post, dev, local):
|
||||
# When we compare a release version, we want to compare it with all of the
|
||||
# trailing zeros removed. So we'll use a reverse the list, drop all the now
|
||||
# leading zeros until we come to something non-zero, then take the rest,
|
||||
# re-reverse it back into the correct order, and make it a tuple and use
|
||||
# that for our sorting key.
|
||||
release = tuple(
|
||||
reversed(list(
|
||||
itertools.dropwhile(
|
||||
lambda x: x == 0,
|
||||
reversed(release),
|
||||
)
|
||||
))
|
||||
)
|
||||
|
||||
# We need to "trick" the sorting algorithm to put 1.0.dev0 before 1.0a0.
|
||||
# We'll do this by abusing the pre-segment, but we _only_ want to do this
|
||||
# if there is no pre- or a post-segment. If we have one of those, then
|
||||
# the normal sorting rules will handle this case correctly.
|
||||
if pre is None and post is None and dev is not None:
|
||||
pre = -Infinity
|
||||
# Versions without a pre-release (except as noted above) should sort after
|
||||
# those with one.
|
||||
elif pre is None:
|
||||
pre = Infinity
|
||||
|
||||
# Versions without a post-segment should sort before those with one.
|
||||
if post is None:
|
||||
post = -Infinity
|
||||
|
||||
# Versions without a development segment should sort after those with one.
|
||||
if dev is None:
|
||||
dev = Infinity
|
||||
|
||||
if local is None:
|
||||
# Versions without a local segment should sort before those with one.
|
||||
local = -Infinity
|
||||
else:
|
||||
# Versions with a local segment need that segment parsed to implement
|
||||
# the sorting rules in PEP440.
|
||||
# - Alphanumeric segments sort before numeric segments
|
||||
# - Alphanumeric segments sort lexicographically
|
||||
# - Numeric segments sort numerically
|
||||
# - Shorter versions sort before longer versions when the prefixes
|
||||
# match exactly
|
||||
local = tuple(
|
||||
(i, "") if isinstance(i, int) else (-Infinity, i)
|
||||
for i in local
|
||||
)
|
||||
|
||||
return epoch, release, pre, post, dev, local
|
||||
@@ -0,0 +1,41 @@
|
||||
from abc import ABC
|
||||
|
||||
__all__ = ["SparseABC", "issparse"]
|
||||
|
||||
|
||||
class SparseABC(ABC):
|
||||
pass
|
||||
|
||||
|
||||
def issparse(x):
|
||||
"""Is `x` of a sparse array or sparse matrix type?
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x
|
||||
object to check for being a sparse array or sparse matrix
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if `x` is a sparse array or a sparse matrix, False otherwise
|
||||
|
||||
Notes
|
||||
-----
|
||||
Use `isinstance(x, sp.sparse.sparray)` to check between an array or matrix.
|
||||
Use `a.format` to check the sparse format, e.g. `a.format == 'csr'`.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import numpy as np
|
||||
>>> from scipy.sparse import csr_array, csr_matrix, issparse
|
||||
>>> issparse(csr_matrix([[5]]))
|
||||
True
|
||||
>>> issparse(csr_array([[5]]))
|
||||
True
|
||||
>>> issparse(np.array([[5]]))
|
||||
False
|
||||
>>> issparse(5)
|
||||
False
|
||||
"""
|
||||
return isinstance(x, SparseABC)
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,373 @@
|
||||
"""
|
||||
Generic test utilities.
|
||||
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import sysconfig
|
||||
import threading
|
||||
from importlib.util import module_from_spec, spec_from_file_location
|
||||
|
||||
import numpy as np
|
||||
import scipy
|
||||
|
||||
try:
|
||||
# Need type: ignore[import-untyped] for mypy >= 1.6
|
||||
import cython # type: ignore[import-untyped]
|
||||
from Cython.Compiler.Version import ( # type: ignore[import-untyped]
|
||||
version as cython_version,
|
||||
)
|
||||
except ImportError:
|
||||
cython = None
|
||||
else:
|
||||
from scipy._lib import _pep440
|
||||
required_version = '3.0.8'
|
||||
if _pep440.parse(cython_version) < _pep440.Version(required_version):
|
||||
# too old or wrong cython, skip Cython API tests
|
||||
cython = None
|
||||
|
||||
|
||||
__all__ = ['PytestTester', 'check_free_memory', '_TestPythranFunc', 'IS_MUSL']
|
||||
|
||||
|
||||
IS_MUSL = False
|
||||
# alternate way is
|
||||
# from packaging.tags import sys_tags
|
||||
# _tags = list(sys_tags())
|
||||
# if 'musllinux' in _tags[0].platform:
|
||||
_v = sysconfig.get_config_var('HOST_GNU_TYPE') or ''
|
||||
if 'musl' in _v:
|
||||
IS_MUSL = True
|
||||
|
||||
|
||||
IS_EDITABLE = 'editable' in scipy.__path__[0]
|
||||
|
||||
|
||||
class FPUModeChangeWarning(RuntimeWarning):
|
||||
"""Warning about FPU mode change"""
|
||||
pass
|
||||
|
||||
|
||||
class PytestTester:
|
||||
"""
|
||||
Run tests for this namespace
|
||||
|
||||
``scipy.test()`` runs tests for all of SciPy, with the default settings.
|
||||
When used from a submodule (e.g., ``scipy.cluster.test()``, only the tests
|
||||
for that namespace are run.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
label : {'fast', 'full'}, optional
|
||||
Whether to run only the fast tests, or also those marked as slow.
|
||||
Default is 'fast'.
|
||||
verbose : int, optional
|
||||
Test output verbosity. Default is 1.
|
||||
extra_argv : list, optional
|
||||
Arguments to pass through to Pytest.
|
||||
doctests : bool, optional
|
||||
Whether to run doctests or not. Default is False.
|
||||
coverage : bool, optional
|
||||
Whether to run tests with code coverage measurements enabled.
|
||||
Default is False.
|
||||
tests : list of str, optional
|
||||
List of module names to run tests for. By default, uses the module
|
||||
from which the ``test`` function is called.
|
||||
parallel : int, optional
|
||||
Run tests in parallel with pytest-xdist, if number given is larger than
|
||||
1. Default is 1.
|
||||
|
||||
"""
|
||||
def __init__(self, module_name):
|
||||
self.module_name = module_name
|
||||
|
||||
def __call__(self, label="fast", verbose=1, extra_argv=None, doctests=False,
|
||||
coverage=False, tests=None, parallel=None):
|
||||
import pytest
|
||||
|
||||
module = sys.modules[self.module_name]
|
||||
module_path = os.path.abspath(module.__path__[0])
|
||||
|
||||
pytest_args = ['--showlocals', '--tb=short']
|
||||
|
||||
if extra_argv is None:
|
||||
extra_argv = []
|
||||
pytest_args += extra_argv
|
||||
if any(arg == "-m" or arg == "--markers" for arg in extra_argv):
|
||||
# Likely conflict with default --mode=fast
|
||||
raise ValueError("Must specify -m before --")
|
||||
|
||||
if verbose and int(verbose) > 1:
|
||||
pytest_args += ["-" + "v"*(int(verbose)-1)]
|
||||
|
||||
if coverage:
|
||||
pytest_args += ["--cov=" + module_path]
|
||||
|
||||
if label == "fast":
|
||||
pytest_args += ["-m", "not slow"]
|
||||
elif label != "full":
|
||||
pytest_args += ["-m", label]
|
||||
|
||||
if tests is None:
|
||||
tests = [self.module_name]
|
||||
|
||||
if parallel is not None and parallel > 1:
|
||||
if _pytest_has_xdist():
|
||||
pytest_args += ['-n', str(parallel)]
|
||||
else:
|
||||
import warnings
|
||||
warnings.warn('Could not run tests in parallel because '
|
||||
'pytest-xdist plugin is not available.',
|
||||
stacklevel=2)
|
||||
|
||||
pytest_args += ['--pyargs'] + list(tests)
|
||||
|
||||
try:
|
||||
code = pytest.main(pytest_args)
|
||||
except SystemExit as exc:
|
||||
code = exc.code
|
||||
|
||||
return (code == 0)
|
||||
|
||||
|
||||
class _TestPythranFunc:
|
||||
'''
|
||||
These are situations that can be tested in our pythran tests:
|
||||
- A function with multiple array arguments and then
|
||||
other positional and keyword arguments.
|
||||
- A function with array-like keywords (e.g. `def somefunc(x0, x1=None)`.
|
||||
Note: list/tuple input is not yet tested!
|
||||
|
||||
`self.arguments`: A dictionary which key is the index of the argument,
|
||||
value is tuple(array value, all supported dtypes)
|
||||
`self.partialfunc`: A function used to freeze some non-array argument
|
||||
that of no interests in the original function
|
||||
'''
|
||||
ALL_INTEGER = [np.int8, np.int16, np.int32, np.int64, np.intc, np.intp]
|
||||
ALL_FLOAT = [np.float32, np.float64]
|
||||
ALL_COMPLEX = [np.complex64, np.complex128]
|
||||
|
||||
def setup_method(self):
|
||||
self.arguments = {}
|
||||
self.partialfunc = None
|
||||
self.expected = None
|
||||
|
||||
def get_optional_args(self, func):
|
||||
# get optional arguments with its default value,
|
||||
# used for testing keywords
|
||||
signature = inspect.signature(func)
|
||||
optional_args = {}
|
||||
for k, v in signature.parameters.items():
|
||||
if v.default is not inspect.Parameter.empty:
|
||||
optional_args[k] = v.default
|
||||
return optional_args
|
||||
|
||||
def get_max_dtype_list_length(self):
|
||||
# get the max supported dtypes list length in all arguments
|
||||
max_len = 0
|
||||
for arg_idx in self.arguments:
|
||||
cur_len = len(self.arguments[arg_idx][1])
|
||||
if cur_len > max_len:
|
||||
max_len = cur_len
|
||||
return max_len
|
||||
|
||||
def get_dtype(self, dtype_list, dtype_idx):
|
||||
# get the dtype from dtype_list via index
|
||||
# if the index is out of range, then return the last dtype
|
||||
if dtype_idx > len(dtype_list)-1:
|
||||
return dtype_list[-1]
|
||||
else:
|
||||
return dtype_list[dtype_idx]
|
||||
|
||||
def test_all_dtypes(self):
|
||||
for type_idx in range(self.get_max_dtype_list_length()):
|
||||
args_array = []
|
||||
for arg_idx in self.arguments:
|
||||
new_dtype = self.get_dtype(self.arguments[arg_idx][1],
|
||||
type_idx)
|
||||
args_array.append(self.arguments[arg_idx][0].astype(new_dtype))
|
||||
self.pythranfunc(*args_array)
|
||||
|
||||
def test_views(self):
|
||||
args_array = []
|
||||
for arg_idx in self.arguments:
|
||||
args_array.append(self.arguments[arg_idx][0][::-1][::-1])
|
||||
self.pythranfunc(*args_array)
|
||||
|
||||
def test_strided(self):
|
||||
args_array = []
|
||||
for arg_idx in self.arguments:
|
||||
args_array.append(np.repeat(self.arguments[arg_idx][0],
|
||||
2, axis=0)[::2])
|
||||
self.pythranfunc(*args_array)
|
||||
|
||||
|
||||
def _pytest_has_xdist():
|
||||
"""
|
||||
Check if the pytest-xdist plugin is installed, providing parallel tests
|
||||
"""
|
||||
# Check xdist exists without importing, otherwise pytests emits warnings
|
||||
from importlib.util import find_spec
|
||||
return find_spec('xdist') is not None
|
||||
|
||||
|
||||
def check_free_memory(free_mb):
|
||||
"""
|
||||
Check *free_mb* of memory is available, otherwise do pytest.skip
|
||||
"""
|
||||
import pytest
|
||||
|
||||
try:
|
||||
mem_free = _parse_size(os.environ['SCIPY_AVAILABLE_MEM'])
|
||||
msg = '{} MB memory required, but environment SCIPY_AVAILABLE_MEM={}'.format(
|
||||
free_mb, os.environ['SCIPY_AVAILABLE_MEM'])
|
||||
except KeyError:
|
||||
mem_free = _get_mem_available()
|
||||
if mem_free is None:
|
||||
pytest.skip("Could not determine available memory; set SCIPY_AVAILABLE_MEM "
|
||||
"variable to free memory in MB to run the test.")
|
||||
msg = f'{free_mb} MB memory required, but {mem_free/1e6} MB available'
|
||||
|
||||
if mem_free < free_mb * 1e6:
|
||||
pytest.skip(msg)
|
||||
|
||||
|
||||
def _parse_size(size_str):
|
||||
suffixes = {'': 1e6,
|
||||
'b': 1.0,
|
||||
'k': 1e3, 'M': 1e6, 'G': 1e9, 'T': 1e12,
|
||||
'kb': 1e3, 'Mb': 1e6, 'Gb': 1e9, 'Tb': 1e12,
|
||||
'kib': 1024.0, 'Mib': 1024.0**2, 'Gib': 1024.0**3, 'Tib': 1024.0**4}
|
||||
m = re.match(r'^\s*(\d+)\s*({})\s*$'.format('|'.join(suffixes.keys())),
|
||||
size_str,
|
||||
re.I)
|
||||
if not m or m.group(2) not in suffixes:
|
||||
raise ValueError("Invalid size string")
|
||||
|
||||
return float(m.group(1)) * suffixes[m.group(2)]
|
||||
|
||||
|
||||
def _get_mem_available():
|
||||
"""
|
||||
Get information about memory available, not counting swap.
|
||||
"""
|
||||
try:
|
||||
import psutil
|
||||
return psutil.virtual_memory().available
|
||||
except (ImportError, AttributeError):
|
||||
pass
|
||||
|
||||
if sys.platform.startswith('linux'):
|
||||
info = {}
|
||||
with open('/proc/meminfo') as f:
|
||||
for line in f:
|
||||
p = line.split()
|
||||
info[p[0].strip(':').lower()] = float(p[1]) * 1e3
|
||||
|
||||
if 'memavailable' in info:
|
||||
# Linux >= 3.14
|
||||
return info['memavailable']
|
||||
else:
|
||||
return info['memfree'] + info['cached']
|
||||
|
||||
return None
|
||||
|
||||
def _test_cython_extension(tmp_path, srcdir):
|
||||
"""
|
||||
Helper function to test building and importing Cython modules that
|
||||
make use of the Cython APIs for BLAS, LAPACK, optimize, and special.
|
||||
"""
|
||||
import pytest
|
||||
try:
|
||||
subprocess.check_call(["meson", "--version"])
|
||||
except FileNotFoundError:
|
||||
pytest.skip("No usable 'meson' found")
|
||||
|
||||
# Make safe for being called by multiple threads within one test
|
||||
tmp_path = tmp_path / str(threading.get_ident())
|
||||
|
||||
# build the examples in a temporary directory
|
||||
mod_name = os.path.split(srcdir)[1]
|
||||
shutil.copytree(srcdir, tmp_path / mod_name)
|
||||
build_dir = tmp_path / mod_name / 'tests' / '_cython_examples'
|
||||
target_dir = build_dir / 'build'
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
|
||||
# Ensure we use the correct Python interpreter even when `meson` is
|
||||
# installed in a different Python environment (see numpy#24956)
|
||||
native_file = str(build_dir / 'interpreter-native-file.ini')
|
||||
with open(native_file, 'w') as f:
|
||||
f.write("[binaries]\n")
|
||||
f.write(f"python = '{sys.executable}'")
|
||||
|
||||
if sys.platform == "win32":
|
||||
subprocess.check_call(["meson", "setup",
|
||||
"--buildtype=release",
|
||||
"--native-file", native_file,
|
||||
"--vsenv", str(build_dir)],
|
||||
cwd=target_dir,
|
||||
)
|
||||
else:
|
||||
subprocess.check_call(["meson", "setup",
|
||||
"--native-file", native_file, str(build_dir)],
|
||||
cwd=target_dir
|
||||
)
|
||||
subprocess.check_call(["meson", "compile", "-vv"], cwd=target_dir)
|
||||
|
||||
# import without adding the directory to sys.path
|
||||
suffix = sysconfig.get_config_var('EXT_SUFFIX')
|
||||
|
||||
def load(modname):
|
||||
so = (target_dir / modname).with_suffix(suffix)
|
||||
spec = spec_from_file_location(modname, so)
|
||||
mod = module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
# test that the module can be imported
|
||||
return load("extending"), load("extending_cpp")
|
||||
|
||||
|
||||
def _run_concurrent_barrier(n_workers, fn, *args, **kwargs):
|
||||
"""
|
||||
Run a given function concurrently across a given number of threads.
|
||||
|
||||
This is equivalent to using a ThreadPoolExecutor, but using the threading
|
||||
primitives instead. This function ensures that the closure passed by
|
||||
parameter gets called concurrently by setting up a barrier before it gets
|
||||
called before any of the threads.
|
||||
|
||||
Arguments
|
||||
---------
|
||||
n_workers: int
|
||||
Number of concurrent threads to spawn.
|
||||
fn: callable
|
||||
Function closure to execute concurrently. Its first argument will
|
||||
be the thread id.
|
||||
*args: tuple
|
||||
Variable number of positional arguments to pass to the function.
|
||||
**kwargs: dict
|
||||
Keyword arguments to pass to the function.
|
||||
"""
|
||||
barrier = threading.Barrier(n_workers)
|
||||
|
||||
def closure(i, *args, **kwargs):
|
||||
barrier.wait()
|
||||
fn(i, *args, **kwargs)
|
||||
|
||||
workers = []
|
||||
for i in range(0, n_workers):
|
||||
workers.append(threading.Thread(
|
||||
target=closure,
|
||||
args=(i,) + args, kwargs=kwargs))
|
||||
|
||||
for worker in workers:
|
||||
worker.start()
|
||||
|
||||
for worker in workers:
|
||||
worker.join()
|
||||
@@ -0,0 +1,58 @@
|
||||
import threading
|
||||
|
||||
import scipy._lib.decorator
|
||||
|
||||
|
||||
__all__ = ['ReentrancyError', 'ReentrancyLock', 'non_reentrant']
|
||||
|
||||
|
||||
class ReentrancyError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class ReentrancyLock:
|
||||
"""
|
||||
Threading lock that raises an exception for reentrant calls.
|
||||
|
||||
Calls from different threads are serialized, and nested calls from the
|
||||
same thread result to an error.
|
||||
|
||||
The object can be used as a context manager or to decorate functions
|
||||
via the decorate() method.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, err_msg):
|
||||
self._rlock = threading.RLock()
|
||||
self._entered = False
|
||||
self._err_msg = err_msg
|
||||
|
||||
def __enter__(self):
|
||||
self._rlock.acquire()
|
||||
if self._entered:
|
||||
self._rlock.release()
|
||||
raise ReentrancyError(self._err_msg)
|
||||
self._entered = True
|
||||
|
||||
def __exit__(self, type, value, traceback):
|
||||
self._entered = False
|
||||
self._rlock.release()
|
||||
|
||||
def decorate(self, func):
|
||||
def caller(func, *a, **kw):
|
||||
with self:
|
||||
return func(*a, **kw)
|
||||
return scipy._lib.decorator.decorate(func, caller)
|
||||
|
||||
|
||||
def non_reentrant(err_msg=None):
|
||||
"""
|
||||
Decorate a function with a threading lock and prevent reentrant calls.
|
||||
"""
|
||||
def decorator(func):
|
||||
msg = err_msg
|
||||
if msg is None:
|
||||
msg = f"{func.__name__} is not re-entrant"
|
||||
lock = ReentrancyLock(msg)
|
||||
return lock.decorate(func)
|
||||
return decorator
|
||||
@@ -0,0 +1,86 @@
|
||||
''' Contexts for *with* statement providing temporary directories
|
||||
'''
|
||||
import os
|
||||
from contextlib import contextmanager
|
||||
from shutil import rmtree
|
||||
from tempfile import mkdtemp
|
||||
|
||||
|
||||
@contextmanager
|
||||
def tempdir():
|
||||
"""Create and return a temporary directory. This has the same
|
||||
behavior as mkdtemp but can be used as a context manager.
|
||||
|
||||
Upon exiting the context, the directory and everything contained
|
||||
in it are removed.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import os
|
||||
>>> with tempdir() as tmpdir:
|
||||
... fname = os.path.join(tmpdir, 'example_file.txt')
|
||||
... with open(fname, 'wt') as fobj:
|
||||
... _ = fobj.write('a string\\n')
|
||||
>>> os.path.exists(tmpdir)
|
||||
False
|
||||
"""
|
||||
d = mkdtemp()
|
||||
yield d
|
||||
rmtree(d)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def in_tempdir():
|
||||
''' Create, return, and change directory to a temporary directory
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import os
|
||||
>>> my_cwd = os.getcwd()
|
||||
>>> with in_tempdir() as tmpdir:
|
||||
... _ = open('test.txt', 'wt').write('some text')
|
||||
... assert os.path.isfile('test.txt')
|
||||
... assert os.path.isfile(os.path.join(tmpdir, 'test.txt'))
|
||||
>>> os.path.exists(tmpdir)
|
||||
False
|
||||
>>> os.getcwd() == my_cwd
|
||||
True
|
||||
'''
|
||||
pwd = os.getcwd()
|
||||
d = mkdtemp()
|
||||
os.chdir(d)
|
||||
yield d
|
||||
os.chdir(pwd)
|
||||
rmtree(d)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def in_dir(dir=None):
|
||||
""" Change directory to given directory for duration of ``with`` block
|
||||
|
||||
Useful when you want to use `in_tempdir` for the final test, but
|
||||
you are still debugging. For example, you may want to do this in the end:
|
||||
|
||||
>>> with in_tempdir() as tmpdir:
|
||||
... # do something complicated which might break
|
||||
... pass
|
||||
|
||||
But, indeed, the complicated thing does break, and meanwhile, the
|
||||
``in_tempdir`` context manager wiped out the directory with the
|
||||
temporary files that you wanted for debugging. So, while debugging, you
|
||||
replace with something like:
|
||||
|
||||
>>> with in_dir() as tmpdir: # Use working directory by default
|
||||
... # do something complicated which might break
|
||||
... pass
|
||||
|
||||
You can then look at the temporary file outputs to debug what is happening,
|
||||
fix, and finally replace ``in_dir`` with ``in_tempdir`` again.
|
||||
"""
|
||||
cwd = os.getcwd()
|
||||
if dir is None:
|
||||
yield cwd
|
||||
return
|
||||
os.chdir(dir)
|
||||
yield dir
|
||||
os.chdir(cwd)
|
||||
@@ -0,0 +1,29 @@
|
||||
BSD 3-Clause License
|
||||
|
||||
Copyright (c) 2018, Quansight-Labs
|
||||
All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are met:
|
||||
|
||||
* Redistributions of source code must retain the above copyright notice, this
|
||||
list of conditions and the following disclaimer.
|
||||
|
||||
* Redistributions in binary form must reproduce the above copyright notice,
|
||||
this list of conditions and the following disclaimer in the documentation
|
||||
and/or other materials provided with the distribution.
|
||||
|
||||
* Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived from
|
||||
this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
@@ -0,0 +1,116 @@
|
||||
"""
|
||||
.. note:
|
||||
If you are looking for overrides for NumPy-specific methods, see the
|
||||
documentation for :obj:`unumpy`. This page explains how to write
|
||||
back-ends and multimethods.
|
||||
|
||||
``uarray`` is built around a back-end protocol, and overridable multimethods.
|
||||
It is necessary to define multimethods for back-ends to be able to override them.
|
||||
See the documentation of :obj:`generate_multimethod` on how to write multimethods.
|
||||
|
||||
|
||||
|
||||
Let's start with the simplest:
|
||||
|
||||
``__ua_domain__`` defines the back-end *domain*. The domain consists of period-
|
||||
separated string consisting of the modules you extend plus the submodule. For
|
||||
example, if a submodule ``module2.submodule`` extends ``module1``
|
||||
(i.e., it exposes dispatchables marked as types available in ``module1``),
|
||||
then the domain string should be ``"module1.module2.submodule"``.
|
||||
|
||||
|
||||
For the purpose of this demonstration, we'll be creating an object and setting
|
||||
its attributes directly. However, note that you can use a module or your own type
|
||||
as a backend as well.
|
||||
|
||||
>>> class Backend: pass
|
||||
>>> be = Backend()
|
||||
>>> be.__ua_domain__ = "ua_examples"
|
||||
|
||||
It might be useful at this point to sidetrack to the documentation of
|
||||
:obj:`generate_multimethod` to find out how to generate a multimethod
|
||||
overridable by :obj:`uarray`. Needless to say, writing a backend and
|
||||
creating multimethods are mostly orthogonal activities, and knowing
|
||||
one doesn't necessarily require knowledge of the other, although it
|
||||
is certainly helpful. We expect core API designers/specifiers to write the
|
||||
multimethods, and implementors to override them. But, as is often the case,
|
||||
similar people write both.
|
||||
|
||||
Without further ado, here's an example multimethod:
|
||||
|
||||
>>> import uarray as ua
|
||||
>>> from uarray import Dispatchable
|
||||
>>> def override_me(a, b):
|
||||
... return Dispatchable(a, int),
|
||||
>>> def override_replacer(args, kwargs, dispatchables):
|
||||
... return (dispatchables[0], args[1]), {}
|
||||
>>> overridden_me = ua.generate_multimethod(
|
||||
... override_me, override_replacer, "ua_examples"
|
||||
... )
|
||||
|
||||
Next comes the part about overriding the multimethod. This requires
|
||||
the ``__ua_function__`` protocol, and the ``__ua_convert__``
|
||||
protocol. The ``__ua_function__`` protocol has the signature
|
||||
``(method, args, kwargs)`` where ``method`` is the passed
|
||||
multimethod, ``args``/``kwargs`` specify the arguments and ``dispatchables``
|
||||
is the list of converted dispatchables passed in.
|
||||
|
||||
>>> def __ua_function__(method, args, kwargs):
|
||||
... return method.__name__, args, kwargs
|
||||
>>> be.__ua_function__ = __ua_function__
|
||||
|
||||
The other protocol of interest is the ``__ua_convert__`` protocol. It has the
|
||||
signature ``(dispatchables, coerce)``. When ``coerce`` is ``False``, conversion
|
||||
between the formats should ideally be an ``O(1)`` operation, but it means that
|
||||
no memory copying should be involved, only views of the existing data.
|
||||
|
||||
>>> def __ua_convert__(dispatchables, coerce):
|
||||
... for d in dispatchables:
|
||||
... if d.type is int:
|
||||
... if coerce and d.coercible:
|
||||
... yield str(d.value)
|
||||
... else:
|
||||
... yield d.value
|
||||
>>> be.__ua_convert__ = __ua_convert__
|
||||
|
||||
Now that we have defined the backend, the next thing to do is to call the multimethod.
|
||||
|
||||
>>> with ua.set_backend(be):
|
||||
... overridden_me(1, "2")
|
||||
('override_me', (1, '2'), {})
|
||||
|
||||
Note that the marked type has no effect on the actual type of the passed object.
|
||||
We can also coerce the type of the input.
|
||||
|
||||
>>> with ua.set_backend(be, coerce=True):
|
||||
... overridden_me(1, "2")
|
||||
... overridden_me(1.0, "2")
|
||||
('override_me', ('1', '2'), {})
|
||||
('override_me', ('1.0', '2'), {})
|
||||
|
||||
Another feature is that if you remove ``__ua_convert__``, the arguments are not
|
||||
converted at all and it's up to the backend to handle that.
|
||||
|
||||
>>> del be.__ua_convert__
|
||||
>>> with ua.set_backend(be):
|
||||
... overridden_me(1, "2")
|
||||
('override_me', (1, '2'), {})
|
||||
|
||||
You also have the option to return ``NotImplemented``, in which case processing moves on
|
||||
to the next back-end, which in this case, doesn't exist. The same applies to
|
||||
``__ua_convert__``.
|
||||
|
||||
>>> be.__ua_function__ = lambda *a, **kw: NotImplemented
|
||||
>>> with ua.set_backend(be):
|
||||
... overridden_me(1, "2")
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
uarray.BackendNotImplementedError: ...
|
||||
|
||||
The last possibility is if we don't have ``__ua_convert__``, in which case the job is
|
||||
left up to ``__ua_function__``, but putting things back into arrays after conversion
|
||||
will not be possible.
|
||||
"""
|
||||
|
||||
from ._backend import *
|
||||
__version__ = '0.8.8.dev0+aa94c5a4.scipy'
|
||||
@@ -0,0 +1,707 @@
|
||||
import typing
|
||||
import types
|
||||
import inspect
|
||||
import functools
|
||||
from . import _uarray
|
||||
import copyreg
|
||||
import pickle
|
||||
import contextlib
|
||||
import threading
|
||||
|
||||
from ._uarray import ( # type: ignore
|
||||
BackendNotImplementedError,
|
||||
_Function,
|
||||
_SkipBackendContext,
|
||||
_SetBackendContext,
|
||||
_BackendState,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"set_backend",
|
||||
"set_global_backend",
|
||||
"skip_backend",
|
||||
"register_backend",
|
||||
"determine_backend",
|
||||
"determine_backend_multi",
|
||||
"clear_backends",
|
||||
"create_multimethod",
|
||||
"generate_multimethod",
|
||||
"_Function",
|
||||
"BackendNotImplementedError",
|
||||
"Dispatchable",
|
||||
"wrap_single_convertor",
|
||||
"wrap_single_convertor_instance",
|
||||
"all_of_type",
|
||||
"mark_as",
|
||||
"set_state",
|
||||
"get_state",
|
||||
"reset_state",
|
||||
"_BackendState",
|
||||
"_SkipBackendContext",
|
||||
"_SetBackendContext",
|
||||
]
|
||||
|
||||
ArgumentExtractorType = typing.Callable[..., tuple["Dispatchable", ...]]
|
||||
ArgumentReplacerType = typing.Callable[
|
||||
[tuple, dict, tuple], tuple[tuple, dict]
|
||||
]
|
||||
|
||||
def unpickle_function(mod_name, qname, self_):
|
||||
import importlib
|
||||
|
||||
try:
|
||||
module = importlib.import_module(mod_name)
|
||||
qname = qname.split(".")
|
||||
func = module
|
||||
for q in qname:
|
||||
func = getattr(func, q)
|
||||
|
||||
if self_ is not None:
|
||||
func = types.MethodType(func, self_)
|
||||
|
||||
return func
|
||||
except (ImportError, AttributeError) as e:
|
||||
from pickle import UnpicklingError
|
||||
|
||||
raise UnpicklingError from e
|
||||
|
||||
|
||||
def pickle_function(func):
|
||||
mod_name = getattr(func, "__module__", None)
|
||||
qname = getattr(func, "__qualname__", None)
|
||||
self_ = getattr(func, "__self__", None)
|
||||
|
||||
try:
|
||||
test = unpickle_function(mod_name, qname, self_)
|
||||
except pickle.UnpicklingError:
|
||||
test = None
|
||||
|
||||
if test is not func:
|
||||
raise pickle.PicklingError(
|
||||
f"Can't pickle {func}: it's not the same object as {test}"
|
||||
)
|
||||
|
||||
return unpickle_function, (mod_name, qname, self_)
|
||||
|
||||
|
||||
def pickle_state(state):
|
||||
return _uarray._BackendState._unpickle, state._pickle()
|
||||
|
||||
|
||||
def pickle_set_backend_context(ctx):
|
||||
return _SetBackendContext, ctx._pickle()
|
||||
|
||||
|
||||
def pickle_skip_backend_context(ctx):
|
||||
return _SkipBackendContext, ctx._pickle()
|
||||
|
||||
|
||||
copyreg.pickle(_Function, pickle_function)
|
||||
copyreg.pickle(_uarray._BackendState, pickle_state)
|
||||
copyreg.pickle(_SetBackendContext, pickle_set_backend_context)
|
||||
copyreg.pickle(_SkipBackendContext, pickle_skip_backend_context)
|
||||
|
||||
|
||||
def get_state():
|
||||
"""
|
||||
Returns an opaque object containing the current state of all the backends.
|
||||
|
||||
Can be used for synchronization between threads/processes.
|
||||
|
||||
See Also
|
||||
--------
|
||||
set_state
|
||||
Sets the state returned by this function.
|
||||
"""
|
||||
return _uarray.get_state()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def reset_state():
|
||||
"""
|
||||
Returns a context manager that resets all state once exited.
|
||||
|
||||
See Also
|
||||
--------
|
||||
set_state
|
||||
Context manager that sets the backend state.
|
||||
get_state
|
||||
Gets a state to be set by this context manager.
|
||||
"""
|
||||
with set_state(get_state()):
|
||||
yield
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def set_state(state):
|
||||
"""
|
||||
A context manager that sets the state of the backends to one returned by :obj:`get_state`.
|
||||
|
||||
See Also
|
||||
--------
|
||||
get_state
|
||||
Gets a state to be set by this context manager.
|
||||
""" # noqa: E501
|
||||
old_state = get_state()
|
||||
_uarray.set_state(state)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_uarray.set_state(old_state, True)
|
||||
|
||||
|
||||
def create_multimethod(*args, **kwargs):
|
||||
"""
|
||||
Creates a decorator for generating multimethods.
|
||||
|
||||
This function creates a decorator that can be used with an argument
|
||||
extractor in order to generate a multimethod. Other than for the
|
||||
argument extractor, all arguments are passed on to
|
||||
:obj:`generate_multimethod`.
|
||||
|
||||
See Also
|
||||
--------
|
||||
generate_multimethod
|
||||
Generates a multimethod.
|
||||
"""
|
||||
|
||||
def wrapper(a):
|
||||
return generate_multimethod(a, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def generate_multimethod(
|
||||
argument_extractor: ArgumentExtractorType,
|
||||
argument_replacer: ArgumentReplacerType,
|
||||
domain: str,
|
||||
default: typing.Callable | None = None,
|
||||
):
|
||||
"""
|
||||
Generates a multimethod.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
argument_extractor : ArgumentExtractorType
|
||||
A callable which extracts the dispatchable arguments. Extracted arguments
|
||||
should be marked by the :obj:`Dispatchable` class. It has the same signature
|
||||
as the desired multimethod.
|
||||
argument_replacer : ArgumentReplacerType
|
||||
A callable with the signature (args, kwargs, dispatchables), which should also
|
||||
return an (args, kwargs) pair with the dispatchables replaced inside the
|
||||
args/kwargs.
|
||||
domain : str
|
||||
A string value indicating the domain of this multimethod.
|
||||
default: Optional[Callable], optional
|
||||
The default implementation of this multimethod, where ``None`` (the default)
|
||||
specifies there is no default implementation.
|
||||
|
||||
Examples
|
||||
--------
|
||||
In this example, ``a`` is to be dispatched over, so we return it, while marking it
|
||||
as an ``int``.
|
||||
The trailing comma is needed because the args have to be returned as an iterable.
|
||||
|
||||
>>> def override_me(a, b):
|
||||
... return Dispatchable(a, int),
|
||||
|
||||
Next, we define the argument replacer that replaces the dispatchables inside
|
||||
args/kwargs with the supplied ones.
|
||||
|
||||
>>> def override_replacer(args, kwargs, dispatchables):
|
||||
... return (dispatchables[0], args[1]), {}
|
||||
|
||||
Next, we define the multimethod.
|
||||
|
||||
>>> overridden_me = generate_multimethod(
|
||||
... override_me, override_replacer, "ua_examples"
|
||||
... )
|
||||
|
||||
Notice that there's no default implementation, unless you supply one.
|
||||
|
||||
>>> overridden_me(1, "a")
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
uarray.BackendNotImplementedError: ...
|
||||
|
||||
>>> overridden_me2 = generate_multimethod(
|
||||
... override_me, override_replacer, "ua_examples", default=lambda x, y: (x, y)
|
||||
... )
|
||||
>>> overridden_me2(1, "a")
|
||||
(1, 'a')
|
||||
|
||||
See Also
|
||||
--------
|
||||
uarray
|
||||
See the module documentation for how to override the method by creating
|
||||
backends.
|
||||
"""
|
||||
kw_defaults, arg_defaults, opts = get_defaults(argument_extractor)
|
||||
ua_func = _Function(
|
||||
argument_extractor,
|
||||
argument_replacer,
|
||||
domain,
|
||||
arg_defaults,
|
||||
kw_defaults,
|
||||
default,
|
||||
)
|
||||
|
||||
return functools.update_wrapper(ua_func, argument_extractor)
|
||||
|
||||
|
||||
def set_backend(backend, coerce=False, only=False):
|
||||
"""
|
||||
A context manager that sets the preferred backend.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
backend
|
||||
The backend to set.
|
||||
coerce
|
||||
Whether or not to coerce to a specific backend's types. Implies ``only``.
|
||||
only
|
||||
Whether or not this should be the last backend to try.
|
||||
|
||||
See Also
|
||||
--------
|
||||
skip_backend: A context manager that allows skipping of backends.
|
||||
set_global_backend: Set a single, global backend for a domain.
|
||||
"""
|
||||
tid = threading.get_native_id()
|
||||
try:
|
||||
return backend.__ua_cache__[tid, "set", coerce, only]
|
||||
except AttributeError:
|
||||
backend.__ua_cache__ = {}
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
ctx = _SetBackendContext(backend, coerce, only)
|
||||
backend.__ua_cache__[tid, "set", coerce, only] = ctx
|
||||
return ctx
|
||||
|
||||
|
||||
def skip_backend(backend):
|
||||
"""
|
||||
A context manager that allows one to skip a given backend from processing
|
||||
entirely. This allows one to use another backend's code in a library that
|
||||
is also a consumer of the same backend.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
backend
|
||||
The backend to skip.
|
||||
|
||||
See Also
|
||||
--------
|
||||
set_backend: A context manager that allows setting of backends.
|
||||
set_global_backend: Set a single, global backend for a domain.
|
||||
"""
|
||||
tid = threading.get_native_id()
|
||||
try:
|
||||
return backend.__ua_cache__[tid, "skip"]
|
||||
except AttributeError:
|
||||
backend.__ua_cache__ = {}
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
ctx = _SkipBackendContext(backend)
|
||||
backend.__ua_cache__[tid, "skip"] = ctx
|
||||
return ctx
|
||||
|
||||
|
||||
def get_defaults(f):
|
||||
sig = inspect.signature(f)
|
||||
kw_defaults = {}
|
||||
arg_defaults = []
|
||||
opts = set()
|
||||
for k, v in sig.parameters.items():
|
||||
if v.default is not inspect.Parameter.empty:
|
||||
kw_defaults[k] = v.default
|
||||
if v.kind in (
|
||||
inspect.Parameter.POSITIONAL_ONLY,
|
||||
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
||||
):
|
||||
arg_defaults.append(v.default)
|
||||
opts.add(k)
|
||||
|
||||
return kw_defaults, tuple(arg_defaults), opts
|
||||
|
||||
|
||||
def set_global_backend(backend, coerce=False, only=False, *, try_last=False):
|
||||
"""
|
||||
This utility method replaces the default backend for permanent use. It
|
||||
will be tried in the list of backends automatically, unless the
|
||||
``only`` flag is set on a backend. This will be the first tried
|
||||
backend outside the :obj:`set_backend` context manager.
|
||||
|
||||
Note that this method is not thread-safe.
|
||||
|
||||
.. warning::
|
||||
We caution library authors against using this function in
|
||||
their code. We do *not* support this use-case. This function
|
||||
is meant to be used only by users themselves, or by a reference
|
||||
implementation, if one exists.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
backend
|
||||
The backend to register.
|
||||
coerce : bool
|
||||
Whether to coerce input types when trying this backend.
|
||||
only : bool
|
||||
If ``True``, no more backends will be tried if this fails.
|
||||
Implied by ``coerce=True``.
|
||||
try_last : bool
|
||||
If ``True``, the global backend is tried after registered backends.
|
||||
|
||||
See Also
|
||||
--------
|
||||
set_backend: A context manager that allows setting of backends.
|
||||
skip_backend: A context manager that allows skipping of backends.
|
||||
"""
|
||||
_uarray.set_global_backend(backend, coerce, only, try_last)
|
||||
|
||||
|
||||
def register_backend(backend):
|
||||
"""
|
||||
This utility method sets registers backend for permanent use. It
|
||||
will be tried in the list of backends automatically, unless the
|
||||
``only`` flag is set on a backend.
|
||||
|
||||
Note that this method is not thread-safe.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
backend
|
||||
The backend to register.
|
||||
"""
|
||||
_uarray.register_backend(backend)
|
||||
|
||||
|
||||
def clear_backends(domain, registered=True, globals=False):
|
||||
"""
|
||||
This utility method clears registered backends.
|
||||
|
||||
.. warning::
|
||||
We caution library authors against using this function in
|
||||
their code. We do *not* support this use-case. This function
|
||||
is meant to be used only by users themselves.
|
||||
|
||||
.. warning::
|
||||
Do NOT use this method inside a multimethod call, or the
|
||||
program is likely to crash.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
domain : Optional[str]
|
||||
The domain for which to de-register backends. ``None`` means
|
||||
de-register for all domains.
|
||||
registered : bool
|
||||
Whether or not to clear registered backends. See :obj:`register_backend`.
|
||||
globals : bool
|
||||
Whether or not to clear global backends. See :obj:`set_global_backend`.
|
||||
|
||||
See Also
|
||||
--------
|
||||
register_backend : Register a backend globally.
|
||||
set_global_backend : Set a global backend.
|
||||
"""
|
||||
_uarray.clear_backends(domain, registered, globals)
|
||||
|
||||
|
||||
class Dispatchable:
|
||||
"""
|
||||
A utility class which marks an argument with a specific dispatch type.
|
||||
|
||||
|
||||
Attributes
|
||||
----------
|
||||
value
|
||||
The value of the Dispatchable.
|
||||
|
||||
type
|
||||
The type of the Dispatchable.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> x = Dispatchable(1, str)
|
||||
>>> x
|
||||
<Dispatchable: type=<class 'str'>, value=1>
|
||||
|
||||
See Also
|
||||
--------
|
||||
all_of_type
|
||||
Marks all unmarked parameters of a function.
|
||||
|
||||
mark_as
|
||||
Allows one to create a utility function to mark as a given type.
|
||||
"""
|
||||
|
||||
def __init__(self, value, dispatch_type, coercible=True):
|
||||
self.value = value
|
||||
self.type = dispatch_type
|
||||
self.coercible = coercible
|
||||
|
||||
def __getitem__(self, index):
|
||||
return (self.type, self.value)[index]
|
||||
|
||||
def __str__(self):
|
||||
return f"<{type(self).__name__}: type={self.type!r}, value={self.value!r}>"
|
||||
|
||||
__repr__ = __str__
|
||||
|
||||
|
||||
def mark_as(dispatch_type):
|
||||
"""
|
||||
Creates a utility function to mark something as a specific type.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> mark_int = mark_as(int)
|
||||
>>> mark_int(1)
|
||||
<Dispatchable: type=<class 'int'>, value=1>
|
||||
"""
|
||||
return functools.partial(Dispatchable, dispatch_type=dispatch_type)
|
||||
|
||||
|
||||
def all_of_type(arg_type):
|
||||
"""
|
||||
Marks all unmarked arguments as a given type.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> @all_of_type(str)
|
||||
... def f(a, b):
|
||||
... return a, Dispatchable(b, int)
|
||||
>>> f('a', 1)
|
||||
(<Dispatchable: type=<class 'str'>, value='a'>,
|
||||
<Dispatchable: type=<class 'int'>, value=1>)
|
||||
"""
|
||||
|
||||
def outer(func):
|
||||
@functools.wraps(func)
|
||||
def inner(*args, **kwargs):
|
||||
extracted_args = func(*args, **kwargs)
|
||||
return tuple(
|
||||
Dispatchable(arg, arg_type)
|
||||
if not isinstance(arg, Dispatchable)
|
||||
else arg
|
||||
for arg in extracted_args
|
||||
)
|
||||
|
||||
return inner
|
||||
|
||||
return outer
|
||||
|
||||
|
||||
def wrap_single_convertor(convert_single):
|
||||
"""
|
||||
Wraps a ``__ua_convert__`` defined for a single element to all elements.
|
||||
If any of them return ``NotImplemented``, the operation is assumed to be
|
||||
undefined.
|
||||
|
||||
Accepts a signature of (value, type, coerce).
|
||||
"""
|
||||
|
||||
@functools.wraps(convert_single)
|
||||
def __ua_convert__(dispatchables, coerce):
|
||||
converted = []
|
||||
for d in dispatchables:
|
||||
c = convert_single(d.value, d.type, coerce and d.coercible)
|
||||
|
||||
if c is NotImplemented:
|
||||
return NotImplemented
|
||||
|
||||
converted.append(c)
|
||||
|
||||
return converted
|
||||
|
||||
return __ua_convert__
|
||||
|
||||
|
||||
def wrap_single_convertor_instance(convert_single):
|
||||
"""
|
||||
Wraps a ``__ua_convert__`` defined for a single element to all elements.
|
||||
If any of them return ``NotImplemented``, the operation is assumed to be
|
||||
undefined.
|
||||
|
||||
Accepts a signature of (value, type, coerce).
|
||||
"""
|
||||
|
||||
@functools.wraps(convert_single)
|
||||
def __ua_convert__(self, dispatchables, coerce):
|
||||
converted = []
|
||||
for d in dispatchables:
|
||||
c = convert_single(self, d.value, d.type, coerce and d.coercible)
|
||||
|
||||
if c is NotImplemented:
|
||||
return NotImplemented
|
||||
|
||||
converted.append(c)
|
||||
|
||||
return converted
|
||||
|
||||
return __ua_convert__
|
||||
|
||||
|
||||
def determine_backend(value, dispatch_type, *, domain, only=True, coerce=False):
|
||||
"""Set the backend to the first active backend that supports ``value``
|
||||
|
||||
This is useful for functions that call multimethods without any dispatchable
|
||||
arguments. You can use :func:`determine_backend` to ensure the same backend
|
||||
is used everywhere in a block of multimethod calls.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value
|
||||
The value being tested
|
||||
dispatch_type
|
||||
The dispatch type associated with ``value``, aka
|
||||
":ref:`marking <MarkingGlossary>`".
|
||||
domain: string
|
||||
The domain to query for backends and set.
|
||||
coerce: bool
|
||||
Whether or not to allow coercion to the backend's types. Implies ``only``.
|
||||
only: bool
|
||||
Whether or not this should be the last backend to try.
|
||||
|
||||
See Also
|
||||
--------
|
||||
set_backend: For when you know which backend to set
|
||||
|
||||
Notes
|
||||
-----
|
||||
|
||||
Support is determined by the ``__ua_convert__`` protocol. Backends not
|
||||
supporting the type must return ``NotImplemented`` from their
|
||||
``__ua_convert__`` if they don't support input of that type.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
Suppose we have two backends ``BackendA`` and ``BackendB`` each supporting
|
||||
different types, ``TypeA`` and ``TypeB``. Neither supporting the other type:
|
||||
|
||||
>>> with ua.set_backend(ex.BackendA):
|
||||
... ex.call_multimethod(ex.TypeB(), ex.TypeB())
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
uarray.BackendNotImplementedError: ...
|
||||
|
||||
Now consider a multimethod that creates a new object of ``TypeA``, or
|
||||
``TypeB`` depending on the active backend.
|
||||
|
||||
>>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB):
|
||||
... res = ex.creation_multimethod()
|
||||
... ex.call_multimethod(res, ex.TypeA())
|
||||
Traceback (most recent call last):
|
||||
...
|
||||
uarray.BackendNotImplementedError: ...
|
||||
|
||||
``res`` is an object of ``TypeB`` because ``BackendB`` is set in the
|
||||
innermost with statement. So, ``call_multimethod`` fails since the types
|
||||
don't match.
|
||||
|
||||
Instead, we need to first find a backend suitable for all of our objects.
|
||||
|
||||
>>> with ua.set_backend(ex.BackendA), ua.set_backend(ex.BackendB):
|
||||
... x = ex.TypeA()
|
||||
... with ua.determine_backend(x, "mark", domain="ua_examples"):
|
||||
... res = ex.creation_multimethod()
|
||||
... ex.call_multimethod(res, x)
|
||||
TypeA
|
||||
|
||||
"""
|
||||
dispatchables = (Dispatchable(value, dispatch_type, coerce),)
|
||||
backend = _uarray.determine_backend(domain, dispatchables, coerce)
|
||||
|
||||
return set_backend(backend, coerce=coerce, only=only)
|
||||
|
||||
|
||||
def determine_backend_multi(
|
||||
dispatchables, *, domain, only=True, coerce=False, **kwargs
|
||||
):
|
||||
"""Set a backend supporting all ``dispatchables``
|
||||
|
||||
This is useful for functions that call multimethods without any dispatchable
|
||||
arguments. You can use :func:`determine_backend_multi` to ensure the same
|
||||
backend is used everywhere in a block of multimethod calls involving
|
||||
multiple arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dispatchables: Sequence[Union[uarray.Dispatchable, Any]]
|
||||
The dispatchables that must be supported
|
||||
domain: string
|
||||
The domain to query for backends and set.
|
||||
coerce: bool
|
||||
Whether or not to allow coercion to the backend's types. Implies ``only``.
|
||||
only: bool
|
||||
Whether or not this should be the last backend to try.
|
||||
dispatch_type: Optional[Any]
|
||||
The default dispatch type associated with ``dispatchables``, aka
|
||||
":ref:`marking <MarkingGlossary>`".
|
||||
|
||||
See Also
|
||||
--------
|
||||
determine_backend: For a single dispatch value
|
||||
set_backend: For when you know which backend to set
|
||||
|
||||
Notes
|
||||
-----
|
||||
|
||||
Support is determined by the ``__ua_convert__`` protocol. Backends not
|
||||
supporting the type must return ``NotImplemented`` from their
|
||||
``__ua_convert__`` if they don't support input of that type.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
:func:`determine_backend` allows the backend to be set from a single
|
||||
object. :func:`determine_backend_multi` allows multiple objects to be
|
||||
checked simultaneously for support in the backend. Suppose we have a
|
||||
``BackendAB`` which supports ``TypeA`` and ``TypeB`` in the same call,
|
||||
and a ``BackendBC`` that doesn't support ``TypeA``.
|
||||
|
||||
>>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC):
|
||||
... a, b = ex.TypeA(), ex.TypeB()
|
||||
... with ua.determine_backend_multi(
|
||||
... [ua.Dispatchable(a, "mark"), ua.Dispatchable(b, "mark")],
|
||||
... domain="ua_examples"
|
||||
... ):
|
||||
... res = ex.creation_multimethod()
|
||||
... ex.call_multimethod(res, a, b)
|
||||
TypeA
|
||||
|
||||
This won't call ``BackendBC`` because it doesn't support ``TypeA``.
|
||||
|
||||
We can also use leave out the ``ua.Dispatchable`` if we specify the
|
||||
default ``dispatch_type`` for the ``dispatchables`` argument.
|
||||
|
||||
>>> with ua.set_backend(ex.BackendAB), ua.set_backend(ex.BackendBC):
|
||||
... a, b = ex.TypeA(), ex.TypeB()
|
||||
... with ua.determine_backend_multi(
|
||||
... [a, b], dispatch_type="mark", domain="ua_examples"
|
||||
... ):
|
||||
... res = ex.creation_multimethod()
|
||||
... ex.call_multimethod(res, a, b)
|
||||
TypeA
|
||||
|
||||
"""
|
||||
if "dispatch_type" in kwargs:
|
||||
disp_type = kwargs.pop("dispatch_type")
|
||||
dispatchables = tuple(
|
||||
d if isinstance(d, Dispatchable) else Dispatchable(d, disp_type)
|
||||
for d in dispatchables
|
||||
)
|
||||
else:
|
||||
dispatchables = tuple(dispatchables)
|
||||
if not all(isinstance(d, Dispatchable) for d in dispatchables):
|
||||
raise TypeError("dispatchables must be instances of uarray.Dispatchable")
|
||||
|
||||
if len(kwargs) != 0:
|
||||
raise TypeError(f"Received unexpected keyword arguments: {kwargs}")
|
||||
|
||||
backend = _uarray.determine_backend(domain, dispatchables, coerce)
|
||||
|
||||
return set_backend(backend, coerce=coerce, only=only)
|
||||
BIN
Binary file not shown.
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
NumPy Array API compatibility library
|
||||
|
||||
This is a small wrapper around NumPy, CuPy, JAX, sparse and others that are
|
||||
compatible with the Array API standard https://data-apis.org/array-api/latest/.
|
||||
See also NEP 47 https://numpy.org/neps/nep-0047-array-api-standard.html.
|
||||
|
||||
Unlike array_api_strict, this is not a strict minimal implementation of the
|
||||
Array API, but rather just an extension of the main NumPy namespace with
|
||||
changes needed to be compliant with the Array API. See
|
||||
https://numpy.org/doc/stable/reference/array_api.html for a full list of
|
||||
changes. In particular, unlike array_api_strict, this package does not use a
|
||||
separate Array object, but rather just uses numpy.ndarray directly.
|
||||
|
||||
Library authors using the Array API may wish to test against array_api_strict
|
||||
to ensure they are not using functionality outside of the standard, but prefer
|
||||
this implementation for the default when working with NumPy arrays.
|
||||
|
||||
"""
|
||||
__version__ = '1.12.0'
|
||||
|
||||
from .common import * # noqa: F401, F403
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
Internal helpers
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from functools import wraps
|
||||
from inspect import signature
|
||||
from types import ModuleType
|
||||
from typing import TypeVar
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def get_xp(xp: ModuleType) -> Callable[[Callable[..., _T]], Callable[..., _T]]:
|
||||
"""
|
||||
Decorator to automatically replace xp with the corresponding array module.
|
||||
|
||||
Use like
|
||||
|
||||
import numpy as np
|
||||
|
||||
@get_xp(np)
|
||||
def func(x, /, xp, kwarg=None):
|
||||
return xp.func(x, kwarg=kwarg)
|
||||
|
||||
Note that xp must be a keyword argument and come after all non-keyword
|
||||
arguments.
|
||||
|
||||
"""
|
||||
|
||||
def inner(f: Callable[..., _T], /) -> Callable[..., _T]:
|
||||
@wraps(f)
|
||||
def wrapped_f(*args: object, **kwargs: object) -> object:
|
||||
return f(*args, xp=xp, **kwargs)
|
||||
|
||||
sig = signature(f)
|
||||
new_sig = sig.replace(
|
||||
parameters=[par for i, par in sig.parameters.items() if i != "xp"]
|
||||
)
|
||||
|
||||
if wrapped_f.__doc__ is None:
|
||||
wrapped_f.__doc__ = f"""\
|
||||
Array API compatibility wrapper for {f.__name__}.
|
||||
|
||||
See the corresponding documentation in NumPy/CuPy and/or the array API
|
||||
specification for more details.
|
||||
|
||||
"""
|
||||
wrapped_f.__signature__ = new_sig # pyright: ignore[reportAttributeAccessIssue]
|
||||
return wrapped_f # pyright: ignore[reportReturnType]
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
__all__ = ["get_xp"]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1 @@
|
||||
from ._helpers import * # noqa: F403
|
||||
@@ -0,0 +1,727 @@
|
||||
"""
|
||||
These are functions that are just aliases of existing functions in NumPy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Sequence, cast
|
||||
|
||||
from ._helpers import _check_device, array_namespace
|
||||
from ._helpers import device as _get_device
|
||||
from ._helpers import is_cupy_namespace as _is_cupy_namespace
|
||||
from ._typing import Array, Device, DType, Namespace
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# TODO: import from typing (requires Python >=3.13)
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
# These functions are modified from the NumPy versions.
|
||||
|
||||
# Creation functions add the device keyword (which does nothing for NumPy and Dask)
|
||||
|
||||
|
||||
def arange(
|
||||
start: float,
|
||||
/,
|
||||
stop: float | None = None,
|
||||
step: float = 1,
|
||||
*,
|
||||
xp: Namespace,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.arange(start, stop=stop, step=step, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def empty(
|
||||
shape: int | tuple[int, ...],
|
||||
xp: Namespace,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.empty(shape, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def empty_like(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.empty_like(x, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def eye(
|
||||
n_rows: int,
|
||||
n_cols: int | None = None,
|
||||
/,
|
||||
*,
|
||||
xp: Namespace,
|
||||
k: int = 0,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.eye(n_rows, M=n_cols, k=k, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def full(
|
||||
shape: int | tuple[int, ...],
|
||||
fill_value: complex,
|
||||
xp: Namespace,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.full(shape, fill_value, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def full_like(
|
||||
x: Array,
|
||||
/,
|
||||
fill_value: complex,
|
||||
*,
|
||||
xp: Namespace,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.full_like(x, fill_value, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def linspace(
|
||||
start: float,
|
||||
stop: float,
|
||||
/,
|
||||
num: int,
|
||||
*,
|
||||
xp: Namespace,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
endpoint: bool = True,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.linspace(start, stop, num, dtype=dtype, endpoint=endpoint, **kwargs)
|
||||
|
||||
|
||||
def ones(
|
||||
shape: int | tuple[int, ...],
|
||||
xp: Namespace,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.ones(shape, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def ones_like(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.ones_like(x, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def zeros(
|
||||
shape: int | tuple[int, ...],
|
||||
xp: Namespace,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.zeros(shape, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
def zeros_like(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
_check_device(xp, device)
|
||||
return xp.zeros_like(x, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
# np.unique() is split into four functions in the array API:
|
||||
# unique_all, unique_counts, unique_inverse, and unique_values (this is done
|
||||
# to remove polymorphic return types).
|
||||
|
||||
# The functions here return namedtuples (np.unique() returns a normal
|
||||
# tuple).
|
||||
|
||||
|
||||
# Note that these named tuples aren't actually part of the standard namespace,
|
||||
# but I don't see any issue with exporting the names here regardless.
|
||||
class UniqueAllResult(NamedTuple):
|
||||
values: Array
|
||||
indices: Array
|
||||
inverse_indices: Array
|
||||
counts: Array
|
||||
|
||||
|
||||
class UniqueCountsResult(NamedTuple):
|
||||
values: Array
|
||||
counts: Array
|
||||
|
||||
|
||||
class UniqueInverseResult(NamedTuple):
|
||||
values: Array
|
||||
inverse_indices: Array
|
||||
|
||||
|
||||
def _unique_kwargs(xp: Namespace) -> dict[str, bool]:
|
||||
# Older versions of NumPy and CuPy do not have equal_nan. Rather than
|
||||
# trying to parse version numbers, just check if equal_nan is in the
|
||||
# signature.
|
||||
s = inspect.signature(xp.unique)
|
||||
if "equal_nan" in s.parameters:
|
||||
return {"equal_nan": False}
|
||||
return {}
|
||||
|
||||
|
||||
def unique_all(x: Array, /, xp: Namespace) -> UniqueAllResult:
|
||||
kwargs = _unique_kwargs(xp)
|
||||
values, indices, inverse_indices, counts = xp.unique(
|
||||
x,
|
||||
return_counts=True,
|
||||
return_index=True,
|
||||
return_inverse=True,
|
||||
**kwargs,
|
||||
)
|
||||
# np.unique() flattens inverse indices, but they need to share x's shape
|
||||
# See https://github.com/numpy/numpy/issues/20638
|
||||
inverse_indices = inverse_indices.reshape(x.shape)
|
||||
return UniqueAllResult(
|
||||
values,
|
||||
indices,
|
||||
inverse_indices,
|
||||
counts,
|
||||
)
|
||||
|
||||
|
||||
def unique_counts(x: Array, /, xp: Namespace) -> UniqueCountsResult:
|
||||
kwargs = _unique_kwargs(xp)
|
||||
res = xp.unique(
|
||||
x, return_counts=True, return_index=False, return_inverse=False, **kwargs
|
||||
)
|
||||
|
||||
return UniqueCountsResult(*res)
|
||||
|
||||
|
||||
def unique_inverse(x: Array, /, xp: Namespace) -> UniqueInverseResult:
|
||||
kwargs = _unique_kwargs(xp)
|
||||
values, inverse_indices = xp.unique(
|
||||
x,
|
||||
return_counts=False,
|
||||
return_index=False,
|
||||
return_inverse=True,
|
||||
**kwargs,
|
||||
)
|
||||
# xp.unique() flattens inverse indices, but they need to share x's shape
|
||||
# See https://github.com/numpy/numpy/issues/20638
|
||||
inverse_indices = inverse_indices.reshape(x.shape)
|
||||
return UniqueInverseResult(values, inverse_indices)
|
||||
|
||||
|
||||
def unique_values(x: Array, /, xp: Namespace) -> Array:
|
||||
kwargs = _unique_kwargs(xp)
|
||||
return xp.unique(
|
||||
x,
|
||||
return_counts=False,
|
||||
return_index=False,
|
||||
return_inverse=False,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
# These functions have different keyword argument names
|
||||
|
||||
|
||||
def std(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int | tuple[int, ...] | None = None,
|
||||
correction: float = 0.0, # correction instead of ddof
|
||||
keepdims: bool = False,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
return xp.std(x, axis=axis, ddof=correction, keepdims=keepdims, **kwargs)
|
||||
|
||||
|
||||
def var(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int | tuple[int, ...] | None = None,
|
||||
correction: float = 0.0, # correction instead of ddof
|
||||
keepdims: bool = False,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
return xp.var(x, axis=axis, ddof=correction, keepdims=keepdims, **kwargs)
|
||||
|
||||
|
||||
# cumulative_sum is renamed from cumsum, and adds the include_initial keyword
|
||||
# argument
|
||||
|
||||
|
||||
def cumulative_sum(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int | None = None,
|
||||
dtype: DType | None = None,
|
||||
include_initial: bool = False,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
wrapped_xp = array_namespace(x)
|
||||
|
||||
# TODO: The standard is not clear about what should happen when x.ndim == 0.
|
||||
if axis is None:
|
||||
if x.ndim > 1:
|
||||
raise ValueError(
|
||||
"axis must be specified in cumulative_sum for more than one dimension"
|
||||
)
|
||||
axis = 0
|
||||
|
||||
res = xp.cumsum(x, axis=axis, dtype=dtype, **kwargs)
|
||||
|
||||
# np.cumsum does not support include_initial
|
||||
if include_initial:
|
||||
initial_shape = list(x.shape)
|
||||
initial_shape[axis] = 1
|
||||
res = xp.concatenate(
|
||||
[
|
||||
wrapped_xp.zeros(
|
||||
shape=initial_shape, dtype=res.dtype, device=_get_device(res)
|
||||
),
|
||||
res,
|
||||
],
|
||||
axis=axis,
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
def cumulative_prod(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int | None = None,
|
||||
dtype: DType | None = None,
|
||||
include_initial: bool = False,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
wrapped_xp = array_namespace(x)
|
||||
|
||||
if axis is None:
|
||||
if x.ndim > 1:
|
||||
raise ValueError(
|
||||
"axis must be specified in cumulative_prod for more than one dimension"
|
||||
)
|
||||
axis = 0
|
||||
|
||||
res = xp.cumprod(x, axis=axis, dtype=dtype, **kwargs)
|
||||
|
||||
# np.cumprod does not support include_initial
|
||||
if include_initial:
|
||||
initial_shape = list(x.shape)
|
||||
initial_shape[axis] = 1
|
||||
res = xp.concatenate(
|
||||
[
|
||||
wrapped_xp.ones(
|
||||
shape=initial_shape, dtype=res.dtype, device=_get_device(res)
|
||||
),
|
||||
res,
|
||||
],
|
||||
axis=axis,
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
# The min and max argument names in clip are different and not optional in numpy, and type
|
||||
# promotion behavior is different.
|
||||
def clip(
|
||||
x: Array,
|
||||
/,
|
||||
min: float | Array | None = None,
|
||||
max: float | Array | None = None,
|
||||
*,
|
||||
xp: Namespace,
|
||||
# TODO: np.clip has other ufunc kwargs
|
||||
out: Array | None = None,
|
||||
) -> Array:
|
||||
def _isscalar(a: object) -> TypeIs[int | float | None]:
|
||||
return isinstance(a, (int, float, type(None)))
|
||||
|
||||
min_shape = () if _isscalar(min) else min.shape
|
||||
max_shape = () if _isscalar(max) else max.shape
|
||||
|
||||
wrapped_xp = array_namespace(x)
|
||||
|
||||
result_shape = xp.broadcast_shapes(x.shape, min_shape, max_shape)
|
||||
|
||||
# np.clip does type promotion but the array API clip requires that the
|
||||
# output have the same dtype as x. We do this instead of just downcasting
|
||||
# the result of xp.clip() to handle some corner cases better (e.g.,
|
||||
# avoiding uint64 -> float64 promotion).
|
||||
|
||||
# Note: cases where min or max overflow (integer) or round (float) in the
|
||||
# wrong direction when downcasting to x.dtype are unspecified. This code
|
||||
# just does whatever NumPy does when it downcasts in the assignment, but
|
||||
# other behavior could be preferred, especially for integers. For example,
|
||||
# this code produces:
|
||||
|
||||
# >>> clip(asarray(0, dtype=int8), asarray(128, dtype=int16), None)
|
||||
# -128
|
||||
|
||||
# but an answer of 0 might be preferred. See
|
||||
# https://github.com/numpy/numpy/issues/24976 for more discussion on this issue.
|
||||
|
||||
# At least handle the case of Python integers correctly (see
|
||||
# https://github.com/numpy/numpy/pull/26892).
|
||||
if wrapped_xp.isdtype(x.dtype, "integral"):
|
||||
if type(min) is int and min <= wrapped_xp.iinfo(x.dtype).min:
|
||||
min = None
|
||||
if type(max) is int and max >= wrapped_xp.iinfo(x.dtype).max:
|
||||
max = None
|
||||
|
||||
dev = _get_device(x)
|
||||
if out is None:
|
||||
out = wrapped_xp.empty(result_shape, dtype=x.dtype, device=dev)
|
||||
assert out is not None # workaround for a type-narrowing issue in pyright
|
||||
out[()] = x
|
||||
|
||||
if min is not None:
|
||||
a = wrapped_xp.asarray(min, dtype=x.dtype, device=dev)
|
||||
a = xp.broadcast_to(a, result_shape)
|
||||
ia = (out < a) | xp.isnan(a)
|
||||
out[ia] = a[ia]
|
||||
|
||||
if max is not None:
|
||||
b = wrapped_xp.asarray(max, dtype=x.dtype, device=dev)
|
||||
b = xp.broadcast_to(b, result_shape)
|
||||
ib = (out > b) | xp.isnan(b)
|
||||
out[ib] = b[ib]
|
||||
|
||||
# Return a scalar for 0-D
|
||||
return out[()]
|
||||
|
||||
|
||||
# Unlike transpose(), the axes argument to permute_dims() is required.
|
||||
def permute_dims(x: Array, /, axes: tuple[int, ...], xp: Namespace) -> Array:
|
||||
return xp.transpose(x, axes)
|
||||
|
||||
|
||||
# np.reshape calls the keyword argument 'newshape' instead of 'shape'
|
||||
def reshape(
|
||||
x: Array,
|
||||
/,
|
||||
shape: tuple[int, ...],
|
||||
xp: Namespace,
|
||||
*,
|
||||
copy: Optional[bool] = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
if copy is True:
|
||||
x = x.copy()
|
||||
elif copy is False:
|
||||
y = x.view()
|
||||
y.shape = shape
|
||||
return y
|
||||
return xp.reshape(x, shape, **kwargs)
|
||||
|
||||
|
||||
# The descending keyword is new in sort and argsort, and 'kind' replaced with
|
||||
# 'stable'
|
||||
def argsort(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int = -1,
|
||||
descending: bool = False,
|
||||
stable: bool = True,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
# Note: this keyword argument is different, and the default is different.
|
||||
# We set it in kwargs like this because numpy.sort uses kind='quicksort'
|
||||
# as the default whereas cupy.sort uses kind=None.
|
||||
if stable:
|
||||
kwargs["kind"] = "stable"
|
||||
if not descending:
|
||||
res = xp.argsort(x, axis=axis, **kwargs)
|
||||
else:
|
||||
# As NumPy has no native descending sort, we imitate it here. Note that
|
||||
# simply flipping the results of xp.argsort(x, ...) would not
|
||||
# respect the relative order like it would in native descending sorts.
|
||||
res = xp.flip(
|
||||
xp.argsort(xp.flip(x, axis=axis), axis=axis, **kwargs),
|
||||
axis=axis,
|
||||
)
|
||||
# Rely on flip()/argsort() to validate axis
|
||||
normalised_axis = axis if axis >= 0 else x.ndim + axis
|
||||
max_i = x.shape[normalised_axis] - 1
|
||||
res = max_i - res
|
||||
return res
|
||||
|
||||
|
||||
def sort(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int = -1,
|
||||
descending: bool = False,
|
||||
stable: bool = True,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
# Note: this keyword argument is different, and the default is different.
|
||||
# We set it in kwargs like this because numpy.sort uses kind='quicksort'
|
||||
# as the default whereas cupy.sort uses kind=None.
|
||||
if stable:
|
||||
kwargs["kind"] = "stable"
|
||||
res = xp.sort(x, axis=axis, **kwargs)
|
||||
if descending:
|
||||
res = xp.flip(res, axis=axis)
|
||||
return res
|
||||
|
||||
|
||||
# nonzero should error for zero-dimensional arrays
|
||||
def nonzero(x: Array, /, xp: Namespace, **kwargs: object) -> tuple[Array, ...]:
|
||||
if x.ndim == 0:
|
||||
raise ValueError("nonzero() does not support zero-dimensional arrays")
|
||||
return xp.nonzero(x, **kwargs)
|
||||
|
||||
|
||||
# ceil, floor, and trunc return integers for integer inputs
|
||||
|
||||
|
||||
def ceil(x: Array, /, xp: Namespace, **kwargs: object) -> Array:
|
||||
if xp.issubdtype(x.dtype, xp.integer):
|
||||
return x
|
||||
return xp.ceil(x, **kwargs)
|
||||
|
||||
|
||||
def floor(x: Array, /, xp: Namespace, **kwargs: object) -> Array:
|
||||
if xp.issubdtype(x.dtype, xp.integer):
|
||||
return x
|
||||
return xp.floor(x, **kwargs)
|
||||
|
||||
|
||||
def trunc(x: Array, /, xp: Namespace, **kwargs: object) -> Array:
|
||||
if xp.issubdtype(x.dtype, xp.integer):
|
||||
return x
|
||||
return xp.trunc(x, **kwargs)
|
||||
|
||||
|
||||
# linear algebra functions
|
||||
|
||||
|
||||
def matmul(x1: Array, x2: Array, /, xp: Namespace, **kwargs: object) -> Array:
|
||||
return xp.matmul(x1, x2, **kwargs)
|
||||
|
||||
|
||||
# Unlike transpose, matrix_transpose only transposes the last two axes.
|
||||
def matrix_transpose(x: Array, /, xp: Namespace) -> Array:
|
||||
if x.ndim < 2:
|
||||
raise ValueError("x must be at least 2-dimensional for matrix_transpose")
|
||||
return xp.swapaxes(x, -1, -2)
|
||||
|
||||
|
||||
def tensordot(
|
||||
x1: Array,
|
||||
x2: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axes: int | tuple[Sequence[int], Sequence[int]] = 2,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
return xp.tensordot(x1, x2, axes=axes, **kwargs)
|
||||
|
||||
|
||||
def vecdot(x1: Array, x2: Array, /, xp: Namespace, *, axis: int = -1) -> Array:
|
||||
if x1.shape[axis] != x2.shape[axis]:
|
||||
raise ValueError("x1 and x2 must have the same size along the given axis")
|
||||
|
||||
if hasattr(xp, "broadcast_tensors"):
|
||||
_broadcast = xp.broadcast_tensors
|
||||
else:
|
||||
_broadcast = xp.broadcast_arrays
|
||||
|
||||
x1_ = xp.moveaxis(x1, axis, -1)
|
||||
x2_ = xp.moveaxis(x2, axis, -1)
|
||||
x1_, x2_ = _broadcast(x1_, x2_)
|
||||
|
||||
res = xp.conj(x1_[..., None, :]) @ x2_[..., None]
|
||||
return res[..., 0, 0]
|
||||
|
||||
|
||||
# isdtype is a new function in the 2022.12 array API specification.
|
||||
|
||||
|
||||
def isdtype(
|
||||
dtype: DType,
|
||||
kind: DType | str | tuple[DType | str, ...],
|
||||
xp: Namespace,
|
||||
*,
|
||||
_tuple: bool = True, # Disallow nested tuples
|
||||
) -> bool:
|
||||
"""
|
||||
Returns a boolean indicating whether a provided dtype is of a specified data type ``kind``.
|
||||
|
||||
Note that outside of this function, this compat library does not yet fully
|
||||
support complex numbers.
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/generated/array_api.isdtype.html
|
||||
for more details
|
||||
"""
|
||||
if isinstance(kind, tuple) and _tuple:
|
||||
return any(
|
||||
isdtype(dtype, k, xp, _tuple=False)
|
||||
for k in cast("tuple[DType | str, ...]", kind)
|
||||
)
|
||||
elif isinstance(kind, str):
|
||||
if kind == "bool":
|
||||
return dtype == xp.bool_
|
||||
elif kind == "signed integer":
|
||||
return xp.issubdtype(dtype, xp.signedinteger)
|
||||
elif kind == "unsigned integer":
|
||||
return xp.issubdtype(dtype, xp.unsignedinteger)
|
||||
elif kind == "integral":
|
||||
return xp.issubdtype(dtype, xp.integer)
|
||||
elif kind == "real floating":
|
||||
return xp.issubdtype(dtype, xp.floating)
|
||||
elif kind == "complex floating":
|
||||
return xp.issubdtype(dtype, xp.complexfloating)
|
||||
elif kind == "numeric":
|
||||
return xp.issubdtype(dtype, xp.number)
|
||||
else:
|
||||
raise ValueError(f"Unrecognized data type kind: {kind!r}")
|
||||
else:
|
||||
# This will allow things that aren't required by the spec, like
|
||||
# isdtype(np.float64, float) or isdtype(np.int64, 'l'). Should we be
|
||||
# more strict here to match the type annotation? Note that the
|
||||
# array_api_strict implementation will be very strict.
|
||||
return dtype == kind
|
||||
|
||||
|
||||
# unstack is a new function in the 2023.12 array API standard
|
||||
def unstack(x: Array, /, xp: Namespace, *, axis: int = 0) -> tuple[Array, ...]:
|
||||
if x.ndim == 0:
|
||||
raise ValueError("Input array must be at least 1-d.")
|
||||
return tuple(xp.moveaxis(x, axis, 0))
|
||||
|
||||
|
||||
# numpy 1.26 does not use the standard definition for sign on complex numbers
|
||||
|
||||
|
||||
def sign(x: Array, /, xp: Namespace, **kwargs: object) -> Array:
|
||||
if isdtype(x.dtype, "complex floating", xp=xp):
|
||||
out = (x / xp.abs(x, **kwargs))[...]
|
||||
# sign(0) = 0 but the above formula would give nan
|
||||
out[x == 0j] = 0j
|
||||
else:
|
||||
out = xp.sign(x, **kwargs)
|
||||
# CuPy sign() does not propagate nans. See
|
||||
# https://github.com/data-apis/array-api-compat/issues/136
|
||||
if _is_cupy_namespace(xp) and isdtype(x.dtype, "real floating", xp=xp):
|
||||
out[xp.isnan(x)] = xp.nan
|
||||
return out[()]
|
||||
|
||||
|
||||
def finfo(type_: DType | Array, /, xp: Namespace) -> Any:
|
||||
# It is surprisingly difficult to recognize a dtype apart from an array.
|
||||
# np.int64 is not the same as np.asarray(1).dtype!
|
||||
try:
|
||||
return xp.finfo(type_)
|
||||
except (ValueError, TypeError):
|
||||
return xp.finfo(type_.dtype)
|
||||
|
||||
|
||||
def iinfo(type_: DType | Array, /, xp: Namespace) -> Any:
|
||||
try:
|
||||
return xp.iinfo(type_)
|
||||
except (ValueError, TypeError):
|
||||
return xp.iinfo(type_.dtype)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"arange",
|
||||
"empty",
|
||||
"empty_like",
|
||||
"eye",
|
||||
"full",
|
||||
"full_like",
|
||||
"linspace",
|
||||
"ones",
|
||||
"ones_like",
|
||||
"zeros",
|
||||
"zeros_like",
|
||||
"UniqueAllResult",
|
||||
"UniqueCountsResult",
|
||||
"UniqueInverseResult",
|
||||
"unique_all",
|
||||
"unique_counts",
|
||||
"unique_inverse",
|
||||
"unique_values",
|
||||
"std",
|
||||
"var",
|
||||
"cumulative_sum",
|
||||
"cumulative_prod",
|
||||
"clip",
|
||||
"permute_dims",
|
||||
"reshape",
|
||||
"argsort",
|
||||
"sort",
|
||||
"nonzero",
|
||||
"ceil",
|
||||
"floor",
|
||||
"trunc",
|
||||
"matmul",
|
||||
"matrix_transpose",
|
||||
"tensordot",
|
||||
"vecdot",
|
||||
"isdtype",
|
||||
"unstack",
|
||||
"sign",
|
||||
"finfo",
|
||||
"iinfo",
|
||||
]
|
||||
_all_ignore = ["inspect", "array_namespace", "NamedTuple"]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,213 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Literal, TypeAlias
|
||||
|
||||
from ._typing import Array, Device, DType, Namespace
|
||||
|
||||
_Norm: TypeAlias = Literal["backward", "ortho", "forward"]
|
||||
|
||||
# Note: NumPy fft functions improperly upcast float32 and complex64 to
|
||||
# complex128, which is why we require wrapping them all here.
|
||||
|
||||
def fft(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
n: int | None = None,
|
||||
axis: int = -1,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.fft(x, n=n, axis=axis, norm=norm)
|
||||
if x.dtype in [xp.float32, xp.complex64]:
|
||||
return res.astype(xp.complex64)
|
||||
return res
|
||||
|
||||
def ifft(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
n: int | None = None,
|
||||
axis: int = -1,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.ifft(x, n=n, axis=axis, norm=norm)
|
||||
if x.dtype in [xp.float32, xp.complex64]:
|
||||
return res.astype(xp.complex64)
|
||||
return res
|
||||
|
||||
def fftn(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
s: Sequence[int] | None = None,
|
||||
axes: Sequence[int] | None = None,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.fftn(x, s=s, axes=axes, norm=norm)
|
||||
if x.dtype in [xp.float32, xp.complex64]:
|
||||
return res.astype(xp.complex64)
|
||||
return res
|
||||
|
||||
def ifftn(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
s: Sequence[int] | None = None,
|
||||
axes: Sequence[int] | None = None,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.ifftn(x, s=s, axes=axes, norm=norm)
|
||||
if x.dtype in [xp.float32, xp.complex64]:
|
||||
return res.astype(xp.complex64)
|
||||
return res
|
||||
|
||||
def rfft(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
n: int | None = None,
|
||||
axis: int = -1,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.rfft(x, n=n, axis=axis, norm=norm)
|
||||
if x.dtype == xp.float32:
|
||||
return res.astype(xp.complex64)
|
||||
return res
|
||||
|
||||
def irfft(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
n: int | None = None,
|
||||
axis: int = -1,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.irfft(x, n=n, axis=axis, norm=norm)
|
||||
if x.dtype == xp.complex64:
|
||||
return res.astype(xp.float32)
|
||||
return res
|
||||
|
||||
def rfftn(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
s: Sequence[int] | None = None,
|
||||
axes: Sequence[int] | None = None,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.rfftn(x, s=s, axes=axes, norm=norm)
|
||||
if x.dtype == xp.float32:
|
||||
return res.astype(xp.complex64)
|
||||
return res
|
||||
|
||||
def irfftn(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
s: Sequence[int] | None = None,
|
||||
axes: Sequence[int] | None = None,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.irfftn(x, s=s, axes=axes, norm=norm)
|
||||
if x.dtype == xp.complex64:
|
||||
return res.astype(xp.float32)
|
||||
return res
|
||||
|
||||
def hfft(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
n: int | None = None,
|
||||
axis: int = -1,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.hfft(x, n=n, axis=axis, norm=norm)
|
||||
if x.dtype in [xp.float32, xp.complex64]:
|
||||
return res.astype(xp.float32)
|
||||
return res
|
||||
|
||||
def ihfft(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
n: int | None = None,
|
||||
axis: int = -1,
|
||||
norm: _Norm = "backward",
|
||||
) -> Array:
|
||||
res = xp.fft.ihfft(x, n=n, axis=axis, norm=norm)
|
||||
if x.dtype in [xp.float32, xp.complex64]:
|
||||
return res.astype(xp.complex64)
|
||||
return res
|
||||
|
||||
def fftfreq(
|
||||
n: int,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
d: float = 1.0,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
) -> Array:
|
||||
if device not in ["cpu", None]:
|
||||
raise ValueError(f"Unsupported device {device!r}")
|
||||
res = xp.fft.fftfreq(n, d=d)
|
||||
if dtype is not None:
|
||||
return res.astype(dtype)
|
||||
return res
|
||||
|
||||
def rfftfreq(
|
||||
n: int,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
d: float = 1.0,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
) -> Array:
|
||||
if device not in ["cpu", None]:
|
||||
raise ValueError(f"Unsupported device {device!r}")
|
||||
res = xp.fft.rfftfreq(n, d=d)
|
||||
if dtype is not None:
|
||||
return res.astype(dtype)
|
||||
return res
|
||||
|
||||
def fftshift(
|
||||
x: Array, /, xp: Namespace, *, axes: int | Sequence[int] | None = None
|
||||
) -> Array:
|
||||
return xp.fft.fftshift(x, axes=axes)
|
||||
|
||||
def ifftshift(
|
||||
x: Array, /, xp: Namespace, *, axes: int | Sequence[int] | None = None
|
||||
) -> Array:
|
||||
return xp.fft.ifftshift(x, axes=axes)
|
||||
|
||||
__all__ = [
|
||||
"fft",
|
||||
"ifft",
|
||||
"fftn",
|
||||
"ifftn",
|
||||
"rfft",
|
||||
"irfft",
|
||||
"rfftn",
|
||||
"irfftn",
|
||||
"hfft",
|
||||
"ihfft",
|
||||
"fftfreq",
|
||||
"rfftfreq",
|
||||
"fftshift",
|
||||
"ifftshift",
|
||||
]
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,232 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Literal, NamedTuple, cast
|
||||
|
||||
import numpy as np
|
||||
|
||||
if np.__version__[0] == "2":
|
||||
from numpy.lib.array_utils import normalize_axis_tuple
|
||||
else:
|
||||
from numpy.core.numeric import normalize_axis_tuple
|
||||
|
||||
from .._internal import get_xp
|
||||
from ._aliases import isdtype, matmul, matrix_transpose, tensordot, vecdot
|
||||
from ._typing import Array, DType, JustFloat, JustInt, Namespace
|
||||
|
||||
|
||||
# These are in the main NumPy namespace but not in numpy.linalg
|
||||
def cross(
|
||||
x1: Array,
|
||||
x2: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int = -1,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
return xp.cross(x1, x2, axis=axis, **kwargs)
|
||||
|
||||
def outer(x1: Array, x2: Array, /, xp: Namespace, **kwargs: object) -> Array:
|
||||
return xp.outer(x1, x2, **kwargs)
|
||||
|
||||
class EighResult(NamedTuple):
|
||||
eigenvalues: Array
|
||||
eigenvectors: Array
|
||||
|
||||
class QRResult(NamedTuple):
|
||||
Q: Array
|
||||
R: Array
|
||||
|
||||
class SlogdetResult(NamedTuple):
|
||||
sign: Array
|
||||
logabsdet: Array
|
||||
|
||||
class SVDResult(NamedTuple):
|
||||
U: Array
|
||||
S: Array
|
||||
Vh: Array
|
||||
|
||||
# These functions are the same as their NumPy counterparts except they return
|
||||
# a namedtuple.
|
||||
def eigh(x: Array, /, xp: Namespace, **kwargs: object) -> EighResult:
|
||||
return EighResult(*xp.linalg.eigh(x, **kwargs))
|
||||
|
||||
def qr(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
mode: Literal["reduced", "complete"] = "reduced",
|
||||
**kwargs: object,
|
||||
) -> QRResult:
|
||||
return QRResult(*xp.linalg.qr(x, mode=mode, **kwargs))
|
||||
|
||||
def slogdet(x: Array, /, xp: Namespace, **kwargs: object) -> SlogdetResult:
|
||||
return SlogdetResult(*xp.linalg.slogdet(x, **kwargs))
|
||||
|
||||
def svd(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
full_matrices: bool = True,
|
||||
**kwargs: object,
|
||||
) -> SVDResult:
|
||||
return SVDResult(*xp.linalg.svd(x, full_matrices=full_matrices, **kwargs))
|
||||
|
||||
# These functions have additional keyword arguments
|
||||
|
||||
# The upper keyword argument is new from NumPy
|
||||
def cholesky(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
upper: bool = False,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
L = xp.linalg.cholesky(x, **kwargs)
|
||||
if upper:
|
||||
U = get_xp(xp)(matrix_transpose)(L)
|
||||
if get_xp(xp)(isdtype)(U.dtype, 'complex floating'):
|
||||
U = xp.conj(U) # pyright: ignore[reportConstantRedefinition]
|
||||
return U
|
||||
return L
|
||||
|
||||
# The rtol keyword argument of matrix_rank() and pinv() is new from NumPy.
|
||||
# Note that it has a different semantic meaning from tol and rcond.
|
||||
def matrix_rank(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
rtol: float | Array | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
# this is different from xp.linalg.matrix_rank, which supports 1
|
||||
# dimensional arrays.
|
||||
if x.ndim < 2:
|
||||
raise xp.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional")
|
||||
S: Array = get_xp(xp)(svdvals)(x, **kwargs)
|
||||
if rtol is None:
|
||||
tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * xp.finfo(S.dtype).eps
|
||||
else:
|
||||
# this is different from xp.linalg.matrix_rank, which does not
|
||||
# multiply the tolerance by the largest singular value.
|
||||
tol = S.max(axis=-1, keepdims=True)*xp.asarray(rtol)[..., xp.newaxis]
|
||||
return xp.count_nonzero(S > tol, axis=-1)
|
||||
|
||||
def pinv(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
rtol: float | Array | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
# this is different from xp.linalg.pinv, which does not multiply the
|
||||
# default tolerance by max(M, N).
|
||||
if rtol is None:
|
||||
rtol = max(x.shape[-2:]) * xp.finfo(x.dtype).eps
|
||||
return xp.linalg.pinv(x, rcond=rtol, **kwargs)
|
||||
|
||||
# These functions are new in the array API spec
|
||||
|
||||
def matrix_norm(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
keepdims: bool = False,
|
||||
ord: Literal[1, 2, -1, -2] | JustFloat | Literal["fro", "nuc"] | None = "fro",
|
||||
) -> Array:
|
||||
return xp.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord)
|
||||
|
||||
# svdvals is not in NumPy (but it is in SciPy). It is equivalent to
|
||||
# xp.linalg.svd(compute_uv=False).
|
||||
def svdvals(x: Array, /, xp: Namespace) -> Array | tuple[Array, ...]:
|
||||
return xp.linalg.svd(x, compute_uv=False)
|
||||
|
||||
def vector_norm(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
axis: int | tuple[int, ...] | None = None,
|
||||
keepdims: bool = False,
|
||||
ord: JustInt | JustFloat = 2,
|
||||
) -> Array:
|
||||
# xp.linalg.norm tries to do a matrix norm whenever axis is a 2-tuple or
|
||||
# when axis=None and the input is 2-D, so to force a vector norm, we make
|
||||
# it so the input is 1-D (for axis=None), or reshape so that norm is done
|
||||
# on a single dimension.
|
||||
if axis is None:
|
||||
# Note: xp.linalg.norm() doesn't handle 0-D arrays
|
||||
_x = x.ravel()
|
||||
_axis = 0
|
||||
elif isinstance(axis, tuple):
|
||||
# Note: The axis argument supports any number of axes, whereas
|
||||
# xp.linalg.norm() only supports a single axis for vector norm.
|
||||
normalized_axis = cast(
|
||||
"tuple[int, ...]",
|
||||
normalize_axis_tuple(axis, x.ndim), # pyright: ignore[reportCallIssue]
|
||||
)
|
||||
rest = tuple(i for i in range(x.ndim) if i not in normalized_axis)
|
||||
newshape = axis + rest
|
||||
_x = xp.transpose(x, newshape).reshape(
|
||||
(math.prod([x.shape[i] for i in axis]), *[x.shape[i] for i in rest]))
|
||||
_axis = 0
|
||||
else:
|
||||
_x = x
|
||||
_axis = axis
|
||||
|
||||
res = xp.linalg.norm(_x, axis=_axis, ord=ord)
|
||||
|
||||
if keepdims:
|
||||
# We can't reuse xp.linalg.norm(keepdims) because of the reshape hacks
|
||||
# above to avoid matrix norm logic.
|
||||
shape = list(x.shape)
|
||||
_axis = cast(
|
||||
"tuple[int, ...]",
|
||||
normalize_axis_tuple( # pyright: ignore[reportCallIssue]
|
||||
range(x.ndim) if axis is None else axis,
|
||||
x.ndim,
|
||||
),
|
||||
)
|
||||
for i in _axis:
|
||||
shape[i] = 1
|
||||
res = xp.reshape(res, tuple(shape))
|
||||
|
||||
return res
|
||||
|
||||
# xp.diagonal and xp.trace operate on the first two axes whereas these
|
||||
# operates on the last two
|
||||
|
||||
def diagonal(x: Array, /, xp: Namespace, *, offset: int = 0, **kwargs: object) -> Array:
|
||||
return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1, **kwargs)
|
||||
|
||||
def trace(
|
||||
x: Array,
|
||||
/,
|
||||
xp: Namespace,
|
||||
*,
|
||||
offset: int = 0,
|
||||
dtype: DType | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
return xp.asarray(
|
||||
xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs)
|
||||
)
|
||||
|
||||
__all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult',
|
||||
'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet',
|
||||
'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm',
|
||||
'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal',
|
||||
'trace']
|
||||
|
||||
_all_ignore = ['math', 'normalize_axis_tuple', 'get_xp', 'np', 'isdtype']
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,192 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import ModuleType as Namespace
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Literal,
|
||||
Protocol,
|
||||
TypeAlias,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
final,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import Incomplete
|
||||
|
||||
SupportsBufferProtocol: TypeAlias = Incomplete
|
||||
Array: TypeAlias = Incomplete
|
||||
Device: TypeAlias = Incomplete
|
||||
DType: TypeAlias = Incomplete
|
||||
else:
|
||||
SupportsBufferProtocol = object
|
||||
Array = object
|
||||
Device = object
|
||||
DType = object
|
||||
|
||||
|
||||
_T_co = TypeVar("_T_co", covariant=True)
|
||||
|
||||
|
||||
# These "Just" types are equivalent to the `Just` type from the `optype` library,
|
||||
# apart from them not being `@runtime_checkable`.
|
||||
# - docs: https://github.com/jorenham/optype/blob/master/README.md#just
|
||||
# - code: https://github.com/jorenham/optype/blob/master/optype/_core/_just.py
|
||||
@final
|
||||
class JustInt(Protocol):
|
||||
@property
|
||||
def __class__(self, /) -> type[int]: ...
|
||||
@__class__.setter
|
||||
def __class__(self, value: type[int], /) -> None: ... # pyright: ignore[reportIncompatibleMethodOverride]
|
||||
|
||||
|
||||
@final
|
||||
class JustFloat(Protocol):
|
||||
@property
|
||||
def __class__(self, /) -> type[float]: ...
|
||||
@__class__.setter
|
||||
def __class__(self, value: type[float], /) -> None: ... # pyright: ignore[reportIncompatibleMethodOverride]
|
||||
|
||||
|
||||
@final
|
||||
class JustComplex(Protocol):
|
||||
@property
|
||||
def __class__(self, /) -> type[complex]: ...
|
||||
@__class__.setter
|
||||
def __class__(self, value: type[complex], /) -> None: ... # pyright: ignore[reportIncompatibleMethodOverride]
|
||||
|
||||
|
||||
#
|
||||
|
||||
|
||||
class NestedSequence(Protocol[_T_co]):
|
||||
def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
|
||||
def __len__(self, /) -> int: ...
|
||||
|
||||
|
||||
class SupportsArrayNamespace(Protocol[_T_co]):
|
||||
def __array_namespace__(self, /, *, api_version: str | None) -> _T_co: ...
|
||||
|
||||
|
||||
class HasShape(Protocol[_T_co]):
|
||||
@property
|
||||
def shape(self, /) -> _T_co: ...
|
||||
|
||||
|
||||
# Return type of `__array_namespace_info__.default_dtypes`
|
||||
Capabilities = TypedDict(
|
||||
"Capabilities",
|
||||
{
|
||||
"boolean indexing": bool,
|
||||
"data-dependent shapes": bool,
|
||||
"max dimensions": int,
|
||||
},
|
||||
)
|
||||
|
||||
# Return type of `__array_namespace_info__.default_dtypes`
|
||||
DefaultDTypes = TypedDict(
|
||||
"DefaultDTypes",
|
||||
{
|
||||
"real floating": DType,
|
||||
"complex floating": DType,
|
||||
"integral": DType,
|
||||
"indexing": DType,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
_DTypeKind: TypeAlias = Literal[
|
||||
"bool",
|
||||
"signed integer",
|
||||
"unsigned integer",
|
||||
"integral",
|
||||
"real floating",
|
||||
"complex floating",
|
||||
"numeric",
|
||||
]
|
||||
# Type of the `kind` parameter in `__array_namespace_info__.dtypes`
|
||||
DTypeKind: TypeAlias = _DTypeKind | tuple[_DTypeKind, ...]
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind="bool")`
|
||||
class DTypesBool(TypedDict):
|
||||
bool: DType
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind="signed integer")`
|
||||
class DTypesSigned(TypedDict):
|
||||
int8: DType
|
||||
int16: DType
|
||||
int32: DType
|
||||
int64: DType
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind="unsigned integer")`
|
||||
class DTypesUnsigned(TypedDict):
|
||||
uint8: DType
|
||||
uint16: DType
|
||||
uint32: DType
|
||||
uint64: DType
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind="integral")`
|
||||
class DTypesIntegral(DTypesSigned, DTypesUnsigned):
|
||||
pass
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind="real floating")`
|
||||
class DTypesReal(TypedDict):
|
||||
float32: DType
|
||||
float64: DType
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind="complex floating")`
|
||||
class DTypesComplex(TypedDict):
|
||||
complex64: DType
|
||||
complex128: DType
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind="numeric")`
|
||||
class DTypesNumeric(DTypesIntegral, DTypesReal, DTypesComplex):
|
||||
pass
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind=None)` (default)
|
||||
class DTypesAll(DTypesBool, DTypesNumeric):
|
||||
pass
|
||||
|
||||
|
||||
# `__array_namespace_info__.dtypes(kind=?)` (fallback)
|
||||
DTypesAny: TypeAlias = Mapping[str, DType]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Array",
|
||||
"Capabilities",
|
||||
"DType",
|
||||
"DTypeKind",
|
||||
"DTypesAny",
|
||||
"DTypesAll",
|
||||
"DTypesBool",
|
||||
"DTypesNumeric",
|
||||
"DTypesIntegral",
|
||||
"DTypesSigned",
|
||||
"DTypesUnsigned",
|
||||
"DTypesReal",
|
||||
"DTypesComplex",
|
||||
"DefaultDTypes",
|
||||
"Device",
|
||||
"HasShape",
|
||||
"Namespace",
|
||||
"JustInt",
|
||||
"JustFloat",
|
||||
"JustComplex",
|
||||
"NestedSequence",
|
||||
"SupportsArrayNamespace",
|
||||
"SupportsBufferProtocol",
|
||||
]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,13 @@
|
||||
from cupy import * # noqa: F403
|
||||
|
||||
# from cupy import * doesn't overwrite these builtin names
|
||||
from cupy import abs, max, min, round # noqa: F401
|
||||
|
||||
# These imports may overwrite names from the import * above.
|
||||
from ._aliases import * # noqa: F403
|
||||
|
||||
# See the comment in the numpy __init__.py
|
||||
__import__(__package__ + '.linalg')
|
||||
__import__(__package__ + '.fft')
|
||||
|
||||
__array_api_version__ = '2024.12'
|
||||
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import cupy as cp
|
||||
|
||||
from ..common import _aliases, _helpers
|
||||
from ..common._typing import NestedSequence, SupportsBufferProtocol
|
||||
from .._internal import get_xp
|
||||
from ._info import __array_namespace_info__
|
||||
from ._typing import Array, Device, DType
|
||||
|
||||
bool = cp.bool_
|
||||
|
||||
# Basic renames
|
||||
acos = cp.arccos
|
||||
acosh = cp.arccosh
|
||||
asin = cp.arcsin
|
||||
asinh = cp.arcsinh
|
||||
atan = cp.arctan
|
||||
atan2 = cp.arctan2
|
||||
atanh = cp.arctanh
|
||||
bitwise_left_shift = cp.left_shift
|
||||
bitwise_invert = cp.invert
|
||||
bitwise_right_shift = cp.right_shift
|
||||
concat = cp.concatenate
|
||||
pow = cp.power
|
||||
|
||||
arange = get_xp(cp)(_aliases.arange)
|
||||
empty = get_xp(cp)(_aliases.empty)
|
||||
empty_like = get_xp(cp)(_aliases.empty_like)
|
||||
eye = get_xp(cp)(_aliases.eye)
|
||||
full = get_xp(cp)(_aliases.full)
|
||||
full_like = get_xp(cp)(_aliases.full_like)
|
||||
linspace = get_xp(cp)(_aliases.linspace)
|
||||
ones = get_xp(cp)(_aliases.ones)
|
||||
ones_like = get_xp(cp)(_aliases.ones_like)
|
||||
zeros = get_xp(cp)(_aliases.zeros)
|
||||
zeros_like = get_xp(cp)(_aliases.zeros_like)
|
||||
UniqueAllResult = get_xp(cp)(_aliases.UniqueAllResult)
|
||||
UniqueCountsResult = get_xp(cp)(_aliases.UniqueCountsResult)
|
||||
UniqueInverseResult = get_xp(cp)(_aliases.UniqueInverseResult)
|
||||
unique_all = get_xp(cp)(_aliases.unique_all)
|
||||
unique_counts = get_xp(cp)(_aliases.unique_counts)
|
||||
unique_inverse = get_xp(cp)(_aliases.unique_inverse)
|
||||
unique_values = get_xp(cp)(_aliases.unique_values)
|
||||
std = get_xp(cp)(_aliases.std)
|
||||
var = get_xp(cp)(_aliases.var)
|
||||
cumulative_sum = get_xp(cp)(_aliases.cumulative_sum)
|
||||
cumulative_prod = get_xp(cp)(_aliases.cumulative_prod)
|
||||
clip = get_xp(cp)(_aliases.clip)
|
||||
permute_dims = get_xp(cp)(_aliases.permute_dims)
|
||||
reshape = get_xp(cp)(_aliases.reshape)
|
||||
argsort = get_xp(cp)(_aliases.argsort)
|
||||
sort = get_xp(cp)(_aliases.sort)
|
||||
nonzero = get_xp(cp)(_aliases.nonzero)
|
||||
ceil = get_xp(cp)(_aliases.ceil)
|
||||
floor = get_xp(cp)(_aliases.floor)
|
||||
trunc = get_xp(cp)(_aliases.trunc)
|
||||
matmul = get_xp(cp)(_aliases.matmul)
|
||||
matrix_transpose = get_xp(cp)(_aliases.matrix_transpose)
|
||||
tensordot = get_xp(cp)(_aliases.tensordot)
|
||||
sign = get_xp(cp)(_aliases.sign)
|
||||
finfo = get_xp(cp)(_aliases.finfo)
|
||||
iinfo = get_xp(cp)(_aliases.iinfo)
|
||||
|
||||
|
||||
# asarray also adds the copy keyword, which is not present in numpy 1.0.
|
||||
def asarray(
|
||||
obj: (
|
||||
Array
|
||||
| bool | int | float | complex
|
||||
| NestedSequence[bool | int | float | complex]
|
||||
| SupportsBufferProtocol
|
||||
),
|
||||
/,
|
||||
*,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
copy: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility wrapper for asarray().
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
"""
|
||||
with cp.cuda.Device(device):
|
||||
if copy is None:
|
||||
return cp.asarray(obj, dtype=dtype, **kwargs)
|
||||
else:
|
||||
res = cp.array(obj, dtype=dtype, copy=copy, **kwargs)
|
||||
if not copy and res is not obj:
|
||||
raise ValueError("Unable to avoid copy while creating an array as requested")
|
||||
return res
|
||||
|
||||
|
||||
def astype(
|
||||
x: Array,
|
||||
dtype: DType,
|
||||
/,
|
||||
*,
|
||||
copy: bool = True,
|
||||
device: Optional[Device] = None,
|
||||
) -> Array:
|
||||
if device is None:
|
||||
return x.astype(dtype=dtype, copy=copy)
|
||||
out = _helpers.to_device(x.astype(dtype=dtype, copy=False), device)
|
||||
return out.copy() if copy and out is x else out
|
||||
|
||||
|
||||
# cupy.count_nonzero does not have keepdims
|
||||
def count_nonzero(
|
||||
x: Array,
|
||||
axis=None,
|
||||
keepdims=False
|
||||
) -> Array:
|
||||
result = cp.count_nonzero(x, axis)
|
||||
if keepdims:
|
||||
if axis is None:
|
||||
return cp.reshape(result, [1]*x.ndim)
|
||||
return cp.expand_dims(result, axis)
|
||||
return result
|
||||
|
||||
|
||||
# take_along_axis: axis defaults to -1 but in cupy (and numpy) axis is a required arg
|
||||
def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
|
||||
return cp.take_along_axis(x, indices, axis=axis)
|
||||
|
||||
|
||||
# These functions are completely new here. If the library already has them
|
||||
# (i.e., numpy 2.0), use the library version instead of our wrapper.
|
||||
if hasattr(cp, 'vecdot'):
|
||||
vecdot = cp.vecdot
|
||||
else:
|
||||
vecdot = get_xp(cp)(_aliases.vecdot)
|
||||
|
||||
if hasattr(cp, 'isdtype'):
|
||||
isdtype = cp.isdtype
|
||||
else:
|
||||
isdtype = get_xp(cp)(_aliases.isdtype)
|
||||
|
||||
if hasattr(cp, 'unstack'):
|
||||
unstack = cp.unstack
|
||||
else:
|
||||
unstack = get_xp(cp)(_aliases.unstack)
|
||||
|
||||
__all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'astype',
|
||||
'acos', 'acosh', 'asin', 'asinh', 'atan',
|
||||
'atan2', 'atanh', 'bitwise_left_shift',
|
||||
'bitwise_invert', 'bitwise_right_shift',
|
||||
'bool', 'concat', 'count_nonzero', 'pow', 'sign',
|
||||
'take_along_axis']
|
||||
|
||||
_all_ignore = ['cp', 'get_xp']
|
||||
@@ -0,0 +1,336 @@
|
||||
"""
|
||||
Array API Inspection namespace
|
||||
|
||||
This is the namespace for inspection functions as defined by the array API
|
||||
standard. See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html for
|
||||
more details.
|
||||
|
||||
"""
|
||||
from cupy import (
|
||||
dtype,
|
||||
cuda,
|
||||
bool_ as bool,
|
||||
intp,
|
||||
int8,
|
||||
int16,
|
||||
int32,
|
||||
int64,
|
||||
uint8,
|
||||
uint16,
|
||||
uint32,
|
||||
uint64,
|
||||
float32,
|
||||
float64,
|
||||
complex64,
|
||||
complex128,
|
||||
)
|
||||
|
||||
|
||||
class __array_namespace_info__:
|
||||
"""
|
||||
Get the array API inspection namespace for CuPy.
|
||||
|
||||
The array API inspection namespace defines the following functions:
|
||||
|
||||
- capabilities()
|
||||
- default_device()
|
||||
- default_dtypes()
|
||||
- dtypes()
|
||||
- devices()
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html
|
||||
for more details.
|
||||
|
||||
Returns
|
||||
-------
|
||||
info : ModuleType
|
||||
The array API inspection namespace for CuPy.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': cupy.float64,
|
||||
'complex floating': cupy.complex128,
|
||||
'integral': cupy.int64,
|
||||
'indexing': cupy.int64}
|
||||
|
||||
"""
|
||||
|
||||
__module__ = 'cupy'
|
||||
|
||||
def capabilities(self):
|
||||
"""
|
||||
Return a dictionary of array API library capabilities.
|
||||
|
||||
The resulting dictionary has the following keys:
|
||||
|
||||
- **"boolean indexing"**: boolean indicating whether an array library
|
||||
supports boolean indexing. Always ``True`` for CuPy.
|
||||
|
||||
- **"data-dependent shapes"**: boolean indicating whether an array
|
||||
library supports data-dependent output shapes. Always ``True`` for
|
||||
CuPy.
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html
|
||||
for more details.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
capabilities : dict
|
||||
A dictionary of array API library capabilities.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.capabilities()
|
||||
{'boolean indexing': True,
|
||||
'data-dependent shapes': True,
|
||||
'max dimensions': 64}
|
||||
|
||||
"""
|
||||
return {
|
||||
"boolean indexing": True,
|
||||
"data-dependent shapes": True,
|
||||
"max dimensions": 64,
|
||||
}
|
||||
|
||||
def default_device(self):
|
||||
"""
|
||||
The default device used for new CuPy arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
device : Device
|
||||
The default device used for new CuPy arrays.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_device()
|
||||
Device(0)
|
||||
|
||||
Notes
|
||||
-----
|
||||
This method returns the static default device when CuPy is initialized.
|
||||
However, the *current* device used by creation functions (``empty`` etc.)
|
||||
can be changed globally or with a context manager.
|
||||
|
||||
See Also
|
||||
--------
|
||||
https://github.com/data-apis/array-api/issues/835
|
||||
"""
|
||||
return cuda.Device(0)
|
||||
|
||||
def default_dtypes(self, *, device=None):
|
||||
"""
|
||||
The default data types used for new CuPy arrays.
|
||||
|
||||
For CuPy, this always returns the following dictionary:
|
||||
|
||||
- **"real floating"**: ``cupy.float64``
|
||||
- **"complex floating"**: ``cupy.complex128``
|
||||
- **"integral"**: ``cupy.intp``
|
||||
- **"indexing"**: ``cupy.intp``
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : str, optional
|
||||
The device to get the default data types for.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary describing the default data types used for new CuPy
|
||||
arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': cupy.float64,
|
||||
'complex floating': cupy.complex128,
|
||||
'integral': cupy.int64,
|
||||
'indexing': cupy.int64}
|
||||
|
||||
"""
|
||||
# TODO: Does this depend on device?
|
||||
return {
|
||||
"real floating": dtype(float64),
|
||||
"complex floating": dtype(complex128),
|
||||
"integral": dtype(intp),
|
||||
"indexing": dtype(intp),
|
||||
}
|
||||
|
||||
def dtypes(self, *, device=None, kind=None):
|
||||
"""
|
||||
The array API data types supported by CuPy.
|
||||
|
||||
Note that this function only returns data types that are defined by
|
||||
the array API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : str, optional
|
||||
The device to get the data types for.
|
||||
kind : str or tuple of str, optional
|
||||
The kind of data types to return. If ``None``, all data types are
|
||||
returned. If a string, only data types of that kind are returned.
|
||||
If a tuple, a dictionary containing the union of the given kinds
|
||||
is returned. The following kinds are supported:
|
||||
|
||||
- ``'bool'``: boolean data types (i.e., ``bool``).
|
||||
- ``'signed integer'``: signed integer data types (i.e., ``int8``,
|
||||
``int16``, ``int32``, ``int64``).
|
||||
- ``'unsigned integer'``: unsigned integer data types (i.e.,
|
||||
``uint8``, ``uint16``, ``uint32``, ``uint64``).
|
||||
- ``'integral'``: integer data types. Shorthand for ``('signed
|
||||
integer', 'unsigned integer')``.
|
||||
- ``'real floating'``: real-valued floating-point data types
|
||||
(i.e., ``float32``, ``float64``).
|
||||
- ``'complex floating'``: complex floating-point data types (i.e.,
|
||||
``complex64``, ``complex128``).
|
||||
- ``'numeric'``: numeric data types. Shorthand for ``('integral',
|
||||
'real floating', 'complex floating')``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary mapping the names of data types to the corresponding
|
||||
CuPy data types.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.dtypes(kind='signed integer')
|
||||
{'int8': cupy.int8,
|
||||
'int16': cupy.int16,
|
||||
'int32': cupy.int32,
|
||||
'int64': cupy.int64}
|
||||
|
||||
"""
|
||||
# TODO: Does this depend on device?
|
||||
if kind is None:
|
||||
return {
|
||||
"bool": dtype(bool),
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if kind == "bool":
|
||||
return {"bool": bool}
|
||||
if kind == "signed integer":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
}
|
||||
if kind == "unsigned integer":
|
||||
return {
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
}
|
||||
if kind == "integral":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
}
|
||||
if kind == "real floating":
|
||||
return {
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
}
|
||||
if kind == "complex floating":
|
||||
return {
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if kind == "numeric":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if isinstance(kind, tuple):
|
||||
res = {}
|
||||
for k in kind:
|
||||
res.update(self.dtypes(kind=k))
|
||||
return res
|
||||
raise ValueError(f"unsupported kind: {kind!r}")
|
||||
|
||||
def devices(self):
|
||||
"""
|
||||
The devices supported by CuPy.
|
||||
|
||||
Returns
|
||||
-------
|
||||
devices : list[Device]
|
||||
The devices supported by CuPy.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes
|
||||
|
||||
"""
|
||||
return [cuda.Device(i) for i in range(cuda.runtime.getDeviceCount())]
|
||||
@@ -0,0 +1,31 @@
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = ["Array", "DType", "Device"]
|
||||
_all_ignore = ["cp"]
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import cupy as cp
|
||||
from cupy import ndarray as Array
|
||||
from cupy.cuda.device import Device
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# NumPy 1.x on Python 3.10 fails to parse np.dtype[]
|
||||
DType = cp.dtype[
|
||||
cp.intp
|
||||
| cp.int8
|
||||
| cp.int16
|
||||
| cp.int32
|
||||
| cp.int64
|
||||
| cp.uint8
|
||||
| cp.uint16
|
||||
| cp.uint32
|
||||
| cp.uint64
|
||||
| cp.float32
|
||||
| cp.float64
|
||||
| cp.complex64
|
||||
| cp.complex128
|
||||
| cp.bool_
|
||||
]
|
||||
else:
|
||||
DType = cp.dtype
|
||||
@@ -0,0 +1,36 @@
|
||||
from cupy.fft import * # noqa: F403
|
||||
# cupy.fft doesn't have __all__. If it is added, replace this with
|
||||
#
|
||||
# from cupy.fft import __all__ as linalg_all
|
||||
_n = {}
|
||||
exec('from cupy.fft import *', _n)
|
||||
del _n['__builtins__']
|
||||
fft_all = list(_n)
|
||||
del _n
|
||||
|
||||
from ..common import _fft
|
||||
from .._internal import get_xp
|
||||
|
||||
import cupy as cp
|
||||
|
||||
fft = get_xp(cp)(_fft.fft)
|
||||
ifft = get_xp(cp)(_fft.ifft)
|
||||
fftn = get_xp(cp)(_fft.fftn)
|
||||
ifftn = get_xp(cp)(_fft.ifftn)
|
||||
rfft = get_xp(cp)(_fft.rfft)
|
||||
irfft = get_xp(cp)(_fft.irfft)
|
||||
rfftn = get_xp(cp)(_fft.rfftn)
|
||||
irfftn = get_xp(cp)(_fft.irfftn)
|
||||
hfft = get_xp(cp)(_fft.hfft)
|
||||
ihfft = get_xp(cp)(_fft.ihfft)
|
||||
fftfreq = get_xp(cp)(_fft.fftfreq)
|
||||
rfftfreq = get_xp(cp)(_fft.rfftfreq)
|
||||
fftshift = get_xp(cp)(_fft.fftshift)
|
||||
ifftshift = get_xp(cp)(_fft.ifftshift)
|
||||
|
||||
__all__ = fft_all + _fft.__all__
|
||||
|
||||
del get_xp
|
||||
del cp
|
||||
del fft_all
|
||||
del _fft
|
||||
@@ -0,0 +1,49 @@
|
||||
from cupy.linalg import * # noqa: F403
|
||||
# cupy.linalg doesn't have __all__. If it is added, replace this with
|
||||
#
|
||||
# from cupy.linalg import __all__ as linalg_all
|
||||
_n = {}
|
||||
exec('from cupy.linalg import *', _n)
|
||||
del _n['__builtins__']
|
||||
linalg_all = list(_n)
|
||||
del _n
|
||||
|
||||
from ..common import _linalg
|
||||
from .._internal import get_xp
|
||||
|
||||
import cupy as cp
|
||||
|
||||
# These functions are in both the main and linalg namespaces
|
||||
from ._aliases import matmul, matrix_transpose, tensordot, vecdot # noqa: F401
|
||||
|
||||
cross = get_xp(cp)(_linalg.cross)
|
||||
outer = get_xp(cp)(_linalg.outer)
|
||||
EighResult = _linalg.EighResult
|
||||
QRResult = _linalg.QRResult
|
||||
SlogdetResult = _linalg.SlogdetResult
|
||||
SVDResult = _linalg.SVDResult
|
||||
eigh = get_xp(cp)(_linalg.eigh)
|
||||
qr = get_xp(cp)(_linalg.qr)
|
||||
slogdet = get_xp(cp)(_linalg.slogdet)
|
||||
svd = get_xp(cp)(_linalg.svd)
|
||||
cholesky = get_xp(cp)(_linalg.cholesky)
|
||||
matrix_rank = get_xp(cp)(_linalg.matrix_rank)
|
||||
pinv = get_xp(cp)(_linalg.pinv)
|
||||
matrix_norm = get_xp(cp)(_linalg.matrix_norm)
|
||||
svdvals = get_xp(cp)(_linalg.svdvals)
|
||||
diagonal = get_xp(cp)(_linalg.diagonal)
|
||||
trace = get_xp(cp)(_linalg.trace)
|
||||
|
||||
# These functions are completely new here. If the library already has them
|
||||
# (i.e., numpy 2.0), use the library version instead of our wrapper.
|
||||
if hasattr(cp.linalg, 'vector_norm'):
|
||||
vector_norm = cp.linalg.vector_norm
|
||||
else:
|
||||
vector_norm = get_xp(cp)(_linalg.vector_norm)
|
||||
|
||||
__all__ = linalg_all + _linalg.__all__
|
||||
|
||||
del get_xp
|
||||
del cp
|
||||
del linalg_all
|
||||
del _linalg
|
||||
@@ -0,0 +1,12 @@
|
||||
from typing import Final
|
||||
|
||||
from dask.array import * # noqa: F403
|
||||
|
||||
# These imports may overwrite names from the import * above.
|
||||
from ._aliases import * # noqa: F403
|
||||
|
||||
__array_api_version__: Final = "2024.12"
|
||||
|
||||
# See the comment in the numpy __init__.py
|
||||
__import__(__package__ + '.linalg')
|
||||
__import__(__package__ + '.fft')
|
||||
@@ -0,0 +1,376 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
# pyright: reportUnknownArgumentType=false
|
||||
# pyright: reportUnknownMemberType=false
|
||||
# pyright: reportUnknownVariableType=false
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from builtins import bool as py_bool
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
import dask.array as da
|
||||
import numpy as np
|
||||
from numpy import bool_ as bool
|
||||
from numpy import (
|
||||
can_cast,
|
||||
complex64,
|
||||
complex128,
|
||||
float32,
|
||||
float64,
|
||||
int8,
|
||||
int16,
|
||||
int32,
|
||||
int64,
|
||||
result_type,
|
||||
uint8,
|
||||
uint16,
|
||||
uint32,
|
||||
uint64,
|
||||
)
|
||||
|
||||
from ..._internal import get_xp
|
||||
from ...common import _aliases, _helpers, array_namespace
|
||||
from ...common._typing import (
|
||||
Array,
|
||||
Device,
|
||||
DType,
|
||||
NestedSequence,
|
||||
SupportsBufferProtocol,
|
||||
)
|
||||
from ._info import __array_namespace_info__
|
||||
|
||||
isdtype = get_xp(np)(_aliases.isdtype)
|
||||
unstack = get_xp(da)(_aliases.unstack)
|
||||
|
||||
|
||||
# da.astype doesn't respect copy=True
|
||||
def astype(
|
||||
x: Array,
|
||||
dtype: DType,
|
||||
/,
|
||||
*,
|
||||
copy: py_bool = True,
|
||||
device: Device | None = None,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility wrapper for astype().
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
"""
|
||||
# TODO: respect device keyword?
|
||||
_helpers._check_device(da, device)
|
||||
|
||||
if not copy and dtype == x.dtype:
|
||||
return x
|
||||
x = x.astype(dtype)
|
||||
return x.copy() if copy else x
|
||||
|
||||
|
||||
# Common aliases
|
||||
|
||||
|
||||
# This arange func is modified from the common one to
|
||||
# not pass stop/step as keyword arguments, which will cause
|
||||
# an error with dask
|
||||
def arange(
|
||||
start: float,
|
||||
/,
|
||||
stop: float | None = None,
|
||||
step: float = 1,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility wrapper for arange().
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
"""
|
||||
# TODO: respect device keyword?
|
||||
_helpers._check_device(da, device)
|
||||
|
||||
args: list[Any] = [start]
|
||||
if stop is not None:
|
||||
args.append(stop)
|
||||
else:
|
||||
# stop is None, so start is actually stop
|
||||
# prepend the default value for start which is 0
|
||||
args.insert(0, 0)
|
||||
args.append(step)
|
||||
|
||||
return da.arange(*args, dtype=dtype, **kwargs)
|
||||
|
||||
|
||||
eye = get_xp(da)(_aliases.eye)
|
||||
linspace = get_xp(da)(_aliases.linspace)
|
||||
UniqueAllResult = get_xp(da)(_aliases.UniqueAllResult)
|
||||
UniqueCountsResult = get_xp(da)(_aliases.UniqueCountsResult)
|
||||
UniqueInverseResult = get_xp(da)(_aliases.UniqueInverseResult)
|
||||
unique_all = get_xp(da)(_aliases.unique_all)
|
||||
unique_counts = get_xp(da)(_aliases.unique_counts)
|
||||
unique_inverse = get_xp(da)(_aliases.unique_inverse)
|
||||
unique_values = get_xp(da)(_aliases.unique_values)
|
||||
permute_dims = get_xp(da)(_aliases.permute_dims)
|
||||
std = get_xp(da)(_aliases.std)
|
||||
var = get_xp(da)(_aliases.var)
|
||||
cumulative_sum = get_xp(da)(_aliases.cumulative_sum)
|
||||
cumulative_prod = get_xp(da)(_aliases.cumulative_prod)
|
||||
empty = get_xp(da)(_aliases.empty)
|
||||
empty_like = get_xp(da)(_aliases.empty_like)
|
||||
full = get_xp(da)(_aliases.full)
|
||||
full_like = get_xp(da)(_aliases.full_like)
|
||||
ones = get_xp(da)(_aliases.ones)
|
||||
ones_like = get_xp(da)(_aliases.ones_like)
|
||||
zeros = get_xp(da)(_aliases.zeros)
|
||||
zeros_like = get_xp(da)(_aliases.zeros_like)
|
||||
reshape = get_xp(da)(_aliases.reshape)
|
||||
matrix_transpose = get_xp(da)(_aliases.matrix_transpose)
|
||||
vecdot = get_xp(da)(_aliases.vecdot)
|
||||
nonzero = get_xp(da)(_aliases.nonzero)
|
||||
ceil = get_xp(np)(_aliases.ceil)
|
||||
floor = get_xp(np)(_aliases.floor)
|
||||
trunc = get_xp(np)(_aliases.trunc)
|
||||
matmul = get_xp(np)(_aliases.matmul)
|
||||
tensordot = get_xp(np)(_aliases.tensordot)
|
||||
sign = get_xp(np)(_aliases.sign)
|
||||
finfo = get_xp(np)(_aliases.finfo)
|
||||
iinfo = get_xp(np)(_aliases.iinfo)
|
||||
|
||||
|
||||
# asarray also adds the copy keyword, which is not present in numpy 1.0.
|
||||
def asarray(
|
||||
obj: complex | NestedSequence[complex] | Array | SupportsBufferProtocol,
|
||||
/,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
copy: py_bool | None = None,
|
||||
**kwargs: object,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility wrapper for asarray().
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
"""
|
||||
# TODO: respect device keyword?
|
||||
_helpers._check_device(da, device)
|
||||
|
||||
if isinstance(obj, da.Array):
|
||||
if dtype is not None and dtype != obj.dtype:
|
||||
if copy is False:
|
||||
raise ValueError("Unable to avoid copy when changing dtype")
|
||||
obj = obj.astype(dtype)
|
||||
return obj.copy() if copy else obj # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
if copy is False:
|
||||
raise ValueError(
|
||||
"Unable to avoid copy when converting a non-dask object to dask"
|
||||
)
|
||||
|
||||
# copy=None to be uniform across dask < 2024.12 and >= 2024.12
|
||||
# see https://github.com/dask/dask/pull/11524/
|
||||
obj = np.array(obj, dtype=dtype, copy=True)
|
||||
return da.from_array(obj)
|
||||
|
||||
|
||||
# Element wise aliases
|
||||
from dask.array import arccos as acos
|
||||
from dask.array import arccosh as acosh
|
||||
from dask.array import arcsin as asin
|
||||
from dask.array import arcsinh as asinh
|
||||
from dask.array import arctan as atan
|
||||
from dask.array import arctan2 as atan2
|
||||
from dask.array import arctanh as atanh
|
||||
|
||||
# Other
|
||||
from dask.array import concatenate as concat
|
||||
from dask.array import invert as bitwise_invert
|
||||
from dask.array import left_shift as bitwise_left_shift
|
||||
from dask.array import power as pow
|
||||
from dask.array import right_shift as bitwise_right_shift
|
||||
|
||||
|
||||
# dask.array.clip does not work unless all three arguments are provided.
|
||||
# Furthermore, the masking workaround in common._aliases.clip cannot work with
|
||||
# dask (meaning uint64 promoting to float64 is going to just be unfixed for
|
||||
# now).
|
||||
def clip(
|
||||
x: Array,
|
||||
/,
|
||||
min: float | Array | None = None,
|
||||
max: float | Array | None = None,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility wrapper for clip().
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
"""
|
||||
|
||||
def _isscalar(a: float | Array | None, /) -> TypeIs[float | None]:
|
||||
return a is None or isinstance(a, (int, float))
|
||||
|
||||
min_shape = () if _isscalar(min) else min.shape
|
||||
max_shape = () if _isscalar(max) else max.shape
|
||||
|
||||
# TODO: This won't handle dask unknown shapes
|
||||
result_shape = np.broadcast_shapes(x.shape, min_shape, max_shape)
|
||||
|
||||
if min is not None:
|
||||
min = da.broadcast_to(da.asarray(min), result_shape)
|
||||
if max is not None:
|
||||
max = da.broadcast_to(da.asarray(max), result_shape)
|
||||
|
||||
if min is None and max is None:
|
||||
return da.positive(x)
|
||||
|
||||
if min is None:
|
||||
return astype(da.minimum(x, max), x.dtype)
|
||||
if max is None:
|
||||
return astype(da.maximum(x, min), x.dtype)
|
||||
|
||||
return astype(da.minimum(da.maximum(x, min), max), x.dtype)
|
||||
|
||||
|
||||
def _ensure_single_chunk(x: Array, axis: int) -> tuple[Array, Callable[[Array], Array]]:
|
||||
"""
|
||||
Make sure that Array is not broken into multiple chunks along axis.
|
||||
|
||||
Returns
|
||||
-------
|
||||
x : Array
|
||||
The input Array with a single chunk along axis.
|
||||
restore : Callable[Array, Array]
|
||||
function to apply to the output to rechunk it back into reasonable chunks
|
||||
"""
|
||||
if axis < 0:
|
||||
axis += x.ndim
|
||||
if x.numblocks[axis] < 2:
|
||||
return x, lambda x: x
|
||||
|
||||
# Break chunks on other axes in an attempt to keep chunk size low
|
||||
x = x.rechunk({i: -1 if i == axis else "auto" for i in range(x.ndim)})
|
||||
|
||||
# Rather than reconstructing the original chunks, which can be a
|
||||
# very expensive affair, just break down oversized chunks without
|
||||
# incurring in any transfers over the network.
|
||||
# This has the downside of a risk of overchunking if the array is
|
||||
# then used in operations against other arrays that match the
|
||||
# original chunking pattern.
|
||||
return x, lambda x: x.rechunk()
|
||||
|
||||
|
||||
def sort(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: int = -1,
|
||||
descending: py_bool = False,
|
||||
stable: py_bool = True,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility layer around the lack of sort() in Dask.
|
||||
|
||||
Warnings
|
||||
--------
|
||||
This function temporarily rechunks the array along `axis` to a single chunk.
|
||||
This can be extremely inefficient and can lead to out-of-memory errors.
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
"""
|
||||
x, restore = _ensure_single_chunk(x, axis)
|
||||
|
||||
meta_xp = array_namespace(x._meta)
|
||||
x = da.map_blocks(
|
||||
meta_xp.sort,
|
||||
x,
|
||||
axis=axis,
|
||||
meta=x._meta,
|
||||
dtype=x.dtype,
|
||||
descending=descending,
|
||||
stable=stable,
|
||||
)
|
||||
|
||||
return restore(x)
|
||||
|
||||
|
||||
def argsort(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: int = -1,
|
||||
descending: py_bool = False,
|
||||
stable: py_bool = True,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility layer around the lack of argsort() in Dask.
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
|
||||
Warnings
|
||||
--------
|
||||
This function temporarily rechunks the array along `axis` into a single chunk.
|
||||
This can be extremely inefficient and can lead to out-of-memory errors.
|
||||
"""
|
||||
x, restore = _ensure_single_chunk(x, axis)
|
||||
|
||||
meta_xp = array_namespace(x._meta)
|
||||
dtype = meta_xp.argsort(x._meta).dtype
|
||||
meta = meta_xp.astype(x._meta, dtype)
|
||||
x = da.map_blocks(
|
||||
meta_xp.argsort,
|
||||
x,
|
||||
axis=axis,
|
||||
meta=meta,
|
||||
dtype=dtype,
|
||||
descending=descending,
|
||||
stable=stable,
|
||||
)
|
||||
|
||||
return restore(x)
|
||||
|
||||
|
||||
# dask.array.count_nonzero does not have keepdims
|
||||
def count_nonzero(
|
||||
x: Array,
|
||||
axis: int | None = None,
|
||||
keepdims: py_bool = False,
|
||||
) -> Array:
|
||||
result = da.count_nonzero(x, axis)
|
||||
if keepdims:
|
||||
if axis is None:
|
||||
return da.reshape(result, [1] * x.ndim)
|
||||
return da.expand_dims(result, axis)
|
||||
return result
|
||||
|
||||
|
||||
__all__ = [
|
||||
"__array_namespace_info__",
|
||||
"count_nonzero",
|
||||
"bool",
|
||||
"int8", "int16", "int32", "int64",
|
||||
"uint8", "uint16", "uint32", "uint64",
|
||||
"float32", "float64",
|
||||
"complex64", "complex128",
|
||||
"asarray", "astype", "can_cast", "result_type",
|
||||
"pow",
|
||||
"concat",
|
||||
"acos", "acosh", "asin", "asinh", "atan", "atan2", "atanh",
|
||||
"bitwise_left_shift", "bitwise_right_shift", "bitwise_invert",
|
||||
] # fmt: skip
|
||||
__all__ += _aliases.__all__
|
||||
_all_ignore = ["array_namespace", "get_xp", "da", "np"]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,416 @@
|
||||
"""
|
||||
Array API Inspection namespace
|
||||
|
||||
This is the namespace for inspection functions as defined by the array API
|
||||
standard. See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html for
|
||||
more details.
|
||||
|
||||
"""
|
||||
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal as L
|
||||
from typing import TypeAlias, overload
|
||||
|
||||
from numpy import bool_ as bool
|
||||
from numpy import (
|
||||
complex64,
|
||||
complex128,
|
||||
dtype,
|
||||
float32,
|
||||
float64,
|
||||
int8,
|
||||
int16,
|
||||
int32,
|
||||
int64,
|
||||
intp,
|
||||
uint8,
|
||||
uint16,
|
||||
uint32,
|
||||
uint64,
|
||||
)
|
||||
|
||||
from ...common._helpers import _DASK_DEVICE, _dask_device
|
||||
from ...common._typing import (
|
||||
Capabilities,
|
||||
DefaultDTypes,
|
||||
DType,
|
||||
DTypeKind,
|
||||
DTypesAll,
|
||||
DTypesAny,
|
||||
DTypesBool,
|
||||
DTypesComplex,
|
||||
DTypesIntegral,
|
||||
DTypesNumeric,
|
||||
DTypesReal,
|
||||
DTypesSigned,
|
||||
DTypesUnsigned,
|
||||
)
|
||||
|
||||
_Device: TypeAlias = L["cpu"] | _dask_device
|
||||
|
||||
|
||||
class __array_namespace_info__:
|
||||
"""
|
||||
Get the array API inspection namespace for Dask.
|
||||
|
||||
The array API inspection namespace defines the following functions:
|
||||
|
||||
- capabilities()
|
||||
- default_device()
|
||||
- default_dtypes()
|
||||
- dtypes()
|
||||
- devices()
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html
|
||||
for more details.
|
||||
|
||||
Returns
|
||||
-------
|
||||
info : ModuleType
|
||||
The array API inspection namespace for Dask.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': dask.float64,
|
||||
'complex floating': dask.complex128,
|
||||
'integral': dask.int64,
|
||||
'indexing': dask.int64}
|
||||
|
||||
"""
|
||||
|
||||
__module__ = "dask.array"
|
||||
|
||||
def capabilities(self) -> Capabilities:
|
||||
"""
|
||||
Return a dictionary of array API library capabilities.
|
||||
|
||||
The resulting dictionary has the following keys:
|
||||
|
||||
- **"boolean indexing"**: boolean indicating whether an array library
|
||||
supports boolean indexing.
|
||||
|
||||
Dask support boolean indexing as long as both the index
|
||||
and the indexed arrays have known shapes.
|
||||
Note however that the output .shape and .size properties
|
||||
will contain a non-compliant math.nan instead of None.
|
||||
|
||||
- **"data-dependent shapes"**: boolean indicating whether an array
|
||||
library supports data-dependent output shapes.
|
||||
|
||||
Dask implements unique_values et.al.
|
||||
Note however that the output .shape and .size properties
|
||||
will contain a non-compliant math.nan instead of None.
|
||||
|
||||
- **"max dimensions"**: integer indicating the maximum number of
|
||||
dimensions supported by the array library.
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html
|
||||
for more details.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
capabilities : dict
|
||||
A dictionary of array API library capabilities.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.capabilities()
|
||||
{'boolean indexing': True,
|
||||
'data-dependent shapes': True,
|
||||
'max dimensions': 64}
|
||||
|
||||
"""
|
||||
return {
|
||||
"boolean indexing": True,
|
||||
"data-dependent shapes": True,
|
||||
"max dimensions": 64,
|
||||
}
|
||||
|
||||
def default_device(self) -> L["cpu"]:
|
||||
"""
|
||||
The default device used for new Dask arrays.
|
||||
|
||||
For Dask, this always returns ``'cpu'``.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
device : Device
|
||||
The default device used for new Dask arrays.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_device()
|
||||
'cpu'
|
||||
|
||||
"""
|
||||
return "cpu"
|
||||
|
||||
def default_dtypes(self, /, *, device: _Device | None = None) -> DefaultDTypes:
|
||||
"""
|
||||
The default data types used for new Dask arrays.
|
||||
|
||||
For Dask, this always returns the following dictionary:
|
||||
|
||||
- **"real floating"**: ``numpy.float64``
|
||||
- **"complex floating"**: ``numpy.complex128``
|
||||
- **"integral"**: ``numpy.intp``
|
||||
- **"indexing"**: ``numpy.intp``
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : str, optional
|
||||
The device to get the default data types for.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary describing the default data types used for new Dask
|
||||
arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': dask.float64,
|
||||
'complex floating': dask.complex128,
|
||||
'integral': dask.int64,
|
||||
'indexing': dask.int64}
|
||||
|
||||
"""
|
||||
if device not in ["cpu", _DASK_DEVICE, None]:
|
||||
raise ValueError(
|
||||
f'Device not understood. Only "cpu" or _DASK_DEVICE is allowed, '
|
||||
f"but received: {device!r}"
|
||||
)
|
||||
return {
|
||||
"real floating": dtype(float64),
|
||||
"complex floating": dtype(complex128),
|
||||
"integral": dtype(intp),
|
||||
"indexing": dtype(intp),
|
||||
}
|
||||
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: None = None
|
||||
) -> DTypesAll: ...
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: L["bool"]
|
||||
) -> DTypesBool: ...
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: L["signed integer"]
|
||||
) -> DTypesSigned: ...
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: L["unsigned integer"]
|
||||
) -> DTypesUnsigned: ...
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: L["integral"]
|
||||
) -> DTypesIntegral: ...
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: L["real floating"]
|
||||
) -> DTypesReal: ...
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: L["complex floating"]
|
||||
) -> DTypesComplex: ...
|
||||
@overload
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: L["numeric"]
|
||||
) -> DTypesNumeric: ...
|
||||
def dtypes(
|
||||
self, /, *, device: _Device | None = None, kind: DTypeKind | None = None
|
||||
) -> DTypesAny:
|
||||
"""
|
||||
The array API data types supported by Dask.
|
||||
|
||||
Note that this function only returns data types that are defined by
|
||||
the array API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : str, optional
|
||||
The device to get the data types for.
|
||||
kind : str or tuple of str, optional
|
||||
The kind of data types to return. If ``None``, all data types are
|
||||
returned. If a string, only data types of that kind are returned.
|
||||
If a tuple, a dictionary containing the union of the given kinds
|
||||
is returned. The following kinds are supported:
|
||||
|
||||
- ``'bool'``: boolean data types (i.e., ``bool``).
|
||||
- ``'signed integer'``: signed integer data types (i.e., ``int8``,
|
||||
``int16``, ``int32``, ``int64``).
|
||||
- ``'unsigned integer'``: unsigned integer data types (i.e.,
|
||||
``uint8``, ``uint16``, ``uint32``, ``uint64``).
|
||||
- ``'integral'``: integer data types. Shorthand for ``('signed
|
||||
integer', 'unsigned integer')``.
|
||||
- ``'real floating'``: real-valued floating-point data types
|
||||
(i.e., ``float32``, ``float64``).
|
||||
- ``'complex floating'``: complex floating-point data types (i.e.,
|
||||
``complex64``, ``complex128``).
|
||||
- ``'numeric'``: numeric data types. Shorthand for ``('integral',
|
||||
'real floating', 'complex floating')``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary mapping the names of data types to the corresponding
|
||||
Dask data types.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.dtypes(kind='signed integer')
|
||||
{'int8': dask.int8,
|
||||
'int16': dask.int16,
|
||||
'int32': dask.int32,
|
||||
'int64': dask.int64}
|
||||
|
||||
"""
|
||||
if device not in ["cpu", _DASK_DEVICE, None]:
|
||||
raise ValueError(
|
||||
'Device not understood. Only "cpu" or _DASK_DEVICE is allowed, but received:'
|
||||
f" {device}"
|
||||
)
|
||||
if kind is None:
|
||||
return {
|
||||
"bool": dtype(bool),
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if kind == "bool":
|
||||
return {"bool": bool}
|
||||
if kind == "signed integer":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
}
|
||||
if kind == "unsigned integer":
|
||||
return {
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
}
|
||||
if kind == "integral":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
}
|
||||
if kind == "real floating":
|
||||
return {
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
}
|
||||
if kind == "complex floating":
|
||||
return {
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if kind == "numeric":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if isinstance(kind, tuple): # type: ignore[reportUnnecessaryIsinstanceCall]
|
||||
res: dict[str, DType] = {}
|
||||
for k in kind:
|
||||
res.update(self.dtypes(kind=k))
|
||||
return res
|
||||
raise ValueError(f"unsupported kind: {kind!r}")
|
||||
|
||||
def devices(self) -> list[_Device]:
|
||||
"""
|
||||
The devices supported by Dask.
|
||||
|
||||
For Dask, this always returns ``['cpu', DASK_DEVICE]``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
devices : list[Device]
|
||||
The devices supported by Dask.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.devices()
|
||||
['cpu', DASK_DEVICE]
|
||||
|
||||
"""
|
||||
return ["cpu", _DASK_DEVICE]
|
||||
@@ -0,0 +1,21 @@
|
||||
from dask.array.fft import * # noqa: F403
|
||||
# dask.array.fft doesn't have __all__. If it is added, replace this with
|
||||
#
|
||||
# from dask.array.fft import __all__ as linalg_all
|
||||
_n = {}
|
||||
exec('from dask.array.fft import *', _n)
|
||||
for k in ("__builtins__", "Sequence", "annotations", "warnings"):
|
||||
_n.pop(k, None)
|
||||
fft_all = list(_n)
|
||||
del _n, k
|
||||
|
||||
from ...common import _fft
|
||||
from ..._internal import get_xp
|
||||
|
||||
import dask.array as da
|
||||
|
||||
fftfreq = get_xp(da)(_fft.fftfreq)
|
||||
rfftfreq = get_xp(da)(_fft.rfftfreq)
|
||||
|
||||
__all__ = fft_all + ["fftfreq", "rfftfreq"]
|
||||
_all_ignore = ["da", "fft_all", "get_xp", "warnings"]
|
||||
@@ -0,0 +1,72 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
import dask.array as da
|
||||
|
||||
# The `matmul` and `tensordot` functions are in both the main and linalg namespaces
|
||||
from dask.array import matmul, outer, tensordot
|
||||
|
||||
# Exports
|
||||
from dask.array.linalg import * # noqa: F403
|
||||
|
||||
from ..._internal import get_xp
|
||||
from ...common import _linalg
|
||||
from ...common._typing import Array as _Array
|
||||
from ._aliases import matrix_transpose, vecdot
|
||||
|
||||
# dask.array.linalg doesn't have __all__. If it is added, replace this with
|
||||
#
|
||||
# from dask.array.linalg import __all__ as linalg_all
|
||||
_n = {}
|
||||
exec('from dask.array.linalg import *', _n)
|
||||
for k in ('__builtins__', 'annotations', 'operator', 'warnings', 'Array'):
|
||||
_n.pop(k, None)
|
||||
linalg_all = list(_n)
|
||||
del _n, k
|
||||
|
||||
EighResult = _linalg.EighResult
|
||||
QRResult = _linalg.QRResult
|
||||
SlogdetResult = _linalg.SlogdetResult
|
||||
SVDResult = _linalg.SVDResult
|
||||
# TODO: use the QR wrapper once dask
|
||||
# supports the mode keyword on QR
|
||||
# https://github.com/dask/dask/issues/10388
|
||||
#qr = get_xp(da)(_linalg.qr)
|
||||
def qr(
|
||||
x: _Array,
|
||||
mode: Literal["reduced", "complete"] = "reduced",
|
||||
**kwargs: object,
|
||||
) -> QRResult:
|
||||
if mode != "reduced":
|
||||
raise ValueError("dask arrays only support using mode='reduced'")
|
||||
return QRResult(*da.linalg.qr(x, **kwargs))
|
||||
trace = get_xp(da)(_linalg.trace)
|
||||
cholesky = get_xp(da)(_linalg.cholesky)
|
||||
matrix_rank = get_xp(da)(_linalg.matrix_rank)
|
||||
matrix_norm = get_xp(da)(_linalg.matrix_norm)
|
||||
|
||||
|
||||
# Wrap the svd functions to not pass full_matrices to dask
|
||||
# when full_matrices=False (as that is the default behavior for dask),
|
||||
# and dask doesn't have the full_matrices keyword
|
||||
def svd(x: _Array, full_matrices: bool = True, **kwargs) -> SVDResult:
|
||||
if full_matrices:
|
||||
raise ValueError("full_matrics=True is not supported by dask.")
|
||||
return da.linalg.svd(x, coerce_signs=False, **kwargs)
|
||||
|
||||
def svdvals(x: _Array) -> _Array:
|
||||
# TODO: can't avoid computing U or V for dask
|
||||
_, s, _ = svd(x)
|
||||
return s
|
||||
|
||||
vector_norm = get_xp(da)(_linalg.vector_norm)
|
||||
diagonal = get_xp(da)(_linalg.diagonal)
|
||||
|
||||
__all__ = linalg_all + ["trace", "outer", "matmul", "tensordot",
|
||||
"matrix_transpose", "vecdot", "EighResult",
|
||||
"QRResult", "SlogdetResult", "SVDResult", "qr",
|
||||
"cholesky", "matrix_rank", "matrix_norm", "svdvals",
|
||||
"vector_norm", "diagonal"]
|
||||
|
||||
_all_ignore = ['get_xp', 'da', 'linalg_all', 'warnings']
|
||||
@@ -0,0 +1,28 @@
|
||||
# ruff: noqa: PLC0414
|
||||
from typing import Final
|
||||
|
||||
from numpy import * # noqa: F403 # pyright: ignore[reportWildcardImportFromLibrary]
|
||||
|
||||
# from numpy import * doesn't overwrite these builtin names
|
||||
from numpy import abs as abs
|
||||
from numpy import max as max
|
||||
from numpy import min as min
|
||||
from numpy import round as round
|
||||
|
||||
# These imports may overwrite names from the import * above.
|
||||
from ._aliases import * # noqa: F403
|
||||
|
||||
# Don't know why, but we have to do an absolute import to import linalg. If we
|
||||
# instead do
|
||||
#
|
||||
# from . import linalg
|
||||
#
|
||||
# It doesn't overwrite np.linalg from above. The import is generated
|
||||
# dynamically so that the library can be vendored.
|
||||
__import__(__package__ + ".linalg")
|
||||
|
||||
__import__(__package__ + ".fft")
|
||||
|
||||
from .linalg import matrix_transpose, vecdot # type: ignore[no-redef] # noqa: F401
|
||||
|
||||
__array_api_version__: Final = "2024.12"
|
||||
@@ -0,0 +1,190 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
from __future__ import annotations
|
||||
|
||||
from builtins import bool as py_bool
|
||||
from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .._internal import get_xp
|
||||
from ..common import _aliases, _helpers
|
||||
from ..common._typing import NestedSequence, SupportsBufferProtocol
|
||||
from ._info import __array_namespace_info__
|
||||
from ._typing import Array, Device, DType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing_extensions import Buffer, TypeIs
|
||||
|
||||
# The values of the `_CopyMode` enum can be either `False`, `True`, or `2`:
|
||||
# https://github.com/numpy/numpy/blob/5a8a6a79d9c2fff8f07dcab5d41e14f8508d673f/numpy/_globals.pyi#L7-L10
|
||||
_Copy: TypeAlias = py_bool | Literal[2] | np._CopyMode
|
||||
|
||||
bool = np.bool_
|
||||
|
||||
# Basic renames
|
||||
acos = np.arccos
|
||||
acosh = np.arccosh
|
||||
asin = np.arcsin
|
||||
asinh = np.arcsinh
|
||||
atan = np.arctan
|
||||
atan2 = np.arctan2
|
||||
atanh = np.arctanh
|
||||
bitwise_left_shift = np.left_shift
|
||||
bitwise_invert = np.invert
|
||||
bitwise_right_shift = np.right_shift
|
||||
concat = np.concatenate
|
||||
pow = np.power
|
||||
|
||||
arange = get_xp(np)(_aliases.arange)
|
||||
empty = get_xp(np)(_aliases.empty)
|
||||
empty_like = get_xp(np)(_aliases.empty_like)
|
||||
eye = get_xp(np)(_aliases.eye)
|
||||
full = get_xp(np)(_aliases.full)
|
||||
full_like = get_xp(np)(_aliases.full_like)
|
||||
linspace = get_xp(np)(_aliases.linspace)
|
||||
ones = get_xp(np)(_aliases.ones)
|
||||
ones_like = get_xp(np)(_aliases.ones_like)
|
||||
zeros = get_xp(np)(_aliases.zeros)
|
||||
zeros_like = get_xp(np)(_aliases.zeros_like)
|
||||
UniqueAllResult = get_xp(np)(_aliases.UniqueAllResult)
|
||||
UniqueCountsResult = get_xp(np)(_aliases.UniqueCountsResult)
|
||||
UniqueInverseResult = get_xp(np)(_aliases.UniqueInverseResult)
|
||||
unique_all = get_xp(np)(_aliases.unique_all)
|
||||
unique_counts = get_xp(np)(_aliases.unique_counts)
|
||||
unique_inverse = get_xp(np)(_aliases.unique_inverse)
|
||||
unique_values = get_xp(np)(_aliases.unique_values)
|
||||
std = get_xp(np)(_aliases.std)
|
||||
var = get_xp(np)(_aliases.var)
|
||||
cumulative_sum = get_xp(np)(_aliases.cumulative_sum)
|
||||
cumulative_prod = get_xp(np)(_aliases.cumulative_prod)
|
||||
clip = get_xp(np)(_aliases.clip)
|
||||
permute_dims = get_xp(np)(_aliases.permute_dims)
|
||||
reshape = get_xp(np)(_aliases.reshape)
|
||||
argsort = get_xp(np)(_aliases.argsort)
|
||||
sort = get_xp(np)(_aliases.sort)
|
||||
nonzero = get_xp(np)(_aliases.nonzero)
|
||||
ceil = get_xp(np)(_aliases.ceil)
|
||||
floor = get_xp(np)(_aliases.floor)
|
||||
trunc = get_xp(np)(_aliases.trunc)
|
||||
matmul = get_xp(np)(_aliases.matmul)
|
||||
matrix_transpose = get_xp(np)(_aliases.matrix_transpose)
|
||||
tensordot = get_xp(np)(_aliases.tensordot)
|
||||
sign = get_xp(np)(_aliases.sign)
|
||||
finfo = get_xp(np)(_aliases.finfo)
|
||||
iinfo = get_xp(np)(_aliases.iinfo)
|
||||
|
||||
|
||||
def _supports_buffer_protocol(obj: object) -> TypeIs[Buffer]: # pyright: ignore[reportUnusedFunction]
|
||||
try:
|
||||
memoryview(obj) # pyright: ignore[reportArgumentType]
|
||||
except TypeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
# asarray also adds the copy keyword, which is not present in numpy 1.0.
|
||||
# asarray() is different enough between numpy, cupy, and dask, the logic
|
||||
# complicated enough that it's easier to define it separately for each module
|
||||
# rather than trying to combine everything into one function in common/
|
||||
def asarray(
|
||||
obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol,
|
||||
/,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
copy: _Copy | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Array:
|
||||
"""
|
||||
Array API compatibility wrapper for asarray().
|
||||
|
||||
See the corresponding documentation in the array library and/or the array API
|
||||
specification for more details.
|
||||
"""
|
||||
_helpers._check_device(np, device)
|
||||
|
||||
if copy is None:
|
||||
copy = np._CopyMode.IF_NEEDED
|
||||
elif copy is False:
|
||||
copy = np._CopyMode.NEVER
|
||||
elif copy is True:
|
||||
copy = np._CopyMode.ALWAYS
|
||||
|
||||
return np.array(obj, copy=copy, dtype=dtype, **kwargs) # pyright: ignore
|
||||
|
||||
|
||||
def astype(
|
||||
x: Array,
|
||||
dtype: DType,
|
||||
/,
|
||||
*,
|
||||
copy: py_bool = True,
|
||||
device: Device | None = None,
|
||||
) -> Array:
|
||||
_helpers._check_device(np, device)
|
||||
return x.astype(dtype=dtype, copy=copy)
|
||||
|
||||
|
||||
# count_nonzero returns a python int for axis=None and keepdims=False
|
||||
# https://github.com/numpy/numpy/issues/17562
|
||||
def count_nonzero(
|
||||
x: Array,
|
||||
axis: int | tuple[int, ...] | None = None,
|
||||
keepdims: py_bool = False,
|
||||
) -> Array:
|
||||
# NOTE: this is currently incorrectly typed in numpy, but will be fixed in
|
||||
# numpy 2.2.5 and 2.3.0: https://github.com/numpy/numpy/pull/28750
|
||||
result = cast("Any", np.count_nonzero(x, axis=axis, keepdims=keepdims)) # pyright: ignore[reportArgumentType, reportCallIssue]
|
||||
if axis is None and not keepdims:
|
||||
return np.asarray(result)
|
||||
return result
|
||||
|
||||
|
||||
# take_along_axis: axis defaults to -1 but in numpy axis is a required arg
|
||||
def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1):
|
||||
return np.take_along_axis(x, indices, axis=axis)
|
||||
|
||||
|
||||
# These functions are completely new here. If the library already has them
|
||||
# (i.e., numpy 2.0), use the library version instead of our wrapper.
|
||||
if hasattr(np, "vecdot"):
|
||||
vecdot = np.vecdot
|
||||
else:
|
||||
vecdot = get_xp(np)(_aliases.vecdot)
|
||||
|
||||
if hasattr(np, "isdtype"):
|
||||
isdtype = np.isdtype
|
||||
else:
|
||||
isdtype = get_xp(np)(_aliases.isdtype)
|
||||
|
||||
if hasattr(np, "unstack"):
|
||||
unstack = np.unstack
|
||||
else:
|
||||
unstack = get_xp(np)(_aliases.unstack)
|
||||
|
||||
__all__ = [
|
||||
"__array_namespace_info__",
|
||||
"asarray",
|
||||
"astype",
|
||||
"acos",
|
||||
"acosh",
|
||||
"asin",
|
||||
"asinh",
|
||||
"atan",
|
||||
"atan2",
|
||||
"atanh",
|
||||
"bitwise_left_shift",
|
||||
"bitwise_invert",
|
||||
"bitwise_right_shift",
|
||||
"bool",
|
||||
"concat",
|
||||
"count_nonzero",
|
||||
"pow",
|
||||
"take_along_axis"
|
||||
]
|
||||
__all__ += _aliases.__all__
|
||||
_all_ignore = ["np", "get_xp"]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,366 @@
|
||||
"""
|
||||
Array API Inspection namespace
|
||||
|
||||
This is the namespace for inspection functions as defined by the array API
|
||||
standard. See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html for
|
||||
more details.
|
||||
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from numpy import bool_ as bool
|
||||
from numpy import (
|
||||
complex64,
|
||||
complex128,
|
||||
dtype,
|
||||
float32,
|
||||
float64,
|
||||
int8,
|
||||
int16,
|
||||
int32,
|
||||
int64,
|
||||
intp,
|
||||
uint8,
|
||||
uint16,
|
||||
uint32,
|
||||
uint64,
|
||||
)
|
||||
|
||||
from ._typing import Device, DType
|
||||
|
||||
|
||||
class __array_namespace_info__:
|
||||
"""
|
||||
Get the array API inspection namespace for NumPy.
|
||||
|
||||
The array API inspection namespace defines the following functions:
|
||||
|
||||
- capabilities()
|
||||
- default_device()
|
||||
- default_dtypes()
|
||||
- dtypes()
|
||||
- devices()
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html
|
||||
for more details.
|
||||
|
||||
Returns
|
||||
-------
|
||||
info : ModuleType
|
||||
The array API inspection namespace for NumPy.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = np.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': numpy.float64,
|
||||
'complex floating': numpy.complex128,
|
||||
'integral': numpy.int64,
|
||||
'indexing': numpy.int64}
|
||||
|
||||
"""
|
||||
|
||||
__module__ = 'numpy'
|
||||
|
||||
def capabilities(self):
|
||||
"""
|
||||
Return a dictionary of array API library capabilities.
|
||||
|
||||
The resulting dictionary has the following keys:
|
||||
|
||||
- **"boolean indexing"**: boolean indicating whether an array library
|
||||
supports boolean indexing. Always ``True`` for NumPy.
|
||||
|
||||
- **"data-dependent shapes"**: boolean indicating whether an array
|
||||
library supports data-dependent output shapes. Always ``True`` for
|
||||
NumPy.
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html
|
||||
for more details.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
capabilities : dict
|
||||
A dictionary of array API library capabilities.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = np.__array_namespace_info__()
|
||||
>>> info.capabilities()
|
||||
{'boolean indexing': True,
|
||||
'data-dependent shapes': True,
|
||||
'max dimensions': 64}
|
||||
|
||||
"""
|
||||
return {
|
||||
"boolean indexing": True,
|
||||
"data-dependent shapes": True,
|
||||
"max dimensions": 64,
|
||||
}
|
||||
|
||||
def default_device(self):
|
||||
"""
|
||||
The default device used for new NumPy arrays.
|
||||
|
||||
For NumPy, this always returns ``'cpu'``.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
device : Device
|
||||
The default device used for new NumPy arrays.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = np.__array_namespace_info__()
|
||||
>>> info.default_device()
|
||||
'cpu'
|
||||
|
||||
"""
|
||||
return "cpu"
|
||||
|
||||
def default_dtypes(
|
||||
self,
|
||||
*,
|
||||
device: Device | None = None,
|
||||
) -> dict[str, dtype[intp | float64 | complex128]]:
|
||||
"""
|
||||
The default data types used for new NumPy arrays.
|
||||
|
||||
For NumPy, this always returns the following dictionary:
|
||||
|
||||
- **"real floating"**: ``numpy.float64``
|
||||
- **"complex floating"**: ``numpy.complex128``
|
||||
- **"integral"**: ``numpy.intp``
|
||||
- **"indexing"**: ``numpy.intp``
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : str, optional
|
||||
The device to get the default data types for. For NumPy, only
|
||||
``'cpu'`` is allowed.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary describing the default data types used for new NumPy
|
||||
arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = np.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': numpy.float64,
|
||||
'complex floating': numpy.complex128,
|
||||
'integral': numpy.int64,
|
||||
'indexing': numpy.int64}
|
||||
|
||||
"""
|
||||
if device not in ["cpu", None]:
|
||||
raise ValueError(
|
||||
'Device not understood. Only "cpu" is allowed, but received:'
|
||||
f' {device}'
|
||||
)
|
||||
return {
|
||||
"real floating": dtype(float64),
|
||||
"complex floating": dtype(complex128),
|
||||
"integral": dtype(intp),
|
||||
"indexing": dtype(intp),
|
||||
}
|
||||
|
||||
def dtypes(
|
||||
self,
|
||||
*,
|
||||
device: Device | None = None,
|
||||
kind: str | tuple[str, ...] | None = None,
|
||||
) -> dict[str, DType]:
|
||||
"""
|
||||
The array API data types supported by NumPy.
|
||||
|
||||
Note that this function only returns data types that are defined by
|
||||
the array API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : str, optional
|
||||
The device to get the data types for. For NumPy, only ``'cpu'`` is
|
||||
allowed.
|
||||
kind : str or tuple of str, optional
|
||||
The kind of data types to return. If ``None``, all data types are
|
||||
returned. If a string, only data types of that kind are returned.
|
||||
If a tuple, a dictionary containing the union of the given kinds
|
||||
is returned. The following kinds are supported:
|
||||
|
||||
- ``'bool'``: boolean data types (i.e., ``bool``).
|
||||
- ``'signed integer'``: signed integer data types (i.e., ``int8``,
|
||||
``int16``, ``int32``, ``int64``).
|
||||
- ``'unsigned integer'``: unsigned integer data types (i.e.,
|
||||
``uint8``, ``uint16``, ``uint32``, ``uint64``).
|
||||
- ``'integral'``: integer data types. Shorthand for ``('signed
|
||||
integer', 'unsigned integer')``.
|
||||
- ``'real floating'``: real-valued floating-point data types
|
||||
(i.e., ``float32``, ``float64``).
|
||||
- ``'complex floating'``: complex floating-point data types (i.e.,
|
||||
``complex64``, ``complex128``).
|
||||
- ``'numeric'``: numeric data types. Shorthand for ``('integral',
|
||||
'real floating', 'complex floating')``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary mapping the names of data types to the corresponding
|
||||
NumPy data types.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = np.__array_namespace_info__()
|
||||
>>> info.dtypes(kind='signed integer')
|
||||
{'int8': numpy.int8,
|
||||
'int16': numpy.int16,
|
||||
'int32': numpy.int32,
|
||||
'int64': numpy.int64}
|
||||
|
||||
"""
|
||||
if device not in ["cpu", None]:
|
||||
raise ValueError(
|
||||
'Device not understood. Only "cpu" is allowed, but received:'
|
||||
f' {device}'
|
||||
)
|
||||
if kind is None:
|
||||
return {
|
||||
"bool": dtype(bool),
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if kind == "bool":
|
||||
return {"bool": dtype(bool)}
|
||||
if kind == "signed integer":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
}
|
||||
if kind == "unsigned integer":
|
||||
return {
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
}
|
||||
if kind == "integral":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
}
|
||||
if kind == "real floating":
|
||||
return {
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
}
|
||||
if kind == "complex floating":
|
||||
return {
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if kind == "numeric":
|
||||
return {
|
||||
"int8": dtype(int8),
|
||||
"int16": dtype(int16),
|
||||
"int32": dtype(int32),
|
||||
"int64": dtype(int64),
|
||||
"uint8": dtype(uint8),
|
||||
"uint16": dtype(uint16),
|
||||
"uint32": dtype(uint32),
|
||||
"uint64": dtype(uint64),
|
||||
"float32": dtype(float32),
|
||||
"float64": dtype(float64),
|
||||
"complex64": dtype(complex64),
|
||||
"complex128": dtype(complex128),
|
||||
}
|
||||
if isinstance(kind, tuple):
|
||||
res: dict[str, DType] = {}
|
||||
for k in kind:
|
||||
res.update(self.dtypes(kind=k))
|
||||
return res
|
||||
raise ValueError(f"unsupported kind: {kind!r}")
|
||||
|
||||
def devices(self) -> list[Device]:
|
||||
"""
|
||||
The devices supported by NumPy.
|
||||
|
||||
For NumPy, this always returns ``['cpu']``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
devices : list[Device]
|
||||
The devices supported by NumPy.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = np.__array_namespace_info__()
|
||||
>>> info.devices()
|
||||
['cpu']
|
||||
|
||||
"""
|
||||
return ["cpu"]
|
||||
|
||||
|
||||
__all__ = ["__array_namespace_info__"]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Literal, TypeAlias
|
||||
|
||||
import numpy as np
|
||||
|
||||
Device: TypeAlias = Literal["cpu"]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
# NumPy 1.x on Python 3.10 fails to parse np.dtype[]
|
||||
DType: TypeAlias = np.dtype[
|
||||
np.bool_
|
||||
| np.integer[Any]
|
||||
| np.float32
|
||||
| np.float64
|
||||
| np.complex64
|
||||
| np.complex128
|
||||
]
|
||||
Array: TypeAlias = np.ndarray[Any, DType]
|
||||
else:
|
||||
DType: TypeAlias = np.dtype
|
||||
Array: TypeAlias = np.ndarray
|
||||
|
||||
__all__ = ["Array", "DType", "Device"]
|
||||
_all_ignore = ["np"]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,35 @@
|
||||
import numpy as np
|
||||
from numpy.fft import __all__ as fft_all
|
||||
from numpy.fft import fft2, ifft2, irfft2, rfft2
|
||||
|
||||
from .._internal import get_xp
|
||||
from ..common import _fft
|
||||
|
||||
fft = get_xp(np)(_fft.fft)
|
||||
ifft = get_xp(np)(_fft.ifft)
|
||||
fftn = get_xp(np)(_fft.fftn)
|
||||
ifftn = get_xp(np)(_fft.ifftn)
|
||||
rfft = get_xp(np)(_fft.rfft)
|
||||
irfft = get_xp(np)(_fft.irfft)
|
||||
rfftn = get_xp(np)(_fft.rfftn)
|
||||
irfftn = get_xp(np)(_fft.irfftn)
|
||||
hfft = get_xp(np)(_fft.hfft)
|
||||
ihfft = get_xp(np)(_fft.ihfft)
|
||||
fftfreq = get_xp(np)(_fft.fftfreq)
|
||||
rfftfreq = get_xp(np)(_fft.rfftfreq)
|
||||
fftshift = get_xp(np)(_fft.fftshift)
|
||||
ifftshift = get_xp(np)(_fft.ifftshift)
|
||||
|
||||
|
||||
__all__ = ["rfft2", "irfft2", "fft2", "ifft2"]
|
||||
__all__ += _fft.__all__
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
|
||||
|
||||
del get_xp
|
||||
del np
|
||||
del fft_all
|
||||
del _fft
|
||||
@@ -0,0 +1,143 @@
|
||||
# pyright: reportAttributeAccessIssue=false
|
||||
# pyright: reportUnknownArgumentType=false
|
||||
# pyright: reportUnknownMemberType=false
|
||||
# pyright: reportUnknownVariableType=false
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
# intersection of `np.linalg.__all__` on numpy 1.22 and 2.2, minus `_linalg.__all__`
|
||||
from numpy.linalg import (
|
||||
LinAlgError,
|
||||
cond,
|
||||
det,
|
||||
eig,
|
||||
eigvals,
|
||||
eigvalsh,
|
||||
inv,
|
||||
lstsq,
|
||||
matrix_power,
|
||||
multi_dot,
|
||||
norm,
|
||||
tensorinv,
|
||||
tensorsolve,
|
||||
)
|
||||
|
||||
from .._internal import get_xp
|
||||
from ..common import _linalg
|
||||
|
||||
# These functions are in both the main and linalg namespaces
|
||||
from ._aliases import matmul, matrix_transpose, tensordot, vecdot # noqa: F401
|
||||
from ._typing import Array
|
||||
|
||||
cross = get_xp(np)(_linalg.cross)
|
||||
outer = get_xp(np)(_linalg.outer)
|
||||
EighResult = _linalg.EighResult
|
||||
QRResult = _linalg.QRResult
|
||||
SlogdetResult = _linalg.SlogdetResult
|
||||
SVDResult = _linalg.SVDResult
|
||||
eigh = get_xp(np)(_linalg.eigh)
|
||||
qr = get_xp(np)(_linalg.qr)
|
||||
slogdet = get_xp(np)(_linalg.slogdet)
|
||||
svd = get_xp(np)(_linalg.svd)
|
||||
cholesky = get_xp(np)(_linalg.cholesky)
|
||||
matrix_rank = get_xp(np)(_linalg.matrix_rank)
|
||||
pinv = get_xp(np)(_linalg.pinv)
|
||||
matrix_norm = get_xp(np)(_linalg.matrix_norm)
|
||||
svdvals = get_xp(np)(_linalg.svdvals)
|
||||
diagonal = get_xp(np)(_linalg.diagonal)
|
||||
trace = get_xp(np)(_linalg.trace)
|
||||
|
||||
# Note: unlike np.linalg.solve, the array API solve() only accepts x2 as a
|
||||
# vector when it is exactly 1-dimensional. All other cases treat x2 as a stack
|
||||
# of matrices. The np.linalg.solve behavior of allowing stacks of both
|
||||
# matrices and vectors is ambiguous c.f.
|
||||
# https://github.com/numpy/numpy/issues/15349 and
|
||||
# https://github.com/data-apis/array-api/issues/285.
|
||||
|
||||
# To workaround this, the below is the code from np.linalg.solve except
|
||||
# only calling solve1 in the exactly 1D case.
|
||||
|
||||
|
||||
# This code is here instead of in common because it is numpy specific. Also
|
||||
# note that CuPy's solve() does not currently support broadcasting (see
|
||||
# https://github.com/cupy/cupy/blob/main/cupy/cublas.py#L43).
|
||||
def solve(x1: Array, x2: Array, /) -> Array:
|
||||
try:
|
||||
from numpy.linalg._linalg import (
|
||||
_assert_stacked_2d,
|
||||
_assert_stacked_square,
|
||||
_commonType,
|
||||
_makearray,
|
||||
_raise_linalgerror_singular,
|
||||
isComplexType,
|
||||
)
|
||||
except ImportError:
|
||||
from numpy.linalg.linalg import (
|
||||
_assert_stacked_2d,
|
||||
_assert_stacked_square,
|
||||
_commonType,
|
||||
_makearray,
|
||||
_raise_linalgerror_singular,
|
||||
isComplexType,
|
||||
)
|
||||
from numpy.linalg import _umath_linalg
|
||||
|
||||
x1, _ = _makearray(x1)
|
||||
_assert_stacked_2d(x1)
|
||||
_assert_stacked_square(x1)
|
||||
x2, wrap = _makearray(x2)
|
||||
t, result_t = _commonType(x1, x2)
|
||||
|
||||
# This part is different from np.linalg.solve
|
||||
gufunc: np.ufunc
|
||||
if x2.ndim == 1:
|
||||
gufunc = _umath_linalg.solve1
|
||||
else:
|
||||
gufunc = _umath_linalg.solve
|
||||
|
||||
# This does nothing currently but is left in because it will be relevant
|
||||
# when complex dtype support is added to the spec in 2022.
|
||||
signature = "DD->D" if isComplexType(t) else "dd->d"
|
||||
with np.errstate(
|
||||
call=_raise_linalgerror_singular,
|
||||
invalid="call",
|
||||
over="ignore",
|
||||
divide="ignore",
|
||||
under="ignore",
|
||||
):
|
||||
r: Array = gufunc(x1, x2, signature=signature)
|
||||
|
||||
return wrap(r.astype(result_t, copy=False))
|
||||
|
||||
|
||||
# These functions are completely new here. If the library already has them
|
||||
# (i.e., numpy 2.0), use the library version instead of our wrapper.
|
||||
if hasattr(np.linalg, "vector_norm"):
|
||||
vector_norm = np.linalg.vector_norm
|
||||
else:
|
||||
vector_norm = get_xp(np)(_linalg.vector_norm)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LinAlgError",
|
||||
"cond",
|
||||
"det",
|
||||
"eig",
|
||||
"eigvals",
|
||||
"eigvalsh",
|
||||
"inv",
|
||||
"lstsq",
|
||||
"matrix_power",
|
||||
"multi_dot",
|
||||
"norm",
|
||||
"tensorinv",
|
||||
"tensorsolve",
|
||||
]
|
||||
__all__ += _linalg.__all__
|
||||
__all__ += ["solve", "vector_norm"]
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,22 @@
|
||||
from torch import * # noqa: F403
|
||||
|
||||
# Several names are not included in the above import *
|
||||
import torch
|
||||
for n in dir(torch):
|
||||
if (n.startswith('_')
|
||||
or n.endswith('_')
|
||||
or 'cuda' in n
|
||||
or 'cpu' in n
|
||||
or 'backward' in n):
|
||||
continue
|
||||
exec(f"{n} = torch.{n}")
|
||||
del n
|
||||
|
||||
# These imports may overwrite names from the import * above.
|
||||
from ._aliases import * # noqa: F403
|
||||
|
||||
# See the comment in the numpy __init__.py
|
||||
__import__(__package__ + '.linalg')
|
||||
__import__(__package__ + '.fft')
|
||||
|
||||
__array_api_version__ = '2024.12'
|
||||
@@ -0,0 +1,855 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import reduce as _reduce, wraps as _wraps
|
||||
from builtins import all as _builtin_all, any as _builtin_any
|
||||
from typing import Any, List, Optional, Sequence, Tuple, Union, Literal
|
||||
|
||||
import torch
|
||||
|
||||
from .._internal import get_xp
|
||||
from ..common import _aliases
|
||||
from ..common._typing import NestedSequence, SupportsBufferProtocol
|
||||
from ._info import __array_namespace_info__
|
||||
from ._typing import Array, Device, DType
|
||||
|
||||
_int_dtypes = {
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
}
|
||||
try:
|
||||
# torch >=2.3
|
||||
_int_dtypes |= {torch.uint16, torch.uint32, torch.uint64}
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
|
||||
_array_api_dtypes = {
|
||||
torch.bool,
|
||||
*_int_dtypes,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
torch.complex64,
|
||||
torch.complex128,
|
||||
}
|
||||
|
||||
_promotion_table = {
|
||||
# ints
|
||||
(torch.int8, torch.int16): torch.int16,
|
||||
(torch.int8, torch.int32): torch.int32,
|
||||
(torch.int8, torch.int64): torch.int64,
|
||||
(torch.int16, torch.int32): torch.int32,
|
||||
(torch.int16, torch.int64): torch.int64,
|
||||
(torch.int32, torch.int64): torch.int64,
|
||||
# ints and uints (mixed sign)
|
||||
(torch.uint8, torch.int8): torch.int16,
|
||||
(torch.uint8, torch.int16): torch.int16,
|
||||
(torch.uint8, torch.int32): torch.int32,
|
||||
(torch.uint8, torch.int64): torch.int64,
|
||||
# floats
|
||||
(torch.float32, torch.float64): torch.float64,
|
||||
# complexes
|
||||
(torch.complex64, torch.complex128): torch.complex128,
|
||||
# Mixed float and complex
|
||||
(torch.float32, torch.complex64): torch.complex64,
|
||||
(torch.float32, torch.complex128): torch.complex128,
|
||||
(torch.float64, torch.complex64): torch.complex128,
|
||||
(torch.float64, torch.complex128): torch.complex128,
|
||||
}
|
||||
|
||||
_promotion_table.update({(b, a): c for (a, b), c in _promotion_table.items()})
|
||||
_promotion_table.update({(a, a): a for a in _array_api_dtypes})
|
||||
|
||||
|
||||
def _two_arg(f):
|
||||
@_wraps(f)
|
||||
def _f(x1, x2, /, **kwargs):
|
||||
x1, x2 = _fix_promotion(x1, x2)
|
||||
return f(x1, x2, **kwargs)
|
||||
if _f.__doc__ is None:
|
||||
_f.__doc__ = f"""\
|
||||
Array API compatibility wrapper for torch.{f.__name__}.
|
||||
|
||||
See the corresponding PyTorch documentation and/or the array API specification
|
||||
for more details.
|
||||
|
||||
"""
|
||||
return _f
|
||||
|
||||
def _fix_promotion(x1, x2, only_scalar=True):
|
||||
if not isinstance(x1, torch.Tensor) or not isinstance(x2, torch.Tensor):
|
||||
return x1, x2
|
||||
if x1.dtype not in _array_api_dtypes or x2.dtype not in _array_api_dtypes:
|
||||
return x1, x2
|
||||
# If an argument is 0-D pytorch downcasts the other argument
|
||||
if not only_scalar or x1.shape == ():
|
||||
dtype = result_type(x1, x2)
|
||||
x2 = x2.to(dtype)
|
||||
if not only_scalar or x2.shape == ():
|
||||
dtype = result_type(x1, x2)
|
||||
x1 = x1.to(dtype)
|
||||
return x1, x2
|
||||
|
||||
|
||||
_py_scalars = (bool, int, float, complex)
|
||||
|
||||
|
||||
def result_type(
|
||||
*arrays_and_dtypes: Array | DType | bool | int | float | complex
|
||||
) -> DType:
|
||||
num = len(arrays_and_dtypes)
|
||||
|
||||
if num == 0:
|
||||
raise ValueError("At least one array or dtype must be provided")
|
||||
|
||||
elif num == 1:
|
||||
x = arrays_and_dtypes[0]
|
||||
if isinstance(x, torch.dtype):
|
||||
return x
|
||||
return x.dtype
|
||||
|
||||
if num == 2:
|
||||
x, y = arrays_and_dtypes
|
||||
return _result_type(x, y)
|
||||
|
||||
else:
|
||||
# sort scalars so that they are treated last
|
||||
scalars, others = [], []
|
||||
for x in arrays_and_dtypes:
|
||||
if isinstance(x, _py_scalars):
|
||||
scalars.append(x)
|
||||
else:
|
||||
others.append(x)
|
||||
if not others:
|
||||
raise ValueError("At least one array or dtype must be provided")
|
||||
|
||||
# combine left-to-right
|
||||
return _reduce(_result_type, others + scalars)
|
||||
|
||||
|
||||
def _result_type(
|
||||
x: Array | DType | bool | int | float | complex,
|
||||
y: Array | DType | bool | int | float | complex,
|
||||
) -> DType:
|
||||
if not (isinstance(x, _py_scalars) or isinstance(y, _py_scalars)):
|
||||
xdt = x if isinstance(x, torch.dtype) else x.dtype
|
||||
ydt = y if isinstance(y, torch.dtype) else y.dtype
|
||||
|
||||
try:
|
||||
return _promotion_table[xdt, ydt]
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
# This doesn't result_type(dtype, dtype) for non-array API dtypes
|
||||
# because torch.result_type only accepts tensors. This does however, allow
|
||||
# cross-kind promotion.
|
||||
x = torch.tensor([], dtype=x) if isinstance(x, torch.dtype) else x
|
||||
y = torch.tensor([], dtype=y) if isinstance(y, torch.dtype) else y
|
||||
return torch.result_type(x, y)
|
||||
|
||||
|
||||
def can_cast(from_: Union[DType, Array], to: DType, /) -> bool:
|
||||
if not isinstance(from_, torch.dtype):
|
||||
from_ = from_.dtype
|
||||
return torch.can_cast(from_, to)
|
||||
|
||||
# Basic renames
|
||||
bitwise_invert = torch.bitwise_not
|
||||
newaxis = None
|
||||
# torch.conj sets the conjugation bit, which breaks conversion to other
|
||||
# libraries. See https://github.com/data-apis/array-api-compat/issues/173
|
||||
conj = torch.conj_physical
|
||||
|
||||
# Two-arg elementwise functions
|
||||
# These require a wrapper to do the correct type promotion on 0-D tensors
|
||||
add = _two_arg(torch.add)
|
||||
atan2 = _two_arg(torch.atan2)
|
||||
bitwise_and = _two_arg(torch.bitwise_and)
|
||||
bitwise_left_shift = _two_arg(torch.bitwise_left_shift)
|
||||
bitwise_or = _two_arg(torch.bitwise_or)
|
||||
bitwise_right_shift = _two_arg(torch.bitwise_right_shift)
|
||||
bitwise_xor = _two_arg(torch.bitwise_xor)
|
||||
copysign = _two_arg(torch.copysign)
|
||||
divide = _two_arg(torch.divide)
|
||||
# Also a rename. torch.equal does not broadcast
|
||||
equal = _two_arg(torch.eq)
|
||||
floor_divide = _two_arg(torch.floor_divide)
|
||||
greater = _two_arg(torch.greater)
|
||||
greater_equal = _two_arg(torch.greater_equal)
|
||||
hypot = _two_arg(torch.hypot)
|
||||
less = _two_arg(torch.less)
|
||||
less_equal = _two_arg(torch.less_equal)
|
||||
logaddexp = _two_arg(torch.logaddexp)
|
||||
# logical functions are not included here because they only accept bool in the
|
||||
# spec, so type promotion is irrelevant.
|
||||
maximum = _two_arg(torch.maximum)
|
||||
minimum = _two_arg(torch.minimum)
|
||||
multiply = _two_arg(torch.multiply)
|
||||
not_equal = _two_arg(torch.not_equal)
|
||||
pow = _two_arg(torch.pow)
|
||||
remainder = _two_arg(torch.remainder)
|
||||
subtract = _two_arg(torch.subtract)
|
||||
|
||||
|
||||
def asarray(
|
||||
obj: (
|
||||
Array
|
||||
| bool | int | float | complex
|
||||
| NestedSequence[bool | int | float | complex]
|
||||
| SupportsBufferProtocol
|
||||
),
|
||||
/,
|
||||
*,
|
||||
dtype: DType | None = None,
|
||||
device: Device | None = None,
|
||||
copy: bool | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Array:
|
||||
# torch.asarray does not respect input->output device propagation
|
||||
# https://github.com/pytorch/pytorch/issues/150199
|
||||
if device is None and isinstance(obj, torch.Tensor):
|
||||
device = obj.device
|
||||
return torch.asarray(obj, dtype=dtype, device=device, copy=copy, **kwargs)
|
||||
|
||||
|
||||
# These wrappers are mostly based on the fact that pytorch uses 'dim' instead
|
||||
# of 'axis'.
|
||||
|
||||
# torch.min and torch.max return a tuple and don't support multiple axes https://github.com/pytorch/pytorch/issues/58745
|
||||
def max(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> Array:
|
||||
# https://github.com/pytorch/pytorch/issues/29137
|
||||
if axis == ():
|
||||
return torch.clone(x)
|
||||
return torch.amax(x, axis, keepdims=keepdims)
|
||||
|
||||
def min(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> Array:
|
||||
# https://github.com/pytorch/pytorch/issues/29137
|
||||
if axis == ():
|
||||
return torch.clone(x)
|
||||
return torch.amin(x, axis, keepdims=keepdims)
|
||||
|
||||
clip = get_xp(torch)(_aliases.clip)
|
||||
unstack = get_xp(torch)(_aliases.unstack)
|
||||
cumulative_sum = get_xp(torch)(_aliases.cumulative_sum)
|
||||
cumulative_prod = get_xp(torch)(_aliases.cumulative_prod)
|
||||
finfo = get_xp(torch)(_aliases.finfo)
|
||||
iinfo = get_xp(torch)(_aliases.iinfo)
|
||||
|
||||
|
||||
# torch.sort also returns a tuple
|
||||
# https://github.com/pytorch/pytorch/issues/70921
|
||||
def sort(x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True, **kwargs) -> Array:
|
||||
return torch.sort(x, dim=axis, descending=descending, stable=stable, **kwargs).values
|
||||
|
||||
def _normalize_axes(axis, ndim):
|
||||
axes = []
|
||||
if ndim == 0 and axis:
|
||||
# Better error message in this case
|
||||
raise IndexError(f"Dimension out of range: {axis[0]}")
|
||||
lower, upper = -ndim, ndim - 1
|
||||
for a in axis:
|
||||
if a < lower or a > upper:
|
||||
# Match torch error message (e.g., from sum())
|
||||
raise IndexError(f"Dimension out of range (expected to be in range of [{lower}, {upper}], but got {a}")
|
||||
if a < 0:
|
||||
a = a + ndim
|
||||
if a in axes:
|
||||
# Use IndexError instead of RuntimeError, and "axis" instead of "dim"
|
||||
raise IndexError(f"Axis {a} appears multiple times in the list of axes")
|
||||
axes.append(a)
|
||||
return sorted(axes)
|
||||
|
||||
def _axis_none_keepdims(x, ndim, keepdims):
|
||||
# Apply keepdims when axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
# Note that this is only valid for the axis=None case.
|
||||
if keepdims:
|
||||
for i in range(ndim):
|
||||
x = torch.unsqueeze(x, 0)
|
||||
return x
|
||||
|
||||
def _reduce_multiple_axes(f, x, axis, keepdims=False, **kwargs):
|
||||
# Some reductions don't support multiple axes
|
||||
# (https://github.com/pytorch/pytorch/issues/56586).
|
||||
axes = _normalize_axes(axis, x.ndim)
|
||||
for a in reversed(axes):
|
||||
x = torch.movedim(x, a, -1)
|
||||
x = torch.flatten(x, -len(axes))
|
||||
|
||||
out = f(x, -1, **kwargs)
|
||||
|
||||
if keepdims:
|
||||
for a in axes:
|
||||
out = torch.unsqueeze(out, a)
|
||||
return out
|
||||
|
||||
|
||||
def _sum_prod_no_axis(x: Array, dtype: DType | None) -> Array:
|
||||
"""
|
||||
Implements `sum(..., axis=())` and `prod(..., axis=())`.
|
||||
|
||||
Works around https://github.com/pytorch/pytorch/issues/29137
|
||||
"""
|
||||
if dtype is not None:
|
||||
return x.clone() if dtype == x.dtype else x.to(dtype)
|
||||
|
||||
# We can't upcast uint8 according to the spec because there is no
|
||||
# torch.uint64, so at least upcast to int64 which is what prod does
|
||||
# when axis=None.
|
||||
if x.dtype in (torch.uint8, torch.int8, torch.int16, torch.int32):
|
||||
return x.to(torch.int64)
|
||||
|
||||
return x.clone()
|
||||
|
||||
|
||||
def prod(x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
dtype: Optional[DType] = None,
|
||||
keepdims: bool = False,
|
||||
**kwargs) -> Array:
|
||||
|
||||
if axis == ():
|
||||
return _sum_prod_no_axis(x, dtype)
|
||||
# torch.prod doesn't support multiple axes
|
||||
# (https://github.com/pytorch/pytorch/issues/56586).
|
||||
if isinstance(axis, tuple):
|
||||
return _reduce_multiple_axes(torch.prod, x, axis, keepdims=keepdims, dtype=dtype, **kwargs)
|
||||
if axis is None:
|
||||
# torch doesn't support keepdims with axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
res = torch.prod(x, dtype=dtype, **kwargs)
|
||||
res = _axis_none_keepdims(res, x.ndim, keepdims)
|
||||
return res
|
||||
|
||||
return torch.prod(x, axis, dtype=dtype, keepdims=keepdims, **kwargs)
|
||||
|
||||
|
||||
def sum(x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
dtype: Optional[DType] = None,
|
||||
keepdims: bool = False,
|
||||
**kwargs) -> Array:
|
||||
|
||||
if axis == ():
|
||||
return _sum_prod_no_axis(x, dtype)
|
||||
if axis is None:
|
||||
# torch doesn't support keepdims with axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
res = torch.sum(x, dtype=dtype, **kwargs)
|
||||
res = _axis_none_keepdims(res, x.ndim, keepdims)
|
||||
return res
|
||||
|
||||
return torch.sum(x, axis, dtype=dtype, keepdims=keepdims, **kwargs)
|
||||
|
||||
def any(x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
keepdims: bool = False,
|
||||
**kwargs) -> Array:
|
||||
|
||||
if axis == ():
|
||||
return x.to(torch.bool)
|
||||
# torch.any doesn't support multiple axes
|
||||
# (https://github.com/pytorch/pytorch/issues/56586).
|
||||
if isinstance(axis, tuple):
|
||||
res = _reduce_multiple_axes(torch.any, x, axis, keepdims=keepdims, **kwargs)
|
||||
return res.to(torch.bool)
|
||||
if axis is None:
|
||||
# torch doesn't support keepdims with axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
res = torch.any(x, **kwargs)
|
||||
res = _axis_none_keepdims(res, x.ndim, keepdims)
|
||||
return res.to(torch.bool)
|
||||
|
||||
# torch.any doesn't return bool for uint8
|
||||
return torch.any(x, axis, keepdims=keepdims).to(torch.bool)
|
||||
|
||||
def all(x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
keepdims: bool = False,
|
||||
**kwargs) -> Array:
|
||||
|
||||
if axis == ():
|
||||
return x.to(torch.bool)
|
||||
# torch.all doesn't support multiple axes
|
||||
# (https://github.com/pytorch/pytorch/issues/56586).
|
||||
if isinstance(axis, tuple):
|
||||
res = _reduce_multiple_axes(torch.all, x, axis, keepdims=keepdims, **kwargs)
|
||||
return res.to(torch.bool)
|
||||
if axis is None:
|
||||
# torch doesn't support keepdims with axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
res = torch.all(x, **kwargs)
|
||||
res = _axis_none_keepdims(res, x.ndim, keepdims)
|
||||
return res.to(torch.bool)
|
||||
|
||||
# torch.all doesn't return bool for uint8
|
||||
return torch.all(x, axis, keepdims=keepdims).to(torch.bool)
|
||||
|
||||
def mean(x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
keepdims: bool = False,
|
||||
**kwargs) -> Array:
|
||||
# https://github.com/pytorch/pytorch/issues/29137
|
||||
if axis == ():
|
||||
return torch.clone(x)
|
||||
if axis is None:
|
||||
# torch doesn't support keepdims with axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
res = torch.mean(x, **kwargs)
|
||||
res = _axis_none_keepdims(res, x.ndim, keepdims)
|
||||
return res
|
||||
return torch.mean(x, axis, keepdims=keepdims, **kwargs)
|
||||
|
||||
def std(x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
correction: Union[int, float] = 0.0,
|
||||
keepdims: bool = False,
|
||||
**kwargs) -> Array:
|
||||
# Note, float correction is not supported
|
||||
# https://github.com/pytorch/pytorch/issues/61492. We don't try to
|
||||
# implement it here for now.
|
||||
|
||||
if isinstance(correction, float):
|
||||
_correction = int(correction)
|
||||
if correction != _correction:
|
||||
raise NotImplementedError("float correction in torch std() is not yet supported")
|
||||
else:
|
||||
_correction = correction
|
||||
|
||||
# https://github.com/pytorch/pytorch/issues/29137
|
||||
if axis == ():
|
||||
return torch.zeros_like(x)
|
||||
if isinstance(axis, int):
|
||||
axis = (axis,)
|
||||
if axis is None:
|
||||
# torch doesn't support keepdims with axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
res = torch.std(x, tuple(range(x.ndim)), correction=_correction, **kwargs)
|
||||
res = _axis_none_keepdims(res, x.ndim, keepdims)
|
||||
return res
|
||||
return torch.std(x, axis, correction=_correction, keepdims=keepdims, **kwargs)
|
||||
|
||||
def var(x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
correction: Union[int, float] = 0.0,
|
||||
keepdims: bool = False,
|
||||
**kwargs) -> Array:
|
||||
# Note, float correction is not supported
|
||||
# https://github.com/pytorch/pytorch/issues/61492. We don't try to
|
||||
# implement it here for now.
|
||||
|
||||
# if isinstance(correction, float):
|
||||
# correction = int(correction)
|
||||
|
||||
# https://github.com/pytorch/pytorch/issues/29137
|
||||
if axis == ():
|
||||
return torch.zeros_like(x)
|
||||
if isinstance(axis, int):
|
||||
axis = (axis,)
|
||||
if axis is None:
|
||||
# torch doesn't support keepdims with axis=None
|
||||
# (https://github.com/pytorch/pytorch/issues/71209)
|
||||
res = torch.var(x, tuple(range(x.ndim)), correction=correction, **kwargs)
|
||||
res = _axis_none_keepdims(res, x.ndim, keepdims)
|
||||
return res
|
||||
return torch.var(x, axis, correction=correction, keepdims=keepdims, **kwargs)
|
||||
|
||||
# torch.concat doesn't support dim=None
|
||||
# https://github.com/pytorch/pytorch/issues/70925
|
||||
def concat(arrays: Union[Tuple[Array, ...], List[Array]],
|
||||
/,
|
||||
*,
|
||||
axis: Optional[int] = 0,
|
||||
**kwargs) -> Array:
|
||||
if axis is None:
|
||||
arrays = tuple(ar.flatten() for ar in arrays)
|
||||
axis = 0
|
||||
return torch.concat(arrays, axis, **kwargs)
|
||||
|
||||
# torch.squeeze only accepts int dim and doesn't require it
|
||||
# https://github.com/pytorch/pytorch/issues/70924. Support for tuple dim was
|
||||
# added at https://github.com/pytorch/pytorch/pull/89017.
|
||||
def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
|
||||
if isinstance(axis, int):
|
||||
axis = (axis,)
|
||||
for a in axis:
|
||||
if x.shape[a] != 1:
|
||||
raise ValueError("squeezed dimensions must be equal to 1")
|
||||
axes = _normalize_axes(axis, x.ndim)
|
||||
# Remove this once pytorch 1.14 is released with the above PR #89017.
|
||||
sequence = [a - i for i, a in enumerate(axes)]
|
||||
for a in sequence:
|
||||
x = torch.squeeze(x, a)
|
||||
return x
|
||||
|
||||
# torch.broadcast_to uses size instead of shape
|
||||
def broadcast_to(x: Array, /, shape: Tuple[int, ...], **kwargs) -> Array:
|
||||
return torch.broadcast_to(x, shape, **kwargs)
|
||||
|
||||
# torch.permute uses dims instead of axes
|
||||
def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
|
||||
return torch.permute(x, axes)
|
||||
|
||||
# The axis parameter doesn't work for flip() and roll()
|
||||
# https://github.com/pytorch/pytorch/issues/71210. Also torch.flip() doesn't
|
||||
# accept axis=None
|
||||
def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, **kwargs) -> Array:
|
||||
if axis is None:
|
||||
axis = tuple(range(x.ndim))
|
||||
# torch.flip doesn't accept dim as an int but the method does
|
||||
# https://github.com/pytorch/pytorch/issues/18095
|
||||
return x.flip(axis, **kwargs)
|
||||
|
||||
def roll(x: Array, /, shift: Union[int, Tuple[int, ...]], *, axis: Optional[Union[int, Tuple[int, ...]]] = None, **kwargs) -> Array:
|
||||
return torch.roll(x, shift, axis, **kwargs)
|
||||
|
||||
def nonzero(x: Array, /, **kwargs) -> Tuple[Array, ...]:
|
||||
if x.ndim == 0:
|
||||
raise ValueError("nonzero() does not support zero-dimensional arrays")
|
||||
return torch.nonzero(x, as_tuple=True, **kwargs)
|
||||
|
||||
|
||||
# torch uses `dim` instead of `axis`
|
||||
def diff(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: int = -1,
|
||||
n: int = 1,
|
||||
prepend: Optional[Array] = None,
|
||||
append: Optional[Array] = None,
|
||||
) -> Array:
|
||||
return torch.diff(x, dim=axis, n=n, prepend=prepend, append=append)
|
||||
|
||||
|
||||
# torch uses `dim` instead of `axis`, does not have keepdims
|
||||
def count_nonzero(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
keepdims: bool = False,
|
||||
) -> Array:
|
||||
result = torch.count_nonzero(x, dim=axis)
|
||||
if keepdims:
|
||||
if isinstance(axis, int):
|
||||
return result.unsqueeze(axis)
|
||||
elif isinstance(axis, tuple):
|
||||
n_axis = [x.ndim + ax if ax < 0 else ax for ax in axis]
|
||||
sh = [1 if i in n_axis else x.shape[i] for i in range(x.ndim)]
|
||||
return torch.reshape(result, sh)
|
||||
return _axis_none_keepdims(result, x.ndim, keepdims)
|
||||
else:
|
||||
return result
|
||||
|
||||
|
||||
# "repeat" is torch.repeat_interleave; also the dim argument
|
||||
def repeat(x: Array, repeats: int | Array, /, *, axis: int | None = None) -> Array:
|
||||
return torch.repeat_interleave(x, repeats, axis)
|
||||
|
||||
|
||||
def where(
|
||||
condition: Array,
|
||||
x1: Array | bool | int | float | complex,
|
||||
x2: Array | bool | int | float | complex,
|
||||
/,
|
||||
) -> Array:
|
||||
x1, x2 = _fix_promotion(x1, x2)
|
||||
return torch.where(condition, x1, x2)
|
||||
|
||||
|
||||
# torch.reshape doesn't have the copy keyword
|
||||
def reshape(x: Array,
|
||||
/,
|
||||
shape: Tuple[int, ...],
|
||||
*,
|
||||
copy: Optional[bool] = None,
|
||||
**kwargs) -> Array:
|
||||
if copy is not None:
|
||||
raise NotImplementedError("torch.reshape doesn't yet support the copy keyword")
|
||||
return torch.reshape(x, shape, **kwargs)
|
||||
|
||||
# torch.arange doesn't support returning empty arrays
|
||||
# (https://github.com/pytorch/pytorch/issues/70915), and doesn't support some
|
||||
# keyword argument combinations
|
||||
# (https://github.com/pytorch/pytorch/issues/70914)
|
||||
def arange(start: Union[int, float],
|
||||
/,
|
||||
stop: Optional[Union[int, float]] = None,
|
||||
step: Union[int, float] = 1,
|
||||
*,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
**kwargs) -> Array:
|
||||
if stop is None:
|
||||
start, stop = 0, start
|
||||
if step > 0 and stop <= start or step < 0 and stop >= start:
|
||||
if dtype is None:
|
||||
if _builtin_all(isinstance(i, int) for i in [start, stop, step]):
|
||||
dtype = torch.int64
|
||||
else:
|
||||
dtype = torch.float32
|
||||
return torch.empty(0, dtype=dtype, device=device, **kwargs)
|
||||
return torch.arange(start, stop, step, dtype=dtype, device=device, **kwargs)
|
||||
|
||||
# torch.eye does not accept None as a default for the second argument and
|
||||
# doesn't support off-diagonals (https://github.com/pytorch/pytorch/issues/70910)
|
||||
def eye(n_rows: int,
|
||||
n_cols: Optional[int] = None,
|
||||
/,
|
||||
*,
|
||||
k: int = 0,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
**kwargs) -> Array:
|
||||
if n_cols is None:
|
||||
n_cols = n_rows
|
||||
z = torch.zeros(n_rows, n_cols, dtype=dtype, device=device, **kwargs)
|
||||
if abs(k) <= n_rows + n_cols:
|
||||
z.diagonal(k).fill_(1)
|
||||
return z
|
||||
|
||||
# torch.linspace doesn't have the endpoint parameter
|
||||
def linspace(start: Union[int, float],
|
||||
stop: Union[int, float],
|
||||
/,
|
||||
num: int,
|
||||
*,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
endpoint: bool = True,
|
||||
**kwargs) -> Array:
|
||||
if not endpoint:
|
||||
return torch.linspace(start, stop, num+1, dtype=dtype, device=device, **kwargs)[:-1]
|
||||
return torch.linspace(start, stop, num, dtype=dtype, device=device, **kwargs)
|
||||
|
||||
# torch.full does not accept an int size
|
||||
# https://github.com/pytorch/pytorch/issues/70906
|
||||
def full(shape: Union[int, Tuple[int, ...]],
|
||||
fill_value: bool | int | float | complex,
|
||||
*,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
**kwargs) -> Array:
|
||||
if isinstance(shape, int):
|
||||
shape = (shape,)
|
||||
|
||||
return torch.full(shape, fill_value, dtype=dtype, device=device, **kwargs)
|
||||
|
||||
# ones, zeros, and empty do not accept shape as a keyword argument
|
||||
def ones(shape: Union[int, Tuple[int, ...]],
|
||||
*,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
**kwargs) -> Array:
|
||||
return torch.ones(shape, dtype=dtype, device=device, **kwargs)
|
||||
|
||||
def zeros(shape: Union[int, Tuple[int, ...]],
|
||||
*,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
**kwargs) -> Array:
|
||||
return torch.zeros(shape, dtype=dtype, device=device, **kwargs)
|
||||
|
||||
def empty(shape: Union[int, Tuple[int, ...]],
|
||||
*,
|
||||
dtype: Optional[DType] = None,
|
||||
device: Optional[Device] = None,
|
||||
**kwargs) -> Array:
|
||||
return torch.empty(shape, dtype=dtype, device=device, **kwargs)
|
||||
|
||||
# tril and triu do not call the keyword argument k
|
||||
|
||||
def tril(x: Array, /, *, k: int = 0) -> Array:
|
||||
return torch.tril(x, k)
|
||||
|
||||
def triu(x: Array, /, *, k: int = 0) -> Array:
|
||||
return torch.triu(x, k)
|
||||
|
||||
# Functions that aren't in torch https://github.com/pytorch/pytorch/issues/58742
|
||||
def expand_dims(x: Array, /, *, axis: int = 0) -> Array:
|
||||
return torch.unsqueeze(x, axis)
|
||||
|
||||
|
||||
def astype(
|
||||
x: Array,
|
||||
dtype: DType,
|
||||
/,
|
||||
*,
|
||||
copy: bool = True,
|
||||
device: Optional[Device] = None,
|
||||
) -> Array:
|
||||
if device is not None:
|
||||
return x.to(device, dtype=dtype, copy=copy)
|
||||
return x.to(dtype=dtype, copy=copy)
|
||||
|
||||
|
||||
def broadcast_arrays(*arrays: Array) -> List[Array]:
|
||||
shape = torch.broadcast_shapes(*[a.shape for a in arrays])
|
||||
return [torch.broadcast_to(a, shape) for a in arrays]
|
||||
|
||||
# Note that these named tuples aren't actually part of the standard namespace,
|
||||
# but I don't see any issue with exporting the names here regardless.
|
||||
from ..common._aliases import (UniqueAllResult, UniqueCountsResult,
|
||||
UniqueInverseResult)
|
||||
|
||||
# https://github.com/pytorch/pytorch/issues/70920
|
||||
def unique_all(x: Array) -> UniqueAllResult:
|
||||
# torch.unique doesn't support returning indices.
|
||||
# https://github.com/pytorch/pytorch/issues/36748. The workaround
|
||||
# suggested in that issue doesn't actually function correctly (it relies
|
||||
# on non-deterministic behavior of scatter()).
|
||||
raise NotImplementedError("unique_all() not yet implemented for pytorch (see https://github.com/pytorch/pytorch/issues/36748)")
|
||||
|
||||
# values, inverse_indices, counts = torch.unique(x, return_counts=True, return_inverse=True)
|
||||
# # torch.unique incorrectly gives a 0 count for nan values.
|
||||
# # https://github.com/pytorch/pytorch/issues/94106
|
||||
# counts[torch.isnan(values)] = 1
|
||||
# return UniqueAllResult(values, indices, inverse_indices, counts)
|
||||
|
||||
def unique_counts(x: Array) -> UniqueCountsResult:
|
||||
values, counts = torch.unique(x, return_counts=True)
|
||||
|
||||
# torch.unique incorrectly gives a 0 count for nan values.
|
||||
# https://github.com/pytorch/pytorch/issues/94106
|
||||
counts[torch.isnan(values)] = 1
|
||||
return UniqueCountsResult(values, counts)
|
||||
|
||||
def unique_inverse(x: Array) -> UniqueInverseResult:
|
||||
values, inverse = torch.unique(x, return_inverse=True)
|
||||
return UniqueInverseResult(values, inverse)
|
||||
|
||||
def unique_values(x: Array) -> Array:
|
||||
return torch.unique(x)
|
||||
|
||||
def matmul(x1: Array, x2: Array, /, **kwargs) -> Array:
|
||||
# torch.matmul doesn't type promote (but differently from _fix_promotion)
|
||||
x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
|
||||
return torch.matmul(x1, x2, **kwargs)
|
||||
|
||||
matrix_transpose = get_xp(torch)(_aliases.matrix_transpose)
|
||||
_vecdot = get_xp(torch)(_aliases.vecdot)
|
||||
|
||||
def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array:
|
||||
x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
|
||||
return _vecdot(x1, x2, axis=axis)
|
||||
|
||||
# torch.tensordot uses dims instead of axes
|
||||
def tensordot(
|
||||
x1: Array,
|
||||
x2: Array,
|
||||
/,
|
||||
*,
|
||||
axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2,
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
# Note: torch.tensordot fails with integer dtypes when there is only 1
|
||||
# element in the axis (https://github.com/pytorch/pytorch/issues/84530).
|
||||
x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
|
||||
return torch.tensordot(x1, x2, dims=axes, **kwargs)
|
||||
|
||||
|
||||
def isdtype(
|
||||
dtype: DType, kind: Union[DType, str, Tuple[Union[DType, str], ...]],
|
||||
*, _tuple=True, # Disallow nested tuples
|
||||
) -> bool:
|
||||
"""
|
||||
Returns a boolean indicating whether a provided dtype is of a specified data type ``kind``.
|
||||
|
||||
Note that outside of this function, this compat library does not yet fully
|
||||
support complex numbers.
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/generated/array_api.isdtype.html
|
||||
for more details
|
||||
"""
|
||||
if isinstance(kind, tuple) and _tuple:
|
||||
return _builtin_any(isdtype(dtype, k, _tuple=False) for k in kind)
|
||||
elif isinstance(kind, str):
|
||||
if kind == 'bool':
|
||||
return dtype == torch.bool
|
||||
elif kind == 'signed integer':
|
||||
return dtype in _int_dtypes and dtype.is_signed
|
||||
elif kind == 'unsigned integer':
|
||||
return dtype in _int_dtypes and not dtype.is_signed
|
||||
elif kind == 'integral':
|
||||
return dtype in _int_dtypes
|
||||
elif kind == 'real floating':
|
||||
return dtype.is_floating_point
|
||||
elif kind == 'complex floating':
|
||||
return dtype.is_complex
|
||||
elif kind == 'numeric':
|
||||
return isdtype(dtype, ('integral', 'real floating', 'complex floating'))
|
||||
else:
|
||||
raise ValueError(f"Unrecognized data type kind: {kind!r}")
|
||||
else:
|
||||
return dtype == kind
|
||||
|
||||
def take(x: Array, indices: Array, /, *, axis: Optional[int] = None, **kwargs) -> Array:
|
||||
if axis is None:
|
||||
if x.ndim != 1:
|
||||
raise ValueError("axis must be specified when ndim > 1")
|
||||
axis = 0
|
||||
return torch.index_select(x, axis, indices, **kwargs)
|
||||
|
||||
|
||||
def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
|
||||
return torch.take_along_dim(x, indices, dim=axis)
|
||||
|
||||
|
||||
def sign(x: Array, /) -> Array:
|
||||
# torch sign() does not support complex numbers and does not propagate
|
||||
# nans. See https://github.com/data-apis/array-api-compat/issues/136
|
||||
if x.dtype.is_complex:
|
||||
out = x/torch.abs(x)
|
||||
# sign(0) = 0 but the above formula would give nan
|
||||
out[x == 0+0j] = 0+0j
|
||||
return out
|
||||
else:
|
||||
out = torch.sign(x)
|
||||
if x.dtype.is_floating_point:
|
||||
out[torch.isnan(x)] = torch.nan
|
||||
return out
|
||||
|
||||
|
||||
def meshgrid(*arrays: Array, indexing: Literal['xy', 'ij'] = 'xy') -> List[Array]:
|
||||
# enforce the default of 'xy'
|
||||
# TODO: is the return type a list or a tuple
|
||||
return list(torch.meshgrid(*arrays, indexing='xy'))
|
||||
|
||||
|
||||
__all__ = ['__array_namespace_info__', 'asarray', 'result_type', 'can_cast',
|
||||
'permute_dims', 'bitwise_invert', 'newaxis', 'conj', 'add',
|
||||
'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
|
||||
'bitwise_right_shift', 'bitwise_xor', 'copysign', 'count_nonzero',
|
||||
'diff', 'divide',
|
||||
'equal', 'floor_divide', 'greater', 'greater_equal', 'hypot',
|
||||
'less', 'less_equal', 'logaddexp', 'maximum', 'minimum',
|
||||
'multiply', 'not_equal', 'pow', 'remainder', 'subtract', 'max',
|
||||
'min', 'clip', 'unstack', 'cumulative_sum', 'cumulative_prod', 'sort', 'prod', 'sum',
|
||||
'any', 'all', 'mean', 'std', 'var', 'concat', 'squeeze',
|
||||
'broadcast_to', 'flip', 'roll', 'nonzero', 'where', 'reshape',
|
||||
'arange', 'eye', 'linspace', 'full', 'ones', 'zeros', 'empty',
|
||||
'tril', 'triu', 'expand_dims', 'astype', 'broadcast_arrays',
|
||||
'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
|
||||
'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
|
||||
'matmul', 'matrix_transpose', 'vecdot', 'tensordot', 'isdtype',
|
||||
'take', 'take_along_axis', 'sign', 'finfo', 'iinfo', 'repeat', 'meshgrid']
|
||||
|
||||
_all_ignore = ['torch', 'get_xp']
|
||||
@@ -0,0 +1,369 @@
|
||||
"""
|
||||
Array API Inspection namespace
|
||||
|
||||
This is the namespace for inspection functions as defined by the array API
|
||||
standard. See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html for
|
||||
more details.
|
||||
|
||||
"""
|
||||
import torch
|
||||
|
||||
from functools import cache
|
||||
|
||||
class __array_namespace_info__:
|
||||
"""
|
||||
Get the array API inspection namespace for PyTorch.
|
||||
|
||||
The array API inspection namespace defines the following functions:
|
||||
|
||||
- capabilities()
|
||||
- default_device()
|
||||
- default_dtypes()
|
||||
- dtypes()
|
||||
- devices()
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/inspection.html
|
||||
for more details.
|
||||
|
||||
Returns
|
||||
-------
|
||||
info : ModuleType
|
||||
The array API inspection namespace for PyTorch.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': numpy.float64,
|
||||
'complex floating': numpy.complex128,
|
||||
'integral': numpy.int64,
|
||||
'indexing': numpy.int64}
|
||||
|
||||
"""
|
||||
|
||||
__module__ = 'torch'
|
||||
|
||||
def capabilities(self):
|
||||
"""
|
||||
Return a dictionary of array API library capabilities.
|
||||
|
||||
The resulting dictionary has the following keys:
|
||||
|
||||
- **"boolean indexing"**: boolean indicating whether an array library
|
||||
supports boolean indexing. Always ``True`` for PyTorch.
|
||||
|
||||
- **"data-dependent shapes"**: boolean indicating whether an array
|
||||
library supports data-dependent output shapes. Always ``True`` for
|
||||
PyTorch.
|
||||
|
||||
See
|
||||
https://data-apis.org/array-api/latest/API_specification/generated/array_api.info.capabilities.html
|
||||
for more details.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
capabilities : dict
|
||||
A dictionary of array API library capabilities.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.capabilities()
|
||||
{'boolean indexing': True,
|
||||
'data-dependent shapes': True,
|
||||
'max dimensions': 64}
|
||||
|
||||
"""
|
||||
return {
|
||||
"boolean indexing": True,
|
||||
"data-dependent shapes": True,
|
||||
"max dimensions": 64,
|
||||
}
|
||||
|
||||
def default_device(self):
|
||||
"""
|
||||
The default device used for new PyTorch arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Returns
|
||||
-------
|
||||
device : Device
|
||||
The default device used for new PyTorch arrays.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_device()
|
||||
device(type='cpu')
|
||||
|
||||
Notes
|
||||
-----
|
||||
This method returns the static default device when PyTorch is initialized.
|
||||
However, the *current* device used by creation functions (``empty`` etc.)
|
||||
can be changed at runtime.
|
||||
|
||||
See Also
|
||||
--------
|
||||
https://github.com/data-apis/array-api/issues/835
|
||||
"""
|
||||
return torch.device("cpu")
|
||||
|
||||
def default_dtypes(self, *, device=None):
|
||||
"""
|
||||
The default data types used for new PyTorch arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : Device, optional
|
||||
The device to get the default data types for.
|
||||
Unused for PyTorch, as all devices use the same default dtypes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary describing the default data types used for new PyTorch
|
||||
arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.default_dtypes()
|
||||
{'real floating': torch.float32,
|
||||
'complex floating': torch.complex64,
|
||||
'integral': torch.int64,
|
||||
'indexing': torch.int64}
|
||||
|
||||
"""
|
||||
# Note: if the default is set to float64, the devices like MPS that
|
||||
# don't support float64 will error. We still return the default_dtype
|
||||
# value here because this error doesn't represent a different default
|
||||
# per-device.
|
||||
default_floating = torch.get_default_dtype()
|
||||
default_complex = torch.complex64 if default_floating == torch.float32 else torch.complex128
|
||||
default_integral = torch.int64
|
||||
return {
|
||||
"real floating": default_floating,
|
||||
"complex floating": default_complex,
|
||||
"integral": default_integral,
|
||||
"indexing": default_integral,
|
||||
}
|
||||
|
||||
|
||||
def _dtypes(self, kind):
|
||||
bool = torch.bool
|
||||
int8 = torch.int8
|
||||
int16 = torch.int16
|
||||
int32 = torch.int32
|
||||
int64 = torch.int64
|
||||
uint8 = torch.uint8
|
||||
# uint16, uint32, and uint64 are present in newer versions of pytorch,
|
||||
# but they aren't generally supported by the array API functions, so
|
||||
# we omit them from this function.
|
||||
float32 = torch.float32
|
||||
float64 = torch.float64
|
||||
complex64 = torch.complex64
|
||||
complex128 = torch.complex128
|
||||
|
||||
if kind is None:
|
||||
return {
|
||||
"bool": bool,
|
||||
"int8": int8,
|
||||
"int16": int16,
|
||||
"int32": int32,
|
||||
"int64": int64,
|
||||
"uint8": uint8,
|
||||
"float32": float32,
|
||||
"float64": float64,
|
||||
"complex64": complex64,
|
||||
"complex128": complex128,
|
||||
}
|
||||
if kind == "bool":
|
||||
return {"bool": bool}
|
||||
if kind == "signed integer":
|
||||
return {
|
||||
"int8": int8,
|
||||
"int16": int16,
|
||||
"int32": int32,
|
||||
"int64": int64,
|
||||
}
|
||||
if kind == "unsigned integer":
|
||||
return {
|
||||
"uint8": uint8,
|
||||
}
|
||||
if kind == "integral":
|
||||
return {
|
||||
"int8": int8,
|
||||
"int16": int16,
|
||||
"int32": int32,
|
||||
"int64": int64,
|
||||
"uint8": uint8,
|
||||
}
|
||||
if kind == "real floating":
|
||||
return {
|
||||
"float32": float32,
|
||||
"float64": float64,
|
||||
}
|
||||
if kind == "complex floating":
|
||||
return {
|
||||
"complex64": complex64,
|
||||
"complex128": complex128,
|
||||
}
|
||||
if kind == "numeric":
|
||||
return {
|
||||
"int8": int8,
|
||||
"int16": int16,
|
||||
"int32": int32,
|
||||
"int64": int64,
|
||||
"uint8": uint8,
|
||||
"float32": float32,
|
||||
"float64": float64,
|
||||
"complex64": complex64,
|
||||
"complex128": complex128,
|
||||
}
|
||||
if isinstance(kind, tuple):
|
||||
res = {}
|
||||
for k in kind:
|
||||
res.update(self.dtypes(kind=k))
|
||||
return res
|
||||
raise ValueError(f"unsupported kind: {kind!r}")
|
||||
|
||||
@cache
|
||||
def dtypes(self, *, device=None, kind=None):
|
||||
"""
|
||||
The array API data types supported by PyTorch.
|
||||
|
||||
Note that this function only returns data types that are defined by
|
||||
the array API.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
device : Device, optional
|
||||
The device to get the data types for.
|
||||
Unused for PyTorch, as all devices use the same dtypes.
|
||||
kind : str or tuple of str, optional
|
||||
The kind of data types to return. If ``None``, all data types are
|
||||
returned. If a string, only data types of that kind are returned.
|
||||
If a tuple, a dictionary containing the union of the given kinds
|
||||
is returned. The following kinds are supported:
|
||||
|
||||
- ``'bool'``: boolean data types (i.e., ``bool``).
|
||||
- ``'signed integer'``: signed integer data types (i.e., ``int8``,
|
||||
``int16``, ``int32``, ``int64``).
|
||||
- ``'unsigned integer'``: unsigned integer data types (i.e.,
|
||||
``uint8``, ``uint16``, ``uint32``, ``uint64``).
|
||||
- ``'integral'``: integer data types. Shorthand for ``('signed
|
||||
integer', 'unsigned integer')``.
|
||||
- ``'real floating'``: real-valued floating-point data types
|
||||
(i.e., ``float32``, ``float64``).
|
||||
- ``'complex floating'``: complex floating-point data types (i.e.,
|
||||
``complex64``, ``complex128``).
|
||||
- ``'numeric'``: numeric data types. Shorthand for ``('integral',
|
||||
'real floating', 'complex floating')``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dtypes : dict
|
||||
A dictionary mapping the names of data types to the corresponding
|
||||
PyTorch data types.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.devices
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.dtypes(kind='signed integer')
|
||||
{'int8': numpy.int8,
|
||||
'int16': numpy.int16,
|
||||
'int32': numpy.int32,
|
||||
'int64': numpy.int64}
|
||||
|
||||
"""
|
||||
res = self._dtypes(kind)
|
||||
for k, v in res.copy().items():
|
||||
try:
|
||||
torch.empty((0,), dtype=v, device=device)
|
||||
except:
|
||||
del res[k]
|
||||
return res
|
||||
|
||||
@cache
|
||||
def devices(self):
|
||||
"""
|
||||
The devices supported by PyTorch.
|
||||
|
||||
Returns
|
||||
-------
|
||||
devices : list[Device]
|
||||
The devices supported by PyTorch.
|
||||
|
||||
See Also
|
||||
--------
|
||||
__array_namespace_info__.capabilities,
|
||||
__array_namespace_info__.default_device,
|
||||
__array_namespace_info__.default_dtypes,
|
||||
__array_namespace_info__.dtypes
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> info = xp.__array_namespace_info__()
|
||||
>>> info.devices()
|
||||
[device(type='cpu'), device(type='mps', index=0), device(type='meta')]
|
||||
|
||||
"""
|
||||
# Torch doesn't have a straightforward way to get the list of all
|
||||
# currently supported devices. To do this, we first parse the error
|
||||
# message of torch.device to get the list of all possible types of
|
||||
# device:
|
||||
try:
|
||||
torch.device('notadevice')
|
||||
raise AssertionError("unreachable") # pragma: nocover
|
||||
except RuntimeError as e:
|
||||
# The error message is something like:
|
||||
# "Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone device type at start of device string: notadevice"
|
||||
devices_names = e.args[0].split('Expected one of ')[1].split(' device type')[0].split(', ')
|
||||
|
||||
# Next we need to check for different indices for different devices.
|
||||
# device(device_name, index=index) doesn't actually check if the
|
||||
# device name or index is valid. We have to try to create a tensor
|
||||
# with it (which is why this function is cached).
|
||||
devices = []
|
||||
for device_name in devices_names:
|
||||
i = 0
|
||||
while True:
|
||||
try:
|
||||
a = torch.empty((0,), device=torch.device(device_name, index=i))
|
||||
if a.device in devices:
|
||||
break
|
||||
devices.append(a.device)
|
||||
except:
|
||||
break
|
||||
i += 1
|
||||
|
||||
return devices
|
||||
@@ -0,0 +1,3 @@
|
||||
__all__ = ["Array", "Device", "DType"]
|
||||
|
||||
from torch import device as Device, dtype as DType, Tensor as Array
|
||||
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Union, Sequence, Literal
|
||||
|
||||
import torch
|
||||
import torch.fft
|
||||
from torch.fft import * # noqa: F403
|
||||
|
||||
from ._typing import Array
|
||||
|
||||
# Several torch fft functions do not map axes to dim
|
||||
|
||||
def fftn(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
s: Sequence[int] = None,
|
||||
axes: Sequence[int] = None,
|
||||
norm: Literal["backward", "ortho", "forward"] = "backward",
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
return torch.fft.fftn(x, s=s, dim=axes, norm=norm, **kwargs)
|
||||
|
||||
def ifftn(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
s: Sequence[int] = None,
|
||||
axes: Sequence[int] = None,
|
||||
norm: Literal["backward", "ortho", "forward"] = "backward",
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
return torch.fft.ifftn(x, s=s, dim=axes, norm=norm, **kwargs)
|
||||
|
||||
def rfftn(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
s: Sequence[int] = None,
|
||||
axes: Sequence[int] = None,
|
||||
norm: Literal["backward", "ortho", "forward"] = "backward",
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
return torch.fft.rfftn(x, s=s, dim=axes, norm=norm, **kwargs)
|
||||
|
||||
def irfftn(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
s: Sequence[int] = None,
|
||||
axes: Sequence[int] = None,
|
||||
norm: Literal["backward", "ortho", "forward"] = "backward",
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
return torch.fft.irfftn(x, s=s, dim=axes, norm=norm, **kwargs)
|
||||
|
||||
def fftshift(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axes: Union[int, Sequence[int]] = None,
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
return torch.fft.fftshift(x, dim=axes, **kwargs)
|
||||
|
||||
def ifftshift(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axes: Union[int, Sequence[int]] = None,
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
return torch.fft.ifftshift(x, dim=axes, **kwargs)
|
||||
|
||||
|
||||
__all__ = torch.fft.__all__ + [
|
||||
"fftn",
|
||||
"ifftn",
|
||||
"rfftn",
|
||||
"irfftn",
|
||||
"fftshift",
|
||||
"ifftshift",
|
||||
]
|
||||
|
||||
_all_ignore = ['torch']
|
||||
@@ -0,0 +1,121 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typing import Optional, Union, Tuple
|
||||
|
||||
from torch.linalg import * # noqa: F403
|
||||
|
||||
# torch.linalg doesn't define __all__
|
||||
# from torch.linalg import __all__ as linalg_all
|
||||
from torch import linalg as torch_linalg
|
||||
linalg_all = [i for i in dir(torch_linalg) if not i.startswith('_')]
|
||||
|
||||
# outer is implemented in torch but aren't in the linalg namespace
|
||||
from torch import outer
|
||||
from ._aliases import _fix_promotion, sum
|
||||
# These functions are in both the main and linalg namespaces
|
||||
from ._aliases import matmul, matrix_transpose, tensordot
|
||||
from ._typing import Array, DType
|
||||
from ..common._typing import JustInt, JustFloat
|
||||
|
||||
# Note: torch.linalg.cross does not default to axis=-1 (it defaults to the
|
||||
# first axis with size 3), see https://github.com/pytorch/pytorch/issues/58743
|
||||
|
||||
# torch.cross also does not support broadcasting when it would add new
|
||||
# dimensions https://github.com/pytorch/pytorch/issues/39656
|
||||
def cross(x1: Array, x2: Array, /, *, axis: int = -1) -> Array:
|
||||
x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
|
||||
if not (-min(x1.ndim, x2.ndim) <= axis < max(x1.ndim, x2.ndim)):
|
||||
raise ValueError(f"axis {axis} out of bounds for cross product of arrays with shapes {x1.shape} and {x2.shape}")
|
||||
if not (x1.shape[axis] == x2.shape[axis] == 3):
|
||||
raise ValueError(f"cross product axis must have size 3, got {x1.shape[axis]} and {x2.shape[axis]}")
|
||||
x1, x2 = torch.broadcast_tensors(x1, x2)
|
||||
return torch_linalg.cross(x1, x2, dim=axis)
|
||||
|
||||
def vecdot(x1: Array, x2: Array, /, *, axis: int = -1, **kwargs) -> Array:
|
||||
from ._aliases import isdtype
|
||||
|
||||
x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
|
||||
|
||||
# torch.linalg.vecdot incorrectly allows broadcasting along the contracted dimension
|
||||
if x1.shape[axis] != x2.shape[axis]:
|
||||
raise ValueError("x1 and x2 must have the same size along the given axis")
|
||||
|
||||
# torch.linalg.vecdot doesn't support integer dtypes
|
||||
if isdtype(x1.dtype, 'integral') or isdtype(x2.dtype, 'integral'):
|
||||
if kwargs:
|
||||
raise RuntimeError("vecdot kwargs not supported for integral dtypes")
|
||||
|
||||
x1_ = torch.moveaxis(x1, axis, -1)
|
||||
x2_ = torch.moveaxis(x2, axis, -1)
|
||||
x1_, x2_ = torch.broadcast_tensors(x1_, x2_)
|
||||
|
||||
res = x1_[..., None, :] @ x2_[..., None]
|
||||
return res[..., 0, 0]
|
||||
return torch.linalg.vecdot(x1, x2, dim=axis, **kwargs)
|
||||
|
||||
def solve(x1: Array, x2: Array, /, **kwargs) -> Array:
|
||||
x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
|
||||
# Torch tries to emulate NumPy 1 solve behavior by using batched 1-D solve
|
||||
# whenever
|
||||
# 1. x1.ndim - 1 == x2.ndim
|
||||
# 2. x1.shape[:-1] == x2.shape
|
||||
#
|
||||
# See linalg_solve_is_vector_rhs in
|
||||
# aten/src/ATen/native/LinearAlgebraUtils.h and
|
||||
# TORCH_META_FUNC(_linalg_solve_ex) in
|
||||
# aten/src/ATen/native/BatchLinearAlgebra.cpp in the PyTorch source code.
|
||||
#
|
||||
# The easiest way to work around this is to prepend a size 1 dimension to
|
||||
# x2, since x2 is already one dimension less than x1.
|
||||
#
|
||||
# See https://github.com/pytorch/pytorch/issues/52915
|
||||
if x2.ndim != 1 and x1.ndim - 1 == x2.ndim and x1.shape[:-1] == x2.shape:
|
||||
x2 = x2[None]
|
||||
return torch.linalg.solve(x1, x2, **kwargs)
|
||||
|
||||
# torch.trace doesn't support the offset argument and doesn't support stacking
|
||||
def trace(x: Array, /, *, offset: int = 0, dtype: Optional[DType] = None) -> Array:
|
||||
# Use our wrapped sum to make sure it does upcasting correctly
|
||||
return sum(torch.diagonal(x, offset=offset, dim1=-2, dim2=-1), axis=-1, dtype=dtype)
|
||||
|
||||
def vector_norm(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: Optional[Union[int, Tuple[int, ...]]] = None,
|
||||
keepdims: bool = False,
|
||||
# JustFloat stands for inf | -inf, which are not valid for Literal
|
||||
ord: JustInt | JustFloat = 2,
|
||||
**kwargs,
|
||||
) -> Array:
|
||||
# torch.vector_norm incorrectly treats axis=() the same as axis=None
|
||||
if axis == ():
|
||||
out = kwargs.get('out')
|
||||
if out is None:
|
||||
dtype = None
|
||||
if x.dtype == torch.complex64:
|
||||
dtype = torch.float32
|
||||
elif x.dtype == torch.complex128:
|
||||
dtype = torch.float64
|
||||
|
||||
out = torch.zeros_like(x, dtype=dtype)
|
||||
|
||||
# The norm of a single scalar works out to abs(x) in every case except
|
||||
# for ord=0, which is x != 0.
|
||||
if ord == 0:
|
||||
out[:] = (x != 0)
|
||||
else:
|
||||
out[:] = torch.abs(x)
|
||||
return out
|
||||
return torch.linalg.vector_norm(x, ord=ord, axis=axis, keepdim=keepdims, **kwargs)
|
||||
|
||||
__all__ = linalg_all + ['outer', 'matmul', 'matrix_transpose', 'tensordot',
|
||||
'cross', 'vecdot', 'solve', 'trace', 'vector_norm']
|
||||
|
||||
_all_ignore = ['torch_linalg', 'sum']
|
||||
|
||||
del linalg_all
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return __all__
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Extra array functions built on top of the array API standard."""
|
||||
|
||||
from ._delegation import isclose, pad
|
||||
from ._lib._at import at
|
||||
from ._lib._funcs import (
|
||||
apply_where,
|
||||
atleast_nd,
|
||||
broadcast_shapes,
|
||||
cov,
|
||||
create_diagonal,
|
||||
expand_dims,
|
||||
kron,
|
||||
nunique,
|
||||
setdiff1d,
|
||||
sinc,
|
||||
)
|
||||
from ._lib._lazy import lazy_apply
|
||||
|
||||
__version__ = "0.8.0.dev0"
|
||||
|
||||
# pylint: disable=duplicate-code
|
||||
__all__ = [
|
||||
"__version__",
|
||||
"apply_where",
|
||||
"at",
|
||||
"atleast_nd",
|
||||
"broadcast_shapes",
|
||||
"cov",
|
||||
"create_diagonal",
|
||||
"expand_dims",
|
||||
"isclose",
|
||||
"kron",
|
||||
"lazy_apply",
|
||||
"nunique",
|
||||
"pad",
|
||||
"setdiff1d",
|
||||
"sinc",
|
||||
]
|
||||
@@ -0,0 +1,171 @@
|
||||
"""Delegation to existing implementations for Public API Functions."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from types import ModuleType
|
||||
from typing import Literal
|
||||
|
||||
from ._lib import _funcs
|
||||
from ._lib._utils._compat import (
|
||||
array_namespace,
|
||||
is_cupy_namespace,
|
||||
is_dask_namespace,
|
||||
is_jax_namespace,
|
||||
is_numpy_namespace,
|
||||
is_pydata_sparse_namespace,
|
||||
is_torch_namespace,
|
||||
)
|
||||
from ._lib._utils._helpers import asarrays
|
||||
from ._lib._utils._typing import Array
|
||||
|
||||
__all__ = ["isclose", "pad"]
|
||||
|
||||
|
||||
def isclose(
|
||||
a: Array | complex,
|
||||
b: Array | complex,
|
||||
*,
|
||||
rtol: float = 1e-05,
|
||||
atol: float = 1e-08,
|
||||
equal_nan: bool = False,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array:
|
||||
"""
|
||||
Return a boolean array where two arrays are element-wise equal within a tolerance.
|
||||
|
||||
The tolerance values are positive, typically very small numbers. The relative
|
||||
difference ``(rtol * abs(b))`` and the absolute difference `atol` are added together
|
||||
to compare against the absolute difference between `a` and `b`.
|
||||
|
||||
NaNs are treated as equal if they are in the same place and if ``equal_nan=True``.
|
||||
Infs are treated as equal if they are in the same place and of the same sign in both
|
||||
arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a, b : Array | int | float | complex | bool
|
||||
Input objects to compare. At least one must be an array.
|
||||
rtol : array_like, optional
|
||||
The relative tolerance parameter (see Notes).
|
||||
atol : array_like, optional
|
||||
The absolute tolerance parameter (see Notes).
|
||||
equal_nan : bool, optional
|
||||
Whether to compare NaN's as equal. If True, NaN's in `a` will be considered
|
||||
equal to NaN's in `b` in the output array.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `a` and `b`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Array
|
||||
A boolean array of shape broadcasted from `a` and `b`, containing ``True`` where
|
||||
`a` is close to `b`, and ``False`` otherwise.
|
||||
|
||||
Warnings
|
||||
--------
|
||||
The default `atol` is not appropriate for comparing numbers with magnitudes much
|
||||
smaller than one (see notes).
|
||||
|
||||
See Also
|
||||
--------
|
||||
math.isclose : Similar function in stdlib for Python scalars.
|
||||
|
||||
Notes
|
||||
-----
|
||||
For finite values, `isclose` uses the following equation to test whether two
|
||||
floating point values are equivalent::
|
||||
|
||||
absolute(a - b) <= (atol + rtol * absolute(b))
|
||||
|
||||
Unlike the built-in `math.isclose`,
|
||||
the above equation is not symmetric in `a` and `b`,
|
||||
so that ``isclose(a, b)`` might be different from ``isclose(b, a)`` in some rare
|
||||
cases.
|
||||
|
||||
The default value of `atol` is not appropriate when the reference value `b` has
|
||||
magnitude smaller than one. For example, it is unlikely that ``a = 1e-9`` and
|
||||
``b = 2e-9`` should be considered "close", yet ``isclose(1e-9, 2e-9)`` is ``True``
|
||||
with default settings. Be sure to select `atol` for the use case at hand, especially
|
||||
for defining the threshold below which a non-zero value in `a` will be considered
|
||||
"close" to a very small or zero value in `b`.
|
||||
|
||||
The comparison of `a` and `b` uses standard broadcasting, which means that `a` and
|
||||
`b` need not have the same shape in order for ``isclose(a, b)`` to evaluate to
|
||||
``True``.
|
||||
|
||||
`isclose` is not defined for non-numeric data types.
|
||||
``bool`` is considered a numeric data-type for this purpose.
|
||||
"""
|
||||
xp = array_namespace(a, b) if xp is None else xp
|
||||
|
||||
if (
|
||||
is_numpy_namespace(xp)
|
||||
or is_cupy_namespace(xp)
|
||||
or is_dask_namespace(xp)
|
||||
or is_jax_namespace(xp)
|
||||
):
|
||||
return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
|
||||
|
||||
if is_torch_namespace(xp):
|
||||
a, b = asarrays(a, b, xp=xp) # Array API 2024.12 support
|
||||
return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
|
||||
|
||||
return _funcs.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp)
|
||||
|
||||
|
||||
def pad(
|
||||
x: Array,
|
||||
pad_width: int | tuple[int, int] | Sequence[tuple[int, int]],
|
||||
mode: Literal["constant"] = "constant",
|
||||
*,
|
||||
constant_values: complex = 0,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array:
|
||||
"""
|
||||
Pad the input array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array
|
||||
Input array.
|
||||
pad_width : int or tuple of ints or sequence of pairs of ints
|
||||
Pad the input array with this many elements from each side.
|
||||
If a sequence of tuples, ``[(before_0, after_0), ... (before_N, after_N)]``,
|
||||
each pair applies to the corresponding axis of ``x``.
|
||||
A single tuple, ``(before, after)``, is equivalent to a list of ``x.ndim``
|
||||
copies of this tuple.
|
||||
mode : str, optional
|
||||
Only "constant" mode is currently supported, which pads with
|
||||
the value passed to `constant_values`.
|
||||
constant_values : python scalar, optional
|
||||
Use this value to pad the input. Default is zero.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
The input array,
|
||||
padded with ``pad_width`` elements equal to ``constant_values``.
|
||||
"""
|
||||
xp = array_namespace(x) if xp is None else xp
|
||||
|
||||
if mode != "constant":
|
||||
msg = "Only `'constant'` mode is currently supported"
|
||||
raise NotImplementedError(msg)
|
||||
|
||||
if (
|
||||
is_numpy_namespace(xp)
|
||||
or is_cupy_namespace(xp)
|
||||
or is_jax_namespace(xp)
|
||||
or is_pydata_sparse_namespace(xp)
|
||||
):
|
||||
return xp.pad(x, pad_width, mode, constant_values=constant_values)
|
||||
|
||||
# https://github.com/pytorch/pytorch/blob/cf76c05b4dc629ac989d1fb8e789d4fac04a095a/torch/_numpy/_funcs_impl.py#L2045-L2056
|
||||
if is_torch_namespace(xp):
|
||||
pad_width = xp.asarray(pad_width)
|
||||
pad_width = xp.broadcast_to(pad_width, (x.ndim, 2))
|
||||
pad_width = xp.flip(pad_width, axis=(0,)).flatten()
|
||||
return xp.nn.functional.pad(x, tuple(pad_width), value=constant_values) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
|
||||
|
||||
return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp)
|
||||
@@ -0,0 +1 @@
|
||||
"""Internals of array-api-extra."""
|
||||
@@ -0,0 +1,463 @@
|
||||
"""Update operations for read-only arrays."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from types import ModuleType
|
||||
from typing import TYPE_CHECKING, ClassVar, cast
|
||||
|
||||
from ._utils import _compat
|
||||
from ._utils._compat import (
|
||||
array_namespace,
|
||||
is_dask_array,
|
||||
is_jax_array,
|
||||
is_torch_array,
|
||||
is_writeable_array,
|
||||
)
|
||||
from ._utils._helpers import meta_namespace
|
||||
from ._utils._typing import Array, SetIndex
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover
|
||||
# TODO import from typing (requires Python >=3.11)
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
class _AtOp(Enum):
|
||||
"""Operations for use in `xpx.at`."""
|
||||
|
||||
SET = "set"
|
||||
ADD = "add"
|
||||
SUBTRACT = "subtract"
|
||||
MULTIPLY = "multiply"
|
||||
DIVIDE = "divide"
|
||||
POWER = "power"
|
||||
MIN = "min"
|
||||
MAX = "max"
|
||||
|
||||
# @override from Python 3.12
|
||||
def __str__(self) -> str: # type: ignore[explicit-override] # pyright: ignore[reportImplicitOverride]
|
||||
"""
|
||||
Return string representation (useful for pytest logs).
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The operation's name.
|
||||
"""
|
||||
return self.value
|
||||
|
||||
|
||||
class Undef(Enum):
|
||||
"""Sentinel for undefined values."""
|
||||
|
||||
UNDEF = 0
|
||||
|
||||
|
||||
_undef = Undef.UNDEF
|
||||
|
||||
|
||||
class at: # pylint: disable=invalid-name # numpydoc ignore=PR02
|
||||
"""
|
||||
Update operations for read-only arrays.
|
||||
|
||||
This implements ``jax.numpy.ndarray.at`` for all writeable
|
||||
backends (those that support ``__setitem__``) and routes
|
||||
to the ``.at[]`` method for JAX arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array
|
||||
Input array.
|
||||
idx : index, optional
|
||||
Only `array API standard compliant indices
|
||||
<https://data-apis.org/array-api/latest/API_specification/indexing.html>`_
|
||||
are supported.
|
||||
|
||||
You may use two alternate syntaxes::
|
||||
|
||||
>>> import array_api_extra as xpx
|
||||
>>> xpx.at(x, idx).set(value) # or add(value), etc.
|
||||
>>> xpx.at(x)[idx].set(value)
|
||||
|
||||
copy : bool, optional
|
||||
None (default)
|
||||
The array parameter *may* be modified in place if it is
|
||||
possible and beneficial for performance.
|
||||
You should not reuse it after calling this function.
|
||||
True
|
||||
Ensure that the inputs are not modified.
|
||||
False
|
||||
Ensure that the update operation writes back to the input.
|
||||
Raise ``ValueError`` if a copy cannot be avoided.
|
||||
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Updated input array.
|
||||
|
||||
Warnings
|
||||
--------
|
||||
(a) When you omit the ``copy`` parameter, you should never reuse the parameter
|
||||
array later on; ideally, you should reassign it immediately::
|
||||
|
||||
>>> import array_api_extra as xpx
|
||||
>>> x = xpx.at(x, 0).set(2)
|
||||
|
||||
The above best practice pattern ensures that the behaviour won't change depending
|
||||
on whether ``x`` is writeable or not, as the original ``x`` object is dereferenced
|
||||
as soon as ``xpx.at`` returns; this way there is no risk to accidentally update it
|
||||
twice.
|
||||
|
||||
On the reverse, the anti-pattern below must be avoided, as it will result in
|
||||
different behaviour on read-only versus writeable arrays::
|
||||
|
||||
>>> x = xp.asarray([0, 0, 0])
|
||||
>>> y = xpx.at(x, 0).set(2)
|
||||
>>> z = xpx.at(x, 1).set(3)
|
||||
|
||||
In the above example, both calls to ``xpx.at`` update ``x`` in place *if possible*.
|
||||
This causes the behaviour to diverge depending on whether ``x`` is writeable or not:
|
||||
|
||||
- If ``x`` is writeable, then after the snippet above you'll have
|
||||
``x == y == z == [2, 3, 0]``
|
||||
- If ``x`` is read-only, then you'll end up with
|
||||
``x == [0, 0, 0]``, ``y == [2, 0, 0]`` and ``z == [0, 3, 0]``.
|
||||
|
||||
The correct pattern to use if you want diverging outputs from the same input is
|
||||
to enforce copies::
|
||||
|
||||
>>> x = xp.asarray([0, 0, 0])
|
||||
>>> y = xpx.at(x, 0).set(2, copy=True) # Never updates x
|
||||
>>> z = xpx.at(x, 1).set(3) # May or may not update x in place
|
||||
>>> del x # avoid accidental reuse of x as we don't know its state anymore
|
||||
|
||||
(b) The array API standard does not support integer array indices.
|
||||
The behaviour of update methods when the index is an array of integers is
|
||||
undefined and will vary between backends; this is particularly true when the
|
||||
index contains multiple occurrences of the same index, e.g.::
|
||||
|
||||
>>> import numpy as np
|
||||
>>> import jax.numpy as jnp
|
||||
>>> import array_api_extra as xpx
|
||||
>>> xpx.at(np.asarray([123]), np.asarray([0, 0])).add(1)
|
||||
array([124])
|
||||
>>> xpx.at(jnp.asarray([123]), jnp.asarray([0, 0])).add(1)
|
||||
Array([125], dtype=int32)
|
||||
|
||||
See Also
|
||||
--------
|
||||
jax.numpy.ndarray.at : Equivalent array method in JAX.
|
||||
|
||||
Notes
|
||||
-----
|
||||
`sparse <https://sparse.pydata.org/>`_, as well as read-only arrays from libraries
|
||||
not explicitly covered by ``array-api-compat``, are not supported by update
|
||||
methods.
|
||||
|
||||
Boolean masks are supported on Dask and jitted JAX arrays exclusively
|
||||
when `idx` has the same shape as `x` and `y` is 0-dimensional.
|
||||
Note that this support is not available in JAX's native
|
||||
``x.at[mask].set(y)``.
|
||||
|
||||
This pattern::
|
||||
|
||||
>>> mask = m(x)
|
||||
>>> x[mask] = f(x[mask])
|
||||
|
||||
Can't be replaced by `at`, as it won't work on Dask and JAX inside jax.jit::
|
||||
|
||||
>>> mask = m(x)
|
||||
>>> x = xpx.at(x, mask).set(f(x[mask]) # Crash on Dask and jax.jit
|
||||
|
||||
You should instead use::
|
||||
|
||||
>>> x = xp.where(m(x), f(x), x)
|
||||
|
||||
Examples
|
||||
--------
|
||||
Given either of these equivalent expressions::
|
||||
|
||||
>>> import array_api_extra as xpx
|
||||
>>> x = xpx.at(x)[1].add(2)
|
||||
>>> x = xpx.at(x, 1).add(2)
|
||||
|
||||
If x is a JAX array, they are the same as::
|
||||
|
||||
>>> x = x.at[1].add(2)
|
||||
|
||||
If x is a read-only NumPy array, they are the same as::
|
||||
|
||||
>>> x = x.copy()
|
||||
>>> x[1] += 2
|
||||
|
||||
For other known backends, they are the same as::
|
||||
|
||||
>>> x[1] += 2
|
||||
"""
|
||||
|
||||
_x: Array
|
||||
_idx: SetIndex | Undef
|
||||
__slots__: ClassVar[tuple[str, ...]] = ("_idx", "_x")
|
||||
|
||||
def __init__(
|
||||
self, x: Array, idx: SetIndex | Undef = _undef, /
|
||||
) -> None: # numpydoc ignore=GL08
|
||||
self._x = x
|
||||
self._idx = idx
|
||||
|
||||
def __getitem__(self, idx: SetIndex, /) -> Self: # numpydoc ignore=PR01,RT01
|
||||
"""
|
||||
Allow for the alternate syntax ``at(x)[start:stop:step]``.
|
||||
|
||||
It looks prettier than ``at(x, slice(start, stop, step))``
|
||||
and feels more intuitive coming from the JAX documentation.
|
||||
"""
|
||||
if self._idx is not _undef:
|
||||
msg = "Index has already been set"
|
||||
raise ValueError(msg)
|
||||
return type(self)(self._x, idx)
|
||||
|
||||
def _op(
|
||||
self,
|
||||
at_op: _AtOp,
|
||||
in_place_op: Callable[[Array, Array | complex], Array] | None,
|
||||
out_of_place_op: Callable[[Array, Array], Array] | None,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None,
|
||||
xp: ModuleType | None,
|
||||
) -> Array:
|
||||
"""
|
||||
Implement all update operations.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
at_op : _AtOp
|
||||
Method of JAX's Array.at[].
|
||||
in_place_op : Callable[[Array, Array | complex], Array] | None
|
||||
In-place operation to apply on mutable backends::
|
||||
|
||||
x[idx] = in_place_op(x[idx], y)
|
||||
|
||||
If None::
|
||||
|
||||
x[idx] = y
|
||||
|
||||
out_of_place_op : Callable[[Array, Array], Array] | None
|
||||
Out-of-place operation to apply when idx is a boolean mask and the backend
|
||||
doesn't support in-place updates::
|
||||
|
||||
x = xp.where(idx, out_of_place_op(x, y), x)
|
||||
|
||||
If None::
|
||||
|
||||
x = xp.where(idx, y, x)
|
||||
|
||||
y : array or complex
|
||||
Right-hand side of the operation.
|
||||
copy : bool or None
|
||||
Whether to copy the input array. See the class docstring for details.
|
||||
xp : array_namespace, optional
|
||||
The array namespace for the input array. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Array
|
||||
Updated `x`.
|
||||
"""
|
||||
from ._funcs import apply_where # pylint: disable=cyclic-import
|
||||
|
||||
x, idx = self._x, self._idx
|
||||
xp = array_namespace(x, y) if xp is None else xp
|
||||
|
||||
if isinstance(idx, Undef):
|
||||
msg = (
|
||||
"Index has not been set.\n"
|
||||
"Usage: either\n"
|
||||
" at(x, idx).set(value)\n"
|
||||
"or\n"
|
||||
" at(x)[idx].set(value)\n"
|
||||
"(same for all other methods)."
|
||||
)
|
||||
raise ValueError(msg)
|
||||
|
||||
if copy not in (True, False, None):
|
||||
msg = f"copy must be True, False, or None; got {copy!r}"
|
||||
raise ValueError(msg)
|
||||
|
||||
writeable = None if copy else is_writeable_array(x)
|
||||
|
||||
# JAX inside jax.jit doesn't support in-place updates with boolean
|
||||
# masks; Dask exclusively supports __setitem__ but not iops.
|
||||
# We can handle the common special case of 0-dimensional y
|
||||
# with where(idx, y, x) instead.
|
||||
if (
|
||||
(is_dask_array(idx) or is_jax_array(idx))
|
||||
and idx.dtype == xp.bool
|
||||
and idx.shape == x.shape
|
||||
):
|
||||
y_xp = xp.asarray(y, dtype=x.dtype, device=_compat.device(x))
|
||||
if y_xp.ndim == 0:
|
||||
if out_of_place_op: # add(), subtract(), ...
|
||||
# suppress inf warnings on Dask
|
||||
out = apply_where(
|
||||
idx, (x, y_xp), out_of_place_op, fill_value=x, xp=xp
|
||||
)
|
||||
# Undo int->float promotion on JAX after _AtOp.DIVIDE
|
||||
out = xp.astype(out, x.dtype, copy=False)
|
||||
else: # set()
|
||||
out = xp.where(idx, y_xp, x)
|
||||
|
||||
if copy is False:
|
||||
x[()] = out
|
||||
return x
|
||||
return out
|
||||
|
||||
# else: this will work on eager JAX and crash on jax.jit and Dask
|
||||
|
||||
if copy or (copy is None and not writeable):
|
||||
if is_jax_array(x):
|
||||
# Use JAX's at[]
|
||||
func = cast(
|
||||
Callable[[Array | complex], Array],
|
||||
getattr(x.at[idx], at_op.value), # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue,reportUnknownArgumentType]
|
||||
)
|
||||
out = func(y)
|
||||
# Undo int->float promotion on JAX after _AtOp.DIVIDE
|
||||
return xp.astype(out, x.dtype, copy=False)
|
||||
|
||||
# Emulate at[] behaviour for non-JAX arrays
|
||||
# with a copy followed by an update
|
||||
|
||||
x = xp.asarray(x, copy=True)
|
||||
# A copy of a read-only numpy array is writeable
|
||||
# Note: this assumes that a copy of a writeable array is writeable
|
||||
assert not writeable
|
||||
writeable = None
|
||||
|
||||
if writeable is None:
|
||||
writeable = is_writeable_array(x)
|
||||
if not writeable:
|
||||
# sparse crashes here
|
||||
msg = f"Can't update read-only array {x}"
|
||||
raise ValueError(msg)
|
||||
|
||||
# Work around bug in PyTorch where __setitem__ doesn't
|
||||
# always support mismatched dtypes
|
||||
# https://github.com/pytorch/pytorch/issues/150017
|
||||
if is_torch_array(y):
|
||||
y = xp.astype(y, x.dtype, copy=False)
|
||||
|
||||
# Backends without boolean indexing (other than JAX) crash here
|
||||
if in_place_op: # add(), subtract(), ...
|
||||
x[idx] = in_place_op(x[idx], y)
|
||||
else: # set()
|
||||
x[idx] = y
|
||||
return x
|
||||
|
||||
def set(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] = y`` and return the update array."""
|
||||
return self._op(_AtOp.SET, None, None, y, copy=copy, xp=xp)
|
||||
|
||||
def add(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] += y`` and return the updated array."""
|
||||
|
||||
# Note for this and all other methods based on _iop:
|
||||
# operator.iadd and operator.add subtly differ in behaviour, as
|
||||
# only iadd will trigger exceptions when y has an incompatible dtype.
|
||||
return self._op(_AtOp.ADD, operator.iadd, operator.add, y, copy=copy, xp=xp)
|
||||
|
||||
def subtract(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] -= y`` and return the updated array."""
|
||||
return self._op(
|
||||
_AtOp.SUBTRACT, operator.isub, operator.sub, y, copy=copy, xp=xp
|
||||
)
|
||||
|
||||
def multiply(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] *= y`` and return the updated array."""
|
||||
return self._op(
|
||||
_AtOp.MULTIPLY, operator.imul, operator.mul, y, copy=copy, xp=xp
|
||||
)
|
||||
|
||||
def divide(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] /= y`` and return the updated array."""
|
||||
return self._op(
|
||||
_AtOp.DIVIDE, operator.itruediv, operator.truediv, y, copy=copy, xp=xp
|
||||
)
|
||||
|
||||
def power(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] **= y`` and return the updated array."""
|
||||
return self._op(_AtOp.POWER, operator.ipow, operator.pow, y, copy=copy, xp=xp)
|
||||
|
||||
def min(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] = minimum(x[idx], y)`` and return the updated array."""
|
||||
# On Dask, this function runs on the chunks, so we need to determine the
|
||||
# namespace that Dask is wrapping.
|
||||
# Note that da.minimum _incidentally_ works on NumPy, CuPy, and sparse
|
||||
# thanks to all these meta-namespaces implementing the __array_ufunc__
|
||||
# interface, but there's no guarantee that it will work for other
|
||||
# wrapped libraries in the future.
|
||||
xp = array_namespace(self._x) if xp is None else xp
|
||||
mxp = meta_namespace(self._x, xp=xp)
|
||||
y = xp.asarray(y)
|
||||
return self._op(_AtOp.MIN, mxp.minimum, mxp.minimum, y, copy=copy, xp=xp)
|
||||
|
||||
def max(
|
||||
self,
|
||||
y: Array | complex,
|
||||
/,
|
||||
copy: bool | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""Apply ``x[idx] = maximum(x[idx], y)`` and return the updated array."""
|
||||
# See note on min()
|
||||
xp = array_namespace(self._x) if xp is None else xp
|
||||
mxp = meta_namespace(self._x, xp=xp)
|
||||
y = xp.asarray(y)
|
||||
return self._op(_AtOp.MAX, mxp.maximum, mxp.maximum, y, copy=copy, xp=xp)
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Backends against which array-api-extra runs its tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
|
||||
__all__ = ["Backend"]
|
||||
|
||||
|
||||
class Backend(Enum): # numpydoc ignore=PR02
|
||||
"""
|
||||
All array library backends explicitly tested by array-api-extra.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value : str
|
||||
Tag of the backend's module, in the format ``<namespace>[:<extra tag>]``.
|
||||
"""
|
||||
|
||||
# Use :<tag> to prevent Enum from deduplicating items with the same value
|
||||
ARRAY_API_STRICT = "array_api_strict"
|
||||
ARRAY_API_STRICTEST = "array_api_strict:strictest"
|
||||
NUMPY = "numpy"
|
||||
NUMPY_READONLY = "numpy:readonly"
|
||||
CUPY = "cupy"
|
||||
TORCH = "torch"
|
||||
TORCH_GPU = "torch:gpu"
|
||||
DASK = "dask.array"
|
||||
SPARSE = "sparse"
|
||||
JAX = "jax.numpy"
|
||||
JAX_GPU = "jax.numpy:gpu"
|
||||
|
||||
def __str__(self) -> str: # type: ignore[explicit-override] # pyright: ignore[reportImplicitOverride] # numpydoc ignore=RT01
|
||||
"""Pretty-print parameterized test names."""
|
||||
return (
|
||||
self.name.lower().replace("_gpu", ":gpu").replace("_readonly", ":readonly")
|
||||
)
|
||||
|
||||
@property
|
||||
def modname(self) -> str: # numpydoc ignore=RT01
|
||||
"""Module name to be imported."""
|
||||
return self.value.split(":")[0]
|
||||
|
||||
def like(self, *others: Backend) -> bool: # numpydoc ignore=PR01,RT01
|
||||
"""Check if this backend uses the same module as others."""
|
||||
return any(self.modname == other.modname for other in others)
|
||||
@@ -0,0 +1,937 @@
|
||||
"""Array-agnostic implementations for the public API."""
|
||||
|
||||
import math
|
||||
import warnings
|
||||
from collections.abc import Callable, Sequence
|
||||
from types import ModuleType, NoneType
|
||||
from typing import cast, overload
|
||||
|
||||
from ._at import at
|
||||
from ._utils import _compat, _helpers
|
||||
from ._utils._compat import array_namespace, is_dask_namespace, is_jax_array
|
||||
from ._utils._helpers import (
|
||||
asarrays,
|
||||
capabilities,
|
||||
eager_shape,
|
||||
meta_namespace,
|
||||
ndindex,
|
||||
)
|
||||
from ._utils._typing import Array
|
||||
|
||||
__all__ = [
|
||||
"apply_where",
|
||||
"atleast_nd",
|
||||
"broadcast_shapes",
|
||||
"cov",
|
||||
"create_diagonal",
|
||||
"expand_dims",
|
||||
"kron",
|
||||
"nunique",
|
||||
"pad",
|
||||
"setdiff1d",
|
||||
"sinc",
|
||||
]
|
||||
|
||||
|
||||
@overload
|
||||
def apply_where( # type: ignore[explicit-any,decorated-any] # numpydoc ignore=GL08
|
||||
cond: Array,
|
||||
args: Array | tuple[Array, ...],
|
||||
f1: Callable[..., Array],
|
||||
f2: Callable[..., Array],
|
||||
/,
|
||||
*,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: ...
|
||||
|
||||
|
||||
@overload
|
||||
def apply_where( # type: ignore[explicit-any,decorated-any] # numpydoc ignore=GL08
|
||||
cond: Array,
|
||||
args: Array | tuple[Array, ...],
|
||||
f1: Callable[..., Array],
|
||||
/,
|
||||
*,
|
||||
fill_value: Array | complex,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: ...
|
||||
|
||||
|
||||
def apply_where( # type: ignore[explicit-any] # numpydoc ignore=PR01,PR02
|
||||
cond: Array,
|
||||
args: Array | tuple[Array, ...],
|
||||
f1: Callable[..., Array],
|
||||
f2: Callable[..., Array] | None = None,
|
||||
/,
|
||||
*,
|
||||
fill_value: Array | complex | None = None,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array:
|
||||
"""
|
||||
Run one of two elementwise functions depending on a condition.
|
||||
|
||||
Equivalent to ``f1(*args) if cond else fill_value`` performed elementwise
|
||||
when `fill_value` is defined, otherwise to ``f1(*args) if cond else f2(*args)``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cond : array
|
||||
The condition, expressed as a boolean array.
|
||||
args : Array or tuple of Arrays
|
||||
Argument(s) to `f1` (and `f2`). Must be broadcastable with `cond`.
|
||||
f1 : callable
|
||||
Elementwise function of `args`, returning a single array.
|
||||
Where `cond` is True, output will be ``f1(arg0[cond], arg1[cond], ...)``.
|
||||
f2 : callable, optional
|
||||
Elementwise function of `args`, returning a single array.
|
||||
Where `cond` is False, output will be ``f2(arg0[cond], arg1[cond], ...)``.
|
||||
Mutually exclusive with `fill_value`.
|
||||
fill_value : Array or scalar, optional
|
||||
If provided, value with which to fill output array where `cond` is False.
|
||||
It does not need to be scalar; it needs however to be broadcastable with
|
||||
`cond` and `args`.
|
||||
Mutually exclusive with `f2`. You must provide one or the other.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `cond` and `args`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Array
|
||||
An array with elements from the output of `f1` where `cond` is True and either
|
||||
the output of `f2` or `fill_value` where `cond` is False. The returned array has
|
||||
data type determined by type promotion rules between the output of `f1` and
|
||||
either `fill_value` or the output of `f2`.
|
||||
|
||||
Notes
|
||||
-----
|
||||
``xp.where(cond, f1(*args), f2(*args))`` requires explicitly evaluating `f1` even
|
||||
when `cond` is False, and `f2` when cond is True. This function evaluates each
|
||||
function only for their matching condition, if the backend allows for it.
|
||||
|
||||
On Dask, `f1` and `f2` are applied to the individual chunks and should use functions
|
||||
from the namespace of the chunks.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
>>> a = xp.asarray([5, 4, 3])
|
||||
>>> b = xp.asarray([0, 2, 2])
|
||||
>>> def f(a, b):
|
||||
... return a // b
|
||||
>>> xpx.apply_where(b != 0, (a, b), f, fill_value=xp.nan)
|
||||
array([ nan, 2., 1.])
|
||||
"""
|
||||
# Parse and normalize arguments
|
||||
if (f2 is None) == (fill_value is None):
|
||||
msg = "Exactly one of `fill_value` or `f2` must be given."
|
||||
raise TypeError(msg)
|
||||
args_ = list(args) if isinstance(args, tuple) else [args]
|
||||
del args
|
||||
|
||||
xp = array_namespace(cond, fill_value, *args_) if xp is None else xp
|
||||
|
||||
if isinstance(fill_value, int | float | complex | NoneType):
|
||||
cond, *args_ = xp.broadcast_arrays(cond, *args_)
|
||||
else:
|
||||
cond, fill_value, *args_ = xp.broadcast_arrays(cond, fill_value, *args_)
|
||||
|
||||
if is_dask_namespace(xp):
|
||||
meta_xp = meta_namespace(cond, fill_value, *args_, xp=xp)
|
||||
# map_blocks doesn't descend into tuples of Arrays
|
||||
return xp.map_blocks(_apply_where, cond, f1, f2, fill_value, *args_, xp=meta_xp)
|
||||
return _apply_where(cond, f1, f2, fill_value, *args_, xp=xp)
|
||||
|
||||
|
||||
def _apply_where( # type: ignore[explicit-any] # numpydoc ignore=PR01,RT01
|
||||
cond: Array,
|
||||
f1: Callable[..., Array],
|
||||
f2: Callable[..., Array] | None,
|
||||
fill_value: Array | int | float | complex | bool | None,
|
||||
*args: Array,
|
||||
xp: ModuleType,
|
||||
) -> Array:
|
||||
"""Helper of `apply_where`. On Dask, this runs on a single chunk."""
|
||||
|
||||
if not capabilities(xp)["boolean indexing"]:
|
||||
# jax.jit does not support assignment by boolean mask
|
||||
return xp.where(cond, f1(*args), f2(*args) if f2 is not None else fill_value)
|
||||
|
||||
temp1 = f1(*(arr[cond] for arr in args))
|
||||
|
||||
if f2 is None:
|
||||
dtype = xp.result_type(temp1, fill_value)
|
||||
if isinstance(fill_value, int | float | complex):
|
||||
out = xp.full_like(cond, dtype=dtype, fill_value=fill_value)
|
||||
else:
|
||||
out = xp.astype(fill_value, dtype, copy=True)
|
||||
else:
|
||||
ncond = ~cond
|
||||
temp2 = f2(*(arr[ncond] for arr in args))
|
||||
dtype = xp.result_type(temp1, temp2)
|
||||
out = xp.empty_like(cond, dtype=dtype)
|
||||
out = at(out, ncond).set(temp2)
|
||||
|
||||
return at(out, cond).set(temp1)
|
||||
|
||||
|
||||
def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array:
|
||||
"""
|
||||
Recursively expand the dimension of an array to at least `ndim`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array
|
||||
Input array.
|
||||
ndim : int
|
||||
The minimum number of dimensions for the result.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
An array with ``res.ndim`` >= `ndim`.
|
||||
If ``x.ndim`` >= `ndim`, `x` is returned.
|
||||
If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes
|
||||
until ``res.ndim`` equals `ndim`.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
>>> x = xp.asarray([1])
|
||||
>>> xpx.atleast_nd(x, ndim=3, xp=xp)
|
||||
Array([[[1]]], dtype=array_api_strict.int64)
|
||||
|
||||
>>> x = xp.asarray([[[1, 2],
|
||||
... [3, 4]]])
|
||||
>>> xpx.atleast_nd(x, ndim=1, xp=xp) is x
|
||||
True
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(x)
|
||||
|
||||
if x.ndim < ndim:
|
||||
x = xp.expand_dims(x, axis=0)
|
||||
x = atleast_nd(x, ndim=ndim, xp=xp)
|
||||
return x
|
||||
|
||||
|
||||
# `float` in signature to accept `math.nan` for Dask.
|
||||
# `int`s are still accepted as `float` is a superclass of `int` in typing
|
||||
def broadcast_shapes(*shapes: tuple[float | None, ...]) -> tuple[int | None, ...]:
|
||||
"""
|
||||
Compute the shape of the broadcasted arrays.
|
||||
|
||||
Duplicates :func:`numpy.broadcast_shapes`, with additional support for
|
||||
None and NaN sizes.
|
||||
|
||||
This is equivalent to ``xp.broadcast_arrays(arr1, arr2, ...)[0].shape``
|
||||
without needing to worry about the backend potentially deep copying
|
||||
the arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
*shapes : tuple[int | None, ...]
|
||||
Shapes of the arrays to broadcast.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple[int | None, ...]
|
||||
The shape of the broadcasted arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
numpy.broadcast_shapes : Equivalent NumPy function.
|
||||
array_api.broadcast_arrays : Function to broadcast actual arrays.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function accepts the Array API's ``None`` for unknown sizes,
|
||||
as well as Dask's non-standard ``math.nan``.
|
||||
Regardless of input, the output always contains ``None`` for unknown sizes.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_extra as xpx
|
||||
>>> xpx.broadcast_shapes((2, 3), (2, 1))
|
||||
(2, 3)
|
||||
>>> xpx.broadcast_shapes((4, 2, 3), (2, 1), (1, 3))
|
||||
(4, 2, 3)
|
||||
"""
|
||||
if not shapes:
|
||||
return () # Match NumPy output
|
||||
|
||||
ndim = max(len(shape) for shape in shapes)
|
||||
out: list[int | None] = []
|
||||
for axis in range(-ndim, 0):
|
||||
sizes = {shape[axis] for shape in shapes if axis >= -len(shape)}
|
||||
# Dask uses NaN for unknown shape, which predates the Array API spec for None
|
||||
none_size = None in sizes or math.nan in sizes
|
||||
sizes -= {1, None, math.nan}
|
||||
if len(sizes) > 1:
|
||||
msg = (
|
||||
"shape mismatch: objects cannot be broadcast to a single shape: "
|
||||
f"{shapes}."
|
||||
)
|
||||
raise ValueError(msg)
|
||||
out.append(None if none_size else cast(int, sizes.pop()) if sizes else 1)
|
||||
|
||||
return tuple(out)
|
||||
|
||||
|
||||
def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
|
||||
"""
|
||||
Estimate a covariance matrix.
|
||||
|
||||
Covariance indicates the level to which two variables vary together.
|
||||
If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`,
|
||||
then the covariance matrix element :math:`C_{ij}` is the covariance of
|
||||
:math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance
|
||||
of :math:`x_i`.
|
||||
|
||||
This provides a subset of the functionality of ``numpy.cov``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
m : array
|
||||
A 1-D or 2-D array containing multiple variables and observations.
|
||||
Each row of `m` represents a variable, and each column a single
|
||||
observation of all those variables.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `m`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
The covariance matrix of the variables.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
|
||||
Consider two variables, :math:`x_0` and :math:`x_1`, which
|
||||
correlate perfectly, but in opposite directions:
|
||||
|
||||
>>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T
|
||||
>>> x
|
||||
Array([[0, 1, 2],
|
||||
[2, 1, 0]], dtype=array_api_strict.int64)
|
||||
|
||||
Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance
|
||||
matrix shows this clearly:
|
||||
|
||||
>>> xpx.cov(x, xp=xp)
|
||||
Array([[ 1., -1.],
|
||||
[-1., 1.]], dtype=array_api_strict.float64)
|
||||
|
||||
Note that element :math:`C_{0,1}`, which shows the correlation between
|
||||
:math:`x_0` and :math:`x_1`, is negative.
|
||||
|
||||
Further, note how `x` and `y` are combined:
|
||||
|
||||
>>> x = xp.asarray([-2.1, -1, 4.3])
|
||||
>>> y = xp.asarray([3, 1.1, 0.12])
|
||||
>>> X = xp.stack((x, y), axis=0)
|
||||
>>> xpx.cov(X, xp=xp)
|
||||
Array([[11.71 , -4.286 ],
|
||||
[-4.286 , 2.14413333]], dtype=array_api_strict.float64)
|
||||
|
||||
>>> xpx.cov(x, xp=xp)
|
||||
Array(11.71, dtype=array_api_strict.float64)
|
||||
|
||||
>>> xpx.cov(y, xp=xp)
|
||||
Array(2.14413333, dtype=array_api_strict.float64)
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(m)
|
||||
|
||||
m = xp.asarray(m, copy=True)
|
||||
dtype = (
|
||||
xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64)
|
||||
)
|
||||
|
||||
m = atleast_nd(m, ndim=2, xp=xp)
|
||||
m = xp.astype(m, dtype)
|
||||
|
||||
avg = _helpers.mean(m, axis=1, xp=xp)
|
||||
|
||||
m_shape = eager_shape(m)
|
||||
fact = m_shape[1] - 1
|
||||
|
||||
if fact <= 0:
|
||||
warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2)
|
||||
fact = 0
|
||||
|
||||
m -= avg[:, None]
|
||||
m_transpose = m.T
|
||||
if xp.isdtype(m_transpose.dtype, "complex floating"):
|
||||
m_transpose = xp.conj(m_transpose)
|
||||
c = m @ m_transpose
|
||||
c /= fact
|
||||
axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
|
||||
return xp.squeeze(c, axis=axes)
|
||||
|
||||
|
||||
def create_diagonal(
|
||||
x: Array, /, *, offset: int = 0, xp: ModuleType | None = None
|
||||
) -> Array:
|
||||
"""
|
||||
Construct a diagonal array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array
|
||||
An array having shape ``(*batch_dims, k)``.
|
||||
offset : int, optional
|
||||
Offset from the leading diagonal (default is ``0``).
|
||||
Use positive ints for diagonals above the leading diagonal,
|
||||
and negative ints for diagonals below the leading diagonal.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
An array having shape ``(*batch_dims, k+abs(offset), k+abs(offset))`` with `x`
|
||||
on the diagonal (offset by `offset`).
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
>>> x = xp.asarray([2, 4, 8])
|
||||
|
||||
>>> xpx.create_diagonal(x, xp=xp)
|
||||
Array([[2, 0, 0],
|
||||
[0, 4, 0],
|
||||
[0, 0, 8]], dtype=array_api_strict.int64)
|
||||
|
||||
>>> xpx.create_diagonal(x, offset=-2, xp=xp)
|
||||
Array([[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[2, 0, 0, 0, 0],
|
||||
[0, 4, 0, 0, 0],
|
||||
[0, 0, 8, 0, 0]], dtype=array_api_strict.int64)
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(x)
|
||||
|
||||
if x.ndim == 0:
|
||||
err_msg = "`x` must be at least 1-dimensional."
|
||||
raise ValueError(err_msg)
|
||||
|
||||
x_shape = eager_shape(x)
|
||||
batch_dims = x_shape[:-1]
|
||||
n = x_shape[-1] + abs(offset)
|
||||
diag = xp.zeros((*batch_dims, n**2), dtype=x.dtype, device=_compat.device(x))
|
||||
|
||||
target_slice = slice(
|
||||
offset if offset >= 0 else abs(offset) * n,
|
||||
min(n * (n - offset), diag.shape[-1]),
|
||||
n + 1,
|
||||
)
|
||||
for index in ndindex(*batch_dims):
|
||||
diag = at(diag)[(*index, target_slice)].set(x[(*index, slice(None))])
|
||||
return xp.reshape(diag, (*batch_dims, n, n))
|
||||
|
||||
|
||||
def expand_dims(
|
||||
a: Array, /, *, axis: int | tuple[int, ...] = (0,), xp: ModuleType | None = None
|
||||
) -> Array:
|
||||
"""
|
||||
Expand the shape of an array.
|
||||
|
||||
Insert (a) new axis/axes that will appear at the position(s) specified by
|
||||
`axis` in the expanded array shape.
|
||||
|
||||
This is ``xp.expand_dims`` for `axis` an int *or a tuple of ints*.
|
||||
Roughly equivalent to ``numpy.expand_dims`` for NumPy arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a : array
|
||||
Array to have its shape expanded.
|
||||
axis : int or tuple of ints, optional
|
||||
Position(s) in the expanded axes where the new axis (or axes) is/are placed.
|
||||
If multiple positions are provided, they should be unique (note that a position
|
||||
given by a positive index could also be referred to by a negative index -
|
||||
that will also result in an error).
|
||||
Default: ``(0,)``.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `a`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
`a` with an expanded shape.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
>>> x = xp.asarray([1, 2])
|
||||
>>> x.shape
|
||||
(2,)
|
||||
|
||||
The following is equivalent to ``x[xp.newaxis, :]`` or ``x[xp.newaxis]``:
|
||||
|
||||
>>> y = xpx.expand_dims(x, axis=0, xp=xp)
|
||||
>>> y
|
||||
Array([[1, 2]], dtype=array_api_strict.int64)
|
||||
>>> y.shape
|
||||
(1, 2)
|
||||
|
||||
The following is equivalent to ``x[:, xp.newaxis]``:
|
||||
|
||||
>>> y = xpx.expand_dims(x, axis=1, xp=xp)
|
||||
>>> y
|
||||
Array([[1],
|
||||
[2]], dtype=array_api_strict.int64)
|
||||
>>> y.shape
|
||||
(2, 1)
|
||||
|
||||
``axis`` may also be a tuple:
|
||||
|
||||
>>> y = xpx.expand_dims(x, axis=(0, 1), xp=xp)
|
||||
>>> y
|
||||
Array([[[1, 2]]], dtype=array_api_strict.int64)
|
||||
|
||||
>>> y = xpx.expand_dims(x, axis=(2, 0), xp=xp)
|
||||
>>> y
|
||||
Array([[[1],
|
||||
[2]]], dtype=array_api_strict.int64)
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(a)
|
||||
|
||||
if not isinstance(axis, tuple):
|
||||
axis = (axis,)
|
||||
ndim = a.ndim + len(axis)
|
||||
if axis != () and (min(axis) < -ndim or max(axis) >= ndim):
|
||||
err_msg = (
|
||||
f"a provided axis position is out of bounds for array of dimension {a.ndim}"
|
||||
)
|
||||
raise IndexError(err_msg)
|
||||
axis = tuple(dim % ndim for dim in axis)
|
||||
if len(set(axis)) != len(axis):
|
||||
err_msg = "Duplicate dimensions specified in `axis`."
|
||||
raise ValueError(err_msg)
|
||||
for i in sorted(axis):
|
||||
a = xp.expand_dims(a, axis=i)
|
||||
return a
|
||||
|
||||
|
||||
def isclose(
|
||||
a: Array | complex,
|
||||
b: Array | complex,
|
||||
*,
|
||||
rtol: float = 1e-05,
|
||||
atol: float = 1e-08,
|
||||
equal_nan: bool = False,
|
||||
xp: ModuleType,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""See docstring in array_api_extra._delegation."""
|
||||
a, b = asarrays(a, b, xp=xp)
|
||||
|
||||
a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating"))
|
||||
b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating"))
|
||||
if a_inexact or b_inexact:
|
||||
# prevent warnings on NumPy and Dask on inf - inf
|
||||
mxp = meta_namespace(a, b, xp=xp)
|
||||
out = apply_where(
|
||||
xp.isinf(a) | xp.isinf(b),
|
||||
(a, b),
|
||||
lambda a, b: mxp.isinf(a) & mxp.isinf(b) & (mxp.sign(a) == mxp.sign(b)), # pyright: ignore[reportUnknownArgumentType]
|
||||
# Note: inf <= inf is True!
|
||||
lambda a, b: mxp.abs(a - b) <= (atol + rtol * mxp.abs(b)), # pyright: ignore[reportUnknownArgumentType]
|
||||
xp=xp,
|
||||
)
|
||||
if equal_nan:
|
||||
out = xp.where(xp.isnan(a) & xp.isnan(b), True, out)
|
||||
return out
|
||||
|
||||
if xp.isdtype(a.dtype, "bool") or xp.isdtype(b.dtype, "bool"):
|
||||
if atol >= 1 or rtol >= 1:
|
||||
return xp.ones_like(a == b)
|
||||
return a == b
|
||||
|
||||
# integer types
|
||||
atol = int(atol)
|
||||
if rtol == 0:
|
||||
return xp.abs(a - b) <= atol
|
||||
|
||||
# Don't rely on OverflowError, as it is not guaranteed by the Array API.
|
||||
nrtol = int(1.0 / rtol)
|
||||
if nrtol > xp.iinfo(b.dtype).max:
|
||||
# rtol * max_int < 1, so it's inconsequential
|
||||
return xp.abs(a - b) <= atol
|
||||
return xp.abs(a - b) <= (atol + xp.abs(b) // nrtol)
|
||||
|
||||
|
||||
def kron(
|
||||
a: Array | complex,
|
||||
b: Array | complex,
|
||||
/,
|
||||
*,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array:
|
||||
"""
|
||||
Kronecker product of two arrays.
|
||||
|
||||
Computes the Kronecker product, a composite array made of blocks of the
|
||||
second array scaled by the first.
|
||||
|
||||
Equivalent to ``numpy.kron`` for NumPy arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a, b : Array | int | float | complex
|
||||
Input arrays or scalars. At least one must be an array.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `a` and `b`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
The Kronecker product of `a` and `b`.
|
||||
|
||||
Notes
|
||||
-----
|
||||
The function assumes that the number of dimensions of `a` and `b`
|
||||
are the same, if necessary prepending the smallest with ones.
|
||||
If ``a.shape = (r0,r1,..,rN)`` and ``b.shape = (s0,s1,...,sN)``,
|
||||
the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``.
|
||||
The elements are products of elements from `a` and `b`, organized
|
||||
explicitly by::
|
||||
|
||||
kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN]
|
||||
|
||||
where::
|
||||
|
||||
kt = it * st + jt, t = 0,...,N
|
||||
|
||||
In the common 2-D case (N=1), the block structure can be visualized::
|
||||
|
||||
[[ a[0,0]*b, a[0,1]*b, ... , a[0,-1]*b ],
|
||||
[ ... ... ],
|
||||
[ a[-1,0]*b, a[-1,1]*b, ... , a[-1,-1]*b ]]
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
>>> xpx.kron(xp.asarray([1, 10, 100]), xp.asarray([5, 6, 7]), xp=xp)
|
||||
Array([ 5, 6, 7, 50, 60, 70, 500,
|
||||
600, 700], dtype=array_api_strict.int64)
|
||||
|
||||
>>> xpx.kron(xp.asarray([5, 6, 7]), xp.asarray([1, 10, 100]), xp=xp)
|
||||
Array([ 5, 50, 500, 6, 60, 600, 7,
|
||||
70, 700], dtype=array_api_strict.int64)
|
||||
|
||||
>>> xpx.kron(xp.eye(2), xp.ones((2, 2)), xp=xp)
|
||||
Array([[1., 1., 0., 0.],
|
||||
[1., 1., 0., 0.],
|
||||
[0., 0., 1., 1.],
|
||||
[0., 0., 1., 1.]], dtype=array_api_strict.float64)
|
||||
|
||||
>>> a = xp.reshape(xp.arange(100), (2, 5, 2, 5))
|
||||
>>> b = xp.reshape(xp.arange(24), (2, 3, 4))
|
||||
>>> c = xpx.kron(a, b, xp=xp)
|
||||
>>> c.shape
|
||||
(2, 10, 6, 20)
|
||||
>>> I = (1, 3, 0, 2)
|
||||
>>> J = (0, 2, 1)
|
||||
>>> J1 = (0,) + J # extend to ndim=4
|
||||
>>> S1 = (1,) + b.shape
|
||||
>>> K = tuple(xp.asarray(I) * xp.asarray(S1) + xp.asarray(J1))
|
||||
>>> c[K] == a[I]*b[J]
|
||||
Array(True, dtype=array_api_strict.bool)
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(a, b)
|
||||
a, b = asarrays(a, b, xp=xp)
|
||||
|
||||
singletons = (1,) * (b.ndim - a.ndim)
|
||||
a = cast(Array, xp.broadcast_to(a, singletons + a.shape))
|
||||
|
||||
nd_b, nd_a = b.ndim, a.ndim
|
||||
nd_max = max(nd_b, nd_a)
|
||||
if nd_a == 0 or nd_b == 0:
|
||||
return xp.multiply(a, b)
|
||||
|
||||
a_shape = eager_shape(a)
|
||||
b_shape = eager_shape(b)
|
||||
|
||||
# Equalise the shapes by prepending smaller one with 1s
|
||||
a_shape = (1,) * max(0, nd_b - nd_a) + a_shape
|
||||
b_shape = (1,) * max(0, nd_a - nd_b) + b_shape
|
||||
|
||||
# Insert empty dimensions
|
||||
a_arr = expand_dims(a, axis=tuple(range(nd_b - nd_a)), xp=xp)
|
||||
b_arr = expand_dims(b, axis=tuple(range(nd_a - nd_b)), xp=xp)
|
||||
|
||||
# Compute the product
|
||||
a_arr = expand_dims(a_arr, axis=tuple(range(1, nd_max * 2, 2)), xp=xp)
|
||||
b_arr = expand_dims(b_arr, axis=tuple(range(0, nd_max * 2, 2)), xp=xp)
|
||||
result = xp.multiply(a_arr, b_arr)
|
||||
|
||||
# Reshape back and return
|
||||
res_shape = tuple(a_s * b_s for a_s, b_s in zip(a_shape, b_shape, strict=True))
|
||||
return xp.reshape(result, res_shape)
|
||||
|
||||
|
||||
def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array:
|
||||
"""
|
||||
Count the number of unique elements in an array.
|
||||
|
||||
Compatible with JAX and Dask, whose laziness would be otherwise
|
||||
problematic.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : Array
|
||||
Input array.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array: 0-dimensional integer array
|
||||
The number of unique elements in `x`. It can be lazy.
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(x)
|
||||
|
||||
if is_jax_array(x):
|
||||
# size= is JAX-specific
|
||||
# https://github.com/data-apis/array-api/issues/883
|
||||
_, counts = xp.unique_counts(x, size=_compat.size(x))
|
||||
return (counts > 0).sum()
|
||||
|
||||
# There are 3 general use cases:
|
||||
# 1. backend has unique_counts and it returns an array with known shape
|
||||
# 2. backend has unique_counts and it returns a None-sized array;
|
||||
# e.g. Dask, ndonnx
|
||||
# 3. backend does not have unique_counts; e.g. wrapped JAX
|
||||
if capabilities(xp)["data-dependent shapes"]:
|
||||
# xp has unique_counts; O(n) complexity
|
||||
_, counts = xp.unique_counts(x)
|
||||
n = _compat.size(counts)
|
||||
if n is None:
|
||||
return xp.sum(xp.ones_like(counts))
|
||||
return xp.asarray(n, device=_compat.device(x))
|
||||
|
||||
# xp does not have unique_counts; O(n*logn) complexity
|
||||
x = xp.reshape(x, (-1,))
|
||||
x = xp.sort(x)
|
||||
mask = x != xp.roll(x, -1)
|
||||
default_int = xp.__array_namespace_info__().default_dtypes(
|
||||
device=_compat.device(x)
|
||||
)["integral"]
|
||||
return xp.maximum(
|
||||
# Special cases:
|
||||
# - array is size 0
|
||||
# - array has all elements equal to each other
|
||||
xp.astype(xp.any(~mask), default_int),
|
||||
xp.sum(xp.astype(mask, default_int)),
|
||||
)
|
||||
|
||||
|
||||
def pad(
|
||||
x: Array,
|
||||
pad_width: int | tuple[int, int] | Sequence[tuple[int, int]],
|
||||
*,
|
||||
constant_values: complex = 0,
|
||||
xp: ModuleType,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""See docstring in `array_api_extra._delegation.py`."""
|
||||
# make pad_width a list of length-2 tuples of ints
|
||||
if isinstance(pad_width, int):
|
||||
pad_width_seq = [(pad_width, pad_width)] * x.ndim
|
||||
elif (
|
||||
isinstance(pad_width, tuple)
|
||||
and len(pad_width) == 2
|
||||
and all(isinstance(i, int) for i in pad_width)
|
||||
):
|
||||
pad_width_seq = [cast(tuple[int, int], pad_width)] * x.ndim
|
||||
else:
|
||||
pad_width_seq = cast(list[tuple[int, int]], list(pad_width))
|
||||
|
||||
# https://github.com/python/typeshed/issues/13376
|
||||
slices: list[slice] = [] # type: ignore[explicit-any]
|
||||
newshape: list[int] = []
|
||||
for ax, w_tpl in enumerate(pad_width_seq):
|
||||
if len(w_tpl) != 2:
|
||||
msg = f"expect a 2-tuple (before, after), got {w_tpl}."
|
||||
raise ValueError(msg)
|
||||
|
||||
sh = eager_shape(x)[ax]
|
||||
|
||||
if w_tpl[0] == 0 and w_tpl[1] == 0:
|
||||
sl = slice(None, None, None)
|
||||
else:
|
||||
start, stop = w_tpl
|
||||
stop = None if stop == 0 else -stop
|
||||
|
||||
sl = slice(start, stop, None)
|
||||
sh += w_tpl[0] + w_tpl[1]
|
||||
|
||||
newshape.append(sh)
|
||||
slices.append(sl)
|
||||
|
||||
padded = xp.full(
|
||||
tuple(newshape),
|
||||
fill_value=constant_values,
|
||||
dtype=x.dtype,
|
||||
device=_compat.device(x),
|
||||
)
|
||||
return at(padded, tuple(slices)).set(x)
|
||||
|
||||
|
||||
def setdiff1d(
|
||||
x1: Array | complex,
|
||||
x2: Array | complex,
|
||||
/,
|
||||
*,
|
||||
assume_unique: bool = False,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array:
|
||||
"""
|
||||
Find the set difference of two arrays.
|
||||
|
||||
Return the unique values in `x1` that are not in `x2`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : array | int | float | complex | bool
|
||||
Input array.
|
||||
x2 : array
|
||||
Input comparison array.
|
||||
assume_unique : bool
|
||||
If ``True``, the input arrays are both assumed to be unique, which
|
||||
can speed up the calculation. Default is ``False``.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x1` and `x2`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
1D array of values in `x1` that are not in `x2`. The result
|
||||
is sorted when `assume_unique` is ``False``, but otherwise only sorted
|
||||
if the input is sorted.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
|
||||
>>> x1 = xp.asarray([1, 2, 3, 2, 4, 1])
|
||||
>>> x2 = xp.asarray([3, 4, 5, 6])
|
||||
>>> xpx.setdiff1d(x1, x2, xp=xp)
|
||||
Array([1, 2], dtype=array_api_strict.int64)
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(x1, x2)
|
||||
# https://github.com/microsoft/pyright/issues/10103
|
||||
x1_, x2_ = asarrays(x1, x2, xp=xp)
|
||||
|
||||
if assume_unique:
|
||||
x1_ = xp.reshape(x1_, (-1,))
|
||||
x2_ = xp.reshape(x2_, (-1,))
|
||||
else:
|
||||
x1_ = xp.unique_values(x1_)
|
||||
x2_ = xp.unique_values(x2_)
|
||||
|
||||
return x1_[_helpers.in1d(x1_, x2_, assume_unique=True, invert=True, xp=xp)]
|
||||
|
||||
|
||||
def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
|
||||
r"""
|
||||
Return the normalized sinc function.
|
||||
|
||||
The sinc function is equal to :math:`\sin(\pi x)/(\pi x)` for any argument
|
||||
:math:`x\ne 0`. ``sinc(0)`` takes the limit value 1, making ``sinc`` not
|
||||
only everywhere continuous but also infinitely differentiable.
|
||||
|
||||
.. note::
|
||||
|
||||
Note the normalization factor of ``pi`` used in the definition.
|
||||
This is the most commonly used definition in signal processing.
|
||||
Use ``sinc(x / xp.pi)`` to obtain the unnormalized sinc function
|
||||
:math:`\sin(x)/x` that is more common in mathematics.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array
|
||||
Array (possibly multi-dimensional) of values for which to calculate
|
||||
``sinc(x)``. Must have a real floating point dtype.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array
|
||||
``sinc(x)`` calculated elementwise, which has the same shape as the input.
|
||||
|
||||
Notes
|
||||
-----
|
||||
The name sinc is short for "sine cardinal" or "sinus cardinalis".
|
||||
|
||||
The sinc function is used in various signal processing applications,
|
||||
including in anti-aliasing, in the construction of a Lanczos resampling
|
||||
filter, and in interpolation.
|
||||
|
||||
For bandlimited interpolation of discrete-time signals, the ideal
|
||||
interpolation kernel is proportional to the sinc function.
|
||||
|
||||
References
|
||||
----------
|
||||
#. Weisstein, Eric W. "Sinc Function." From MathWorld--A Wolfram Web
|
||||
Resource. https://mathworld.wolfram.com/SincFunction.html
|
||||
#. Wikipedia, "Sinc function",
|
||||
https://en.wikipedia.org/wiki/Sinc_function
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import array_api_strict as xp
|
||||
>>> import array_api_extra as xpx
|
||||
>>> x = xp.linspace(-4, 4, 41)
|
||||
>>> xpx.sinc(x, xp=xp)
|
||||
Array([-3.89817183e-17, -4.92362781e-02,
|
||||
-8.40918587e-02, -8.90384387e-02,
|
||||
-5.84680802e-02, 3.89817183e-17,
|
||||
6.68206631e-02, 1.16434881e-01,
|
||||
1.26137788e-01, 8.50444803e-02,
|
||||
-3.89817183e-17, -1.03943254e-01,
|
||||
-1.89206682e-01, -2.16236208e-01,
|
||||
-1.55914881e-01, 3.89817183e-17,
|
||||
2.33872321e-01, 5.04551152e-01,
|
||||
7.56826729e-01, 9.35489284e-01,
|
||||
1.00000000e+00, 9.35489284e-01,
|
||||
7.56826729e-01, 5.04551152e-01,
|
||||
2.33872321e-01, 3.89817183e-17,
|
||||
-1.55914881e-01, -2.16236208e-01,
|
||||
-1.89206682e-01, -1.03943254e-01,
|
||||
-3.89817183e-17, 8.50444803e-02,
|
||||
1.26137788e-01, 1.16434881e-01,
|
||||
6.68206631e-02, 3.89817183e-17,
|
||||
-5.84680802e-02, -8.90384387e-02,
|
||||
-8.40918587e-02, -4.92362781e-02,
|
||||
-3.89817183e-17], dtype=array_api_strict.float64)
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(x)
|
||||
|
||||
if not xp.isdtype(x.dtype, "real floating"):
|
||||
err_msg = "`x` must have a real floating data type."
|
||||
raise ValueError(err_msg)
|
||||
# no scalars in `where` - array-api#807
|
||||
y = xp.pi * xp.where(
|
||||
xp.astype(x, xp.bool),
|
||||
x,
|
||||
xp.asarray(xp.finfo(x.dtype).eps, dtype=x.dtype, device=_compat.device(x)),
|
||||
)
|
||||
return xp.sin(y) / y
|
||||
@@ -0,0 +1,357 @@
|
||||
"""Public API Functions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Callable, Sequence
|
||||
from functools import partial, wraps
|
||||
from types import ModuleType
|
||||
from typing import TYPE_CHECKING, Any, ParamSpec, TypeAlias, cast, overload
|
||||
|
||||
from ._funcs import broadcast_shapes
|
||||
from ._utils import _compat
|
||||
from ._utils._compat import (
|
||||
array_namespace,
|
||||
is_dask_namespace,
|
||||
is_jax_namespace,
|
||||
)
|
||||
from ._utils._helpers import is_python_scalar
|
||||
from ._utils._typing import Array, DType
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover
|
||||
import numpy as np
|
||||
from numpy.typing import ArrayLike
|
||||
|
||||
NumPyObject: TypeAlias = np.ndarray[Any, Any] | np.generic # type: ignore[explicit-any]
|
||||
else:
|
||||
# Sphinx hack
|
||||
NumPyObject = Any
|
||||
|
||||
P = ParamSpec("P")
|
||||
|
||||
|
||||
@overload
|
||||
def lazy_apply( # type: ignore[decorated-any, valid-type]
|
||||
func: Callable[P, Array | ArrayLike],
|
||||
*args: Array | complex | None,
|
||||
shape: tuple[int | None, ...] | None = None,
|
||||
dtype: DType | None = None,
|
||||
as_numpy: bool = False,
|
||||
xp: ModuleType | None = None,
|
||||
**kwargs: P.kwargs, # pyright: ignore[reportGeneralTypeIssues]
|
||||
) -> Array: ... # numpydoc ignore=GL08
|
||||
|
||||
|
||||
@overload
|
||||
def lazy_apply( # type: ignore[decorated-any, valid-type]
|
||||
func: Callable[P, Sequence[Array | ArrayLike]],
|
||||
*args: Array | complex | None,
|
||||
shape: Sequence[tuple[int | None, ...]],
|
||||
dtype: Sequence[DType] | None = None,
|
||||
as_numpy: bool = False,
|
||||
xp: ModuleType | None = None,
|
||||
**kwargs: P.kwargs, # pyright: ignore[reportGeneralTypeIssues]
|
||||
) -> tuple[Array, ...]: ... # numpydoc ignore=GL08
|
||||
|
||||
|
||||
def lazy_apply( # type: ignore[valid-type] # numpydoc ignore=GL07,SA04
|
||||
func: Callable[P, Array | ArrayLike | Sequence[Array | ArrayLike]],
|
||||
*args: Array | complex | None,
|
||||
shape: tuple[int | None, ...] | Sequence[tuple[int | None, ...]] | None = None,
|
||||
dtype: DType | Sequence[DType] | None = None,
|
||||
as_numpy: bool = False,
|
||||
xp: ModuleType | None = None,
|
||||
**kwargs: P.kwargs, # pyright: ignore[reportGeneralTypeIssues]
|
||||
) -> Array | tuple[Array, ...]:
|
||||
"""
|
||||
Lazily apply an eager function.
|
||||
|
||||
If the backend of the input arrays is lazy, e.g. Dask or jitted JAX, the execution
|
||||
of the function is delayed until the graph is materialized; if it's eager, the
|
||||
function is executed immediately.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : callable
|
||||
The function to apply.
|
||||
|
||||
It must accept one or more array API compliant arrays as positional arguments.
|
||||
If `as_numpy=True`, inputs are converted to NumPy before they are passed to
|
||||
`func`.
|
||||
It must return either a single array-like or a sequence of array-likes.
|
||||
|
||||
`func` must be a pure function, i.e. without side effects, as depending on the
|
||||
backend it may be executed more than once or never.
|
||||
*args : Array | int | float | complex | bool | None
|
||||
One or more Array API compliant arrays, Python scalars, or None's.
|
||||
|
||||
If `as_numpy=True`, you need to be able to apply :func:`numpy.asarray` to
|
||||
non-None args to convert them to NumPy; read notes below about specific
|
||||
backends.
|
||||
shape : tuple[int | None, ...] | Sequence[tuple[int | None, ...]], optional
|
||||
Output shape or sequence of output shapes, one for each output of `func`.
|
||||
Default: assume single output and broadcast shapes of the input arrays.
|
||||
dtype : DType | Sequence[DType], optional
|
||||
Output dtype or sequence of output dtypes, one for each output of `func`.
|
||||
dtype(s) must belong to the same array namespace as the input arrays.
|
||||
Default: infer the result type(s) from the input arrays.
|
||||
as_numpy : bool, optional
|
||||
If True, convert the input arrays to NumPy before passing them to `func`.
|
||||
This is particularly useful to make NumPy-only functions, e.g. written in Cython
|
||||
or Numba, work transparently with array API-compliant arrays.
|
||||
Default: False.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `args`. Default: infer.
|
||||
**kwargs : Any, optional
|
||||
Additional keyword arguments to pass verbatim to `func`.
|
||||
They cannot contain Array objects.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Array | tuple[Array, ...]
|
||||
The result(s) of `func` applied to the input arrays, wrapped in the same
|
||||
array namespace as the inputs.
|
||||
If shape is omitted or a single `tuple[int | None, ...]`, return a single array.
|
||||
Otherwise, return a tuple of arrays.
|
||||
|
||||
Notes
|
||||
-----
|
||||
JAX
|
||||
This allows applying eager functions to jitted JAX arrays, which are lazy.
|
||||
The function won't be applied until the JAX array is materialized.
|
||||
When running inside ``jax.jit``, `shape` must be fully known, i.e. it cannot
|
||||
contain any `None` elements.
|
||||
|
||||
.. warning::
|
||||
|
||||
`func` must never raise inside ``jax.jit``, as the resulting behavior is
|
||||
undefined.
|
||||
|
||||
Using this with `as_numpy=False` is particularly useful to apply non-jittable
|
||||
JAX functions to arrays on GPU devices.
|
||||
If ``as_numpy=True``, the :doc:`jax:transfer_guard` may prevent arrays on a GPU
|
||||
device from being transferred back to CPU. This is treated as an implicit
|
||||
transfer.
|
||||
|
||||
PyTorch, CuPy
|
||||
If ``as_numpy=True``, these backends raise by default if you attempt to convert
|
||||
arrays on a GPU device to NumPy.
|
||||
|
||||
Sparse
|
||||
If ``as_numpy=True``, by default sparse prevents implicit densification through
|
||||
:func:`numpy.asarray`. `This safety mechanism can be disabled
|
||||
<https://sparse.pydata.org/en/stable/operations.html#package-configuration>`_.
|
||||
|
||||
Dask
|
||||
This allows applying eager functions to Dask arrays.
|
||||
The Dask graph won't be computed until the user calls ``compute()`` or
|
||||
``persist()`` down the line.
|
||||
|
||||
The function name will be prominently visible on the user-facing Dask
|
||||
dashboard and on Prometheus metrics, so it is recommended for it to be
|
||||
meaningful.
|
||||
|
||||
`lazy_apply` doesn't know if `func` reduces along any axes; also, shape
|
||||
changes are non-trivial in chunked Dask arrays. For these reasons, all inputs
|
||||
will be rechunked into a single chunk.
|
||||
|
||||
.. warning::
|
||||
|
||||
The whole operation needs to fit in memory all at once on a single worker.
|
||||
|
||||
The outputs will also be returned as a single chunk and you should consider
|
||||
rechunking them into smaller chunks afterwards.
|
||||
|
||||
If you want to distribute the calculation across multiple workers, you
|
||||
should use :func:`dask.array.map_blocks`, :func:`dask.array.map_overlap`,
|
||||
:func:`dask.array.blockwise`, or a native Dask wrapper instead of
|
||||
`lazy_apply`.
|
||||
|
||||
Dask wrapping around other backends
|
||||
If ``as_numpy=False``, `func` will receive in input eager arrays of the meta
|
||||
namespace, as defined by the ``._meta`` attribute of the input Dask arrays.
|
||||
The outputs of `func` will be wrapped by the meta namespace, and then wrapped
|
||||
again by Dask.
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
When ``xp=jax.numpy``, the output `shape` is unknown (it contains ``None`` on
|
||||
one or more axes) and this function was called inside ``jax.jit``.
|
||||
RuntimeError
|
||||
When ``xp=sparse`` and auto-densification is disabled.
|
||||
Exception (backend-specific)
|
||||
When the backend disallows implicit device to host transfers and the input
|
||||
arrays are on a non-CPU device, e.g. on GPU.
|
||||
|
||||
See Also
|
||||
--------
|
||||
jax.transfer_guard
|
||||
jax.pure_callback
|
||||
dask.array.map_blocks
|
||||
dask.array.map_overlap
|
||||
dask.array.blockwise
|
||||
"""
|
||||
args_not_none = [arg for arg in args if arg is not None]
|
||||
array_args = [arg for arg in args_not_none if not is_python_scalar(arg)]
|
||||
if not array_args:
|
||||
msg = "Must have at least one argument array"
|
||||
raise ValueError(msg)
|
||||
if xp is None:
|
||||
xp = array_namespace(*args)
|
||||
|
||||
# Normalize and validate shape and dtype
|
||||
shapes: list[tuple[int | None, ...]]
|
||||
dtypes: list[DType]
|
||||
multi_output = False
|
||||
|
||||
if shape is None:
|
||||
shapes = [broadcast_shapes(*(arg.shape for arg in array_args))]
|
||||
elif all(isinstance(s, int | None) for s in shape):
|
||||
# Do not test for shape to be a tuple
|
||||
# https://github.com/data-apis/array-api/issues/891#issuecomment-2637430522
|
||||
shapes = [cast(tuple[int | None, ...], shape)]
|
||||
else:
|
||||
shapes = list(shape) # type: ignore[arg-type] # pyright: ignore[reportAssignmentType]
|
||||
multi_output = True
|
||||
|
||||
if dtype is None:
|
||||
dtypes = [xp.result_type(*args_not_none)] * len(shapes)
|
||||
elif multi_output:
|
||||
if not isinstance(dtype, Sequence):
|
||||
msg = "Got multiple shapes but only one dtype"
|
||||
raise ValueError(msg)
|
||||
dtypes = list(dtype) # pyright: ignore[reportUnknownArgumentType]
|
||||
else:
|
||||
if isinstance(dtype, Sequence):
|
||||
msg = "Got single shape but multiple dtypes"
|
||||
raise ValueError(msg)
|
||||
|
||||
dtypes = [dtype]
|
||||
|
||||
if len(shapes) != len(dtypes):
|
||||
msg = f"Got {len(shapes)} shapes and {len(dtypes)} dtypes"
|
||||
raise ValueError(msg)
|
||||
del shape
|
||||
del dtype
|
||||
# End of shape and dtype parsing
|
||||
|
||||
# Backend-specific branches
|
||||
if is_dask_namespace(xp):
|
||||
import dask
|
||||
|
||||
metas: list[Array] = [arg._meta for arg in array_args] # pylint: disable=protected-access # pyright: ignore[reportAttributeAccessIssue]
|
||||
meta_xp = array_namespace(*metas)
|
||||
|
||||
wrapped = dask.delayed( # type: ignore[attr-defined] # pyright: ignore[reportPrivateImportUsage]
|
||||
_lazy_apply_wrapper(func, as_numpy, multi_output, meta_xp),
|
||||
pure=True,
|
||||
)
|
||||
# This finalizes each arg, which is the same as arg.rechunk(-1).
|
||||
# Please read docstring above for why we're not using
|
||||
# dask.array.map_blocks or dask.array.blockwise!
|
||||
delayed_out = wrapped(*args, **kwargs)
|
||||
|
||||
out = tuple(
|
||||
xp.from_delayed(
|
||||
delayed_out[i], # pyright: ignore[reportIndexIssue]
|
||||
# Dask's unknown shapes diverge from the Array API specification
|
||||
shape=tuple(math.nan if s is None else s for s in shape),
|
||||
dtype=dtype,
|
||||
meta=metas[0],
|
||||
)
|
||||
for i, (shape, dtype) in enumerate(zip(shapes, dtypes, strict=True))
|
||||
)
|
||||
|
||||
elif is_jax_namespace(xp) and _is_jax_jit_enabled(xp):
|
||||
# Delay calling func with jax.pure_callback, which will forward to func eager
|
||||
# JAX arrays. Do not use jax.pure_callback when running outside of the JIT,
|
||||
# as it does not support raising exceptions:
|
||||
# https://github.com/jax-ml/jax/issues/26102
|
||||
import jax
|
||||
|
||||
if any(None in shape for shape in shapes):
|
||||
msg = "Output shape must be fully known when running inside jax.jit"
|
||||
raise ValueError(msg)
|
||||
|
||||
# Shield kwargs from being coerced into JAX arrays.
|
||||
# jax.pure_callback calls jax.jit under the hood, but without the chance of
|
||||
# passing static_argnames / static_argnums.
|
||||
wrapped = _lazy_apply_wrapper(
|
||||
partial(func, **kwargs), as_numpy, multi_output, xp
|
||||
)
|
||||
|
||||
# suppress unused-ignore to run mypy in -e lint as well as -e dev
|
||||
out = cast( # type: ignore[bad-cast,unused-ignore]
|
||||
tuple[Array, ...],
|
||||
jax.pure_callback(
|
||||
wrapped,
|
||||
tuple(
|
||||
jax.ShapeDtypeStruct(shape, dtype) # pyright: ignore[reportUnknownArgumentType]
|
||||
for shape, dtype in zip(shapes, dtypes, strict=True)
|
||||
),
|
||||
*args,
|
||||
),
|
||||
)
|
||||
|
||||
else:
|
||||
# Eager backends, including non-jitted JAX
|
||||
wrapped = _lazy_apply_wrapper(func, as_numpy, multi_output, xp)
|
||||
out = wrapped(*args, **kwargs)
|
||||
|
||||
return out if multi_output else out[0]
|
||||
|
||||
|
||||
def _is_jax_jit_enabled(xp: ModuleType) -> bool: # numpydoc ignore=PR01,RT01
|
||||
"""Return True if this function is being called inside ``jax.jit``."""
|
||||
import jax # pylint: disable=import-outside-toplevel
|
||||
|
||||
x = xp.asarray(False)
|
||||
try:
|
||||
return bool(x)
|
||||
except jax.errors.TracerBoolConversionError:
|
||||
return True
|
||||
|
||||
|
||||
def _lazy_apply_wrapper( # type: ignore[explicit-any] # numpydoc ignore=PR01,RT01
|
||||
func: Callable[..., Array | ArrayLike | Sequence[Array | ArrayLike]],
|
||||
as_numpy: bool,
|
||||
multi_output: bool,
|
||||
xp: ModuleType,
|
||||
) -> Callable[..., tuple[Array, ...]]:
|
||||
"""
|
||||
Helper of `lazy_apply`.
|
||||
|
||||
Given a function that accepts one or more arrays as positional arguments and returns
|
||||
a single array-like or a sequence of array-likes, return a function that accepts the
|
||||
same number of Array API arrays and always returns a tuple of Array API array.
|
||||
|
||||
Any keyword arguments are passed through verbatim to the wrapped function.
|
||||
"""
|
||||
|
||||
# On Dask, @wraps causes the graph key to contain the wrapped function's name
|
||||
@wraps(func)
|
||||
def wrapper( # type: ignore[decorated-any,explicit-any]
|
||||
*args: Array | complex | None, **kwargs: Any
|
||||
) -> tuple[Array, ...]: # numpydoc ignore=GL08
|
||||
args_list = []
|
||||
device = None
|
||||
for arg in args:
|
||||
if arg is not None and not is_python_scalar(arg):
|
||||
if device is None:
|
||||
device = _compat.device(arg)
|
||||
if as_numpy:
|
||||
import numpy as np
|
||||
|
||||
arg = cast(Array, np.asarray(arg)) # type: ignore[bad-cast] # noqa: PLW2901
|
||||
args_list.append(arg)
|
||||
assert device is not None
|
||||
|
||||
out = func(*args_list, **kwargs)
|
||||
|
||||
if multi_output:
|
||||
assert isinstance(out, Sequence)
|
||||
return tuple(xp.asarray(o, device=device) for o in out)
|
||||
return (xp.asarray(out, device=device),)
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,278 @@
|
||||
"""
|
||||
Testing utilities.
|
||||
|
||||
Note that this is private API; don't expect it to be stable.
|
||||
See also ..testing for public testing utilities.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from types import ModuleType
|
||||
from typing import Any, cast
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from ._utils._compat import (
|
||||
array_namespace,
|
||||
is_array_api_strict_namespace,
|
||||
is_cupy_namespace,
|
||||
is_dask_namespace,
|
||||
is_jax_namespace,
|
||||
is_numpy_namespace,
|
||||
is_pydata_sparse_namespace,
|
||||
is_torch_namespace,
|
||||
to_device,
|
||||
)
|
||||
from ._utils._typing import Array, Device
|
||||
|
||||
__all__ = ["as_numpy_array", "xp_assert_close", "xp_assert_equal", "xp_assert_less"]
|
||||
|
||||
|
||||
def _check_ns_shape_dtype(
|
||||
actual: Array,
|
||||
desired: Array,
|
||||
check_dtype: bool,
|
||||
check_shape: bool,
|
||||
check_scalar: bool,
|
||||
) -> ModuleType: # numpydoc ignore=RT03
|
||||
"""
|
||||
Assert that namespace, shape and dtype of the two arrays match.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
actual : Array
|
||||
The array produced by the tested function.
|
||||
desired : Array
|
||||
The expected array (typically hardcoded).
|
||||
check_dtype, check_shape : bool, default: True
|
||||
Whether to check agreement between actual and desired dtypes and shapes
|
||||
check_scalar : bool, default: False
|
||||
NumPy only: whether to check agreement between actual and desired types -
|
||||
0d array vs scalar.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Arrays namespace.
|
||||
"""
|
||||
actual_xp = array_namespace(actual) # Raises on scalars and lists
|
||||
desired_xp = array_namespace(desired)
|
||||
|
||||
msg = f"namespaces do not match: {actual_xp} != f{desired_xp}"
|
||||
assert actual_xp == desired_xp, msg
|
||||
|
||||
if check_shape:
|
||||
actual_shape = actual.shape
|
||||
desired_shape = desired.shape
|
||||
if is_dask_namespace(desired_xp):
|
||||
# Dask uses nan instead of None for unknown shapes
|
||||
if any(math.isnan(i) for i in cast(tuple[float, ...], actual_shape)):
|
||||
actual_shape = actual.compute().shape # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue]
|
||||
if any(math.isnan(i) for i in cast(tuple[float, ...], desired_shape)):
|
||||
desired_shape = desired.compute().shape # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
msg = f"shapes do not match: {actual_shape} != f{desired_shape}"
|
||||
assert actual_shape == desired_shape, msg
|
||||
|
||||
if check_dtype:
|
||||
msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
|
||||
assert actual.dtype == desired.dtype, msg
|
||||
|
||||
if is_numpy_namespace(actual_xp) and check_scalar:
|
||||
# only NumPy distinguishes between scalars and arrays; we do if check_scalar.
|
||||
_msg = (
|
||||
"array-ness does not match:\n Actual: "
|
||||
f"{type(actual)}\n Desired: {type(desired)}"
|
||||
)
|
||||
assert np.isscalar(actual) == np.isscalar(desired), _msg
|
||||
|
||||
return desired_xp
|
||||
|
||||
|
||||
def as_numpy_array(array: Array, *, xp: ModuleType) -> np.typing.NDArray[Any]: # type: ignore[explicit-any]
|
||||
"""
|
||||
Convert array to NumPy, bypassing GPU-CPU transfer guards and densification guards.
|
||||
"""
|
||||
if is_cupy_namespace(xp):
|
||||
return xp.asnumpy(array)
|
||||
if is_pydata_sparse_namespace(xp):
|
||||
return array.todense() # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
if is_torch_namespace(xp):
|
||||
array = to_device(array, "cpu")
|
||||
if is_array_api_strict_namespace(xp):
|
||||
cpu: Device = xp.Device("CPU_DEVICE")
|
||||
array = to_device(array, cpu)
|
||||
if is_jax_namespace(xp):
|
||||
import jax
|
||||
|
||||
# Note: only needed if the transfer guard is enabled
|
||||
cpu = cast(Device, jax.devices("cpu")[0])
|
||||
array = to_device(array, cpu)
|
||||
|
||||
return np.asarray(array)
|
||||
|
||||
|
||||
def xp_assert_equal(
|
||||
actual: Array,
|
||||
desired: Array,
|
||||
*,
|
||||
err_msg: str = "",
|
||||
check_dtype: bool = True,
|
||||
check_shape: bool = True,
|
||||
check_scalar: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Array-API compatible version of `np.testing.assert_array_equal`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
actual : Array
|
||||
The array produced by the tested function.
|
||||
desired : Array
|
||||
The expected array (typically hardcoded).
|
||||
err_msg : str, optional
|
||||
Error message to display on failure.
|
||||
check_dtype, check_shape : bool, default: True
|
||||
Whether to check agreement between actual and desired dtypes and shapes
|
||||
check_scalar : bool, default: False
|
||||
NumPy only: whether to check agreement between actual and desired types -
|
||||
0d array vs scalar.
|
||||
|
||||
See Also
|
||||
--------
|
||||
xp_assert_close : Similar function for inexact equality checks.
|
||||
numpy.testing.assert_array_equal : Similar function for NumPy arrays.
|
||||
"""
|
||||
xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
|
||||
actual_np = as_numpy_array(actual, xp=xp)
|
||||
desired_np = as_numpy_array(desired, xp=xp)
|
||||
np.testing.assert_array_equal(actual_np, desired_np, err_msg=err_msg)
|
||||
|
||||
|
||||
def xp_assert_less(
|
||||
x: Array,
|
||||
y: Array,
|
||||
*,
|
||||
err_msg: str = "",
|
||||
check_dtype: bool = True,
|
||||
check_shape: bool = True,
|
||||
check_scalar: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Array-API compatible version of `np.testing.assert_array_less`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x, y : Array
|
||||
The arrays to compare according to ``x < y`` (elementwise).
|
||||
err_msg : str, optional
|
||||
Error message to display on failure.
|
||||
check_dtype, check_shape : bool, default: True
|
||||
Whether to check agreement between actual and desired dtypes and shapes
|
||||
check_scalar : bool, default: False
|
||||
NumPy only: whether to check agreement between actual and desired types -
|
||||
0d array vs scalar.
|
||||
|
||||
See Also
|
||||
--------
|
||||
xp_assert_close : Similar function for inexact equality checks.
|
||||
numpy.testing.assert_array_equal : Similar function for NumPy arrays.
|
||||
"""
|
||||
xp = _check_ns_shape_dtype(x, y, check_dtype, check_shape, check_scalar)
|
||||
x_np = as_numpy_array(x, xp=xp)
|
||||
y_np = as_numpy_array(y, xp=xp)
|
||||
np.testing.assert_array_less(x_np, y_np, err_msg=err_msg)
|
||||
|
||||
|
||||
def xp_assert_close(
|
||||
actual: Array,
|
||||
desired: Array,
|
||||
*,
|
||||
rtol: float | None = None,
|
||||
atol: float = 0,
|
||||
err_msg: str = "",
|
||||
check_dtype: bool = True,
|
||||
check_shape: bool = True,
|
||||
check_scalar: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Array-API compatible version of `np.testing.assert_allclose`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
actual : Array
|
||||
The array produced by the tested function.
|
||||
desired : Array
|
||||
The expected array (typically hardcoded).
|
||||
rtol : float, optional
|
||||
Relative tolerance. Default: dtype-dependent.
|
||||
atol : float, optional
|
||||
Absolute tolerance. Default: 0.
|
||||
err_msg : str, optional
|
||||
Error message to display on failure.
|
||||
check_dtype, check_shape : bool, default: True
|
||||
Whether to check agreement between actual and desired dtypes and shapes
|
||||
check_scalar : bool, default: False
|
||||
NumPy only: whether to check agreement between actual and desired types -
|
||||
0d array vs scalar.
|
||||
|
||||
See Also
|
||||
--------
|
||||
xp_assert_equal : Similar function for exact equality checks.
|
||||
isclose : Public function for checking closeness.
|
||||
numpy.testing.assert_allclose : Similar function for NumPy arrays.
|
||||
|
||||
Notes
|
||||
-----
|
||||
The default `atol` and `rtol` differ from `xp.all(xpx.isclose(a, b))`.
|
||||
"""
|
||||
xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
|
||||
|
||||
if rtol is None:
|
||||
if xp.isdtype(actual.dtype, ("real floating", "complex floating")):
|
||||
# multiplier of 4 is used as for `np.float64` this puts the default `rtol`
|
||||
# roughly half way between sqrt(eps) and the default for
|
||||
# `numpy.testing.assert_allclose`, 1e-7
|
||||
rtol = xp.finfo(actual.dtype).eps ** 0.5 * 4
|
||||
else:
|
||||
rtol = 1e-7
|
||||
|
||||
actual_np = as_numpy_array(actual, xp=xp)
|
||||
desired_np = as_numpy_array(desired, xp=xp)
|
||||
np.testing.assert_allclose( # pyright: ignore[reportCallIssue]
|
||||
actual_np,
|
||||
desired_np,
|
||||
rtol=rtol, # pyright: ignore[reportArgumentType]
|
||||
atol=atol,
|
||||
err_msg=err_msg,
|
||||
)
|
||||
|
||||
|
||||
def xfail(
|
||||
request: pytest.FixtureRequest, *, reason: str, strict: bool | None = None
|
||||
) -> None:
|
||||
"""
|
||||
XFAIL the currently running test.
|
||||
|
||||
Unlike ``pytest.xfail``, allow rest of test to execute instead of immediately
|
||||
halting it, so that it may result in a XPASS.
|
||||
xref https://github.com/pandas-dev/pandas/issues/38902
|
||||
|
||||
Parameters
|
||||
----------
|
||||
request : pytest.FixtureRequest
|
||||
``request`` argument of the test function.
|
||||
reason : str
|
||||
Reason for the expected failure.
|
||||
strict: bool, optional
|
||||
If True, the test will be marked as failed if it passes.
|
||||
If False, the test will be marked as passed if it fails.
|
||||
Default: ``xfail_strict`` value in ``pyproject.toml``, or False if absent.
|
||||
"""
|
||||
if strict is not None:
|
||||
marker = pytest.mark.xfail(reason=reason, strict=strict)
|
||||
else:
|
||||
marker = pytest.mark.xfail(reason=reason)
|
||||
request.node.add_marker(marker)
|
||||
@@ -0,0 +1 @@
|
||||
"""Modules housing private utility functions."""
|
||||
@@ -0,0 +1,74 @@
|
||||
"""Acquire helpers from array-api-compat."""
|
||||
# Allow packages that vendor both `array-api-extra` and
|
||||
# `array-api-compat` to override the import location
|
||||
|
||||
# pylint: disable=duplicate-code
|
||||
try:
|
||||
from ...._array_api_compat_vendor import (
|
||||
array_namespace,
|
||||
device,
|
||||
is_array_api_obj,
|
||||
is_array_api_strict_namespace,
|
||||
is_cupy_array,
|
||||
is_cupy_namespace,
|
||||
is_dask_array,
|
||||
is_dask_namespace,
|
||||
is_jax_array,
|
||||
is_jax_namespace,
|
||||
is_lazy_array,
|
||||
is_numpy_array,
|
||||
is_numpy_namespace,
|
||||
is_pydata_sparse_array,
|
||||
is_pydata_sparse_namespace,
|
||||
is_torch_array,
|
||||
is_torch_namespace,
|
||||
is_writeable_array,
|
||||
size,
|
||||
to_device,
|
||||
)
|
||||
except ImportError:
|
||||
from array_api_compat import (
|
||||
array_namespace,
|
||||
device,
|
||||
is_array_api_obj,
|
||||
is_array_api_strict_namespace,
|
||||
is_cupy_array,
|
||||
is_cupy_namespace,
|
||||
is_dask_array,
|
||||
is_dask_namespace,
|
||||
is_jax_array,
|
||||
is_jax_namespace,
|
||||
is_lazy_array,
|
||||
is_numpy_array,
|
||||
is_numpy_namespace,
|
||||
is_pydata_sparse_array,
|
||||
is_pydata_sparse_namespace,
|
||||
is_torch_array,
|
||||
is_torch_namespace,
|
||||
is_writeable_array,
|
||||
size,
|
||||
to_device,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"array_namespace",
|
||||
"device",
|
||||
"is_array_api_obj",
|
||||
"is_array_api_strict_namespace",
|
||||
"is_cupy_array",
|
||||
"is_cupy_namespace",
|
||||
"is_dask_array",
|
||||
"is_dask_namespace",
|
||||
"is_jax_array",
|
||||
"is_jax_namespace",
|
||||
"is_lazy_array",
|
||||
"is_numpy_array",
|
||||
"is_numpy_namespace",
|
||||
"is_pydata_sparse_array",
|
||||
"is_pydata_sparse_namespace",
|
||||
"is_torch_array",
|
||||
"is_torch_namespace",
|
||||
"is_writeable_array",
|
||||
"size",
|
||||
"to_device",
|
||||
]
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Static type stubs for `_compat.py`."""
|
||||
|
||||
# https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
|
||||
from __future__ import annotations
|
||||
|
||||
from types import ModuleType
|
||||
from typing import Any, TypeGuard
|
||||
|
||||
# TODO import from typing (requires Python >=3.13)
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
from ._typing import Array, Device
|
||||
|
||||
# pylint: disable=missing-class-docstring,unused-argument
|
||||
|
||||
def array_namespace(
|
||||
*xs: Array | complex | None,
|
||||
api_version: str | None = None,
|
||||
use_compat: bool | None = None,
|
||||
) -> ModuleType: ...
|
||||
def device(x: Array, /) -> Device: ...
|
||||
def is_array_api_obj(x: object, /) -> TypeIs[Array]: ...
|
||||
def is_array_api_strict_namespace(xp: ModuleType, /) -> bool: ...
|
||||
def is_cupy_namespace(xp: ModuleType, /) -> bool: ...
|
||||
def is_dask_namespace(xp: ModuleType, /) -> bool: ...
|
||||
def is_jax_namespace(xp: ModuleType, /) -> bool: ...
|
||||
def is_numpy_namespace(xp: ModuleType, /) -> bool: ...
|
||||
def is_pydata_sparse_namespace(xp: ModuleType, /) -> bool: ...
|
||||
def is_torch_namespace(xp: ModuleType, /) -> bool: ...
|
||||
def is_cupy_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def is_dask_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def is_jax_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def is_numpy_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def is_pydata_sparse_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def is_torch_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def is_lazy_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def is_writeable_array(x: object, /) -> TypeGuard[Array]: ...
|
||||
def size(x: Array, /) -> int | None: ...
|
||||
def to_device( # type: ignore[explicit-any]
|
||||
x: Array,
|
||||
device: Device, # pylint: disable=redefined-outer-name
|
||||
/,
|
||||
*,
|
||||
stream: int | Any | None = None,
|
||||
) -> Array: ...
|
||||
@@ -0,0 +1,559 @@
|
||||
"""Helper functions used by `array_api_extra/_funcs.py`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import math
|
||||
import pickle
|
||||
import types
|
||||
from collections.abc import Callable, Generator, Iterable
|
||||
from functools import wraps
|
||||
from types import ModuleType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
ClassVar,
|
||||
Generic,
|
||||
Literal,
|
||||
ParamSpec,
|
||||
TypeAlias,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
from . import _compat
|
||||
from ._compat import (
|
||||
array_namespace,
|
||||
is_array_api_obj,
|
||||
is_dask_namespace,
|
||||
is_jax_namespace,
|
||||
is_numpy_array,
|
||||
is_pydata_sparse_namespace,
|
||||
)
|
||||
from ._typing import Array
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover
|
||||
# TODO import from typing (requires Python >=3.12 and >=3.13)
|
||||
from typing_extensions import TypeIs, override
|
||||
else:
|
||||
|
||||
def override(func):
|
||||
return func
|
||||
|
||||
|
||||
P = ParamSpec("P")
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"asarrays",
|
||||
"capabilities",
|
||||
"eager_shape",
|
||||
"in1d",
|
||||
"is_python_scalar",
|
||||
"jax_autojit",
|
||||
"mean",
|
||||
"meta_namespace",
|
||||
"pickle_flatten",
|
||||
"pickle_unflatten",
|
||||
]
|
||||
|
||||
|
||||
def in1d(
|
||||
x1: Array,
|
||||
x2: Array,
|
||||
/,
|
||||
*,
|
||||
assume_unique: bool = False,
|
||||
invert: bool = False,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""
|
||||
Check whether each element of an array is also present in a second array.
|
||||
|
||||
Returns a boolean array the same length as `x1` that is True
|
||||
where an element of `x1` is in `x2` and False otherwise.
|
||||
|
||||
This function has been adapted using the original implementation
|
||||
present in numpy:
|
||||
https://github.com/numpy/numpy/blob/v1.26.0/numpy/lib/arraysetops.py#L524-L758
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(x1, x2)
|
||||
|
||||
x1_shape = eager_shape(x1)
|
||||
x2_shape = eager_shape(x2)
|
||||
|
||||
# This code is run to make the code significantly faster
|
||||
if x2_shape[0] < 10 * x1_shape[0] ** 0.145 and isinstance(x2, Iterable):
|
||||
if invert:
|
||||
mask = xp.ones(x1_shape[0], dtype=xp.bool, device=_compat.device(x1))
|
||||
for a in x2:
|
||||
mask &= x1 != a
|
||||
else:
|
||||
mask = xp.zeros(x1_shape[0], dtype=xp.bool, device=_compat.device(x1))
|
||||
for a in x2:
|
||||
mask |= x1 == a
|
||||
return mask
|
||||
|
||||
rev_idx = xp.empty(0) # placeholder
|
||||
if not assume_unique:
|
||||
x1, rev_idx = xp.unique_inverse(x1)
|
||||
x2 = xp.unique_values(x2)
|
||||
|
||||
ar = xp.concat((x1, x2))
|
||||
device_ = _compat.device(ar)
|
||||
# We need this to be a stable sort.
|
||||
order = xp.argsort(ar, stable=True)
|
||||
reverse_order = xp.argsort(order, stable=True)
|
||||
sar = xp.take(ar, order, axis=0)
|
||||
ar_size = _compat.size(sar)
|
||||
assert ar_size is not None, "xp.unique*() on lazy backends raises"
|
||||
if ar_size >= 1:
|
||||
bool_ar = sar[1:] != sar[:-1] if invert else sar[1:] == sar[:-1]
|
||||
else:
|
||||
bool_ar = xp.asarray([False]) if invert else xp.asarray([True])
|
||||
flag = xp.concat((bool_ar, xp.asarray([invert], device=device_)))
|
||||
ret = xp.take(flag, reverse_order, axis=0)
|
||||
|
||||
if assume_unique:
|
||||
return ret[: x1.shape[0]]
|
||||
return xp.take(ret, rev_idx, axis=0)
|
||||
|
||||
|
||||
def mean(
|
||||
x: Array,
|
||||
/,
|
||||
*,
|
||||
axis: int | tuple[int, ...] | None = None,
|
||||
keepdims: bool = False,
|
||||
xp: ModuleType | None = None,
|
||||
) -> Array: # numpydoc ignore=PR01,RT01
|
||||
"""
|
||||
Complex mean, https://github.com/data-apis/array-api/issues/846.
|
||||
"""
|
||||
if xp is None:
|
||||
xp = array_namespace(x)
|
||||
|
||||
if xp.isdtype(x.dtype, "complex floating"):
|
||||
x_real = xp.real(x)
|
||||
x_imag = xp.imag(x)
|
||||
mean_real = xp.mean(x_real, axis=axis, keepdims=keepdims)
|
||||
mean_imag = xp.mean(x_imag, axis=axis, keepdims=keepdims)
|
||||
return mean_real + (mean_imag * xp.asarray(1j))
|
||||
return xp.mean(x, axis=axis, keepdims=keepdims)
|
||||
|
||||
|
||||
def is_python_scalar(x: object) -> TypeIs[complex]: # numpydoc ignore=PR01,RT01
|
||||
"""Return True if `x` is a Python scalar, False otherwise."""
|
||||
# isinstance(x, float) returns True for np.float64
|
||||
# isinstance(x, complex) returns True for np.complex128
|
||||
# bool is a subclass of int
|
||||
return isinstance(x, int | float | complex) and not is_numpy_array(x)
|
||||
|
||||
|
||||
def asarrays(
|
||||
a: Array | complex,
|
||||
b: Array | complex,
|
||||
xp: ModuleType,
|
||||
) -> tuple[Array, Array]:
|
||||
"""
|
||||
Ensure both `a` and `b` are arrays.
|
||||
|
||||
If `b` is a python scalar, it is converted to the same dtype as `a`, and vice versa.
|
||||
|
||||
Behavior is not specified when mixing a Python ``float`` and an array with an
|
||||
integer data type; this may give ``float32``, ``float64``, or raise an exception.
|
||||
Behavior is implementation-specific.
|
||||
|
||||
Similarly, behavior is not specified when mixing a Python ``complex`` and an array
|
||||
with a real-valued data type; this may give ``complex64``, ``complex128``, or raise
|
||||
an exception. Behavior is implementation-specific.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
a, b : Array | int | float | complex | bool
|
||||
Input arrays or scalars. At least one must be an array.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for `x`. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Array, Array
|
||||
The input arrays, possibly converted to arrays if they were scalars.
|
||||
|
||||
See Also
|
||||
--------
|
||||
mixing-arrays-with-python-scalars : Array API specification for the behavior.
|
||||
"""
|
||||
a_scalar = is_python_scalar(a)
|
||||
b_scalar = is_python_scalar(b)
|
||||
if not a_scalar and not b_scalar:
|
||||
# This includes misc. malformed input e.g. str
|
||||
return a, b # type: ignore[return-value]
|
||||
|
||||
swap = False
|
||||
if a_scalar:
|
||||
swap = True
|
||||
b, a = a, b
|
||||
|
||||
if is_array_api_obj(a):
|
||||
# a is an Array API object
|
||||
# b is a int | float | complex | bool
|
||||
xa = a
|
||||
|
||||
# https://data-apis.org/array-api/draft/API_specification/type_promotion.html#mixing-arrays-with-python-scalars
|
||||
same_dtype = {
|
||||
bool: "bool",
|
||||
int: ("integral", "real floating", "complex floating"),
|
||||
float: ("real floating", "complex floating"),
|
||||
complex: "complex floating",
|
||||
}
|
||||
kind = same_dtype[type(cast(complex, b))] # type: ignore[index]
|
||||
if xp.isdtype(a.dtype, kind):
|
||||
xb = xp.asarray(b, dtype=a.dtype)
|
||||
else:
|
||||
# Undefined behaviour. Let the function deal with it, if it can.
|
||||
xb = xp.asarray(b)
|
||||
|
||||
else:
|
||||
# Neither a nor b are Array API objects.
|
||||
# Note: we can only reach this point when one explicitly passes
|
||||
# xp=xp to the calling function; otherwise we fail earlier on
|
||||
# array_namespace(a, b).
|
||||
xa, xb = xp.asarray(a), xp.asarray(b)
|
||||
|
||||
return (xb, xa) if swap else (xa, xb)
|
||||
|
||||
|
||||
def ndindex(*x: int) -> Generator[tuple[int, ...]]:
|
||||
"""
|
||||
Generate all N-dimensional indices for a given array shape.
|
||||
|
||||
Given the shape of an array, an ndindex instance iterates over the N-dimensional
|
||||
index of the array. At each iteration a tuple of indices is returned, the last
|
||||
dimension is iterated over first.
|
||||
|
||||
This has an identical API to numpy.ndindex.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
*x : int
|
||||
The shape of the array.
|
||||
"""
|
||||
if not x:
|
||||
yield ()
|
||||
return
|
||||
for i in ndindex(*x[:-1]):
|
||||
for j in range(x[-1]):
|
||||
yield *i, j
|
||||
|
||||
|
||||
def eager_shape(x: Array, /) -> tuple[int, ...]:
|
||||
"""
|
||||
Return shape of an array. Raise if shape is not fully defined.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : Array
|
||||
Input array.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tuple[int, ...]
|
||||
Shape of the array.
|
||||
"""
|
||||
shape = x.shape
|
||||
# Dask arrays uses non-standard NaN instead of None
|
||||
if any(s is None or math.isnan(s) for s in shape):
|
||||
msg = "Unsupported lazy shape"
|
||||
raise TypeError(msg)
|
||||
return cast(tuple[int, ...], shape)
|
||||
|
||||
|
||||
def meta_namespace(
|
||||
*arrays: Array | complex | None, xp: ModuleType | None = None
|
||||
) -> ModuleType:
|
||||
"""
|
||||
Get the namespace of Dask chunks.
|
||||
|
||||
On all other backends, just return the namespace of the arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
*arrays : Array | int | float | complex | bool | None
|
||||
Input arrays.
|
||||
xp : array_namespace, optional
|
||||
The standard-compatible namespace for the input arrays. Default: infer.
|
||||
|
||||
Returns
|
||||
-------
|
||||
array_namespace
|
||||
If xp is Dask, the namespace of the Dask chunks;
|
||||
otherwise, the namespace of the arrays.
|
||||
"""
|
||||
xp = array_namespace(*arrays) if xp is None else xp
|
||||
if not is_dask_namespace(xp):
|
||||
return xp
|
||||
# Quietly skip scalars and None's
|
||||
metas = [cast(Array | None, getattr(a, "_meta", None)) for a in arrays]
|
||||
return array_namespace(*metas)
|
||||
|
||||
|
||||
def capabilities(xp: ModuleType) -> dict[str, int]:
|
||||
"""
|
||||
Return patched ``xp.__array_namespace_info__().capabilities()``.
|
||||
|
||||
TODO this helper should be eventually removed once all the special cases
|
||||
it handles are fixed in the respective backends.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
xp : array_namespace
|
||||
The standard-compatible namespace.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
Capabilities of the namespace.
|
||||
"""
|
||||
if is_pydata_sparse_namespace(xp):
|
||||
# No __array_namespace_info__(); no indexing by sparse arrays
|
||||
return {"boolean indexing": False, "data-dependent shapes": True}
|
||||
out = xp.__array_namespace_info__().capabilities()
|
||||
if is_jax_namespace(xp) and out["boolean indexing"]:
|
||||
# FIXME https://github.com/jax-ml/jax/issues/27418
|
||||
# Fixed in jax >=0.6.0
|
||||
out = out.copy()
|
||||
out["boolean indexing"] = False
|
||||
return out
|
||||
|
||||
|
||||
_BASIC_PICKLED_TYPES = frozenset((
|
||||
bool, int, float, complex, str, bytes, bytearray,
|
||||
list, tuple, dict, set, frozenset, range, slice,
|
||||
types.NoneType, types.EllipsisType,
|
||||
)) # fmt: skip
|
||||
_BASIC_REST_TYPES = frozenset((
|
||||
type, types.BuiltinFunctionType, types.FunctionType, types.ModuleType
|
||||
)) # fmt: skip
|
||||
|
||||
FlattenRest: TypeAlias = tuple[object, ...]
|
||||
|
||||
|
||||
def pickle_flatten(
|
||||
obj: object, cls: type[T] | tuple[type[T], ...]
|
||||
) -> tuple[list[T], FlattenRest]:
|
||||
"""
|
||||
Use the pickle machinery to extract objects out of an arbitrary container.
|
||||
|
||||
Unlike regular ``pickle.dumps``, this function always succeeds.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
obj : object
|
||||
The object to pickle.
|
||||
cls : type | tuple[type, ...]
|
||||
One or multiple classes to extract from the object.
|
||||
The instances of these classes inside ``obj`` will not be pickled.
|
||||
|
||||
Returns
|
||||
-------
|
||||
instances : list[cls]
|
||||
All instances of ``cls`` found inside ``obj`` (not pickled).
|
||||
rest
|
||||
Opaque object containing the pickled bytes plus all other objects where
|
||||
``__reduce__`` / ``__reduce_ex__`` is either not implemented or raised.
|
||||
These are unpickleable objects, types, modules, and functions.
|
||||
|
||||
This object is *typically* hashable save for fairly exotic objects
|
||||
that are neither pickleable nor hashable.
|
||||
|
||||
This object is pickleable if everything except ``instances`` was pickleable
|
||||
in the input object.
|
||||
|
||||
See Also
|
||||
--------
|
||||
pickle_unflatten : Reverse function.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> class A:
|
||||
... def __repr__(self):
|
||||
... return "<A>"
|
||||
>>> class NS:
|
||||
... def __repr__(self):
|
||||
... return "<NS>"
|
||||
... def __reduce__(self):
|
||||
... assert False, "not serializable"
|
||||
>>> obj = {1: A(), 2: [A(), NS(), A()]}
|
||||
>>> instances, rest = pickle_flatten(obj, A)
|
||||
>>> instances
|
||||
[<A>, <A>, <A>]
|
||||
>>> pickle_unflatten(instances, rest)
|
||||
{1: <A>, 2: [<A>, <NS>, <A>]}
|
||||
|
||||
This can be also used to swap inner objects; the only constraint is that
|
||||
the number of objects in and out must be the same:
|
||||
|
||||
>>> pickle_unflatten(["foo", "bar", "baz"], rest)
|
||||
{1: "foo", 2: ["bar", <NS>, "baz"]}
|
||||
"""
|
||||
instances: list[T] = []
|
||||
rest: list[object] = []
|
||||
|
||||
class Pickler(pickle.Pickler): # numpydoc ignore=GL08
|
||||
"""
|
||||
Use the `pickle.Pickler.persistent_id` hook to extract objects.
|
||||
"""
|
||||
|
||||
@override
|
||||
def persistent_id(self, obj: object) -> Literal[0, 1, None]: # pyright: ignore[reportIncompatibleMethodOverride] # numpydoc ignore=GL08
|
||||
if isinstance(obj, cls):
|
||||
instances.append(obj) # type: ignore[arg-type]
|
||||
return 0
|
||||
|
||||
typ_ = type(obj)
|
||||
if typ_ in _BASIC_PICKLED_TYPES: # No subclasses!
|
||||
# If obj is a collection, recursively descend inside it
|
||||
return None
|
||||
if typ_ in _BASIC_REST_TYPES:
|
||||
rest.append(obj)
|
||||
return 1
|
||||
|
||||
try:
|
||||
# Note: a class that defines __slots__ without defining __getstate__
|
||||
# cannot be pickled with __reduce__(), but can with __reduce_ex__(5)
|
||||
_ = obj.__reduce_ex__(pickle.HIGHEST_PROTOCOL)
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
rest.append(obj)
|
||||
return 1
|
||||
|
||||
# Object can be pickled. Let the Pickler recursively descend inside it.
|
||||
return None
|
||||
|
||||
f = io.BytesIO()
|
||||
p = Pickler(f, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
p.dump(obj)
|
||||
return instances, (f.getvalue(), *rest)
|
||||
|
||||
|
||||
def pickle_unflatten(instances: Iterable[object], rest: FlattenRest) -> Any: # type: ignore[explicit-any]
|
||||
"""
|
||||
Reverse of ``pickle_flatten``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
instances : Iterable
|
||||
Inner objects to be reinserted into the flattened container.
|
||||
rest : FlattenRest
|
||||
Extra bits, as returned by ``pickle_flatten``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
object
|
||||
The outer object originally passed to ``pickle_flatten`` after a
|
||||
pickle->unpickle round-trip.
|
||||
|
||||
See Also
|
||||
--------
|
||||
pickle_flatten : Serializing function.
|
||||
pickle.loads : Standard unpickle function.
|
||||
|
||||
Notes
|
||||
-----
|
||||
The `instances` iterable must yield at least the same number of elements as the ones
|
||||
returned by ``pickle_without``, but the elements do not need to be the same objects
|
||||
or even the same types of objects. Excess elements, if any, will be left untouched.
|
||||
"""
|
||||
iters = iter(instances), iter(rest)
|
||||
pik = cast(bytes, next(iters[1]))
|
||||
|
||||
class Unpickler(pickle.Unpickler): # numpydoc ignore=GL08
|
||||
"""Mirror of the overridden Pickler in pickle_flatten."""
|
||||
|
||||
@override
|
||||
def persistent_load(self, pid: Literal[0, 1]) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # numpydoc ignore=GL08
|
||||
try:
|
||||
return next(iters[pid])
|
||||
except StopIteration as e:
|
||||
msg = "Not enough objects to unpickle"
|
||||
raise ValueError(msg) from e
|
||||
|
||||
f = io.BytesIO(pik)
|
||||
return Unpickler(f).load()
|
||||
|
||||
|
||||
class _AutoJITWrapper(Generic[T]): # numpydoc ignore=PR01
|
||||
"""
|
||||
Helper of :func:`jax_autojit`.
|
||||
|
||||
Wrap arbitrary inputs and outputs of the jitted function and
|
||||
convert them to/from PyTrees.
|
||||
"""
|
||||
|
||||
obj: T
|
||||
_registered: ClassVar[bool] = False
|
||||
__slots__: tuple[str, ...] = ("obj",)
|
||||
|
||||
def __init__(self, obj: T) -> None: # numpydoc ignore=GL08
|
||||
self._register()
|
||||
self.obj = obj
|
||||
|
||||
@classmethod
|
||||
def _register(cls): # numpydoc ignore=SS06
|
||||
"""
|
||||
Register upon first use instead of at import time, to avoid
|
||||
globally importing JAX.
|
||||
"""
|
||||
if not cls._registered:
|
||||
import jax
|
||||
|
||||
jax.tree_util.register_pytree_node(
|
||||
cls,
|
||||
lambda obj: pickle_flatten(obj, jax.Array), # pyright: ignore[reportUnknownArgumentType]
|
||||
lambda aux_data, children: pickle_unflatten(children, aux_data), # pyright: ignore[reportUnknownArgumentType]
|
||||
)
|
||||
cls._registered = True
|
||||
|
||||
|
||||
def jax_autojit(
|
||||
func: Callable[P, T],
|
||||
) -> Callable[P, T]: # numpydoc ignore=PR01,RT01,SS03
|
||||
"""
|
||||
Wrap `func` with ``jax.jit``, with the following differences:
|
||||
|
||||
- Python scalar arguments and return values are not automatically converted to
|
||||
``jax.Array`` objects.
|
||||
- All non-array arguments are automatically treated as static.
|
||||
Unlike ``jax.jit``, static arguments must be either hashable or serializable with
|
||||
``pickle``.
|
||||
- Unlike ``jax.jit``, non-array arguments and return values are not limited to
|
||||
tuple/list/dict, but can be any object serializable with ``pickle``.
|
||||
- Automatically descend into non-array arguments and find ``jax.Array`` objects
|
||||
inside them, then rebuild the arguments when entering `func`, swapping the JAX
|
||||
concrete arrays with tracer objects.
|
||||
- Automatically descend into non-array return values and find ``jax.Array`` objects
|
||||
inside them, then rebuild them downstream of exiting the JIT, swapping the JAX
|
||||
tracer objects with concrete arrays.
|
||||
|
||||
See Also
|
||||
--------
|
||||
jax.jit : JAX JIT compilation function.
|
||||
"""
|
||||
import jax
|
||||
|
||||
@jax.jit # type: ignore[misc] # pyright: ignore[reportUntypedFunctionDecorator]
|
||||
def inner( # type: ignore[decorated-any,explicit-any] # numpydoc ignore=GL08
|
||||
wargs: _AutoJITWrapper[Any],
|
||||
) -> _AutoJITWrapper[T]:
|
||||
args, kwargs = wargs.obj
|
||||
res = func(*args, **kwargs) # pyright: ignore[reportCallIssue]
|
||||
return _AutoJITWrapper(res)
|
||||
|
||||
@wraps(func)
|
||||
def outer(*args: P.args, **kwargs: P.kwargs) -> T: # numpydoc ignore=GL08
|
||||
wargs = _AutoJITWrapper((args, kwargs))
|
||||
return inner(wargs).obj
|
||||
|
||||
return outer
|
||||
@@ -0,0 +1,10 @@
|
||||
# numpydoc ignore=GL08
|
||||
# pylint: disable=missing-module-docstring,duplicate-code
|
||||
|
||||
Array = object
|
||||
DType = object
|
||||
Device = object
|
||||
GetIndex = object
|
||||
SetIndex = object
|
||||
|
||||
__all__ = ["Array", "DType", "Device", "GetIndex", "SetIndex"]
|
||||
@@ -0,0 +1,105 @@
|
||||
"""Static typing helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import EllipsisType
|
||||
from typing import Protocol, TypeAlias
|
||||
|
||||
# TODO import from typing (requires Python >=3.12)
|
||||
from typing_extensions import override
|
||||
|
||||
# TODO: use array-api-typing once it is available
|
||||
|
||||
class Array(Protocol): # pylint: disable=missing-class-docstring
|
||||
# Unary operations
|
||||
def __abs__(self) -> Array: ...
|
||||
def __pos__(self) -> Array: ...
|
||||
def __neg__(self) -> Array: ...
|
||||
def __invert__(self) -> Array: ...
|
||||
# Binary operations
|
||||
def __add__(self, other: Array | complex, /) -> Array: ...
|
||||
def __sub__(self, other: Array | complex, /) -> Array: ...
|
||||
def __mul__(self, other: Array | complex, /) -> Array: ...
|
||||
def __truediv__(self, other: Array | complex, /) -> Array: ...
|
||||
def __floordiv__(self, other: Array | complex, /) -> Array: ...
|
||||
def __mod__(self, other: Array | complex, /) -> Array: ...
|
||||
def __pow__(self, other: Array | complex, /) -> Array: ...
|
||||
def __matmul__(self, other: Array, /) -> Array: ...
|
||||
def __and__(self, other: Array | int, /) -> Array: ...
|
||||
def __or__(self, other: Array | int, /) -> Array: ...
|
||||
def __xor__(self, other: Array | int, /) -> Array: ...
|
||||
def __lshift__(self, other: Array | int, /) -> Array: ...
|
||||
def __rshift__(self, other: Array | int, /) -> Array: ...
|
||||
def __lt__(self, other: Array | complex, /) -> Array: ...
|
||||
def __le__(self, other: Array | complex, /) -> Array: ...
|
||||
def __gt__(self, other: Array | complex, /) -> Array: ...
|
||||
def __ge__(self, other: Array | complex, /) -> Array: ...
|
||||
@override
|
||||
def __eq__(self, other: Array | complex, /) -> Array: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
|
||||
@override
|
||||
def __ne__(self, other: Array | complex, /) -> Array: ... # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
|
||||
# Reflected operations
|
||||
def __radd__(self, other: Array | complex, /) -> Array: ...
|
||||
def __rsub__(self, other: Array | complex, /) -> Array: ...
|
||||
def __rmul__(self, other: Array | complex, /) -> Array: ...
|
||||
def __rtruediv__(self, other: Array | complex, /) -> Array: ...
|
||||
def __rfloordiv__(self, other: Array | complex, /) -> Array: ...
|
||||
def __rmod__(self, other: Array | complex, /) -> Array: ...
|
||||
def __rpow__(self, other: Array | complex, /) -> Array: ...
|
||||
def __rmatmul__(self, other: Array, /) -> Array: ...
|
||||
def __rand__(self, other: Array | int, /) -> Array: ...
|
||||
def __ror__(self, other: Array | int, /) -> Array: ...
|
||||
def __rxor__(self, other: Array | int, /) -> Array: ...
|
||||
def __rlshift__(self, other: Array | int, /) -> Array: ...
|
||||
def __rrshift__(self, other: Array | int, /) -> Array: ...
|
||||
# Attributes
|
||||
@property
|
||||
def dtype(self) -> DType: ...
|
||||
@property
|
||||
def device(self) -> Device: ...
|
||||
@property
|
||||
def mT(self) -> Array: ... # pylint: disable=invalid-name
|
||||
@property
|
||||
def ndim(self) -> int: ...
|
||||
@property
|
||||
def shape(self) -> tuple[int | None, ...]: ...
|
||||
@property
|
||||
def size(self) -> int | None: ...
|
||||
@property
|
||||
def T(self) -> Array: ... # pylint: disable=invalid-name
|
||||
# Collection operations (note: an Array does not have to be Sized or Iterable)
|
||||
def __getitem__(self, key: GetIndex, /) -> Array: ...
|
||||
def __setitem__(self, key: SetIndex, value: Array | complex, /) -> None: ...
|
||||
# Materialization methods (may raise on lazy arrays)
|
||||
def __bool__(self) -> bool: ...
|
||||
def __complex__(self) -> complex: ...
|
||||
def __float__(self) -> float: ...
|
||||
def __index__(self) -> int: ...
|
||||
def __int__(self) -> int: ...
|
||||
|
||||
# Misc methods (frequently not implemented in Arrays wrapped by array-api-compat)
|
||||
# def __array_namespace__(*, api_version: str | None) -> ModuleType: ...
|
||||
# def __dlpack__(
|
||||
# *,
|
||||
# stream: int | Any | None = None,
|
||||
# max_version: tuple[int, int] | None = None,
|
||||
# dl_device: tuple[int, int] | None = None, # tuple[Enum, int]
|
||||
# copy: bool | None = None,
|
||||
# ) -> Any: ...
|
||||
# def __dlpack_device__() -> tuple[int, int]: ... # tuple[Enum, int]
|
||||
# def to_device(device: Device, /, *, stream: int | Any | None = None) -> Array: ...
|
||||
|
||||
class DType(Protocol): # pylint: disable=missing-class-docstring
|
||||
pass
|
||||
|
||||
class Device(Protocol): # pylint: disable=missing-class-docstring
|
||||
pass
|
||||
|
||||
SetIndex: TypeAlias = ( # type: ignore[explicit-any]
|
||||
int | slice | EllipsisType | Array | tuple[int | slice | EllipsisType | Array, ...]
|
||||
)
|
||||
GetIndex: TypeAlias = ( # type: ignore[explicit-any]
|
||||
SetIndex | None | tuple[int | slice | EllipsisType | None | Array, ...]
|
||||
)
|
||||
|
||||
__all__ = ["Array", "DType", "Device", "GetIndex", "SetIndex"]
|
||||
@@ -0,0 +1,359 @@
|
||||
"""
|
||||
Public testing utilities.
|
||||
|
||||
See also _lib._testing for additional private testing utilities.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import enum
|
||||
import warnings
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from functools import wraps
|
||||
from types import ModuleType
|
||||
from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast
|
||||
|
||||
from ._lib._utils._compat import is_dask_namespace, is_jax_namespace
|
||||
from ._lib._utils._helpers import jax_autojit, pickle_flatten, pickle_unflatten
|
||||
|
||||
__all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover
|
||||
# TODO import override from typing (requires Python >=3.12)
|
||||
import pytest
|
||||
from dask.typing import Graph, Key, SchedulerGetCallable
|
||||
from typing_extensions import override
|
||||
|
||||
else:
|
||||
# Sphinx hacks
|
||||
SchedulerGetCallable = object
|
||||
|
||||
def override(func):
|
||||
return func
|
||||
|
||||
|
||||
P = ParamSpec("P")
|
||||
T = TypeVar("T")
|
||||
|
||||
_ufuncs_tags: dict[object, dict[str, Any]] = {} # type: ignore[explicit-any]
|
||||
|
||||
|
||||
class Deprecated(enum.Enum):
|
||||
"""Unique type for deprecated parameters."""
|
||||
|
||||
DEPRECATED = 1
|
||||
|
||||
|
||||
DEPRECATED = Deprecated.DEPRECATED
|
||||
|
||||
|
||||
def lazy_xp_function( # type: ignore[explicit-any]
|
||||
func: Callable[..., Any],
|
||||
*,
|
||||
allow_dask_compute: bool | int = False,
|
||||
jax_jit: bool = True,
|
||||
static_argnums: Deprecated = DEPRECATED,
|
||||
static_argnames: Deprecated = DEPRECATED,
|
||||
) -> None: # numpydoc ignore=GL07
|
||||
"""
|
||||
Tag a function to be tested on lazy backends.
|
||||
|
||||
Tag a function so that when any tests are executed with ``xp=jax.numpy`` the
|
||||
function is replaced with a jitted version of itself, and when it is executed with
|
||||
``xp=dask.array`` the function will raise if it attempts to materialize the graph.
|
||||
This will be later expanded to provide test coverage for other lazy backends.
|
||||
|
||||
In order for the tag to be effective, the test or a fixture must call
|
||||
:func:`patch_lazy_xp_functions`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : callable
|
||||
Function to be tested.
|
||||
allow_dask_compute : bool | int, optional
|
||||
Whether `func` is allowed to internally materialize the Dask graph, or maximum
|
||||
number of times it is allowed to do so. This is typically triggered by
|
||||
``bool()``, ``float()``, or ``np.asarray()``.
|
||||
|
||||
Set to 1 if you are aware that `func` converts the input parameters to NumPy and
|
||||
want to let it do so at least for the time being, knowing that it is going to be
|
||||
extremely detrimental for performance.
|
||||
|
||||
If a test needs values higher than 1 to pass, it is a canary that the conversion
|
||||
to NumPy/bool/float is happening multiple times, which translates to multiple
|
||||
computations of the whole graph. Short of making the function fully lazy, you
|
||||
should at least add explicit calls to ``np.asarray()`` early in the function.
|
||||
*Note:* the counter of `allow_dask_compute` resets after each call to `func`, so
|
||||
a test function that invokes `func` multiple times should still work with this
|
||||
parameter set to 1.
|
||||
|
||||
Set to True to allow `func` to materialize the graph an unlimited number
|
||||
of times.
|
||||
|
||||
Default: False, meaning that `func` must be fully lazy and never materialize the
|
||||
graph.
|
||||
jax_jit : bool, optional
|
||||
Set to True to replace `func` with a smart variant of ``jax.jit(func)`` after
|
||||
calling the :func:`patch_lazy_xp_functions` test helper with ``xp=jax.numpy``.
|
||||
Set to False if `func` is only compatible with eager (non-jitted) JAX.
|
||||
|
||||
Unlike with vanilla ``jax.jit``, all arguments and return types that are not JAX
|
||||
arrays are treated as static; the function can accept and return arbitrary
|
||||
wrappers around JAX arrays. This difference is because, in real life, most users
|
||||
won't wrap the function directly with ``jax.jit`` but rather they will use it
|
||||
within their own code, which is itself then wrapped by ``jax.jit``, and
|
||||
internally consume the function's outputs.
|
||||
|
||||
In other words, the pattern that is being tested is::
|
||||
|
||||
>>> @jax.jit
|
||||
... def user_func(x):
|
||||
... y = user_prepares_inputs(x)
|
||||
... z = func(y, some_static_arg=True)
|
||||
... return user_consumes(z)
|
||||
|
||||
Default: True.
|
||||
static_argnums :
|
||||
Deprecated; ignored
|
||||
static_argnames :
|
||||
Deprecated; ignored
|
||||
|
||||
See Also
|
||||
--------
|
||||
patch_lazy_xp_functions : Companion function to call from the test or fixture.
|
||||
jax.jit : JAX function to compile a function for performance.
|
||||
|
||||
Examples
|
||||
--------
|
||||
In ``test_mymodule.py``::
|
||||
|
||||
from array_api_extra.testing import lazy_xp_function from mymodule import myfunc
|
||||
|
||||
lazy_xp_function(myfunc)
|
||||
|
||||
def test_myfunc(xp):
|
||||
a = xp.asarray([1, 2])
|
||||
# When xp=jax.numpy, this is similar to `b = jax.jit(myfunc)(a)`
|
||||
# When xp=dask.array, crash on compute() or persist()
|
||||
b = myfunc(a)
|
||||
|
||||
Notes
|
||||
-----
|
||||
In order for this tag to be effective, the test function must be imported into the
|
||||
test module globals without its namespace; alternatively its namespace must be
|
||||
declared in a ``lazy_xp_modules`` list in the test module globals.
|
||||
|
||||
Example 1::
|
||||
|
||||
from mymodule import myfunc
|
||||
|
||||
lazy_xp_function(myfunc)
|
||||
|
||||
def test_myfunc(xp):
|
||||
x = myfunc(xp.asarray([1, 2]))
|
||||
|
||||
Example 2::
|
||||
|
||||
import mymodule
|
||||
|
||||
lazy_xp_modules = [mymodule]
|
||||
lazy_xp_function(mymodule.myfunc)
|
||||
|
||||
def test_myfunc(xp):
|
||||
x = mymodule.myfunc(xp.asarray([1, 2]))
|
||||
|
||||
A test function can circumvent this monkey-patching system by using a namespace
|
||||
outside of the two above patterns. You need to sanitize your code to make sure this
|
||||
only happens intentionally.
|
||||
|
||||
Example 1::
|
||||
|
||||
import mymodule
|
||||
from mymodule import myfunc
|
||||
|
||||
lazy_xp_function(myfunc)
|
||||
|
||||
def test_myfunc(xp):
|
||||
a = xp.asarray([1, 2])
|
||||
b = myfunc(a) # This is wrapped when xp=jax.numpy or xp=dask.array
|
||||
c = mymodule.myfunc(a) # This is not
|
||||
|
||||
Example 2::
|
||||
|
||||
import mymodule
|
||||
|
||||
class naked:
|
||||
myfunc = mymodule.myfunc
|
||||
|
||||
lazy_xp_modules = [mymodule]
|
||||
lazy_xp_function(mymodule.myfunc)
|
||||
|
||||
def test_myfunc(xp):
|
||||
a = xp.asarray([1, 2])
|
||||
b = mymodule.myfunc(a) # This is wrapped when xp=jax.numpy or xp=dask.array
|
||||
c = naked.myfunc(a) # This is not
|
||||
"""
|
||||
if static_argnums is not DEPRECATED or static_argnames is not DEPRECATED:
|
||||
warnings.warn(
|
||||
(
|
||||
"The `static_argnums` and `static_argnames` parameters are deprecated "
|
||||
"and ignored. They will be removed in a future version."
|
||||
),
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
tags = {
|
||||
"allow_dask_compute": allow_dask_compute,
|
||||
"jax_jit": jax_jit,
|
||||
}
|
||||
|
||||
try:
|
||||
func._lazy_xp_function = tags # type: ignore[attr-defined] # pylint: disable=protected-access # pyright: ignore[reportFunctionMemberAccess]
|
||||
except AttributeError: # @cython.vectorize
|
||||
_ufuncs_tags[func] = tags
|
||||
|
||||
|
||||
def patch_lazy_xp_functions(
|
||||
request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, *, xp: ModuleType
|
||||
) -> None:
|
||||
"""
|
||||
Test lazy execution of functions tagged with :func:`lazy_xp_function`.
|
||||
|
||||
If ``xp==jax.numpy``, search for all functions which have been tagged with
|
||||
:func:`lazy_xp_function` in the globals of the module that defines the current test,
|
||||
as well as in the ``lazy_xp_modules`` list in the globals of the same module,
|
||||
and wrap them with :func:`jax.jit`. Unwrap them at the end of the test.
|
||||
|
||||
If ``xp==dask.array``, wrap the functions with a decorator that disables
|
||||
``compute()`` and ``persist()`` and ensures that exceptions and warnings are raised
|
||||
eagerly.
|
||||
|
||||
This function should be typically called by your library's `xp` fixture that runs
|
||||
tests on multiple backends::
|
||||
|
||||
@pytest.fixture(params=[numpy, array_api_strict, jax.numpy, dask.array])
|
||||
def xp(request, monkeypatch):
|
||||
patch_lazy_xp_functions(request, monkeypatch, xp=request.param)
|
||||
return request.param
|
||||
|
||||
but it can be otherwise be called by the test itself too.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
request : pytest.FixtureRequest
|
||||
Pytest fixture, as acquired by the test itself or by one of its fixtures.
|
||||
monkeypatch : pytest.MonkeyPatch
|
||||
Pytest fixture, as acquired by the test itself or by one of its fixtures.
|
||||
xp : array_namespace
|
||||
Array namespace to be tested.
|
||||
|
||||
See Also
|
||||
--------
|
||||
lazy_xp_function : Tag a function to be tested on lazy backends.
|
||||
pytest.FixtureRequest : `request` test function parameter.
|
||||
"""
|
||||
mod = cast(ModuleType, request.module)
|
||||
mods = [mod, *cast(list[ModuleType], getattr(mod, "lazy_xp_modules", []))]
|
||||
|
||||
def iter_tagged() -> ( # type: ignore[explicit-any]
|
||||
Iterator[tuple[ModuleType, str, Callable[..., Any], dict[str, Any]]]
|
||||
):
|
||||
for mod in mods:
|
||||
for name, func in mod.__dict__.items():
|
||||
tags: dict[str, Any] | None = None # type: ignore[explicit-any]
|
||||
with contextlib.suppress(AttributeError):
|
||||
tags = func._lazy_xp_function # pylint: disable=protected-access
|
||||
if tags is None:
|
||||
with contextlib.suppress(KeyError, TypeError):
|
||||
tags = _ufuncs_tags[func]
|
||||
if tags is not None:
|
||||
yield mod, name, func, tags
|
||||
|
||||
if is_dask_namespace(xp):
|
||||
for mod, name, func, tags in iter_tagged():
|
||||
n = tags["allow_dask_compute"]
|
||||
if n is True:
|
||||
n = 1_000_000
|
||||
elif n is False:
|
||||
n = 0
|
||||
wrapped = _dask_wrap(func, n)
|
||||
monkeypatch.setattr(mod, name, wrapped)
|
||||
|
||||
elif is_jax_namespace(xp):
|
||||
for mod, name, func, tags in iter_tagged():
|
||||
if tags["jax_jit"]:
|
||||
wrapped = jax_autojit(func)
|
||||
monkeypatch.setattr(mod, name, wrapped)
|
||||
|
||||
|
||||
class CountingDaskScheduler(SchedulerGetCallable):
|
||||
"""
|
||||
Dask scheduler that counts how many times `dask.compute` is called.
|
||||
|
||||
If the number of times exceeds 'max_count', it raises an error.
|
||||
This is a wrapper around Dask's own 'synchronous' scheduler.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
max_count : int
|
||||
Maximum number of allowed calls to `dask.compute`.
|
||||
msg : str
|
||||
Assertion to raise when the count exceeds `max_count`.
|
||||
"""
|
||||
|
||||
count: int
|
||||
max_count: int
|
||||
msg: str
|
||||
|
||||
def __init__(self, max_count: int, msg: str): # numpydoc ignore=GL08
|
||||
self.count = 0
|
||||
self.max_count = max_count
|
||||
self.msg = msg
|
||||
|
||||
@override
|
||||
def __call__(self, dsk: Graph, keys: Sequence[Key] | Key, **kwargs: Any) -> Any: # type: ignore[decorated-any,explicit-any] # numpydoc ignore=GL08
|
||||
import dask
|
||||
|
||||
self.count += 1
|
||||
# This should yield a nice traceback to the
|
||||
# offending line in the user's code
|
||||
assert self.count <= self.max_count, self.msg
|
||||
|
||||
return dask.get(dsk, keys, **kwargs) # type: ignore[attr-defined,no-untyped-call] # pyright: ignore[reportPrivateImportUsage]
|
||||
|
||||
|
||||
def _dask_wrap(
|
||||
func: Callable[P, T], n: int
|
||||
) -> Callable[P, T]: # numpydoc ignore=PR01,RT01
|
||||
"""
|
||||
Wrap `func` to raise if it attempts to call `dask.compute` more than `n` times.
|
||||
|
||||
After the function returns, materialize the graph in order to re-raise exceptions.
|
||||
"""
|
||||
import dask
|
||||
import dask.array as da
|
||||
|
||||
func_name = getattr(func, "__name__", str(func))
|
||||
n_str = f"only up to {n}" if n else "no"
|
||||
msg = (
|
||||
f"Called `dask.compute()` or `dask.persist()` {n + 1} times, "
|
||||
f"but {n_str} calls are allowed. Set "
|
||||
f"`lazy_xp_function({func_name}, allow_dask_compute={n + 1})` "
|
||||
"to allow for more (but note that this will harm performance). "
|
||||
)
|
||||
|
||||
@wraps(func)
|
||||
def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: # numpydoc ignore=GL08
|
||||
scheduler = CountingDaskScheduler(n, msg)
|
||||
with dask.config.set({"scheduler": scheduler}): # pyright: ignore[reportPrivateImportUsage]
|
||||
out = func(*args, **kwargs)
|
||||
|
||||
# Block until the graph materializes and reraise exceptions. This allows
|
||||
# `pytest.raises` and `pytest.warns` to work as expected. Note that this would
|
||||
# not work on scheduler='distributed', as it would not block.
|
||||
arrays, rest = pickle_flatten(out, da.Array)
|
||||
arrays = dask.persist(arrays, scheduler="threads")[0] # type: ignore[attr-defined,no-untyped-call,func-returns-value,index] # pyright: ignore[reportPrivateImportUsage]
|
||||
return pickle_unflatten(arrays, rest) # pyright: ignore[reportUnknownArgumentType]
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,20 @@
|
||||
from .main import minimize
|
||||
from .utils import show_versions
|
||||
|
||||
# PEP0440 compatible formatted version, see:
|
||||
# https://www.python.org/dev/peps/pep-0440/
|
||||
#
|
||||
# Final release markers:
|
||||
# X.Y.0 # For first release after an increment in Y
|
||||
# X.Y.Z # For bugfix releases
|
||||
#
|
||||
# Admissible pre-release markers:
|
||||
# X.YaN # Alpha release
|
||||
# X.YbN # Beta release
|
||||
# X.YrcN # Release Candidate
|
||||
#
|
||||
# Dev branch marker is: 'X.Y.dev' or 'X.Y.devN' where N is an integer.
|
||||
# 'X.Y.dev0' is the canonical version of 'X.Y.dev'.
|
||||
__version__ = "1.1.2"
|
||||
|
||||
__all__ = ["minimize", "show_versions"]
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,132 @@
|
||||
import sys
|
||||
from enum import Enum
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
# Exit status.
|
||||
class ExitStatus(Enum):
|
||||
"""
|
||||
Exit statuses.
|
||||
"""
|
||||
|
||||
RADIUS_SUCCESS = 0
|
||||
TARGET_SUCCESS = 1
|
||||
FIXED_SUCCESS = 2
|
||||
CALLBACK_SUCCESS = 3
|
||||
FEASIBLE_SUCCESS = 4
|
||||
MAX_EVAL_WARNING = 5
|
||||
MAX_ITER_WARNING = 6
|
||||
INFEASIBLE_ERROR = -1
|
||||
LINALG_ERROR = -2
|
||||
|
||||
|
||||
class Options(str, Enum):
|
||||
"""
|
||||
Options.
|
||||
"""
|
||||
|
||||
DEBUG = "debug"
|
||||
FEASIBILITY_TOL = "feasibility_tol"
|
||||
FILTER_SIZE = "filter_size"
|
||||
HISTORY_SIZE = "history_size"
|
||||
MAX_EVAL = "maxfev"
|
||||
MAX_ITER = "maxiter"
|
||||
NPT = "nb_points"
|
||||
RHOBEG = "radius_init"
|
||||
RHOEND = "radius_final"
|
||||
SCALE = "scale"
|
||||
STORE_HISTORY = "store_history"
|
||||
TARGET = "target"
|
||||
VERBOSE = "disp"
|
||||
|
||||
|
||||
class Constants(str, Enum):
|
||||
"""
|
||||
Constants.
|
||||
"""
|
||||
|
||||
DECREASE_RADIUS_FACTOR = "decrease_radius_factor"
|
||||
INCREASE_RADIUS_FACTOR = "increase_radius_factor"
|
||||
INCREASE_RADIUS_THRESHOLD = "increase_radius_threshold"
|
||||
DECREASE_RADIUS_THRESHOLD = "decrease_radius_threshold"
|
||||
DECREASE_RESOLUTION_FACTOR = "decrease_resolution_factor"
|
||||
LARGE_RESOLUTION_THRESHOLD = "large_resolution_threshold"
|
||||
MODERATE_RESOLUTION_THRESHOLD = "moderate_resolution_threshold"
|
||||
LOW_RATIO = "low_ratio"
|
||||
HIGH_RATIO = "high_ratio"
|
||||
VERY_LOW_RATIO = "very_low_ratio"
|
||||
PENALTY_INCREASE_THRESHOLD = "penalty_increase_threshold"
|
||||
PENALTY_INCREASE_FACTOR = "penalty_increase_factor"
|
||||
SHORT_STEP_THRESHOLD = "short_step_threshold"
|
||||
LOW_RADIUS_FACTOR = "low_radius_factor"
|
||||
BYRD_OMOJOKUN_FACTOR = "byrd_omojokun_factor"
|
||||
THRESHOLD_RATIO_CONSTRAINTS = "threshold_ratio_constraints"
|
||||
LARGE_SHIFT_FACTOR = "large_shift_factor"
|
||||
LARGE_GRADIENT_FACTOR = "large_gradient_factor"
|
||||
RESOLUTION_FACTOR = "resolution_factor"
|
||||
IMPROVE_TCG = "improve_tcg"
|
||||
|
||||
|
||||
# Default options.
|
||||
DEFAULT_OPTIONS = {
|
||||
Options.DEBUG.value: False,
|
||||
Options.FEASIBILITY_TOL.value: np.sqrt(np.finfo(float).eps),
|
||||
Options.FILTER_SIZE.value: sys.maxsize,
|
||||
Options.HISTORY_SIZE.value: sys.maxsize,
|
||||
Options.MAX_EVAL.value: lambda n: 500 * n,
|
||||
Options.MAX_ITER.value: lambda n: 1000 * n,
|
||||
Options.NPT.value: lambda n: 2 * n + 1,
|
||||
Options.RHOBEG.value: 1.0,
|
||||
Options.RHOEND.value: 1e-6,
|
||||
Options.SCALE.value: False,
|
||||
Options.STORE_HISTORY.value: False,
|
||||
Options.TARGET.value: -np.inf,
|
||||
Options.VERBOSE.value: False,
|
||||
}
|
||||
|
||||
# Default constants.
|
||||
DEFAULT_CONSTANTS = {
|
||||
Constants.DECREASE_RADIUS_FACTOR.value: 0.5,
|
||||
Constants.INCREASE_RADIUS_FACTOR.value: np.sqrt(2.0),
|
||||
Constants.INCREASE_RADIUS_THRESHOLD.value: 2.0,
|
||||
Constants.DECREASE_RADIUS_THRESHOLD.value: 1.4,
|
||||
Constants.DECREASE_RESOLUTION_FACTOR.value: 0.1,
|
||||
Constants.LARGE_RESOLUTION_THRESHOLD.value: 250.0,
|
||||
Constants.MODERATE_RESOLUTION_THRESHOLD.value: 16.0,
|
||||
Constants.LOW_RATIO.value: 0.1,
|
||||
Constants.HIGH_RATIO.value: 0.7,
|
||||
Constants.VERY_LOW_RATIO.value: 0.01,
|
||||
Constants.PENALTY_INCREASE_THRESHOLD.value: 1.5,
|
||||
Constants.PENALTY_INCREASE_FACTOR.value: 2.0,
|
||||
Constants.SHORT_STEP_THRESHOLD.value: 0.5,
|
||||
Constants.LOW_RADIUS_FACTOR.value: 0.1,
|
||||
Constants.BYRD_OMOJOKUN_FACTOR.value: 0.8,
|
||||
Constants.THRESHOLD_RATIO_CONSTRAINTS.value: 2.0,
|
||||
Constants.LARGE_SHIFT_FACTOR.value: 10.0,
|
||||
Constants.LARGE_GRADIENT_FACTOR.value: 10.0,
|
||||
Constants.RESOLUTION_FACTOR.value: 2.0,
|
||||
Constants.IMPROVE_TCG.value: True,
|
||||
}
|
||||
|
||||
# Printing options.
|
||||
PRINT_OPTIONS = {
|
||||
"threshold": 6,
|
||||
"edgeitems": 2,
|
||||
"linewidth": sys.maxsize,
|
||||
"formatter": {
|
||||
"float_kind": lambda x: np.format_float_scientific(
|
||||
x,
|
||||
precision=3,
|
||||
unique=False,
|
||||
pad_left=2,
|
||||
)
|
||||
},
|
||||
}
|
||||
|
||||
# Constants.
|
||||
BARRIER = 2.0 ** min(
|
||||
100,
|
||||
np.finfo(float).maxexp // 2,
|
||||
-np.finfo(float).minexp // 2,
|
||||
)
|
||||
@@ -0,0 +1,14 @@
|
||||
from .geometry import cauchy_geometry, spider_geometry
|
||||
from .optim import (
|
||||
tangential_byrd_omojokun,
|
||||
constrained_tangential_byrd_omojokun,
|
||||
normal_byrd_omojokun,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"cauchy_geometry",
|
||||
"spider_geometry",
|
||||
"tangential_byrd_omojokun",
|
||||
"constrained_tangential_byrd_omojokun",
|
||||
"normal_byrd_omojokun",
|
||||
]
|
||||
@@ -0,0 +1,387 @@
|
||||
import inspect
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ..utils import get_arrays_tol
|
||||
|
||||
|
||||
TINY = np.finfo(float).tiny
|
||||
|
||||
|
||||
def cauchy_geometry(const, grad, curv, xl, xu, delta, debug):
|
||||
r"""
|
||||
Maximize approximately the absolute value of a quadratic function subject
|
||||
to bound constraints in a trust region.
|
||||
|
||||
This function solves approximately
|
||||
|
||||
.. math::
|
||||
|
||||
\max_{s \in \mathbb{R}^n} \quad \bigg\lvert c + g^{\mathsf{T}} s +
|
||||
\frac{1}{2} s^{\mathsf{T}} H s \bigg\rvert \quad \text{s.t.} \quad
|
||||
\left\{ \begin{array}{l}
|
||||
l \le s \le u,\\
|
||||
\lVert s \rVert \le \Delta,
|
||||
\end{array} \right.
|
||||
|
||||
by maximizing the objective function along the constrained Cauchy
|
||||
direction.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
const : float
|
||||
Constant :math:`c` as shown above.
|
||||
grad : `numpy.ndarray`, shape (n,)
|
||||
Gradient :math:`g` as shown above.
|
||||
curv : callable
|
||||
Curvature of :math:`H` along any vector.
|
||||
|
||||
``curv(s) -> float``
|
||||
|
||||
returns :math:`s^{\mathsf{T}} H s`.
|
||||
xl : `numpy.ndarray`, shape (n,)
|
||||
Lower bounds :math:`l` as shown above.
|
||||
xu : `numpy.ndarray`, shape (n,)
|
||||
Upper bounds :math:`u` as shown above.
|
||||
delta : float
|
||||
Trust-region radius :math:`\Delta` as shown above.
|
||||
debug : bool
|
||||
Whether to make debugging tests during the execution.
|
||||
|
||||
Returns
|
||||
-------
|
||||
`numpy.ndarray`, shape (n,)
|
||||
Approximate solution :math:`s`.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function is described as the first alternative in Section 6.5 of [1]_.
|
||||
It is assumed that the origin is feasible with respect to the bound
|
||||
constraints and that `delta` is finite and positive.
|
||||
|
||||
References
|
||||
----------
|
||||
.. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods
|
||||
and Software*. PhD thesis, Department of Applied Mathematics, The Hong
|
||||
Kong Polytechnic University, Hong Kong, China, 2022. URL:
|
||||
https://theses.lib.polyu.edu.hk/handle/200/12294.
|
||||
"""
|
||||
if debug:
|
||||
assert isinstance(const, float)
|
||||
assert isinstance(grad, np.ndarray) and grad.ndim == 1
|
||||
assert inspect.signature(curv).bind(grad)
|
||||
assert isinstance(xl, np.ndarray) and xl.shape == grad.shape
|
||||
assert isinstance(xu, np.ndarray) and xu.shape == grad.shape
|
||||
assert isinstance(delta, float)
|
||||
assert isinstance(debug, bool)
|
||||
tol = get_arrays_tol(xl, xu)
|
||||
assert np.all(xl <= tol)
|
||||
assert np.all(xu >= -tol)
|
||||
assert np.isfinite(delta) and delta > 0.0
|
||||
xl = np.minimum(xl, 0.0)
|
||||
xu = np.maximum(xu, 0.0)
|
||||
|
||||
# To maximize the absolute value of a quadratic function, we maximize the
|
||||
# function itself or its negative, and we choose the solution that provides
|
||||
# the largest function value.
|
||||
step1, q_val1 = _cauchy_geom(const, grad, curv, xl, xu, delta, debug)
|
||||
step2, q_val2 = _cauchy_geom(
|
||||
-const,
|
||||
-grad,
|
||||
lambda x: -curv(x),
|
||||
xl,
|
||||
xu,
|
||||
delta,
|
||||
debug,
|
||||
)
|
||||
step = step1 if abs(q_val1) >= abs(q_val2) else step2
|
||||
|
||||
if debug:
|
||||
assert np.all(xl <= step)
|
||||
assert np.all(step <= xu)
|
||||
assert np.linalg.norm(step) < 1.1 * delta
|
||||
return step
|
||||
|
||||
|
||||
def spider_geometry(const, grad, curv, xpt, xl, xu, delta, debug):
|
||||
r"""
|
||||
Maximize approximately the absolute value of a quadratic function subject
|
||||
to bound constraints in a trust region.
|
||||
|
||||
This function solves approximately
|
||||
|
||||
.. math::
|
||||
|
||||
\max_{s \in \mathbb{R}^n} \quad \bigg\lvert c + g^{\mathsf{T}} s +
|
||||
\frac{1}{2} s^{\mathsf{T}} H s \bigg\rvert \quad \text{s.t.} \quad
|
||||
\left\{ \begin{array}{l}
|
||||
l \le s \le u,\\
|
||||
\lVert s \rVert \le \Delta,
|
||||
\end{array} \right.
|
||||
|
||||
by maximizing the objective function along given straight lines.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
const : float
|
||||
Constant :math:`c` as shown above.
|
||||
grad : `numpy.ndarray`, shape (n,)
|
||||
Gradient :math:`g` as shown above.
|
||||
curv : callable
|
||||
Curvature of :math:`H` along any vector.
|
||||
|
||||
``curv(s) -> float``
|
||||
|
||||
returns :math:`s^{\mathsf{T}} H s`.
|
||||
xpt : `numpy.ndarray`, shape (n, npt)
|
||||
Points defining the straight lines. The straight lines considered are
|
||||
the ones passing through the origin and the points in `xpt`.
|
||||
xl : `numpy.ndarray`, shape (n,)
|
||||
Lower bounds :math:`l` as shown above.
|
||||
xu : `numpy.ndarray`, shape (n,)
|
||||
Upper bounds :math:`u` as shown above.
|
||||
delta : float
|
||||
Trust-region radius :math:`\Delta` as shown above.
|
||||
debug : bool
|
||||
Whether to make debugging tests during the execution.
|
||||
|
||||
Returns
|
||||
-------
|
||||
`numpy.ndarray`, shape (n,)
|
||||
Approximate solution :math:`s`.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This function is described as the second alternative in Section 6.5 of
|
||||
[1]_. It is assumed that the origin is feasible with respect to the bound
|
||||
constraints and that `delta` is finite and positive.
|
||||
|
||||
References
|
||||
----------
|
||||
.. [1] T. M. Ragonneau. *Model-Based Derivative-Free Optimization Methods
|
||||
and Software*. PhD thesis, Department of Applied Mathematics, The Hong
|
||||
Kong Polytechnic University, Hong Kong, China, 2022. URL:
|
||||
https://theses.lib.polyu.edu.hk/handle/200/12294.
|
||||
"""
|
||||
if debug:
|
||||
assert isinstance(const, float)
|
||||
assert isinstance(grad, np.ndarray) and grad.ndim == 1
|
||||
assert inspect.signature(curv).bind(grad)
|
||||
assert (
|
||||
isinstance(xpt, np.ndarray)
|
||||
and xpt.ndim == 2
|
||||
and xpt.shape[0] == grad.size
|
||||
)
|
||||
assert isinstance(xl, np.ndarray) and xl.shape == grad.shape
|
||||
assert isinstance(xu, np.ndarray) and xu.shape == grad.shape
|
||||
assert isinstance(delta, float)
|
||||
assert isinstance(debug, bool)
|
||||
tol = get_arrays_tol(xl, xu)
|
||||
assert np.all(xl <= tol)
|
||||
assert np.all(xu >= -tol)
|
||||
assert np.isfinite(delta) and delta > 0.0
|
||||
xl = np.minimum(xl, 0.0)
|
||||
xu = np.maximum(xu, 0.0)
|
||||
|
||||
# Iterate through the straight lines.
|
||||
step = np.zeros_like(grad)
|
||||
q_val = const
|
||||
s_norm = np.linalg.norm(xpt, axis=0)
|
||||
|
||||
# Set alpha_xl to the step size for the lower-bound constraint and
|
||||
# alpha_xu to the step size for the upper-bound constraint.
|
||||
|
||||
# xl.shape = (N,)
|
||||
# xpt.shape = (N, M)
|
||||
# i_xl_pos.shape = (M, N)
|
||||
i_xl_pos = (xl > -np.inf) & (xpt.T > -TINY * xl)
|
||||
i_xl_neg = (xl > -np.inf) & (xpt.T < TINY * xl)
|
||||
i_xu_pos = (xu < np.inf) & (xpt.T > TINY * xu)
|
||||
i_xu_neg = (xu < np.inf) & (xpt.T < -TINY * xu)
|
||||
|
||||
# (M, N)
|
||||
alpha_xl_pos = np.atleast_2d(
|
||||
np.broadcast_to(xl, i_xl_pos.shape)[i_xl_pos] / xpt.T[i_xl_pos]
|
||||
)
|
||||
# (M,)
|
||||
alpha_xl_pos = np.max(alpha_xl_pos, axis=1, initial=-np.inf)
|
||||
# make sure it's (M,)
|
||||
alpha_xl_pos = np.broadcast_to(np.atleast_1d(alpha_xl_pos), xpt.shape[1])
|
||||
|
||||
alpha_xl_neg = np.atleast_2d(
|
||||
np.broadcast_to(xl, i_xl_neg.shape)[i_xl_neg] / xpt.T[i_xl_neg]
|
||||
)
|
||||
alpha_xl_neg = np.max(alpha_xl_neg, axis=1, initial=np.inf)
|
||||
alpha_xl_neg = np.broadcast_to(np.atleast_1d(alpha_xl_neg), xpt.shape[1])
|
||||
|
||||
alpha_xu_neg = np.atleast_2d(
|
||||
np.broadcast_to(xu, i_xu_neg.shape)[i_xu_neg] / xpt.T[i_xu_neg]
|
||||
)
|
||||
alpha_xu_neg = np.max(alpha_xu_neg, axis=1, initial=-np.inf)
|
||||
alpha_xu_neg = np.broadcast_to(np.atleast_1d(alpha_xu_neg), xpt.shape[1])
|
||||
|
||||
alpha_xu_pos = np.atleast_2d(
|
||||
np.broadcast_to(xu, i_xu_pos.shape)[i_xu_pos] / xpt.T[i_xu_pos]
|
||||
)
|
||||
alpha_xu_pos = np.max(alpha_xu_pos, axis=1, initial=np.inf)
|
||||
alpha_xu_pos = np.broadcast_to(np.atleast_1d(alpha_xu_pos), xpt.shape[1])
|
||||
|
||||
for k in range(xpt.shape[1]):
|
||||
# Set alpha_tr to the step size for the trust-region constraint.
|
||||
if s_norm[k] > TINY * delta:
|
||||
alpha_tr = max(delta / s_norm[k], 0.0)
|
||||
else:
|
||||
# The current straight line is basically zero.
|
||||
continue
|
||||
|
||||
alpha_bd_pos = max(min(alpha_xu_pos[k], alpha_xl_neg[k]), 0.0)
|
||||
alpha_bd_neg = min(max(alpha_xl_pos[k], alpha_xu_neg[k]), 0.0)
|
||||
|
||||
# Set alpha_quad_pos and alpha_quad_neg to the step size to the extrema
|
||||
# of the quadratic function along the positive and negative directions.
|
||||
grad_step = grad @ xpt[:, k]
|
||||
curv_step = curv(xpt[:, k])
|
||||
if (
|
||||
grad_step >= 0.0
|
||||
and curv_step < -TINY * grad_step
|
||||
or grad_step <= 0.0
|
||||
and curv_step > -TINY * grad_step
|
||||
):
|
||||
alpha_quad_pos = max(-grad_step / curv_step, 0.0)
|
||||
else:
|
||||
alpha_quad_pos = np.inf
|
||||
if (
|
||||
grad_step >= 0.0
|
||||
and curv_step > TINY * grad_step
|
||||
or grad_step <= 0.0
|
||||
and curv_step < TINY * grad_step
|
||||
):
|
||||
alpha_quad_neg = min(-grad_step / curv_step, 0.0)
|
||||
else:
|
||||
alpha_quad_neg = -np.inf
|
||||
|
||||
# Select the step that provides the largest value of the objective
|
||||
# function if it improves the current best. The best positive step is
|
||||
# either the one that reaches the constraints or the one that reaches
|
||||
# the extremum of the objective function along the current direction
|
||||
# (only possible if the resulting step is feasible). We test both, and
|
||||
# we perform similar calculations along the negative step.
|
||||
# N.B.: we select the largest possible step among all the ones that
|
||||
# maximize the objective function. This is to avoid returning the zero
|
||||
# step in some extreme cases.
|
||||
alpha_pos = min(alpha_tr, alpha_bd_pos)
|
||||
alpha_neg = max(-alpha_tr, alpha_bd_neg)
|
||||
q_val_pos = (
|
||||
const + alpha_pos * grad_step + 0.5 * alpha_pos**2.0 * curv_step
|
||||
)
|
||||
q_val_neg = (
|
||||
const + alpha_neg * grad_step + 0.5 * alpha_neg**2.0 * curv_step
|
||||
)
|
||||
if alpha_quad_pos < alpha_pos:
|
||||
q_val_quad_pos = (
|
||||
const
|
||||
+ alpha_quad_pos * grad_step
|
||||
+ 0.5 * alpha_quad_pos**2.0 * curv_step
|
||||
)
|
||||
if abs(q_val_quad_pos) > abs(q_val_pos):
|
||||
alpha_pos = alpha_quad_pos
|
||||
q_val_pos = q_val_quad_pos
|
||||
if alpha_quad_neg > alpha_neg:
|
||||
q_val_quad_neg = (
|
||||
const
|
||||
+ alpha_quad_neg * grad_step
|
||||
+ 0.5 * alpha_quad_neg**2.0 * curv_step
|
||||
)
|
||||
if abs(q_val_quad_neg) > abs(q_val_neg):
|
||||
alpha_neg = alpha_quad_neg
|
||||
q_val_neg = q_val_quad_neg
|
||||
if abs(q_val_pos) >= abs(q_val_neg) and abs(q_val_pos) > abs(q_val):
|
||||
step = np.clip(alpha_pos * xpt[:, k], xl, xu)
|
||||
q_val = q_val_pos
|
||||
elif abs(q_val_neg) > abs(q_val_pos) and abs(q_val_neg) > abs(q_val):
|
||||
step = np.clip(alpha_neg * xpt[:, k], xl, xu)
|
||||
q_val = q_val_neg
|
||||
|
||||
if debug:
|
||||
assert np.all(xl <= step)
|
||||
assert np.all(step <= xu)
|
||||
assert np.linalg.norm(step) < 1.1 * delta
|
||||
return step
|
||||
|
||||
|
||||
def _cauchy_geom(const, grad, curv, xl, xu, delta, debug):
|
||||
"""
|
||||
Same as `bound_constrained_cauchy_step` without the absolute value.
|
||||
"""
|
||||
# Calculate the initial active set.
|
||||
fixed_xl = (xl < 0.0) & (grad > 0.0)
|
||||
fixed_xu = (xu > 0.0) & (grad < 0.0)
|
||||
|
||||
# Calculate the Cauchy step.
|
||||
cauchy_step = np.zeros_like(grad)
|
||||
cauchy_step[fixed_xl] = xl[fixed_xl]
|
||||
cauchy_step[fixed_xu] = xu[fixed_xu]
|
||||
if np.linalg.norm(cauchy_step) > delta:
|
||||
working = fixed_xl | fixed_xu
|
||||
while True:
|
||||
# Calculate the Cauchy step for the directions in the working set.
|
||||
g_norm = np.linalg.norm(grad[working])
|
||||
delta_reduced = np.sqrt(
|
||||
delta**2.0 - cauchy_step[~working] @ cauchy_step[~working]
|
||||
)
|
||||
if g_norm > TINY * abs(delta_reduced):
|
||||
mu = max(delta_reduced / g_norm, 0.0)
|
||||
else:
|
||||
break
|
||||
cauchy_step[working] = mu * grad[working]
|
||||
|
||||
# Update the working set.
|
||||
fixed_xl = working & (cauchy_step < xl)
|
||||
fixed_xu = working & (cauchy_step > xu)
|
||||
if not np.any(fixed_xl) and not np.any(fixed_xu):
|
||||
# Stop the calculations as the Cauchy step is now feasible.
|
||||
break
|
||||
cauchy_step[fixed_xl] = xl[fixed_xl]
|
||||
cauchy_step[fixed_xu] = xu[fixed_xu]
|
||||
working = working & ~(fixed_xl | fixed_xu)
|
||||
|
||||
# Calculate the step that maximizes the quadratic along the Cauchy step.
|
||||
grad_step = grad @ cauchy_step
|
||||
if grad_step >= 0.0:
|
||||
# Set alpha_tr to the step size for the trust-region constraint.
|
||||
s_norm = np.linalg.norm(cauchy_step)
|
||||
if s_norm > TINY * delta:
|
||||
alpha_tr = max(delta / s_norm, 0.0)
|
||||
else:
|
||||
# The Cauchy step is basically zero.
|
||||
alpha_tr = 0.0
|
||||
|
||||
# Set alpha_quad to the step size for the maximization problem.
|
||||
curv_step = curv(cauchy_step)
|
||||
if curv_step < -TINY * grad_step:
|
||||
alpha_quad = max(-grad_step / curv_step, 0.0)
|
||||
else:
|
||||
alpha_quad = np.inf
|
||||
|
||||
# Set alpha_bd to the step size for the bound constraints.
|
||||
i_xl = (xl > -np.inf) & (cauchy_step < TINY * xl)
|
||||
i_xu = (xu < np.inf) & (cauchy_step > TINY * xu)
|
||||
alpha_xl = np.min(xl[i_xl] / cauchy_step[i_xl], initial=np.inf)
|
||||
alpha_xu = np.min(xu[i_xu] / cauchy_step[i_xu], initial=np.inf)
|
||||
alpha_bd = min(alpha_xl, alpha_xu)
|
||||
|
||||
# Calculate the solution and the corresponding function value.
|
||||
alpha = min(alpha_tr, alpha_quad, alpha_bd)
|
||||
step = np.clip(alpha * cauchy_step, xl, xu)
|
||||
q_val = const + alpha * grad_step + 0.5 * alpha**2.0 * curv_step
|
||||
else:
|
||||
# This case is never reached in exact arithmetic. It prevents this
|
||||
# function to return a step that decreases the objective function.
|
||||
step = np.zeros_like(grad)
|
||||
q_val = const
|
||||
|
||||
if debug:
|
||||
assert np.all(xl <= step)
|
||||
assert np.all(step <= xu)
|
||||
assert np.linalg.norm(step) < 1.1 * delta
|
||||
return step, q_val
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,18 @@
|
||||
from .exceptions import (
|
||||
MaxEvalError,
|
||||
TargetSuccess,
|
||||
CallbackSuccess,
|
||||
FeasibleSuccess,
|
||||
)
|
||||
from .math import get_arrays_tol, exact_1d_array
|
||||
from .versions import show_versions
|
||||
|
||||
__all__ = [
|
||||
"MaxEvalError",
|
||||
"TargetSuccess",
|
||||
"CallbackSuccess",
|
||||
"FeasibleSuccess",
|
||||
"get_arrays_tol",
|
||||
"exact_1d_array",
|
||||
"show_versions",
|
||||
]
|
||||
@@ -0,0 +1,22 @@
|
||||
class MaxEvalError(Exception):
|
||||
"""
|
||||
Exception raised when the maximum number of evaluations is reached.
|
||||
"""
|
||||
|
||||
|
||||
class TargetSuccess(Exception):
|
||||
"""
|
||||
Exception raised when the target value is reached.
|
||||
"""
|
||||
|
||||
|
||||
class CallbackSuccess(StopIteration):
|
||||
"""
|
||||
Exception raised when the callback function raises a ``StopIteration``.
|
||||
"""
|
||||
|
||||
|
||||
class FeasibleSuccess(Exception):
|
||||
"""
|
||||
Exception raised when a feasible point of a feasible problem is found.
|
||||
"""
|
||||
@@ -0,0 +1,77 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
EPS = np.finfo(float).eps
|
||||
|
||||
|
||||
def get_arrays_tol(*arrays):
|
||||
"""
|
||||
Get a relative tolerance for a set of arrays.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
*arrays: tuple
|
||||
Set of `numpy.ndarray` to get the tolerance for.
|
||||
|
||||
Returns
|
||||
-------
|
||||
float
|
||||
Relative tolerance for the set of arrays.
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If no array is provided.
|
||||
"""
|
||||
if len(arrays) == 0:
|
||||
raise ValueError("At least one array must be provided.")
|
||||
size = max(array.size for array in arrays)
|
||||
weight = max(
|
||||
np.max(np.abs(array[np.isfinite(array)]), initial=1.0)
|
||||
for array in arrays
|
||||
)
|
||||
return 10.0 * EPS * max(size, 1.0) * weight
|
||||
|
||||
|
||||
def exact_1d_array(x, message):
|
||||
"""
|
||||
Preprocess a 1-dimensional array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array_like
|
||||
Array to be preprocessed.
|
||||
message : str
|
||||
Error message if `x` cannot be interpreter as a 1-dimensional array.
|
||||
|
||||
Returns
|
||||
-------
|
||||
`numpy.ndarray`
|
||||
Preprocessed array.
|
||||
"""
|
||||
x = np.atleast_1d(np.squeeze(x)).astype(float)
|
||||
if x.ndim != 1:
|
||||
raise ValueError(message)
|
||||
return x
|
||||
|
||||
|
||||
def exact_2d_array(x, message):
|
||||
"""
|
||||
Preprocess a 2-dimensional array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array_like
|
||||
Array to be preprocessed.
|
||||
message : str
|
||||
Error message if `x` cannot be interpreter as a 2-dimensional array.
|
||||
|
||||
Returns
|
||||
-------
|
||||
`numpy.ndarray`
|
||||
Preprocessed array.
|
||||
"""
|
||||
x = np.atleast_2d(x).astype(float)
|
||||
if x.ndim != 2:
|
||||
raise ValueError(message)
|
||||
return x
|
||||
@@ -0,0 +1,67 @@
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
|
||||
|
||||
def _get_sys_info():
|
||||
"""
|
||||
Get useful system information.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
Useful system information.
|
||||
"""
|
||||
return {
|
||||
"python": sys.version.replace(os.linesep, " "),
|
||||
"executable": sys.executable,
|
||||
"machine": platform.platform(),
|
||||
}
|
||||
|
||||
|
||||
def _get_deps_info():
|
||||
"""
|
||||
Get the versions of the dependencies.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
Versions of the dependencies.
|
||||
"""
|
||||
deps = ["cobyqa", "numpy", "scipy", "setuptools", "pip"]
|
||||
deps_info = {}
|
||||
for module in deps:
|
||||
try:
|
||||
deps_info[module] = version(module)
|
||||
except PackageNotFoundError:
|
||||
deps_info[module] = None
|
||||
return deps_info
|
||||
|
||||
|
||||
def show_versions():
|
||||
"""
|
||||
Display useful system and dependencies information.
|
||||
|
||||
When reporting issues, please include this information.
|
||||
"""
|
||||
print("System settings")
|
||||
print("---------------")
|
||||
sys_info = _get_sys_info()
|
||||
print(
|
||||
"\n".join(
|
||||
f"{k:>{max(map(len, sys_info.keys())) + 1}}: {v}"
|
||||
for k, v in sys_info.items()
|
||||
)
|
||||
)
|
||||
|
||||
print()
|
||||
print("Python dependencies")
|
||||
print("-------------------")
|
||||
deps_info = _get_deps_info()
|
||||
print(
|
||||
"\n".join(
|
||||
f"{k:>{max(map(len, deps_info.keys())) + 1}}: {v}"
|
||||
for k, v in deps_info.items()
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,399 @@
|
||||
# ######################### LICENSE ############################ #
|
||||
|
||||
# Copyright (c) 2005-2015, Michele Simionato
|
||||
# All rights reserved.
|
||||
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are
|
||||
# met:
|
||||
|
||||
# Redistributions of source code must retain the above copyright
|
||||
# notice, this list of conditions and the following disclaimer.
|
||||
# Redistributions in bytecode form must reproduce the above copyright
|
||||
# notice, this list of conditions and the following disclaimer in
|
||||
# the documentation and/or other materials provided with the
|
||||
# distribution.
|
||||
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
||||
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
||||
# HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
|
||||
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
|
||||
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
|
||||
# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
|
||||
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
|
||||
# TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
|
||||
# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
|
||||
# DAMAGE.
|
||||
|
||||
"""
|
||||
Decorator module, see https://pypi.python.org/pypi/decorator
|
||||
for the documentation.
|
||||
"""
|
||||
import re
|
||||
import sys
|
||||
import inspect
|
||||
import operator
|
||||
import itertools
|
||||
import collections
|
||||
|
||||
from inspect import getfullargspec
|
||||
|
||||
__version__ = '4.0.5'
|
||||
|
||||
|
||||
def get_init(cls):
|
||||
return cls.__init__
|
||||
|
||||
|
||||
# getargspec has been deprecated in Python 3.5
|
||||
ArgSpec = collections.namedtuple(
|
||||
'ArgSpec', 'args varargs varkw defaults')
|
||||
|
||||
|
||||
def getargspec(f):
|
||||
"""A replacement for inspect.getargspec"""
|
||||
spec = getfullargspec(f)
|
||||
return ArgSpec(spec.args, spec.varargs, spec.varkw, spec.defaults)
|
||||
|
||||
|
||||
DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(')
|
||||
|
||||
|
||||
# basic functionality
|
||||
class FunctionMaker:
|
||||
"""
|
||||
An object with the ability to create functions with a given signature.
|
||||
It has attributes name, doc, module, signature, defaults, dict, and
|
||||
methods update and make.
|
||||
"""
|
||||
|
||||
# Atomic get-and-increment provided by the GIL
|
||||
_compile_count = itertools.count()
|
||||
|
||||
def __init__(self, func=None, name=None, signature=None,
|
||||
defaults=None, doc=None, module=None, funcdict=None):
|
||||
self.shortsignature = signature
|
||||
if func:
|
||||
# func can be a class or a callable, but not an instance method
|
||||
self.name = func.__name__
|
||||
if self.name == '<lambda>': # small hack for lambda functions
|
||||
self.name = '_lambda_'
|
||||
self.doc = func.__doc__
|
||||
self.module = func.__module__
|
||||
if inspect.isfunction(func):
|
||||
argspec = getfullargspec(func)
|
||||
self.annotations = getattr(func, '__annotations__', {})
|
||||
for a in ('args', 'varargs', 'varkw', 'defaults', 'kwonlyargs',
|
||||
'kwonlydefaults'):
|
||||
setattr(self, a, getattr(argspec, a))
|
||||
for i, arg in enumerate(self.args):
|
||||
setattr(self, f'arg{i}', arg)
|
||||
allargs = list(self.args)
|
||||
allshortargs = list(self.args)
|
||||
if self.varargs:
|
||||
allargs.append('*' + self.varargs)
|
||||
allshortargs.append('*' + self.varargs)
|
||||
elif self.kwonlyargs:
|
||||
allargs.append('*') # single star syntax
|
||||
for a in self.kwonlyargs:
|
||||
allargs.append(f'{a}=None')
|
||||
allshortargs.append(f'{a}={a}')
|
||||
if self.varkw:
|
||||
allargs.append('**' + self.varkw)
|
||||
allshortargs.append('**' + self.varkw)
|
||||
self.signature = ', '.join(allargs)
|
||||
self.shortsignature = ', '.join(allshortargs)
|
||||
self.dict = func.__dict__.copy()
|
||||
# func=None happens when decorating a caller
|
||||
if name:
|
||||
self.name = name
|
||||
if signature is not None:
|
||||
self.signature = signature
|
||||
if defaults:
|
||||
self.defaults = defaults
|
||||
if doc:
|
||||
self.doc = doc
|
||||
if module:
|
||||
self.module = module
|
||||
if funcdict:
|
||||
self.dict = funcdict
|
||||
# check existence required attributes
|
||||
assert hasattr(self, 'name')
|
||||
if not hasattr(self, 'signature'):
|
||||
raise TypeError(f'You are decorating a non-function: {func}')
|
||||
|
||||
def update(self, func, **kw):
|
||||
"Update the signature of func with the data in self"
|
||||
func.__name__ = self.name
|
||||
func.__doc__ = getattr(self, 'doc', None)
|
||||
func.__dict__ = getattr(self, 'dict', {})
|
||||
func.__defaults__ = getattr(self, 'defaults', ())
|
||||
func.__kwdefaults__ = getattr(self, 'kwonlydefaults', None)
|
||||
func.__annotations__ = getattr(self, 'annotations', None)
|
||||
try:
|
||||
frame = sys._getframe(3)
|
||||
except AttributeError: # for IronPython and similar implementations
|
||||
callermodule = '?'
|
||||
else:
|
||||
callermodule = frame.f_globals.get('__name__', '?')
|
||||
func.__module__ = getattr(self, 'module', callermodule)
|
||||
func.__dict__.update(kw)
|
||||
|
||||
def make(self, src_templ, evaldict=None, addsource=False, **attrs):
|
||||
"Make a new function from a given template and update the signature"
|
||||
src = src_templ % vars(self) # expand name and signature
|
||||
evaldict = evaldict or {}
|
||||
mo = DEF.match(src)
|
||||
if mo is None:
|
||||
raise SyntaxError(f'not a valid function template\n{src}')
|
||||
name = mo.group(1) # extract the function name
|
||||
names = set([name] + [arg.strip(' *') for arg in
|
||||
self.shortsignature.split(',')])
|
||||
for n in names:
|
||||
if n in ('_func_', '_call_'):
|
||||
raise NameError(f'{n} is overridden in\n{src}')
|
||||
if not src.endswith('\n'): # add a newline just for safety
|
||||
src += '\n' # this is needed in old versions of Python
|
||||
|
||||
# Ensure each generated function has a unique filename for profilers
|
||||
# (such as cProfile) that depend on the tuple of (<filename>,
|
||||
# <definition line>, <function name>) being unique.
|
||||
filename = f'<decorator-gen-{next(self._compile_count)}>'
|
||||
try:
|
||||
code = compile(src, filename, 'single')
|
||||
exec(code, evaldict)
|
||||
except: # noqa: E722
|
||||
print('Error in generated code:', file=sys.stderr)
|
||||
print(src, file=sys.stderr)
|
||||
raise
|
||||
func = evaldict[name]
|
||||
if addsource:
|
||||
attrs['__source__'] = src
|
||||
self.update(func, **attrs)
|
||||
return func
|
||||
|
||||
@classmethod
|
||||
def create(cls, obj, body, evaldict, defaults=None,
|
||||
doc=None, module=None, addsource=True, **attrs):
|
||||
"""
|
||||
Create a function from the strings name, signature, and body.
|
||||
evaldict is the evaluation dictionary. If addsource is true, an
|
||||
attribute __source__ is added to the result. The attributes attrs
|
||||
are added, if any.
|
||||
"""
|
||||
if isinstance(obj, str): # "name(signature)"
|
||||
name, rest = obj.strip().split('(', 1)
|
||||
signature = rest[:-1] # strip a right parens
|
||||
func = None
|
||||
else: # a function
|
||||
name = None
|
||||
signature = None
|
||||
func = obj
|
||||
self = cls(func, name, signature, defaults, doc, module)
|
||||
ibody = '\n'.join(' ' + line for line in body.splitlines())
|
||||
return self.make('def %(name)s(%(signature)s):\n' + ibody,
|
||||
evaldict, addsource, **attrs)
|
||||
|
||||
|
||||
def decorate(func, caller):
|
||||
"""
|
||||
decorate(func, caller) decorates a function using a caller.
|
||||
"""
|
||||
evaldict = func.__globals__.copy()
|
||||
evaldict['_call_'] = caller
|
||||
evaldict['_func_'] = func
|
||||
fun = FunctionMaker.create(
|
||||
func, "return _call_(_func_, %(shortsignature)s)",
|
||||
evaldict, __wrapped__=func)
|
||||
if hasattr(func, '__qualname__'):
|
||||
fun.__qualname__ = func.__qualname__
|
||||
return fun
|
||||
|
||||
|
||||
def decorator(caller, _func=None):
|
||||
"""decorator(caller) converts a caller function into a decorator"""
|
||||
if _func is not None: # return a decorated function
|
||||
# this is obsolete behavior; you should use decorate instead
|
||||
return decorate(_func, caller)
|
||||
# else return a decorator function
|
||||
if inspect.isclass(caller):
|
||||
name = caller.__name__.lower()
|
||||
callerfunc = get_init(caller)
|
||||
doc = (f'decorator({caller.__name__}) converts functions/generators into '
|
||||
f'factories of {caller.__name__} objects')
|
||||
elif inspect.isfunction(caller):
|
||||
if caller.__name__ == '<lambda>':
|
||||
name = '_lambda_'
|
||||
else:
|
||||
name = caller.__name__
|
||||
callerfunc = caller
|
||||
doc = caller.__doc__
|
||||
else: # assume caller is an object with a __call__ method
|
||||
name = caller.__class__.__name__.lower()
|
||||
callerfunc = caller.__call__.__func__
|
||||
doc = caller.__call__.__doc__
|
||||
evaldict = callerfunc.__globals__.copy()
|
||||
evaldict['_call_'] = caller
|
||||
evaldict['_decorate_'] = decorate
|
||||
return FunctionMaker.create(
|
||||
f'{name}(func)', 'return _decorate_(func, _call_)',
|
||||
evaldict, doc=doc, module=caller.__module__,
|
||||
__wrapped__=caller)
|
||||
|
||||
|
||||
# ####################### contextmanager ####################### #
|
||||
|
||||
try: # Python >= 3.2
|
||||
from contextlib import _GeneratorContextManager
|
||||
except ImportError: # Python >= 2.5
|
||||
from contextlib import GeneratorContextManager as _GeneratorContextManager
|
||||
|
||||
|
||||
class ContextManager(_GeneratorContextManager):
|
||||
def __call__(self, func):
|
||||
"""Context manager decorator"""
|
||||
return FunctionMaker.create(
|
||||
func, "with _self_: return _func_(%(shortsignature)s)",
|
||||
dict(_self_=self, _func_=func), __wrapped__=func)
|
||||
|
||||
|
||||
init = getfullargspec(_GeneratorContextManager.__init__)
|
||||
n_args = len(init.args)
|
||||
if n_args == 2 and not init.varargs: # (self, genobj) Python 2.7
|
||||
def __init__(self, g, *a, **k):
|
||||
return _GeneratorContextManager.__init__(self, g(*a, **k))
|
||||
ContextManager.__init__ = __init__
|
||||
elif n_args == 2 and init.varargs: # (self, gen, *a, **k) Python 3.4
|
||||
pass
|
||||
elif n_args == 4: # (self, gen, args, kwds) Python 3.5
|
||||
def __init__(self, g, *a, **k):
|
||||
return _GeneratorContextManager.__init__(self, g, a, k)
|
||||
ContextManager.__init__ = __init__
|
||||
|
||||
contextmanager = decorator(ContextManager)
|
||||
|
||||
|
||||
# ############################ dispatch_on ############################ #
|
||||
|
||||
def append(a, vancestors):
|
||||
"""
|
||||
Append ``a`` to the list of the virtual ancestors, unless it is already
|
||||
included.
|
||||
"""
|
||||
add = True
|
||||
for j, va in enumerate(vancestors):
|
||||
if issubclass(va, a):
|
||||
add = False
|
||||
break
|
||||
if issubclass(a, va):
|
||||
vancestors[j] = a
|
||||
add = False
|
||||
if add:
|
||||
vancestors.append(a)
|
||||
|
||||
|
||||
# inspired from simplegeneric by P.J. Eby and functools.singledispatch
|
||||
def dispatch_on(*dispatch_args):
|
||||
"""
|
||||
Factory of decorators turning a function into a generic function
|
||||
dispatching on the given arguments.
|
||||
"""
|
||||
assert dispatch_args, 'No dispatch args passed'
|
||||
dispatch_str = f"({', '.join(dispatch_args)},)"
|
||||
|
||||
def check(arguments, wrong=operator.ne, msg=''):
|
||||
"""Make sure one passes the expected number of arguments"""
|
||||
if wrong(len(arguments), len(dispatch_args)):
|
||||
raise TypeError(f'Expected {len(dispatch_args)} arguments, '
|
||||
'got {len(arguments)}{msg}')
|
||||
|
||||
def gen_func_dec(func):
|
||||
"""Decorator turning a function into a generic function"""
|
||||
|
||||
# first check the dispatch arguments
|
||||
argset = set(getfullargspec(func).args)
|
||||
if not set(dispatch_args) <= argset:
|
||||
raise NameError(f'Unknown dispatch arguments {dispatch_str}')
|
||||
|
||||
typemap = {}
|
||||
|
||||
def vancestors(*types):
|
||||
"""
|
||||
Get a list of sets of virtual ancestors for the given types
|
||||
"""
|
||||
check(types)
|
||||
ras = [[] for _ in range(len(dispatch_args))]
|
||||
for types_ in typemap:
|
||||
for t, type_, ra in zip(types, types_, ras):
|
||||
if issubclass(t, type_) and type_ not in t.__mro__:
|
||||
append(type_, ra)
|
||||
return [set(ra) for ra in ras]
|
||||
|
||||
def ancestors(*types):
|
||||
"""
|
||||
Get a list of virtual MROs, one for each type
|
||||
"""
|
||||
check(types)
|
||||
lists = []
|
||||
for t, vas in zip(types, vancestors(*types)):
|
||||
n_vas = len(vas)
|
||||
if n_vas > 1:
|
||||
raise RuntimeError(
|
||||
f'Ambiguous dispatch for {t}: {vas}')
|
||||
elif n_vas == 1:
|
||||
va, = vas
|
||||
mro = type('t', (t, va), {}).__mro__[1:]
|
||||
else:
|
||||
mro = t.__mro__
|
||||
lists.append(mro[:-1]) # discard t and object
|
||||
return lists
|
||||
|
||||
def register(*types):
|
||||
"""
|
||||
Decorator to register an implementation for the given types
|
||||
"""
|
||||
check(types)
|
||||
|
||||
def dec(f):
|
||||
check(getfullargspec(f).args, operator.lt, ' in ' + f.__name__)
|
||||
typemap[types] = f
|
||||
return f
|
||||
return dec
|
||||
|
||||
def dispatch_info(*types):
|
||||
"""
|
||||
An utility to introspect the dispatch algorithm
|
||||
"""
|
||||
check(types)
|
||||
lst = [tuple(a.__name__ for a in anc)
|
||||
for anc in itertools.product(*ancestors(*types))]
|
||||
return lst
|
||||
|
||||
def _dispatch(dispatch_args, *args, **kw):
|
||||
types = tuple(type(arg) for arg in dispatch_args)
|
||||
try: # fast path
|
||||
f = typemap[types]
|
||||
except KeyError:
|
||||
pass
|
||||
else:
|
||||
return f(*args, **kw)
|
||||
combinations = itertools.product(*ancestors(*types))
|
||||
next(combinations) # the first one has been already tried
|
||||
for types_ in combinations:
|
||||
f = typemap.get(types_)
|
||||
if f is not None:
|
||||
return f(*args, **kw)
|
||||
|
||||
# else call the default implementation
|
||||
return func(*args, **kw)
|
||||
|
||||
return FunctionMaker.create(
|
||||
func, f'return _f_({dispatch_str}, %%(shortsignature)s)',
|
||||
dict(_f_=_dispatch), register=register, default=func,
|
||||
typemap=typemap, vancestors=vancestors, ancestors=ancestors,
|
||||
dispatch_info=dispatch_info, __wrapped__=func)
|
||||
|
||||
gen_func_dec.__name__ = 'dispatch_on' + dispatch_str
|
||||
return gen_func_dec
|
||||
@@ -0,0 +1,274 @@
|
||||
from inspect import Parameter, signature
|
||||
import functools
|
||||
import warnings
|
||||
from importlib import import_module
|
||||
from scipy._lib._docscrape import FunctionDoc
|
||||
|
||||
|
||||
__all__ = ["_deprecated"]
|
||||
|
||||
|
||||
# Object to use as default value for arguments to be deprecated. This should
|
||||
# be used over 'None' as the user could parse 'None' as a positional argument
|
||||
_NoValue = object()
|
||||
|
||||
def _sub_module_deprecation(*, sub_package, module, private_modules, all,
|
||||
attribute, correct_module=None, dep_version="1.16.0"):
|
||||
"""Helper function for deprecating modules that are public but were
|
||||
intended to be private.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
sub_package : str
|
||||
Subpackage the module belongs to eg. stats
|
||||
module : str
|
||||
Public but intended private module to deprecate
|
||||
private_modules : list
|
||||
Private replacement(s) for `module`; should contain the
|
||||
content of ``all``, possibly spread over several modules.
|
||||
all : list
|
||||
``__all__`` belonging to `module`
|
||||
attribute : str
|
||||
The attribute in `module` being accessed
|
||||
correct_module : str, optional
|
||||
Module in `sub_package` that `attribute` should be imported from.
|
||||
Default is that `attribute` should be imported from ``scipy.sub_package``.
|
||||
dep_version : str, optional
|
||||
Version in which deprecated attributes will be removed.
|
||||
"""
|
||||
if correct_module is not None:
|
||||
correct_import = f"scipy.{sub_package}.{correct_module}"
|
||||
else:
|
||||
correct_import = f"scipy.{sub_package}"
|
||||
|
||||
if attribute not in all:
|
||||
raise AttributeError(
|
||||
f"`scipy.{sub_package}.{module}` has no attribute `{attribute}`; "
|
||||
f"furthermore, `scipy.{sub_package}.{module}` is deprecated "
|
||||
f"and will be removed in SciPy 2.0.0."
|
||||
)
|
||||
|
||||
attr = getattr(import_module(correct_import), attribute, None)
|
||||
|
||||
if attr is not None:
|
||||
message = (
|
||||
f"Please import `{attribute}` from the `{correct_import}` namespace; "
|
||||
f"the `scipy.{sub_package}.{module}` namespace is deprecated "
|
||||
f"and will be removed in SciPy 2.0.0."
|
||||
)
|
||||
else:
|
||||
message = (
|
||||
f"`scipy.{sub_package}.{module}.{attribute}` is deprecated along with "
|
||||
f"the `scipy.{sub_package}.{module}` namespace. "
|
||||
f"`scipy.{sub_package}.{module}.{attribute}` will be removed "
|
||||
f"in SciPy {dep_version}, and the `scipy.{sub_package}.{module}` namespace "
|
||||
f"will be removed in SciPy 2.0.0."
|
||||
)
|
||||
|
||||
warnings.warn(message, category=DeprecationWarning, stacklevel=3)
|
||||
|
||||
for module in private_modules:
|
||||
try:
|
||||
return getattr(import_module(f"scipy.{sub_package}.{module}"), attribute)
|
||||
except AttributeError as e:
|
||||
# still raise an error if the attribute isn't in any of the expected
|
||||
# private modules
|
||||
if module == private_modules[-1]:
|
||||
raise e
|
||||
continue
|
||||
|
||||
|
||||
def _deprecated(msg, stacklevel=2):
|
||||
"""Deprecate a function by emitting a warning on use."""
|
||||
def wrap(fun):
|
||||
if isinstance(fun, type):
|
||||
warnings.warn(
|
||||
f"Trying to deprecate class {fun!r}",
|
||||
category=RuntimeWarning, stacklevel=2)
|
||||
return fun
|
||||
|
||||
@functools.wraps(fun)
|
||||
def call(*args, **kwargs):
|
||||
warnings.warn(msg, category=DeprecationWarning,
|
||||
stacklevel=stacklevel)
|
||||
return fun(*args, **kwargs)
|
||||
call.__doc__ = fun.__doc__
|
||||
return call
|
||||
|
||||
return wrap
|
||||
|
||||
|
||||
class _DeprecationHelperStr:
|
||||
"""
|
||||
Helper class used by deprecate_cython_api
|
||||
"""
|
||||
def __init__(self, content, message):
|
||||
self._content = content
|
||||
self._message = message
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self._content)
|
||||
|
||||
def __eq__(self, other):
|
||||
res = (self._content == other)
|
||||
if res:
|
||||
warnings.warn(self._message, category=DeprecationWarning,
|
||||
stacklevel=2)
|
||||
return res
|
||||
|
||||
|
||||
def deprecate_cython_api(module, routine_name, new_name=None, message=None):
|
||||
"""
|
||||
Deprecate an exported cdef function in a public Cython API module.
|
||||
|
||||
Only functions can be deprecated; typedefs etc. cannot.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
module : module
|
||||
Public Cython API module (e.g. scipy.linalg.cython_blas).
|
||||
routine_name : str
|
||||
Name of the routine to deprecate. May also be a fused-type
|
||||
routine (in which case its all specializations are deprecated).
|
||||
new_name : str
|
||||
New name to include in the deprecation warning message
|
||||
message : str
|
||||
Additional text in the deprecation warning message
|
||||
|
||||
Examples
|
||||
--------
|
||||
Usually, this function would be used in the top-level of the
|
||||
module ``.pyx`` file:
|
||||
|
||||
>>> from scipy._lib.deprecation import deprecate_cython_api
|
||||
>>> import scipy.linalg.cython_blas as mod
|
||||
>>> deprecate_cython_api(mod, "dgemm", "dgemm_new",
|
||||
... message="Deprecated in Scipy 1.5.0")
|
||||
>>> del deprecate_cython_api, mod
|
||||
|
||||
After this, Cython modules that use the deprecated function emit a
|
||||
deprecation warning when they are imported.
|
||||
|
||||
"""
|
||||
old_name = f"{module.__name__}.{routine_name}"
|
||||
|
||||
if new_name is None:
|
||||
depdoc = f"`{old_name}` is deprecated!"
|
||||
else:
|
||||
depdoc = f"`{old_name}` is deprecated, use `{new_name}` instead!"
|
||||
|
||||
if message is not None:
|
||||
depdoc += "\n" + message
|
||||
|
||||
d = module.__pyx_capi__
|
||||
|
||||
# Check if the function is a fused-type function with a mangled name
|
||||
j = 0
|
||||
has_fused = False
|
||||
while True:
|
||||
fused_name = f"__pyx_fuse_{j}{routine_name}"
|
||||
if fused_name in d:
|
||||
has_fused = True
|
||||
d[_DeprecationHelperStr(fused_name, depdoc)] = d.pop(fused_name)
|
||||
j += 1
|
||||
else:
|
||||
break
|
||||
|
||||
# If not, apply deprecation to the named routine
|
||||
if not has_fused:
|
||||
d[_DeprecationHelperStr(routine_name, depdoc)] = d.pop(routine_name)
|
||||
|
||||
|
||||
# taken from scikit-learn, see
|
||||
# https://github.com/scikit-learn/scikit-learn/blob/1.3.0/sklearn/utils/validation.py#L38
|
||||
def _deprecate_positional_args(func=None, *, version=None,
|
||||
deprecated_args=None, custom_message=""):
|
||||
"""Decorator for methods that issues warnings for positional arguments.
|
||||
|
||||
Using the keyword-only argument syntax in pep 3102, arguments after the
|
||||
* will issue a warning when passed as a positional argument.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : callable, default=None
|
||||
Function to check arguments on.
|
||||
version : callable, default=None
|
||||
The version when positional arguments will result in error.
|
||||
deprecated_args : set of str, optional
|
||||
Arguments to deprecate - whether passed by position or keyword.
|
||||
custom_message : str, optional
|
||||
Custom message to add to deprecation warning and documentation.
|
||||
"""
|
||||
if version is None:
|
||||
msg = "Need to specify a version where signature will be changed"
|
||||
raise ValueError(msg)
|
||||
|
||||
deprecated_args = set() if deprecated_args is None else set(deprecated_args)
|
||||
|
||||
def _inner_deprecate_positional_args(f):
|
||||
sig = signature(f)
|
||||
kwonly_args = []
|
||||
all_args = []
|
||||
|
||||
for name, param in sig.parameters.items():
|
||||
if param.kind == Parameter.POSITIONAL_OR_KEYWORD:
|
||||
all_args.append(name)
|
||||
elif param.kind == Parameter.KEYWORD_ONLY:
|
||||
kwonly_args.append(name)
|
||||
|
||||
def warn_deprecated_args(kwargs):
|
||||
intersection = deprecated_args.intersection(kwargs)
|
||||
if intersection:
|
||||
message = (f"Arguments {intersection} are deprecated, whether passed "
|
||||
"by position or keyword. They will be removed in SciPy "
|
||||
f"{version}. ")
|
||||
message += custom_message
|
||||
warnings.warn(message, category=DeprecationWarning, stacklevel=3)
|
||||
|
||||
@functools.wraps(f)
|
||||
def inner_f(*args, **kwargs):
|
||||
|
||||
extra_args = len(args) - len(all_args)
|
||||
if extra_args <= 0:
|
||||
warn_deprecated_args(kwargs)
|
||||
return f(*args, **kwargs)
|
||||
|
||||
# extra_args > 0
|
||||
kwonly_extra_args = set(kwonly_args[:extra_args]) - deprecated_args
|
||||
args_msg = ", ".join(kwonly_extra_args)
|
||||
warnings.warn(
|
||||
(
|
||||
f"You are passing as positional arguments: {args_msg}. "
|
||||
"Please change your invocation to use keyword arguments. "
|
||||
f"From SciPy {version}, passing these as positional "
|
||||
"arguments will result in an error."
|
||||
),
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
kwargs.update(zip(sig.parameters, args))
|
||||
warn_deprecated_args(kwargs)
|
||||
return f(**kwargs)
|
||||
|
||||
doc = FunctionDoc(inner_f)
|
||||
kwonly_extra_args = set(kwonly_args) - deprecated_args
|
||||
admonition = f"""
|
||||
.. deprecated:: {version}
|
||||
Use of argument(s) ``{kwonly_extra_args}`` by position is deprecated; beginning in
|
||||
SciPy {version}, these will be keyword-only. """
|
||||
if deprecated_args:
|
||||
admonition += (f"Argument(s) ``{deprecated_args}`` are deprecated, whether "
|
||||
"passed by position or keyword; they will be removed in "
|
||||
f"SciPy {version}. ")
|
||||
admonition += custom_message
|
||||
doc['Extended Summary'] += [admonition]
|
||||
|
||||
doc = str(doc).split("\n", 1)[1] # remove signature
|
||||
inner_f.__doc__ = str(doc)
|
||||
|
||||
return inner_f
|
||||
|
||||
if func is not None:
|
||||
return _inner_deprecate_positional_args(func)
|
||||
|
||||
return _inner_deprecate_positional_args
|
||||
@@ -0,0 +1,366 @@
|
||||
"""Utilities to allow inserting docstring fragments for common
|
||||
parameters into function and method docstrings."""
|
||||
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from typing import Protocol, TypeVar
|
||||
import sys
|
||||
|
||||
__all__ = [
|
||||
"docformat",
|
||||
"inherit_docstring_from",
|
||||
"indentcount_lines",
|
||||
"filldoc",
|
||||
"unindent_dict",
|
||||
"unindent_string",
|
||||
"extend_notes_in_docstring",
|
||||
"replace_notes_in_docstring",
|
||||
"doc_replace",
|
||||
]
|
||||
|
||||
_F = TypeVar("_F", bound=Callable[..., object])
|
||||
|
||||
|
||||
class Decorator(Protocol):
|
||||
"""A decorator of a function."""
|
||||
|
||||
def __call__(self, func: _F, /) -> _F: ...
|
||||
|
||||
|
||||
def docformat(docstring: str, docdict: Mapping[str, str] | None = None) -> str:
|
||||
"""Fill a function docstring from variables in dictionary.
|
||||
|
||||
Adapt the indent of the inserted docs
|
||||
|
||||
Parameters
|
||||
----------
|
||||
docstring : str
|
||||
A docstring from a function, possibly with dict formatting strings.
|
||||
docdict : dict[str, str], optional
|
||||
A dictionary with keys that match the dict formatting strings
|
||||
and values that are docstring fragments to be inserted. The
|
||||
indentation of the inserted docstrings is set to match the
|
||||
minimum indentation of the ``docstring`` by adding this
|
||||
indentation to all lines of the inserted string, except the
|
||||
first.
|
||||
|
||||
Returns
|
||||
-------
|
||||
docstring : str
|
||||
string with requested ``docdict`` strings inserted.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> docformat(' Test string with %(value)s', {'value':'inserted value'})
|
||||
' Test string with inserted value'
|
||||
>>> docstring = 'First line\\n Second line\\n %(value)s'
|
||||
>>> inserted_string = "indented\\nstring"
|
||||
>>> docdict = {'value': inserted_string}
|
||||
>>> docformat(docstring, docdict)
|
||||
'First line\\n Second line\\n indented\\n string'
|
||||
"""
|
||||
if not docstring:
|
||||
return docstring
|
||||
if docdict is None:
|
||||
docdict = {}
|
||||
if not docdict:
|
||||
return docstring
|
||||
lines = docstring.expandtabs().splitlines()
|
||||
# Find the minimum indent of the main docstring, after first line
|
||||
if len(lines) < 2:
|
||||
icount = 0
|
||||
else:
|
||||
icount = indentcount_lines(lines[1:])
|
||||
indent = " " * icount
|
||||
# Insert this indent to dictionary docstrings
|
||||
indented = {}
|
||||
for name, dstr in docdict.items():
|
||||
lines = dstr.expandtabs().splitlines()
|
||||
indented[name] = ("\n" + indent).join(lines)
|
||||
return docstring % indented
|
||||
|
||||
|
||||
def inherit_docstring_from(cls: object) -> Decorator:
|
||||
"""This decorator modifies the decorated function's docstring by
|
||||
replacing occurrences of '%(super)s' with the docstring of the
|
||||
method of the same name from the class `cls`.
|
||||
|
||||
If the decorated method has no docstring, it is simply given the
|
||||
docstring of `cls`s method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cls : type or object
|
||||
A class with a method with the same name as the decorated method.
|
||||
The docstring of the method in this class replaces '%(super)s' in the
|
||||
docstring of the decorated method.
|
||||
|
||||
Returns
|
||||
-------
|
||||
decfunc : function
|
||||
The decorator function that modifies the __doc__ attribute
|
||||
of its argument.
|
||||
|
||||
Examples
|
||||
--------
|
||||
In the following, the docstring for Bar.func created using the
|
||||
docstring of `Foo.func`.
|
||||
|
||||
>>> class Foo:
|
||||
... def func(self):
|
||||
... '''Do something useful.'''
|
||||
... return
|
||||
...
|
||||
>>> class Bar(Foo):
|
||||
... @inherit_docstring_from(Foo)
|
||||
... def func(self):
|
||||
... '''%(super)s
|
||||
... Do it fast.
|
||||
... '''
|
||||
... return
|
||||
...
|
||||
>>> b = Bar()
|
||||
>>> b.func.__doc__
|
||||
'Do something useful.\n Do it fast.\n '
|
||||
"""
|
||||
|
||||
def _doc(func: _F) -> _F:
|
||||
cls_docstring = getattr(cls, func.__name__).__doc__
|
||||
func_docstring = func.__doc__
|
||||
if func_docstring is None:
|
||||
func.__doc__ = cls_docstring
|
||||
else:
|
||||
new_docstring = func_docstring % dict(super=cls_docstring)
|
||||
func.__doc__ = new_docstring
|
||||
return func
|
||||
|
||||
return _doc
|
||||
|
||||
|
||||
def extend_notes_in_docstring(cls: object, notes: str) -> Decorator:
|
||||
"""This decorator replaces the decorated function's docstring
|
||||
with the docstring from corresponding method in `cls`.
|
||||
It extends the 'Notes' section of that docstring to include
|
||||
the given `notes`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cls : type or object
|
||||
A class with a method with the same name as the decorated method.
|
||||
The docstring of the method in this class replaces the docstring of the
|
||||
decorated method.
|
||||
notes : str
|
||||
Additional notes to append to the 'Notes' section of the docstring.
|
||||
|
||||
Returns
|
||||
-------
|
||||
decfunc : function
|
||||
The decorator function that modifies the __doc__ attribute
|
||||
of its argument.
|
||||
"""
|
||||
|
||||
def _doc(func: _F) -> _F:
|
||||
cls_docstring = getattr(cls, func.__name__).__doc__
|
||||
# If python is called with -OO option,
|
||||
# there is no docstring
|
||||
if cls_docstring is None:
|
||||
return func
|
||||
end_of_notes = cls_docstring.find(" References\n")
|
||||
if end_of_notes == -1:
|
||||
end_of_notes = cls_docstring.find(" Examples\n")
|
||||
if end_of_notes == -1:
|
||||
end_of_notes = len(cls_docstring)
|
||||
func.__doc__ = (
|
||||
cls_docstring[:end_of_notes] + notes + cls_docstring[end_of_notes:]
|
||||
)
|
||||
return func
|
||||
|
||||
return _doc
|
||||
|
||||
|
||||
def replace_notes_in_docstring(cls: object, notes: str) -> Decorator:
|
||||
"""This decorator replaces the decorated function's docstring
|
||||
with the docstring from corresponding method in `cls`.
|
||||
It replaces the 'Notes' section of that docstring with
|
||||
the given `notes`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cls : type or object
|
||||
A class with a method with the same name as the decorated method.
|
||||
The docstring of the method in this class replaces the docstring of the
|
||||
decorated method.
|
||||
notes : str
|
||||
The notes to replace the existing 'Notes' section with.
|
||||
|
||||
Returns
|
||||
-------
|
||||
decfunc : function
|
||||
The decorator function that modifies the __doc__ attribute
|
||||
of its argument.
|
||||
"""
|
||||
|
||||
def _doc(func: _F) -> _F:
|
||||
cls_docstring = getattr(cls, func.__name__).__doc__
|
||||
notes_header = " Notes\n -----\n"
|
||||
# If python is called with -OO option,
|
||||
# there is no docstring
|
||||
if cls_docstring is None:
|
||||
return func
|
||||
start_of_notes = cls_docstring.find(notes_header)
|
||||
end_of_notes = cls_docstring.find(" References\n")
|
||||
if end_of_notes == -1:
|
||||
end_of_notes = cls_docstring.find(" Examples\n")
|
||||
if end_of_notes == -1:
|
||||
end_of_notes = len(cls_docstring)
|
||||
func.__doc__ = (
|
||||
cls_docstring[: start_of_notes + len(notes_header)]
|
||||
+ notes
|
||||
+ cls_docstring[end_of_notes:]
|
||||
)
|
||||
return func
|
||||
|
||||
return _doc
|
||||
|
||||
|
||||
def indentcount_lines(lines: Iterable[str]) -> int:
|
||||
"""Minimum indent for all lines in line list
|
||||
|
||||
Parameters
|
||||
----------
|
||||
lines : Iterable[str]
|
||||
The lines to find the minimum indent of.
|
||||
|
||||
Returns
|
||||
-------
|
||||
indent : int
|
||||
The minimum indent.
|
||||
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> lines = [' one', ' two', ' three']
|
||||
>>> indentcount_lines(lines)
|
||||
1
|
||||
>>> lines = []
|
||||
>>> indentcount_lines(lines)
|
||||
0
|
||||
>>> lines = [' one']
|
||||
>>> indentcount_lines(lines)
|
||||
1
|
||||
>>> indentcount_lines([' '])
|
||||
0
|
||||
"""
|
||||
indentno = sys.maxsize
|
||||
for line in lines:
|
||||
stripped = line.lstrip()
|
||||
if stripped:
|
||||
indentno = min(indentno, len(line) - len(stripped))
|
||||
if indentno == sys.maxsize:
|
||||
return 0
|
||||
return indentno
|
||||
|
||||
|
||||
def filldoc(docdict: Mapping[str, str], unindent_params: bool = True) -> Decorator:
|
||||
"""Return docstring decorator using docdict variable dictionary.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
docdict : dict[str, str]
|
||||
A dictionary containing name, docstring fragment pairs.
|
||||
unindent_params : bool, optional
|
||||
If True, strip common indentation from all parameters in docdict.
|
||||
Default is False.
|
||||
|
||||
Returns
|
||||
-------
|
||||
decfunc : function
|
||||
The decorator function that applies dictionary to its
|
||||
argument's __doc__ attribute.
|
||||
"""
|
||||
if unindent_params:
|
||||
docdict = unindent_dict(docdict)
|
||||
|
||||
def decorate(func: _F) -> _F:
|
||||
# __doc__ may be None for optimized Python (-OO)
|
||||
doc = func.__doc__ or ""
|
||||
func.__doc__ = docformat(doc, docdict)
|
||||
return func
|
||||
|
||||
return decorate
|
||||
|
||||
|
||||
def unindent_dict(docdict: Mapping[str, str]) -> dict[str, str]:
|
||||
"""Unindent all strings in a docdict.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
docdict : dict[str, str]
|
||||
A dictionary with string values to unindent.
|
||||
|
||||
Returns
|
||||
-------
|
||||
docdict : dict[str, str]
|
||||
The `docdict` dictionary but each of its string values are unindented.
|
||||
"""
|
||||
can_dict: dict[str, str] = {}
|
||||
for name, dstr in docdict.items():
|
||||
can_dict[name] = unindent_string(dstr)
|
||||
return can_dict
|
||||
|
||||
|
||||
def unindent_string(docstring: str) -> str:
|
||||
"""Set docstring to minimum indent for all lines, including first.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
docstring : str
|
||||
The input docstring to unindent.
|
||||
|
||||
Returns
|
||||
-------
|
||||
docstring : str
|
||||
The unindented docstring.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> unindent_string(' two')
|
||||
'two'
|
||||
>>> unindent_string(' two\\n three')
|
||||
'two\\n three'
|
||||
"""
|
||||
lines = docstring.expandtabs().splitlines()
|
||||
icount = indentcount_lines(lines)
|
||||
if icount == 0:
|
||||
return docstring
|
||||
return "\n".join([line[icount:] for line in lines])
|
||||
|
||||
|
||||
def doc_replace(obj: object, oldval: str, newval: str) -> Decorator:
|
||||
"""Decorator to take the docstring from obj, with oldval replaced by newval
|
||||
|
||||
Equivalent to ``func.__doc__ = obj.__doc__.replace(oldval, newval)``
|
||||
|
||||
Parameters
|
||||
----------
|
||||
obj : object
|
||||
A class or object whose docstring will be used as the basis for the
|
||||
replacement operation.
|
||||
oldval : str
|
||||
The string to search for in the docstring.
|
||||
newval : str
|
||||
The string to replace `oldval` with in the docstring.
|
||||
|
||||
Returns
|
||||
-------
|
||||
decfunc : function
|
||||
A decorator function that replaces occurrences of `oldval` with `newval`
|
||||
in the docstring of the decorated function.
|
||||
"""
|
||||
# __doc__ may be None for optimized Python (-OO)
|
||||
doc = (obj.__doc__ or "").replace(oldval, newval)
|
||||
|
||||
def inner(func: _F) -> _F:
|
||||
func.__doc__ = doc
|
||||
return func
|
||||
|
||||
return inner
|
||||
BIN
Binary file not shown.
@@ -0,0 +1,212 @@
|
||||
# Bounds may appear unused in this file but we need to import it to make it available to the user
|
||||
from scipy.optimize import NonlinearConstraint, LinearConstraint, Bounds
|
||||
from .common._nonlinear_constraints import process_nl_constraints
|
||||
from .common._linear_constraints import (
|
||||
combine_multiple_linear_constraints,
|
||||
separate_LC_into_eq_and_ineq,
|
||||
)
|
||||
from .common._bounds import process_bounds
|
||||
from enum import Enum
|
||||
from .common._project import _project
|
||||
from .common.linalg import get_arrays_tol
|
||||
from .cobyla.cobyla import cobyla
|
||||
import numpy as np
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
class ConstraintType(Enum):
|
||||
LINEAR_OBJECT = 5
|
||||
NONLINEAR_OBJECT = 10
|
||||
LINEAR_DICT = 15
|
||||
NONLINEAR_DICT = 20
|
||||
|
||||
|
||||
def get_constraint_type(constraint):
|
||||
if isinstance(constraint, dict) and ("A" in constraint) and ("lb" in constraint) and ("ub" in constraint):
|
||||
return ConstraintType.LINEAR_DICT
|
||||
elif isinstance(constraint, dict) and ("fun" in constraint) and ("lb" in constraint) and ("ub" in constraint):
|
||||
return ConstraintType.NONLINEAR_DICT
|
||||
elif hasattr(constraint, "A") and hasattr(constraint, "lb") and hasattr(constraint, "ub"):
|
||||
return ConstraintType.LINEAR_OBJECT
|
||||
elif hasattr(constraint, "fun") and hasattr(constraint, "lb") and hasattr(constraint, "ub"):
|
||||
return ConstraintType.NONLINEAR_OBJECT
|
||||
else:
|
||||
raise ValueError(f"Constraint type {type(constraint)} not recognized")
|
||||
|
||||
|
||||
def process_constraints(constraints):
|
||||
# First throw it back if it's an empty tuple
|
||||
if not constraints:
|
||||
return None, None
|
||||
# Next figure out if it's a list of constraints or a single constraint
|
||||
# If it's a single constraint, make it a list, and then the remaining logic
|
||||
# doesn't have to change
|
||||
if not isinstance(constraints, Iterable):
|
||||
constraints = [constraints]
|
||||
|
||||
# Separate out the linear and nonlinear constraints
|
||||
linear_constraints = []
|
||||
nonlinear_constraints = []
|
||||
for constraint in constraints:
|
||||
constraint_type = get_constraint_type(constraint)
|
||||
if constraint_type is ConstraintType.LINEAR_OBJECT:
|
||||
linear_constraints.append(constraint)
|
||||
elif constraint_type is ConstraintType.NONLINEAR_OBJECT:
|
||||
nonlinear_constraints.append(constraint)
|
||||
elif constraint_type == ConstraintType.LINEAR_DICT:
|
||||
linear_constraints.append(LinearConstraint(constraint["A"], constraint["lb"], constraint["ub"]))
|
||||
elif constraint_type == ConstraintType.NONLINEAR_DICT:
|
||||
nonlinear_constraints.append(NonlinearConstraint(constraint["fun"], constraint["lb"], constraint["ub"]))
|
||||
else:
|
||||
raise ValueError("Constraint type not recognized")
|
||||
|
||||
if len(nonlinear_constraints) > 0:
|
||||
nonlinear_constraint_function = process_nl_constraints(nonlinear_constraints)
|
||||
else:
|
||||
nonlinear_constraint_function = None
|
||||
|
||||
# Determine if we have multiple linear constraints, just 1, or none, and process accordingly
|
||||
if len(linear_constraints) > 1:
|
||||
linear_constraint = combine_multiple_linear_constraints(linear_constraints)
|
||||
elif len(linear_constraints) == 1:
|
||||
linear_constraint = linear_constraints[0]
|
||||
else:
|
||||
linear_constraint = None
|
||||
|
||||
return linear_constraint, nonlinear_constraint_function
|
||||
|
||||
|
||||
def minimize(fun, x0, args=(), method=None, bounds=None, constraints=(), callback=None, options=None):
|
||||
|
||||
linear_constraint, nonlinear_constraint_function = process_constraints(constraints)
|
||||
|
||||
options = {'quiet': True} if options is None else options
|
||||
quiet = options.get("quiet", True)
|
||||
|
||||
if method is None:
|
||||
if nonlinear_constraint_function is not None:
|
||||
if not quiet: print("Nonlinear constraints detected, applying COBYLA")
|
||||
method = "cobyla"
|
||||
elif linear_constraint is not None:
|
||||
if not quiet: print("Linear constraints detected without nonlinear constraints, applying LINCOA")
|
||||
method = "lincoa"
|
||||
elif bounds is not None:
|
||||
if not quiet: print("Bounds without linear or nonlinear constraints detected, applying BOBYQA")
|
||||
method = "bobyqa"
|
||||
else:
|
||||
if not quiet: print("No bounds or constraints detected, applying NEWUOA")
|
||||
method = "newuoa"
|
||||
else:
|
||||
# Raise some errors if methods were called with inappropriate options
|
||||
method = method.lower()
|
||||
if method not in ('newuoa', 'uobyqa', 'bobyqa', 'cobyla', 'lincoa'):
|
||||
raise ValueError(f"Method must be one of NEWUOA, UOBYQA, BOBYQA, COBYLA, or LINCOA, not '{method}'")
|
||||
if method != "cobyla" and nonlinear_constraint_function is not None:
|
||||
raise ValueError("Nonlinear constraints were provided for an algorithm that cannot handle them")
|
||||
if method not in ("cobyla", "lincoa") and linear_constraint is not None:
|
||||
raise ValueError("Linear constraints were provided for an algorithm that cannot handle them")
|
||||
if method not in ("cobyla", "bobyqa", "lincoa") and bounds is not None:
|
||||
raise ValueError("Bounds were provided for an algorithm that cannot handle them")
|
||||
|
||||
# Try to get the length of x0. If we can't that likely means it's a scalar, and
|
||||
# in that case we turn it into an array and wrap the original function so that it
|
||||
# can accept an array and return a scalar.
|
||||
try:
|
||||
lenx0 = len(x0)
|
||||
except TypeError:
|
||||
x0 = np.array([x0])
|
||||
original_scalar_fun = fun
|
||||
def scalar_fun(x):
|
||||
return original_scalar_fun(x[0], *args)
|
||||
fun = scalar_fun
|
||||
lenx0 = 1
|
||||
|
||||
lb, ub = process_bounds(bounds, lenx0)
|
||||
|
||||
# Check which variables are fixed and eliminate them from the problem.
|
||||
# Save the indices and values so that we can call the original function with
|
||||
# an array of the appropriate size, and so that we can add the fixed values to the
|
||||
# result when COBYLA returns.
|
||||
tol = get_arrays_tol(lb, ub)
|
||||
_fixed_idx = (
|
||||
(lb <= ub)
|
||||
& (np.abs(lb - ub) < tol)
|
||||
)
|
||||
if any(_fixed_idx):
|
||||
_fixed_values = 0.5 * (
|
||||
lb[_fixed_idx] + ub[_fixed_idx]
|
||||
)
|
||||
_fixed_values = np.clip(
|
||||
_fixed_values,
|
||||
lb[_fixed_idx],
|
||||
ub[_fixed_idx],
|
||||
)
|
||||
x0 = x0[~_fixed_idx]
|
||||
lb = lb[~_fixed_idx]
|
||||
ub = ub[~_fixed_idx]
|
||||
original_fun = fun
|
||||
def fixed_fun(x):
|
||||
newx = np.zeros(lenx0)
|
||||
newx[_fixed_idx] = _fixed_values
|
||||
newx[~_fixed_idx] = x
|
||||
return original_fun(newx, *args)
|
||||
fun = fixed_fun
|
||||
|
||||
|
||||
# Project x0 onto the feasible set
|
||||
if nonlinear_constraint_function is None:
|
||||
result = _project(x0, lb, ub, {"linear": linear_constraint, "nonlinear": None})
|
||||
x0 = result.x
|
||||
|
||||
if linear_constraint is not None:
|
||||
A_eq, b_eq, A_ineq, b_ineq = separate_LC_into_eq_and_ineq(linear_constraint)
|
||||
else:
|
||||
A_eq, b_eq, A_ineq, b_ineq = None, None, None, None
|
||||
|
||||
if nonlinear_constraint_function is not None:
|
||||
# If there is a nonlinear constraint function, we will call COBYLA, which needs the number of nonlinear
|
||||
# constraints (m_nlcon). In order to get this number we need to evaluate the constraint function at x0.
|
||||
# The constraint value at x0 (nlconstr0) is not discarded but passed down to the Fortran backend, as its
|
||||
# evaluation is assumed to be expensive. We also evaluate the objective function at x0 and pass the result
|
||||
# (f0) down to the Fortran backend, which expects nlconstr0 and f0 to be provided in sync.
|
||||
def calcfc(x):
|
||||
f = fun(x, *args)
|
||||
nlconstr = nonlinear_constraint_function(x)
|
||||
return f, nlconstr
|
||||
else:
|
||||
def calcfc(x):
|
||||
f = fun(x, *args)
|
||||
constr = np.zeros(0)
|
||||
return f, constr
|
||||
|
||||
f0, nlconstr0 = calcfc(x0)
|
||||
|
||||
if 'quiet' in options:
|
||||
del options['quiet']
|
||||
|
||||
if 'maxfev' in options:
|
||||
options['maxfun'] = options['maxfev']
|
||||
del options['maxfev']
|
||||
|
||||
result = cobyla(
|
||||
calcfc,
|
||||
len(nlconstr0),
|
||||
x0,
|
||||
A_ineq,
|
||||
b_ineq,
|
||||
A_eq,
|
||||
b_eq,
|
||||
lb,
|
||||
ub,
|
||||
f0=f0,
|
||||
nlconstr0=nlconstr0,
|
||||
callback=callback,
|
||||
**options
|
||||
)
|
||||
|
||||
if any(_fixed_idx):
|
||||
newx = np.zeros(lenx0)
|
||||
newx[_fixed_idx] = _fixed_values
|
||||
newx[~_fixed_idx] = result.x
|
||||
result.x = newx
|
||||
return result
|
||||
@@ -0,0 +1,559 @@
|
||||
'''
|
||||
This module provides Powell's COBYLA algorithm.
|
||||
|
||||
Translated from Zaikun Zhang's modern-Fortran reference implementation in PRIMA.
|
||||
|
||||
Dedicated to late Professor M. J. D. Powell FRS (1936--2015).
|
||||
|
||||
Python translation by Nickolai Belakovski.
|
||||
|
||||
N.B.:
|
||||
|
||||
1. The modern-Fortran reference implementation in PRIMA contains bug fixes and improvements over the
|
||||
original Fortran 77 implementation by Powell. Consequently, the PRIMA implementation behaves differently
|
||||
from the original Fortran 77 implementation by Powell. Therefore, it is important to point out that
|
||||
you are using PRIMA rather than the original solvers if you want your results to be reproducible.
|
||||
|
||||
2. Compared to Powell's Fortran 77 implementation, the modern-Fortran implementation and hence any
|
||||
faithful translation like this one generally produce better solutions with fewer function evaluations,
|
||||
making them preferable for applications with expensive function evaluations. However, if function
|
||||
evaluations are not the dominant cost in your application, the Fortran 77 solvers are likely to be
|
||||
faster, as they are more efficient in terms of memory usage and flops thanks to the careful and
|
||||
ingenious (but unmaintained and unmaintainable) implementation by Powell.
|
||||
|
||||
See the PRIMA documentation (www.libprima.net) for more information.
|
||||
'''
|
||||
|
||||
from ..common.evaluate import evaluate, moderatex, moderatef, moderatec
|
||||
from ..common.consts import (EPS, RHOBEG_DEFAULT, RHOEND_DEFAULT, CTOL_DEFAULT,
|
||||
CWEIGHT_DEFAULT, FTARGET_DEFAULT, IPRINT_DEFAULT,
|
||||
MAXFUN_DIM_DEFAULT, DEBUGGING, BOUNDMAX,
|
||||
ETA1_DEFAULT, ETA2_DEFAULT, GAMMA1_DEFAULT,
|
||||
GAMMA2_DEFAULT)
|
||||
from ..common.preproc import preproc
|
||||
from ..common.present import present
|
||||
from ..common.linalg import matprod
|
||||
from .cobylb import cobylb
|
||||
import numpy as np
|
||||
from dataclasses import dataclass
|
||||
from copy import copy
|
||||
|
||||
|
||||
@dataclass
|
||||
class COBYLAResult:
|
||||
x: np.ndarray
|
||||
f: float
|
||||
constr: np.ndarray
|
||||
cstrv: float
|
||||
nf: int
|
||||
xhist: np.ndarray | None
|
||||
fhist: np.ndarray | None
|
||||
chist: np.ndarray | None
|
||||
conhist: np.ndarray | None
|
||||
info: int
|
||||
|
||||
|
||||
def cobyla(calcfc, m_nlcon, x, Aineq=None, bineq=None, Aeq=None, beq=None,
|
||||
xl=None, xu=None, f0=None, nlconstr0=None, rhobeg=None, rhoend=None,
|
||||
ftarget=FTARGET_DEFAULT, ctol=CTOL_DEFAULT, cweight=CWEIGHT_DEFAULT,
|
||||
maxfun=None, iprint=IPRINT_DEFAULT, eta1=None, eta2=None,
|
||||
gamma1=GAMMA1_DEFAULT, gamma2=GAMMA2_DEFAULT, maxhist=None, maxfilt=2000,
|
||||
callback=None):
|
||||
"""
|
||||
Among all the arguments, only CALCFC, M_NLCON, and X are obligatory. The others are
|
||||
OPTIONAL and you can neglect them unless you are familiar with the algorithm. Any
|
||||
unspecified optional input will take the default value detailed below. For
|
||||
instance, we may invoke the solver as follows.
|
||||
|
||||
# First define CALCFC, M_NLCON, and X, and then do the following.
|
||||
result = cobyla(calcfc, m_nlcon, x)
|
||||
|
||||
or
|
||||
|
||||
# First define CALCFC, M_NLCON, X, Aineq, and Bineq, and then do the following.
|
||||
result = cobyla(calcfc, m_nlcon, x, Aineq=Aineq, bineq=bineq, rhobeg=1.0e0,
|
||||
rhoend=1.0e-6)
|
||||
|
||||
####################################################################################
|
||||
# IMPORTANT NOTICE: The user must set M_NLCON correctly to the number of nonlinear
|
||||
# constraints, namely the size of NLCONSTR introduced below. Set it to 0 if there
|
||||
# is no nonlinear constraint.
|
||||
####################################################################################
|
||||
|
||||
See examples/cobyla/cobyla_example.py for a concrete example.
|
||||
|
||||
A detailed introduction to the arguments is as follows.
|
||||
|
||||
####################################################################################
|
||||
# INPUTS
|
||||
####################################################################################
|
||||
|
||||
CALCFC
|
||||
Input, function.
|
||||
f, nlconstr = CALCFC(X) should evaluate the objective function and nonlinear
|
||||
constraints at the given vector X; it should return a tuple consisting of the
|
||||
objective function value and the nonlinear constraint value. It must be provided
|
||||
by the user, and its definition must conform to the following interface:
|
||||
#-------------------------------------------------------------------------#
|
||||
def calcfc(x):
|
||||
f = 0.0
|
||||
nlconstr = np.zeros(m_nlcon)
|
||||
return f, nlconstr
|
||||
#-------------------------------------------------------------------------#
|
||||
|
||||
M_NLCON
|
||||
Input, scalar.
|
||||
M_NLCON must be set to the number of nonlinear constraints, namely the size of
|
||||
NLCONSTR(X).
|
||||
N.B.:
|
||||
1. Why don't we define M_NLCON as optional and default it to 0 when it is absent?
|
||||
This is because we need to allocate memory for CONSTR_LOC using M_NLCON. To
|
||||
ensure that the size of CONSTR_LOC is correct, we require the user to specify
|
||||
M_NLCON explicitly.
|
||||
|
||||
X
|
||||
Input, vector.
|
||||
As an input, X should be an N-dimensional vector that contains the starting
|
||||
point, N being the dimension of the problem.
|
||||
|
||||
Aineq, Bineq
|
||||
Input, matrix of size [Mineq, N] and vector of size Mineq unless they are both
|
||||
empty, default: None and None.
|
||||
Aineq and Bineq represent the linear inequality constraints: Aineq*X <= Bineq.
|
||||
|
||||
Aeq, Beq
|
||||
Input, matrix of size [Meq, N] and vector of size Meq unless they are both
|
||||
empty, default: None and None.
|
||||
Aeq and Beq represent the linear equality constraints: Aeq*X = Beq.
|
||||
|
||||
XL, XU
|
||||
Input, vectors of size N unless they are both None, default: None and None.
|
||||
XL is the lower bound for X. If XL is None, X has no
|
||||
lower bound. Any entry of XL that is NaN or below -BOUNDMAX will be taken as
|
||||
-BOUNDMAX, which effectively means there is no lower bound for the corresponding
|
||||
entry of X. The value of BOUNDMAX is 0.25*HUGE(X), which is about 8.6E37 for
|
||||
single precision and 4.5E307 for double precision. XU is similar.
|
||||
|
||||
F0
|
||||
Input, scalar.
|
||||
F0, if present, should be set to the objective function value of the starting X.
|
||||
|
||||
NLCONSTR0
|
||||
Input, vector.
|
||||
NLCONSTR0, if present, should be set to the nonlinear constraint value at the
|
||||
starting X; in addition, SIZE(NLCONSTR0) must be M_NLCON, or the solver will
|
||||
abort.
|
||||
|
||||
RHOBEG, RHOEND
|
||||
Inputs, scalars, default: RHOBEG = 1, RHOEND = 10^-6. RHOBEG and RHOEND must be
|
||||
set to the initial and final values of a trust-region radius, both being positive
|
||||
and RHOEND <= RHOBEG. Typically RHOBEG should be about one tenth of the greatest
|
||||
expected change to a variable, and RHOEND should indicate the accuracy that is
|
||||
required in the final values of the variables.
|
||||
|
||||
FTARGET
|
||||
Input, scalar, default: -Inf.
|
||||
FTARGET is the target function value. The algorithm will terminate when a
|
||||
feasible point with a function value <= FTARGET is found.
|
||||
|
||||
CTOL
|
||||
Input, scalar, default: sqrt(machine epsilon).
|
||||
CTOL is the tolerance of constraint violation. X is considered feasible if
|
||||
CSTRV(X) <= CTOL.
|
||||
N.B.:
|
||||
1. CTOL is absolute, not relative.
|
||||
2. CTOL is used only when selecting the returned X. It does not affect the
|
||||
iterations of the algorithm.
|
||||
|
||||
CWEIGHT
|
||||
Input, scalar, default: CWEIGHT_DFT defined in common/consts.py.
|
||||
CWEIGHT is the weight that the constraint violation takes in the selection of the
|
||||
returned X.
|
||||
|
||||
MAXFUN
|
||||
Input, integer scalar, default: MAXFUN_DIM_DFT*N with MAXFUN_DIM_DFT defined in
|
||||
common/consts.py. MAXFUN is the maximal number of calls of CALCFC.
|
||||
|
||||
IPRINT
|
||||
Input, integer scalar, default: 0.
|
||||
The value of IPRINT should be set to 0, 1, -1, 2, -2, 3, or -3, which controls
|
||||
how much information will be printed during the computation:
|
||||
0: there will be no printing;
|
||||
1: a message will be printed to the screen at the return, showing the best vector
|
||||
of variables found and its objective function value;
|
||||
2: in addition to 1, each new value of RHO is printed to the screen, with the
|
||||
best vector of variables so far and its objective function value; each new
|
||||
value of CPEN is also printed;
|
||||
3: in addition to 2, each function evaluation with its variables will be printed
|
||||
to the screen; -1, -2, -3: the same information as 1, 2, 3 will be printed,
|
||||
not to the screen but to a file named COBYLA_output.txt; the file will be
|
||||
created if it does not exist; the new output will be appended to the end of
|
||||
this file if it already exists.
|
||||
Note that IPRINT = +/-3 can be costly in terms of time and/or space.
|
||||
|
||||
ETA1, ETA2, GAMMA1, GAMMA2
|
||||
Input, scalars, default: ETA1 = 0.1, ETA2 = 0.7, GAMMA1 = 0.5, and GAMMA2 = 2.
|
||||
ETA1, ETA2, GAMMA1, and GAMMA2 are parameters in the updating scheme of the
|
||||
trust-region radius detailed in the subroutine TRRAD in trustregion.py. Roughly
|
||||
speaking, the trust-region radius is contracted by a factor of GAMMA1 when the
|
||||
reduction ratio is below ETA1, and enlarged by a factor of GAMMA2 when the
|
||||
reduction ratio is above ETA2. It is required that 0 < ETA1 <= ETA2 < 1 and
|
||||
0 < GAMMA1 < 1 < GAMMA2. Normally, ETA1 <= 0.25. It is NOT advised to set
|
||||
ETA1 >= 0.5.
|
||||
|
||||
MAXFILT
|
||||
Input, scalar.
|
||||
MAXFILT is a nonnegative integer indicating the maximal length of the filter used
|
||||
for selecting the returned solution; default: MAXFILT_DFT (a value lower than
|
||||
MIN_MAXFILT is not recommended);
|
||||
see common/consts.py for the definitions of MAXFILT_DFT and MIN_MAXFILT.
|
||||
|
||||
CALLBACK
|
||||
Input, function to report progress and optionally request termination.
|
||||
|
||||
|
||||
####################################################################################
|
||||
# OUTPUTS
|
||||
####################################################################################
|
||||
|
||||
The output is a single data structure, COBYLAResult, with the following fields:
|
||||
|
||||
X
|
||||
Output, vector.
|
||||
As an output, X will be set to an approximate minimizer.
|
||||
|
||||
F
|
||||
Output, scalar.
|
||||
F will be set to the objective function value of X at exit.
|
||||
|
||||
CONSTR
|
||||
Output, vector.
|
||||
CONSTR will be set to the constraint value of X at exit.
|
||||
|
||||
CSTRV
|
||||
Output, scalar.
|
||||
CSTRV will be set to the constraint violation of X at exit, i.e.,
|
||||
max([0, XL - X, X - XU, Aineq*X - Bineq, ABS(Aeq*X -Beq), NLCONSTR(X)]).
|
||||
|
||||
NF
|
||||
Output, scalar.
|
||||
NF will be set to the number of calls of CALCFC at exit.
|
||||
|
||||
XHIST, FHIST, CHIST, CONHIST, MAXHIST
|
||||
XHIST: Output, rank 2 array;
|
||||
FHIST: Output, rank 1 array;
|
||||
CHIST: Output, rank 1 array;
|
||||
CONHIST: Output, rank 2 array;
|
||||
MAXHIST: Input, scalar, default: MAXFUN
|
||||
XHIST, if present, will output the history of iterates; FHIST, if present, will
|
||||
output the history function values; CHIST, if present, will output the history of
|
||||
constraint violations; CONHIST, if present, will output the history of constraint
|
||||
values; MAXHIST should be a nonnegative integer, and XHIST/FHIST/CHIST/CONHIST
|
||||
will output only the history of the last MAXHIST iterations.
|
||||
Therefore, MAXHIST= 0 means XHIST/FHIST/CONHIST/CHIST will output
|
||||
nothing, while setting MAXHIST = MAXFUN requests XHIST/FHIST/CHIST/CONHIST to
|
||||
output all the history. If XHIST is present, its size at exit will be
|
||||
(N, min(NF, MAXHIST)); if FHIST is present, its size at exit will be
|
||||
min(NF, MAXHIST); if CHIST is present, its size at exit will be min(NF, MAXHIST);
|
||||
if CONHIST is present, its size at exit will be (M, min(NF, MAXHIST)).
|
||||
|
||||
IMPORTANT NOTICE:
|
||||
Setting MAXHIST to a large value can be costly in terms of memory for large
|
||||
problems.
|
||||
MAXHIST will be reset to a smaller value if the memory needed exceeds MAXHISTMEM
|
||||
defined in common/consts.py
|
||||
Use *HIST with caution!!! (N.B.: the algorithm is NOT designed for large
|
||||
problems).
|
||||
|
||||
INFO
|
||||
Output, scalar.
|
||||
INFO is the exit flag. It will be set to one of the following values defined in
|
||||
common/infos.py:
|
||||
SMALL_TR_RADIUS: the lower bound for the trust region radius is reached;
|
||||
FTARGET_ACHIEVED: the target function value is reached;
|
||||
MAXFUN_REACHED: the objective function has been evaluated MAXFUN times;
|
||||
MAXTR_REACHED: the trust region iteration has been performed MAXTR times (MAXTR = 2*MAXFUN);
|
||||
NAN_INF_X: NaN or Inf occurs in X;
|
||||
DAMAGING_ROUNDING: rounding errors are becoming damaging.
|
||||
#--------------------------------------------------------------------------#
|
||||
The following case(s) should NEVER occur unless there is a bug.
|
||||
NAN_INF_F: the objective function returns NaN or +Inf;
|
||||
NAN_INF_MODEL: NaN or Inf occurs in the model;
|
||||
TRSUBP_FAILED: a trust region step failed to reduce the model
|
||||
#--------------------------------------------------------------------------#
|
||||
"""
|
||||
|
||||
# Local variables
|
||||
solver = "COBYLA"
|
||||
srname = "COBYLA"
|
||||
|
||||
# Sizes
|
||||
mineq = len(bineq) if present(bineq) else 0
|
||||
meq = len(beq) if present(beq) else 0
|
||||
mxl = sum(xl > -BOUNDMAX) if present(xl) else 0
|
||||
mxu = sum(xu < BOUNDMAX) if present(xu) else 0
|
||||
mmm = mxu + mxl + 2*meq + mineq + m_nlcon
|
||||
num_vars = len(x)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert m_nlcon >= 0, f'{srname} M_NLCON >= 0'
|
||||
assert num_vars >= 1, f'{srname} N >= 1'
|
||||
|
||||
assert present(Aineq) == present(bineq), \
|
||||
f'{srname} Aineq and Bineq are both present or both absent'
|
||||
if (present(Aineq)):
|
||||
assert Aineq.shape == (mineq, num_vars), f'{srname} SIZE(Aineq) == [Mineq, N]'
|
||||
|
||||
assert present(Aeq) == present(beq), \
|
||||
f'{srname} Aeq and Beq are both present or both absent'
|
||||
if (present(Aeq)):
|
||||
assert Aeq.shape == (meq, num_vars), f'{srname} SIZE(Aeq) == [Meq, N]'
|
||||
|
||||
if (present(xl)):
|
||||
assert len(xl) == num_vars, f'{srname} SIZE(XL) == N'
|
||||
if (present(xu)):
|
||||
assert len(xu) == num_vars, f'{srname} SIZE(XU) == N'
|
||||
|
||||
|
||||
# N.B.: If NLCONSTR0 is present, then F0 must be present, and we assume that
|
||||
# F(X0) = F0 even if F0 is NaN; if NLCONSTR0 is absent, then F0 must be either
|
||||
# absent or NaN, both of which will be interpreted as F(X0) is not provided.
|
||||
if present(nlconstr0):
|
||||
assert present(f0), f'{srname} If NLCONSTR0 is present, then F0 is present'
|
||||
if present(f0):
|
||||
assert np.isnan(f0) or present(nlconstr0), \
|
||||
f'{srname} If F0 is present and not NaN, then NLCONSTR0 is present'
|
||||
|
||||
|
||||
|
||||
# Exit if the size of NLCONSTR0 is inconsistent with M_NLCON.
|
||||
if present(nlconstr0):
|
||||
assert np.size(nlconstr0) == m_nlcon
|
||||
|
||||
# Read the inputs.
|
||||
|
||||
if xl is not None:
|
||||
xl = copy(xl)
|
||||
xl[np.isnan(xl)] = -BOUNDMAX
|
||||
xl[xl < -BOUNDMAX] = -BOUNDMAX
|
||||
|
||||
if xu is not None:
|
||||
xu = copy(xu)
|
||||
xu[np.isnan(xu)] = BOUNDMAX
|
||||
xu[xu > BOUNDMAX] = BOUNDMAX
|
||||
|
||||
# Wrap the linear and bound constraints into a single constraint: AMAT@X <= BVEC.
|
||||
amat, bvec = get_lincon(Aeq, Aineq, beq, bineq, xl, xu)
|
||||
|
||||
# Create constraint vector
|
||||
constr = np.zeros(mmm)
|
||||
|
||||
# Set [F_LOC, CONSTR_LOC] to [F(X0), CONSTR(X0)] after evaluating the latter if
|
||||
# needed. In this way, COBYLB only needs one interface.
|
||||
# N.B.: Due to the preconditions above, there are two possibilities for F0 and
|
||||
# NLCONSTR0.
|
||||
# If NLCONSTR0 is present, then F0 must be present, and we assume that F(X0) = F0
|
||||
# even if F0 is NaN.
|
||||
# If NLCONSTR0 is absent, then F0 must be either absent or NaN, both of which will
|
||||
# be interpreted as F(X0) is not provided and we have to evaluate F(X0) and
|
||||
# NLCONSTR(X0) now.
|
||||
if (present(f0) and present(nlconstr0) and all(np.isfinite(x))):
|
||||
f = moderatef(f0)
|
||||
if amat is not None:
|
||||
constr[:mmm - m_nlcon] = moderatec(matprod(amat, x) - bvec)
|
||||
constr[mmm - m_nlcon:] = moderatec(nlconstr0)
|
||||
else:
|
||||
x = moderatex(x)
|
||||
f, constr = evaluate(calcfc, x, m_nlcon, amat, bvec)
|
||||
constr[:mmm - m_nlcon] = moderatec(constr[:mmm - m_nlcon])
|
||||
# N.B.: Do NOT call FMSG, SAVEHIST, or SAVEFILT for the function/constraint evaluation at X0.
|
||||
# They will be called during the initialization, which will read the function/constraint at X0.
|
||||
cstrv = max(np.append(0, constr))
|
||||
|
||||
|
||||
# If RHOBEG is present, use it; otherwise, RHOBEG takes the default value for
|
||||
# RHOBEG, taking the value of RHOEND into account. Note that RHOEND is considered
|
||||
# only if it is present and it is VALID (i.e., finite and positive). The other
|
||||
# inputs are read similarly.
|
||||
if present(rhobeg):
|
||||
rhobeg = rhobeg
|
||||
elif present(rhoend) and np.isfinite(rhoend) and rhoend > 0:
|
||||
rhobeg = max(10 * rhoend, RHOBEG_DEFAULT)
|
||||
else:
|
||||
rhobeg = RHOBEG_DEFAULT
|
||||
|
||||
if present(rhoend):
|
||||
rhoend = rhoend
|
||||
elif rhobeg > 0:
|
||||
rhoend = max(EPS, min(RHOEND_DEFAULT/RHOBEG_DEFAULT * rhobeg, RHOEND_DEFAULT))
|
||||
else:
|
||||
rhoend = RHOEND_DEFAULT
|
||||
|
||||
maxfun = maxfun if present(maxfun) else MAXFUN_DIM_DEFAULT * num_vars
|
||||
|
||||
if present(eta1):
|
||||
eta1 = eta1
|
||||
elif present(eta2) and 0 < eta2 < 1:
|
||||
eta1 = max(EPS, eta2 / 7)
|
||||
else:
|
||||
eta1 = ETA1_DEFAULT
|
||||
|
||||
if present(eta2):
|
||||
eta2 = eta2
|
||||
elif 0 < eta1 < 1:
|
||||
eta2 = (eta1 + 2) / 3
|
||||
else:
|
||||
eta2 = ETA2_DEFAULT
|
||||
|
||||
maxhist = (
|
||||
maxhist
|
||||
if present(maxhist)
|
||||
else max(maxfun, num_vars + 2, MAXFUN_DIM_DEFAULT * num_vars)
|
||||
)
|
||||
|
||||
# Preprocess the inputs in case some of them are invalid. It does nothing if all
|
||||
# inputs are valid.
|
||||
(
|
||||
iprint,
|
||||
maxfun,
|
||||
maxhist,
|
||||
ftarget,
|
||||
rhobeg,
|
||||
rhoend,
|
||||
npt, # Unused in COBYLA
|
||||
maxfilt,
|
||||
ctol,
|
||||
cweight,
|
||||
eta1,
|
||||
eta2,
|
||||
gamma1,
|
||||
gamma2,
|
||||
_x0, # Unused in COBYLA
|
||||
) = preproc(
|
||||
solver,
|
||||
num_vars,
|
||||
iprint,
|
||||
maxfun,
|
||||
maxhist,
|
||||
ftarget,
|
||||
rhobeg,
|
||||
rhoend,
|
||||
num_constraints=mmm,
|
||||
maxfilt=maxfilt,
|
||||
ctol=ctol,
|
||||
cweight=cweight,
|
||||
eta1=eta1,
|
||||
eta2=eta2,
|
||||
gamma1=gamma1,
|
||||
gamma2=gamma2,
|
||||
is_constrained=(mmm > 0),
|
||||
)
|
||||
|
||||
# Further revise MAXHIST according to MAXHISTMEM, and allocate memory for the history.
|
||||
# In MATLAB/Python/Julia/R implementation, we should simply set MAXHIST = MAXFUN and initialize
|
||||
# CHIST = NaN(1, MAXFUN), CONHIST = NaN(M, MAXFUN), FHIST = NaN(1, MAXFUN), XHIST = NaN(N, MAXFUN)
|
||||
# if they are requested; replace MAXFUN with 0 for the history that is not requested.
|
||||
# prehist(maxhist, num_vars, present(xhist), xhist_loc, present(fhist), fhist_loc, &
|
||||
# & present(chist), chist_loc, m, present(conhist), conhist_loc)
|
||||
|
||||
# call cobylb, which performs the real calculations
|
||||
x, f, constr, cstrv, nf, xhist, fhist, chist, conhist, info = cobylb(
|
||||
calcfc,
|
||||
iprint,
|
||||
maxfilt,
|
||||
maxfun,
|
||||
amat,
|
||||
bvec,
|
||||
ctol,
|
||||
cweight,
|
||||
eta1,
|
||||
eta2,
|
||||
ftarget,
|
||||
gamma1,
|
||||
gamma2,
|
||||
rhobeg,
|
||||
rhoend,
|
||||
constr,
|
||||
f,
|
||||
x,
|
||||
maxhist,
|
||||
callback
|
||||
)
|
||||
|
||||
return COBYLAResult(x, f, constr, cstrv, nf, xhist, fhist, chist, conhist, info)
|
||||
|
||||
|
||||
def get_lincon(Aeq=None, Aineq=None, beq=None, bineq=None, xl=None, xu=None):
|
||||
"""
|
||||
This subroutine wraps the linear and bound constraints into a single constraint:
|
||||
AMAT*X <= BVEC.
|
||||
|
||||
N.B.:
|
||||
|
||||
LINCOA normalizes the linear constraints so that each constraint has a gradient
|
||||
of norm 1. However, COBYLA does not do this.
|
||||
"""
|
||||
|
||||
# Sizes
|
||||
if Aeq is not None:
|
||||
num_vars = Aeq.shape[1]
|
||||
elif Aineq is not None:
|
||||
num_vars = Aineq.shape[1]
|
||||
elif xl is not None:
|
||||
num_vars = len(xl)
|
||||
elif xu is not None:
|
||||
num_vars = len(xu)
|
||||
else:
|
||||
return None, None
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert Aineq is None or Aineq.shape == (len(bineq), num_vars)
|
||||
assert Aeq is None or Aeq.shape == (len(beq), num_vars)
|
||||
assert (xl is None or xu is None) or len(xl) == len(xu) == num_vars
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
# Define the indices of the nontrivial bound constraints.
|
||||
ixl = np.where(xl > -BOUNDMAX)[0] if xl is not None else None
|
||||
ixu = np.where(xu < BOUNDMAX)[0] if xu is not None else None
|
||||
|
||||
# Wrap the linear constraints.
|
||||
# The bound constraint XL <= X <= XU is handled as two constraints:
|
||||
# -X <= -XL, X <= XU.
|
||||
# The equality constraint Aeq*X = Beq is handled as two constraints:
|
||||
# -Aeq*X <= -Beq, Aeq*X <= Beq.
|
||||
# N.B.:
|
||||
# 1. The treatment of the equality constraints is naive. One may choose to
|
||||
# eliminate them instead.
|
||||
idmat = np.eye(num_vars)
|
||||
amat = np.vstack([
|
||||
-idmat[ixl, :] if ixl is not None else np.empty((0, num_vars)),
|
||||
idmat[ixu, :] if ixu is not None else np.empty((0, num_vars)),
|
||||
-Aeq if Aeq is not None else np.empty((0, num_vars)),
|
||||
Aeq if Aeq is not None else np.empty((0, num_vars)),
|
||||
Aineq if Aineq is not None else np.empty((0, num_vars))
|
||||
])
|
||||
bvec = np.hstack([
|
||||
-xl[ixl] if ixl is not None else np.empty(0),
|
||||
xu[ixu] if ixu is not None else np.empty(0),
|
||||
-beq if beq is not None else np.empty(0),
|
||||
beq if beq is not None else np.empty(0),
|
||||
bineq if bineq is not None else np.empty(0)
|
||||
])
|
||||
|
||||
amat = amat if amat.shape[0] > 0 else None
|
||||
bvec = bvec if bvec.shape[0] > 0 else None
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert (amat is None and bvec is None) or amat.shape == (len(bvec), num_vars)
|
||||
|
||||
return amat, bvec
|
||||
@@ -0,0 +1,714 @@
|
||||
'''
|
||||
This module performs the major calculations of COBYLA.
|
||||
|
||||
Translated from Zaikun Zhang's modern-Fortran reference implementation in PRIMA.
|
||||
|
||||
Dedicated to late Professor M. J. D. Powell FRS (1936--2015).
|
||||
|
||||
Python translation by Nickolai Belakovski.
|
||||
'''
|
||||
|
||||
import numpy as np
|
||||
from ..common.checkbreak import checkbreak_con
|
||||
from ..common.consts import REALMAX, EPS, DEBUGGING, MIN_MAXFILT
|
||||
from ..common.infos import INFO_DEFAULT, MAXTR_REACHED, DAMAGING_ROUNDING, \
|
||||
SMALL_TR_RADIUS, CALLBACK_TERMINATE
|
||||
from ..common.evaluate import evaluate
|
||||
from ..common.history import savehist
|
||||
from ..common.linalg import isinv, matprod, inprod, norm, primasum, primapow2
|
||||
from ..common.message import fmsg, retmsg, rhomsg
|
||||
from ..common.ratio import redrat
|
||||
from ..common.redrho import redrho
|
||||
from ..common.selectx import savefilt, selectx
|
||||
from .update import updatepole, findpole, updatexfc
|
||||
from .geometry import setdrop_tr, geostep
|
||||
from .trustregion import trstlp, trrad
|
||||
from .initialize import initxfc, initfilt
|
||||
|
||||
|
||||
def cobylb(calcfc, iprint, maxfilt, maxfun, amat, bvec, ctol, cweight, eta1, eta2,
|
||||
ftarget, gamma1, gamma2, rhobeg, rhoend, constr, f, x, maxhist, callback):
|
||||
'''
|
||||
This subroutine performs the actual computations of COBYLA.
|
||||
'''
|
||||
|
||||
# Outputs
|
||||
xhist = []
|
||||
fhist = []
|
||||
chist = []
|
||||
conhist = []
|
||||
|
||||
# Local variables
|
||||
solver = 'COBYLA'
|
||||
A = np.zeros((np.size(x), np.size(constr))) # A contains the approximate gradient for the constraints
|
||||
distsq = np.zeros(np.size(x) + 1)
|
||||
# CPENMIN is the minimum of the penalty parameter CPEN for the L-infinity
|
||||
# constraint violation in the merit function. Note that CPENMIN = 0 in Powell's
|
||||
# implementation, which allows CPEN to be 0. Here, we take CPENMIN > 0 so that CPEN
|
||||
# is always positive. This avoids the situation where PREREM becomes 0 when
|
||||
# PREREF = 0 = CPEN. It brings two advantages as follows.
|
||||
# 1. If the trust-region subproblem solver works correctly and the trust-region
|
||||
# center is not optimal for the subproblem, then PREREM > 0 is guaranteed. This
|
||||
# is because, in theory, PREREC >= 0 and MAX(PREREC, PREREF) > 0, and the
|
||||
# definition of CPEN in GETCPEN ensures that PREREM > 0.
|
||||
# 2. There is no need to revise ACTREM and PREREM when CPEN = 0 and F = FVAL(N+1)
|
||||
# as in lines 312--314 of Powell's cobylb.f code. Powell's code revises ACTREM
|
||||
# to CVAL(N + 1) - CSTRV and PREREM to PREREC in this case, which is crucial for
|
||||
# feasibility problems.
|
||||
cpenmin = EPS
|
||||
|
||||
# Sizes
|
||||
m_lcon = np.size(bvec) if bvec is not None else 0
|
||||
num_constraints = np.size(constr)
|
||||
m_nlcon = num_constraints - m_lcon
|
||||
num_vars = np.size(x)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert abs(iprint) <= 3
|
||||
assert num_constraints >= m_lcon and m_lcon >= 0
|
||||
assert num_vars >= 1
|
||||
assert maxfun >= num_vars + 2
|
||||
assert rhobeg >= rhoend and rhoend > 0
|
||||
assert all(np.isfinite(x))
|
||||
assert 0 <= eta1 <= eta2 < 1
|
||||
assert 0 < gamma1 < 1 < gamma2
|
||||
assert 0 <= ctol
|
||||
assert 0 <= cweight
|
||||
assert 0 <= maxhist <= maxfun
|
||||
assert amat is None or np.shape(amat) == (m_lcon, num_vars)
|
||||
assert min(MIN_MAXFILT, maxfun) <= maxfilt <= maxfun
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
# Initialize SIM, FVAL, CONMAT, and CVAL, together with the history.
|
||||
# After the initialization, SIM[:, NUM_VARS] holds the vertex of the initial
|
||||
# simplex with the smallest function value (regardless of the constraint
|
||||
# violation), and SIM[:, :NUM_VARS] holds the displacements from the other vertices
|
||||
# to SIM[:, NUM_VARS]. FVAL, CONMAT, and CVAL hold the function values, constraint
|
||||
# values, and constraint violations on the vertices in the order corresponding to
|
||||
# SIM.
|
||||
evaluated, conmat, cval, sim, simi, fval, nf, subinfo = initxfc(calcfc, iprint,
|
||||
maxfun, constr, amat, bvec, ctol, f, ftarget, rhobeg, x,
|
||||
xhist, fhist, chist, conhist, maxhist)
|
||||
|
||||
# Initialize the filter, including xfilt, ffilt, confilt, cfilt, and nfilt.
|
||||
# N.B.: The filter is used only when selecting which iterate to return. It does not
|
||||
# interfere with the iterations. COBYLA is NOT a filter method but a trust-region
|
||||
# method based on an L-infinity merit function. Powell's implementation does not
|
||||
# use a filter to select the iterate, possibly returning a suboptimal iterate.
|
||||
cfilt = np.zeros(np.minimum(np.maximum(maxfilt, 1), maxfun))
|
||||
confilt = np.zeros((np.size(constr), np.size(cfilt)))
|
||||
ffilt = np.zeros(np.size(cfilt))
|
||||
xfilt = np.zeros((np.size(x), np.size(cfilt)))
|
||||
nfilt = initfilt(conmat, ctol, cweight, cval, fval, sim, evaluated, cfilt, confilt,
|
||||
ffilt, xfilt)
|
||||
|
||||
# Check whether to return due to abnormal cases that may occur during the initialization.
|
||||
if subinfo != INFO_DEFAULT:
|
||||
info = subinfo
|
||||
# Return the best calculated values of the variables
|
||||
# N.B: Selectx and findpole choose X by different standards, one cannot replace the other
|
||||
kopt = selectx(ffilt[:nfilt], cfilt[:nfilt], cweight, ctol)
|
||||
x = xfilt[:, kopt]
|
||||
f = ffilt[kopt]
|
||||
constr = confilt[:, kopt]
|
||||
cstrv = cfilt[kopt]
|
||||
# print a return message according to IPRINT.
|
||||
retmsg(solver, info, iprint, nf, f, x, cstrv, constr)
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert nf <= maxfun
|
||||
assert np.size(x) == num_vars and not any(np.isnan(x))
|
||||
assert not (np.isnan(f) or np.isposinf(f))
|
||||
# assert np.size(xhist, 0) == n and np.size(xhist, 1) == maxxhist
|
||||
# assert not any(np.isnan(xhist(:, 1:min(nf, maxxhist))))
|
||||
# The last calculated X can be Inf (finite + finite can be Inf numerically).
|
||||
# assert np.size(fhist) == maxfhist
|
||||
# assert not any(np.isnan(fhist(1:min(nf, maxfhist))) or np.isposinf(fhist(1:min(nf, maxfhist))))
|
||||
# assert np.size(conhist, 0) == m and np.size(conhist, 1) == maxconhist
|
||||
# assert not any(np.isnan(conhist(:, 1:min(nf, maxconhist))) or np.isneginf(conhist(:, 1:min(nf, maxconhist))))
|
||||
# assert np.size(chist) == maxchist
|
||||
# assert not any(chist(1:min(nf, maxchist)) < 0 or np.isnan(chist(1:min(nf, maxchist))) or np.isposinf(chist(1:min(nf, maxchist))))
|
||||
# nhist = minval([nf, maxfhist, maxchist])
|
||||
# assert not any(isbetter(fhist(1:nhist), chist(1:nhist), f, cstrv, ctol))
|
||||
return x, f, constr, cstrv, nf, xhist, fhist, chist, conhist, info
|
||||
|
||||
|
||||
# Set some more initial values.
|
||||
# We must initialize shortd, ratio, and jdrop_tr because these get defined on
|
||||
# branches that are not guaranteed to be executed, but their values are used later.
|
||||
# Our initialization of CPEN differs from Powell's in two ways. First, we use the
|
||||
# ratio defined in (13) of Powell's COBYLA paper to initialize CPEN. Second, we
|
||||
# impose CPEN >= CPENMIN > 0. Powell's code simply initializes CPEN to 0.
|
||||
rho = rhobeg
|
||||
delta = rhobeg
|
||||
cpen = np.maximum(cpenmin, np.minimum(1.0E3, fcratio(conmat, fval))) # Powell's code: CPEN = ZERO
|
||||
shortd = False
|
||||
ratio = -1
|
||||
jdrop_tr = 0
|
||||
|
||||
# If DELTA <= GAMMA3*RHO after an update, we set DELTA to RHO. GAMMA3 must be less
|
||||
# than GAMMA2. The reason is as follows. Imagine a very successful step with
|
||||
# DNORM = the un-updated DELTA = RHO. The TRRAD will update DELTA to GAMMA2*RHO.
|
||||
# If GAMMA3 >= GAMMA2, then DELTA will be reset to RHO, which is not reasonable as
|
||||
# D is very successful. See paragraph two of Sec 5.2.5 in T. M. Ragonneau's thesis:
|
||||
# "Model-Based Derivative-Free Optimization Methods and Software." According to
|
||||
# test on 20230613, for COBYLA, this Powellful updating scheme of DELTA works
|
||||
# slightly better than setting directly DELTA = max(NEW_DELTA, RHO).
|
||||
gamma3 = np.maximum(1, np.minimum(0.75 * gamma2, 1.5))
|
||||
|
||||
# MAXTR is the maximal number of trust region iterations. Each trust-region
|
||||
# iteration takes 1 or 2 function evaluations unless the trust-region step is short
|
||||
# or the trust-region subproblem solver fails but the geometry step is not invoked.
|
||||
# Thus the following MAXTR is unlikely to be reached.
|
||||
maxtr = 10 * maxfun
|
||||
info = MAXTR_REACHED
|
||||
|
||||
# Begin the iterative procedure
|
||||
# After solving a trust-region subproblem, we use three boolean variables to
|
||||
# control the workflow.
|
||||
# SHORTD - Is the trust-region trial step too short to invoke # a function
|
||||
# evaluation?
|
||||
# IMPROVE_GEO - Will we improve the model after the trust-region iteration? If yes,
|
||||
# a geometry step will be taken, corresponding to the "Branch (Delta)"
|
||||
# in the COBYLA paper.
|
||||
# REDUCE_RHO - Will we reduce rho after the trust-region iteration?
|
||||
# COBYLA never sets IMPROVE_GEO and REDUCE_RHO to True simultaneously.
|
||||
for tr in range(maxtr):
|
||||
# Increase the penalty parameter CPEN, if needed, so that
|
||||
# PREREM = PREREF + CPEN * PREREC > 0.
|
||||
# This is the first (out of two) update of CPEN, where CPEN increases or
|
||||
# remains the same.
|
||||
# N.B.: CPEN and the merit function PHI = FVAL + CPEN*CVAL are used in three
|
||||
# places only.
|
||||
# 1. In FINDPOLE/UPDATEPOLE, deciding the optimal vertex of the current simplex.
|
||||
# 2. After the trust-region trial step, calculating the reduction ratio.
|
||||
# 3. In GEOSTEP, deciding the direction of the geometry step.
|
||||
# They do not appear explicitly in the trust-region subproblem, though the
|
||||
# trust-region center (i.e. the current optimal vertex) is defined by them.
|
||||
cpen = getcpen(amat, bvec, conmat, cpen, cval, delta, fval, rho, sim, simi)
|
||||
|
||||
# Switch the best vertex of the current simplex to SIM[:, NUM_VARS].
|
||||
conmat, cval, fval, sim, simi, subinfo = updatepole(cpen, conmat, cval, fval,
|
||||
sim, simi)
|
||||
# Check whether to exit due to damaging rounding in UPDATEPOLE.
|
||||
if subinfo == DAMAGING_ROUNDING:
|
||||
info = subinfo
|
||||
break # Better action to take? Geometry step, or simply continue?
|
||||
|
||||
# Does the interpolation set have adequate geometry? It affects improve_geo and
|
||||
# reduce_rho.
|
||||
adequate_geo = all(primasum(primapow2(sim[:, :num_vars]), axis=0) <= 4 * primapow2(delta))
|
||||
|
||||
# Calculate the linear approximations to the objective and constraint functions.
|
||||
# N.B.: TRSTLP accesses A mostly by columns, so it is more reasonable to save A
|
||||
# instead of A^T.
|
||||
# Zaikun 2023108: According to a test on 2023108, calculating G and
|
||||
# A(:, M_LCON+1:M) by solving the linear systems SIM^T*G = FVAL(1:N)-FVAL(N+1)
|
||||
# and SIM^T*A = CONMAT(:, 1:N)-CONMAT(:, N+1) does not seem to improve or worsen
|
||||
# the performance of COBYLA in terms of the number of function evaluations. The
|
||||
# system was solved by SOLVE in LINALG_MOD based on a QR factorization of SIM
|
||||
# (not necessarily a good algorithm). No preconditioning or scaling was used.
|
||||
g = matprod((fval[:num_vars] - fval[num_vars]), simi)
|
||||
A[:, :m_lcon] = amat.T if amat is not None else amat
|
||||
A[:, m_lcon:] = matprod((conmat[m_lcon:, :num_vars] -
|
||||
np.tile(conmat[m_lcon:, num_vars], (num_vars, 1)).T), simi).T
|
||||
|
||||
# Calculate the trust-region trial step d. Note that d does NOT depend on cpen.
|
||||
d = trstlp(A, -conmat[:, num_vars], delta, g)
|
||||
dnorm = min(delta, norm(d))
|
||||
|
||||
# Is the trust-region trial step short? N.B.: we compare DNORM with RHO, not
|
||||
# DELTA. Powell's code especially defines SHORTD by SHORTD = (DNORM < 0.5 *
|
||||
# RHO). In our tests 1/10 seems to work better than 1/2 or 1/4, especially for
|
||||
# linearly constrained problems. Note that LINCOA has a slightly more
|
||||
# sophisticated way of defining SHORTD, taking into account whether D causes a
|
||||
# change to the active set. Should we try the same here?
|
||||
shortd = (dnorm <= 0.1 * rho)
|
||||
|
||||
# Predict the change to F (PREREF) and to the constraint violation (PREREC) due
|
||||
# to D. We have the following in precise arithmetic. They may fail to hold due
|
||||
# to rounding errors.
|
||||
# 1. B[:NUM_CONSTRAINTS] = -CONMAT[:, NUM_VARS] and hence
|
||||
# np.max(np.append(B[:NUM_CONSTRAINTS] - D@A[:, :NUM_CONSTRAINTS], 0)) is the
|
||||
# L-infinity violation of the linearized constraints corresponding to D. When
|
||||
# D=0, the violation is np.max(np.append(B[:NUM_CONSTRAINTS], 0)) =
|
||||
# CVAL[NUM_VARS]. PREREC is the reduction of this violation achieved by D,
|
||||
# which is nonnegative in theory; PREREC = 0 iff B[:NUM_CONSTRAINTS] <= 0, i.e.
|
||||
# the trust-region center satisfies the linearized constraints.
|
||||
# 2. PREREF may be negative or 0, but it is positive when PREREC = 0 and shortd
|
||||
# is False
|
||||
# 3. Due to 2, in theory, max(PREREC, PREREF) > 0 if shortd is False.
|
||||
preref = -inprod(d, g) # Can be negative
|
||||
prerec = cval[num_vars] - np.max(np.append(0, conmat[:, num_vars] + matprod(d, A)))
|
||||
|
||||
# Evaluate PREREM, which is the predicted reduction in the merit function.
|
||||
# In theory, PREREM >= 0 and it is 0 iff CPEN = 0 = PREREF. This may not be true
|
||||
# numerically.
|
||||
prerem = preref + cpen * prerec
|
||||
trfail = not (prerem > 1.0E-6 * min(cpen, 1) * rho)
|
||||
|
||||
if shortd or trfail:
|
||||
# Reduce DELTA if D is short or if D fails to render PREREM > 0. The latter
|
||||
# can only happen due to rounding errors. This seems quite important for
|
||||
# performance
|
||||
delta *= 0.1
|
||||
if delta <= gamma3 * rho:
|
||||
delta = rho # set delta to rho when it is close to or below
|
||||
else:
|
||||
# Calculate the next value of the objective and constraint functions.
|
||||
# If X is close to one of the points in the interpolation set, then we do
|
||||
# not evaluate the objective and constraints at X, assuming them to have
|
||||
# the values at the closest point.
|
||||
# N.B.: If this happens, do NOT include X into the filter, as F and CONSTR
|
||||
# are inaccurate.
|
||||
x = sim[:, num_vars] + d
|
||||
distsq[num_vars] = primasum(primapow2(x - sim[:, num_vars]))
|
||||
distsq[:num_vars] = primasum(primapow2(x.reshape(num_vars, 1) -
|
||||
(sim[:, num_vars].reshape(num_vars, 1) + sim[:, :num_vars])), axis=0)
|
||||
j = np.argmin(distsq)
|
||||
if distsq[j] <= primapow2(1e-4 * rhoend):
|
||||
f = fval[j]
|
||||
constr = conmat[:, j]
|
||||
cstrv = cval[j]
|
||||
else:
|
||||
# Evaluate the objective and constraints at X, taking care of possible
|
||||
# inf/nan values.
|
||||
f, constr = evaluate(calcfc, x, m_nlcon, amat, bvec)
|
||||
cstrv = np.max(np.append(0, constr))
|
||||
nf += 1
|
||||
# Save X, F, CONSTR, CSTRV into the history.
|
||||
savehist(maxhist, x, xhist, f, fhist, cstrv, chist, constr, conhist)
|
||||
# Save X, F, CONSTR, CSTRV into the filter.
|
||||
nfilt, cfilt, ffilt, xfilt, confilt = savefilt(cstrv, ctol, cweight, f,
|
||||
x, nfilt, cfilt, ffilt,
|
||||
xfilt, constr, confilt)
|
||||
|
||||
# Print a message about the function/constraint evaluation according to
|
||||
# iprint
|
||||
fmsg(solver, 'Trust region', iprint, nf, delta, f, x, cstrv, constr)
|
||||
|
||||
# Evaluate ACTREM, which is the actual reduction in the merit function
|
||||
actrem = (fval[num_vars] + cpen * cval[num_vars]) - (f + cpen * cstrv)
|
||||
|
||||
# Calculate the reduction ratio by redrat, which hands inf/nan carefully
|
||||
ratio = redrat(actrem, prerem, eta1)
|
||||
|
||||
# Update DELTA. After this, DELTA < DNORM may hold.
|
||||
# N.B.:
|
||||
# 1. Powell's code uses RHO as the trust-region radius and updates it as
|
||||
# follows.
|
||||
# Reduce RHO to GAMMA1*RHO if ADEQUATE_GEO is TRUE and either SHORTD is
|
||||
# TRUE or RATIO < ETA1, and then revise RHO to RHOEND if its new value is
|
||||
# not more than GAMMA3*RHOEND; RHO remains unchanged in all other cases;
|
||||
# in particular, RHO is never increased.
|
||||
# 2. Our implementation uses DELTA as the trust-region radius, while using
|
||||
# RHO as a lower bound for DELTA. DELTA is updated in a way that is
|
||||
# typical for trust-region methods, and it is revised to RHO if its new
|
||||
# value is not more than GAMMA3*RHO. RHO reflects the current resolution
|
||||
# of the algorithm; its update is essentially the same as the update of
|
||||
# RHO in Powell's code (see the definition of REDUCE_RHO below). Our
|
||||
# implementation aligns with UOBYQA/NEWUOA/BOBYQA/LINCOA and improves the
|
||||
# performance of COBYLA.
|
||||
# 3. The same as Powell's code, we do not reduce RHO unless ADEQUATE_GEO is
|
||||
# TRUE. This is also how Powell updated RHO in
|
||||
# UOBYQA/NEWUOA/BOBYQA/LINCOA. What about we also use ADEQUATE_GEO ==
|
||||
# TRUE as a prerequisite for reducing DELTA? The argument would be that
|
||||
# the bad (small) value of RATIO may be because of a bad geometry (and
|
||||
# hence a bad model) rather than an improperly large DELTA, and it might
|
||||
# be good to try improving the geometry first without reducing DELTA.
|
||||
# However, according to a test on 20230206, it does not improve the
|
||||
# performance if we skip the update of DELTA when ADEQUATE_GEO is FALSE
|
||||
# and RATIO < 0.1. Therefore, we choose to update DELTA without checking
|
||||
# ADEQUATE_GEO.
|
||||
|
||||
delta = trrad(delta, dnorm, eta1, eta2, gamma1, gamma2, ratio)
|
||||
if delta <= gamma3*rho:
|
||||
delta = rho # Set delta to rho when it is close to or below.
|
||||
|
||||
# Is the newly generated X better than the current best point?
|
||||
ximproved = actrem > 0 # If ACTREM is NaN, then XIMPROVED should and will be False
|
||||
|
||||
# Set JDROP_TR to the index of the vertex to be replaced with X. JDROP_TR = 0 means there
|
||||
# is no good point to replace, and X will not be included into the simplex; in this case,
|
||||
# the geometry of the simplex likely needs improvement, which will be handled below.
|
||||
jdrop_tr = setdrop_tr(ximproved, d, delta, rho, sim, simi)
|
||||
|
||||
# Update SIM, SIMI, FVAL, CONMAT, and CVAL so that SIM[:, JDROP_TR] is replaced with D.
|
||||
# UPDATEXFC does nothing if JDROP_TR is None, as the algorithm decides to discard X.
|
||||
sim, simi, fval, conmat, cval, subinfo = updatexfc(jdrop_tr, constr, cpen, cstrv, d, f, conmat, cval, fval, sim, simi)
|
||||
# Check whether to break due to damaging rounding in UPDATEXFC
|
||||
if subinfo == DAMAGING_ROUNDING:
|
||||
info = subinfo
|
||||
break # Better action to take? Geometry step, or a RESCUE as in BOBYQA?
|
||||
|
||||
# Check whether to break due to maxfun, ftarget, etc.
|
||||
subinfo = checkbreak_con(maxfun, nf, cstrv, ctol, f, ftarget, x)
|
||||
if subinfo != INFO_DEFAULT:
|
||||
info = subinfo
|
||||
break
|
||||
# End of if SHORTD or TRFAIL. The normal trust-region calculation ends.
|
||||
|
||||
# Before the next trust-region iteration, we possibly improve the geometry of the simplex or
|
||||
# reduce RHO according to IMPROVE_GEO and REDUCE_RHO. Now we decide these indicators.
|
||||
# N.B.: We must ensure that the algorithm does not set IMPROVE_GEO = True at infinitely many
|
||||
# consecutive iterations without moving SIM[:, NUM_VARS] or reducing RHO. Otherwise, the algorithm
|
||||
# will get stuck in repetitive invocations of GEOSTEP. This is ensured by the following facts:
|
||||
# 1. If an iteration sets IMPROVE_GEO to True, it must also reduce DELTA or set DELTA to RHO.
|
||||
# 2. If SIM[:, NUM_VARS] and RHO remain unchanged, then ADEQUATE_GEO will become True after at
|
||||
# most NUM_VARS invocations of GEOSTEP.
|
||||
|
||||
# BAD_TRSTEP: Is the last trust-region step bad?
|
||||
bad_trstep = shortd or trfail or ratio <= 0 or jdrop_tr is None
|
||||
# IMPROVE_GEO: Should we take a geometry step to improve the geometry of the interpolation set?
|
||||
improve_geo = bad_trstep and not adequate_geo
|
||||
# REDUCE_RHO: Should we enhance the resolution by reducing rho?
|
||||
reduce_rho = bad_trstep and adequate_geo and max(delta, dnorm) <= rho
|
||||
|
||||
# COBYLA never sets IMPROVE_GEO and REDUCE_RHO to True simultaneously.
|
||||
# assert not (IMPROVE_GEO and REDUCE_RHO), 'IMPROVE_GEO or REDUCE_RHO are not both TRUE, COBYLA'
|
||||
|
||||
# If SHORTD or TRFAIL is True, then either IMPROVE_GEO or REDUCE_RHO is True unless ADEQUATE_GEO
|
||||
# is True and max(DELTA, DNORM) > RHO.
|
||||
# assert not (shortd or trfail) or (improve_geo or reduce_rho or (adequate_geo and max(delta, dnorm) > rho)), \
|
||||
# 'If SHORTD or TRFAIL is TRUE, then either IMPROVE_GEO or REDUCE_RHO is TRUE unless ADEQUATE_GEO is TRUE and MAX(DELTA, DNORM) > RHO'
|
||||
|
||||
# Comments on BAD_TRSTEP:
|
||||
# 1. Powell's definition of BAD_TRSTEP is as follows. The one used above seems to work better,
|
||||
# especially for linearly constrained problems due to the factor TENTH (= ETA1).
|
||||
# !bad_trstep = (shortd .or. actrem <= 0 .or. actrem < TENTH * prerem .or. jdrop_tr == 0)
|
||||
# Besides, Powell did not check PREREM > 0 in BAD_TRSTEP, which is reasonable to do but has
|
||||
# little impact upon the performance.
|
||||
# 2. NEWUOA/BOBYQA/LINCOA would define BAD_TRSTEP, IMPROVE_GEO, and REDUCE_RHO as follows. Two
|
||||
# different thresholds are used in BAD_TRSTEP. It outperforms Powell's version.
|
||||
# !bad_trstep = (shortd .or. trfail .or. ratio <= eta1 .or. jdrop_tr == 0)
|
||||
# !improve_geo = bad_trstep .and. .not. adequate_geo
|
||||
# !bad_trstep = (shortd .or. trfail .or. ratio <= 0 .or. jdrop_tr == 0)
|
||||
# !reduce_rho = bad_trstep .and. adequate_geo .and. max(delta, dnorm) <= rho
|
||||
# 3. Theoretically, JDROP_TR > 0 when ACTREM > 0 (guaranteed by RATIO > 0). However, in Powell's
|
||||
# implementation, JDROP_TR may be 0 even RATIO > 0 due to NaN. The modernized code has rectified
|
||||
# this in the function SETDROP_TR. After this rectification, we can indeed simplify the
|
||||
# definition of BAD_TRSTEP by removing the condition JDROP_TR == 0. We retain it for robustness.
|
||||
|
||||
# Comments on REDUCE_RHO:
|
||||
# When SHORTD is TRUE, UOBYQA/NEWUOA/BOBYQA/LINCOA all set REDUCE_RHO to TRUE if the recent
|
||||
# models are sufficiently accurate according to certain criteria. See the paragraph around (37)
|
||||
# in the UOBYQA paper and the discussions about Box 14 in the NEWUOA paper. This strategy is
|
||||
# crucial for the performance of the solvers. However, as of 20221111, we have not managed to
|
||||
# make it work in COBYLA. As in NEWUOA, we recorded the errors of the recent models, and set
|
||||
# REDUCE_RHO to true if they are small (e.g., ALL(ABS(MODERR_REC) <= 0.1 * MAXVAL(ABS(A))*RHO) or
|
||||
# ALL(ABS(MODERR_REC) <= RHO**2)) when SHORTD is TRUE. It made little impact on the performance.
|
||||
|
||||
|
||||
# Since COBYLA never sets IMPROVE_GEO and REDUCE_RHO to TRUE simultaneously, the following
|
||||
# two blocks are exchangeable: IF (IMPROVE_GEO) ... END IF and IF (REDUCE_RHO) ... END IF.
|
||||
|
||||
# Improve the geometry of the simplex by removing a point and adding a new one.
|
||||
# If the current interpolation set has acceptable geometry, then we skip the geometry step.
|
||||
# The code has a small difference from Powell's original code here: If the current geometry
|
||||
# is acceptable, then we will continue with a new trust-region iteration; however, at the
|
||||
# beginning of the iteration, CPEN may be updated, which may alter the pole point SIM(:, N+1)
|
||||
# by UPDATEPOLE; the quality of the interpolation point depends on SIM(:, N + 1), meaning
|
||||
# that the same interpolation set may have good or bad geometry with respect to different
|
||||
# "poles"; if the geometry turns out bad with the new pole, the original COBYLA code will
|
||||
# take a geometry step, but our code here will NOT do it but continue to take a trust-region
|
||||
# step. The argument is this: even if the geometry step is not skipped in the first place, the
|
||||
# geometry may turn out bad again after the pole is altered due to an update to CPEN; should
|
||||
# we take another geometry step in that case? If no, why should we do it here? Indeed, this
|
||||
# distinction makes no practical difference for CUTEst problems with at most 100 variables
|
||||
# and 5000 constraints, while the algorithm framework is simplified.
|
||||
if improve_geo and not all(primasum(primapow2(sim[:, :num_vars]), axis=0) <= 4 * primapow2(delta)):
|
||||
# Before the geometry step, updatepole has been called either implicitly by UPDATEXFC or
|
||||
# explicitly after CPEN is updated, so that SIM[:, :NUM_VARS] is the optimal vertex.
|
||||
|
||||
# Decide a vertex to drop from the simplex. It will be replaced with SIM[:, NUM_VARS] + D to
|
||||
# improve the geometry of the simplex.
|
||||
# N.B.:
|
||||
# 1. COBYLA never sets JDROP_GEO = num_vars.
|
||||
# 2. The following JDROP_GEO comes from UOBYQA/NEWUOA/BOBYQA/LINCOA.
|
||||
# 3. In Powell's original algorithm, the geometry of the simplex is considered acceptable
|
||||
# iff the distance between any vertex and the pole is at most 2.1*DELTA, and the distance
|
||||
# between any vertex and the opposite face of the simplex is at least 0.25*DELTA, as
|
||||
# specified in (14) of the COBYLA paper. Correspondingly, JDROP_GEO is set to the index of
|
||||
# the vertex with the largest distance to the pole provided that the distance is larger than
|
||||
# 2.1*DELTA, or the vertex with the smallest distance to the opposite face of the simplex,
|
||||
# in which case the distance must be less than 0.25*DELTA, as the current simplex does not
|
||||
# have acceptable geometry (see (15)--(16) of the COBYLA paper). Once JDROP_GEO is set, the
|
||||
# algorithm replaces SIM(:, JDROP_GEO) with D specified in (17) of the COBYLA paper, which
|
||||
# is orthogonal to the face opposite to SIM(:, JDROP_GEO) and has a length of 0.5*DELTA,
|
||||
# intending to improve the geometry of the simplex as per (14).
|
||||
# 4. Powell's geometry-improving procedure outlined above has an intrinsic flaw: it may lead
|
||||
# to infinite cycling, as was observed in a test on 20240320. In this test, the geometry-
|
||||
# improving point introduced in the previous iteration was replaced with the trust-region
|
||||
# trial point in the current iteration, which was then replaced with the same geometry-
|
||||
# improving point in the next iteration, and so on. In this process, the simplex alternated
|
||||
# between two configurations, neither of which had acceptable geometry. Thus RHO was never
|
||||
# reduced, leading to infinite cycling. (N.B.: Our implementation uses DELTA as the trust
|
||||
# region radius, with RHO being its lower bound. When the infinite cycling occurred in this
|
||||
# test, DELTA = RHO and it could not be reduced due to the requirement that DELTA >= RHO.)
|
||||
jdrop_geo = np.argmax(primasum(primapow2(sim[:, :num_vars]), axis=0), axis=0)
|
||||
|
||||
# Calculate the geometry step D.
|
||||
delbar = delta/2
|
||||
d = geostep(jdrop_geo, amat, bvec, conmat, cpen, cval, delbar, fval, simi)
|
||||
|
||||
# Calculate the next value of the objective and constraint functions.
|
||||
# If X is close to one of the points in the interpolation set, then we do not evaluate the
|
||||
# objective and constraints at X, assuming them to have the values at the closest point.
|
||||
# N.B.:
|
||||
# 1. If this happens, do NOT include X into the filter, as F and CONSTR are inaccurate.
|
||||
# 2. In precise arithmetic, the geometry improving step ensures that the distance between X
|
||||
# and any interpolation point is at least DELBAR, yet X may be close to them due to
|
||||
# rounding. In an experiment with single precision on 20240317, X = SIM(:, N+1) occurred.
|
||||
x = sim[:, num_vars] + d
|
||||
distsq[num_vars] = primasum(primapow2(x - sim[:, num_vars]))
|
||||
distsq[:num_vars] = primasum(primapow2(x.reshape(num_vars, 1) -
|
||||
(sim[:, num_vars].reshape(num_vars, 1) + sim[:, :num_vars])), axis=0)
|
||||
j = np.argmin(distsq)
|
||||
if distsq[j] <= primapow2(1e-4 * rhoend):
|
||||
f = fval[j]
|
||||
constr = conmat[:, j]
|
||||
cstrv = cval[j]
|
||||
else:
|
||||
# Evaluate the objective and constraints at X, taking care of possible
|
||||
# inf/nan values.
|
||||
f, constr = evaluate(calcfc, x, m_nlcon, amat, bvec)
|
||||
cstrv = np.max(np.append(0, constr))
|
||||
nf += 1
|
||||
# Save X, F, CONSTR, CSTRV into the history.
|
||||
savehist(maxhist, x, xhist, f, fhist, cstrv, chist, constr, conhist)
|
||||
# Save X, F, CONSTR, CSTRV into the filter.
|
||||
nfilt, cfilt, ffilt, xfilt, confilt = savefilt(cstrv, ctol, cweight, f,
|
||||
x, nfilt, cfilt, ffilt,
|
||||
xfilt, constr, confilt)
|
||||
|
||||
# Print a message about the function/constraint evaluation according to iprint
|
||||
fmsg(solver, 'Geometry', iprint, nf, delta, f, x, cstrv, constr)
|
||||
# Update SIM, SIMI, FVAL, CONMAT, and CVAL so that SIM(:, JDROP_GEO) is replaced with D.
|
||||
sim, simi, fval, conmat, cval, subinfo = updatexfc(jdrop_geo, constr, cpen, cstrv, d, f, conmat, cval, fval, sim, simi)
|
||||
# Check whether to break due to damaging rounding in UPDATEXFC
|
||||
if subinfo == DAMAGING_ROUNDING:
|
||||
info = subinfo
|
||||
break # Better action to take? Geometry step, or simply continue?
|
||||
|
||||
# Check whether to break due to maxfun, ftarget, etc.
|
||||
subinfo = checkbreak_con(maxfun, nf, cstrv, ctol, f, ftarget, x)
|
||||
if subinfo != INFO_DEFAULT:
|
||||
info = subinfo
|
||||
break
|
||||
# end of if improve_geo. The procedure of improving the geometry ends.
|
||||
|
||||
# The calculations with the current RHO are complete. Enhance the resolution of the algorithm
|
||||
# by reducing RHO; update DELTA and CPEN at the same time.
|
||||
if reduce_rho:
|
||||
if rho <= rhoend:
|
||||
info = SMALL_TR_RADIUS
|
||||
break
|
||||
delta = max(0.5 * rho, redrho(rho, rhoend))
|
||||
rho = redrho(rho, rhoend)
|
||||
# THe second (out of two) updates of CPEN, where CPEN decreases or remains the same.
|
||||
# Powell's code: cpen = min(cpen, fcratio(fval, conmat)), which may set CPEN to 0.
|
||||
cpen = np.maximum(cpenmin, np.minimum(cpen, fcratio(conmat, fval)))
|
||||
# Print a message about the reduction of rho according to iprint
|
||||
rhomsg(solver, iprint, nf, fval[num_vars], rho, sim[:, num_vars], cval[num_vars], conmat[:, num_vars], cpen)
|
||||
conmat, cval, fval, sim, simi, subinfo = updatepole(cpen, conmat, cval, fval, sim, simi)
|
||||
# Check whether to break due to damaging rounding detected in updatepole
|
||||
if subinfo == DAMAGING_ROUNDING:
|
||||
info = subinfo
|
||||
break # Better action to take? Geometry step, or simply continue?
|
||||
# End of if reduce_rho. The procedure of reducing RHO ends.
|
||||
# Report the current best value, and check if user asks for early termination.
|
||||
if callback:
|
||||
terminate = callback(sim[:, num_vars], fval[num_vars], nf, tr, cval[num_vars], conmat[:, num_vars])
|
||||
if terminate:
|
||||
info = CALLBACK_TERMINATE
|
||||
break
|
||||
# End of for loop. The iterative procedure ends
|
||||
|
||||
# Return from the calculation, after trying the last trust-region step if it has not been tried yet.
|
||||
# Ensure that D has not been updated after SHORTD == TRUE occurred, or the code below is incorrect.
|
||||
x = sim[:, num_vars] + d
|
||||
if (info == SMALL_TR_RADIUS and
|
||||
shortd and
|
||||
norm(x - sim[:, num_vars]) > 1.0E-3 * rhoend and
|
||||
nf < maxfun):
|
||||
# Zaikun 20230615: UPDATEXFC or UPDATEPOLE is not called since the last trust-region step. Hence
|
||||
# SIM[:, NUM_VARS] remains unchanged. Otherwise SIM[:, NUM_VARS] + D would not make sense.
|
||||
f, constr = evaluate(calcfc, x, m_nlcon, amat, bvec)
|
||||
cstrv = np.max(np.append(0, constr))
|
||||
nf += 1
|
||||
savehist(maxhist, x, xhist, f, fhist, cstrv, chist, constr, conhist)
|
||||
nfilt, cfilt, ffilt, xfilt, confilt = savefilt(cstrv, ctol, cweight, f, x, nfilt, cfilt, ffilt, xfilt, constr, confilt)
|
||||
# Zaikun 20230512: DELTA has been updated. RHO is only indicative here. TO BE IMPROVED.
|
||||
fmsg(solver, 'Trust region', iprint, nf, rho, f, x, cstrv, constr)
|
||||
|
||||
# Return the best calculated values of the variables
|
||||
# N.B.: SELECTX and FINDPOLE choose X by different standards, one cannot replace the other.
|
||||
kopt = selectx(ffilt[:nfilt], cfilt[:nfilt], max(cpen, cweight), ctol)
|
||||
x = xfilt[:, kopt]
|
||||
f = ffilt[kopt]
|
||||
constr = confilt[:, kopt]
|
||||
cstrv = cfilt[kopt]
|
||||
|
||||
# Print a return message according to IPRINT.
|
||||
retmsg(solver, info, iprint, nf, f, x, cstrv, constr)
|
||||
return x, f, constr, cstrv, nf, xhist, fhist, chist, conhist, info
|
||||
|
||||
|
||||
|
||||
def getcpen(amat, bvec, conmat, cpen, cval, delta, fval, rho, sim, simi):
|
||||
'''
|
||||
This function gets the penalty parameter CPEN so that PREREM = PREREF + CPEN * PREREC > 0.
|
||||
See the discussions around equation (9) of the COBYLA paper.
|
||||
'''
|
||||
|
||||
# Even after nearly all of the pycutest problems were showing nearly bit for bit
|
||||
# identical results between Python and the Fortran bindings, HS102 was still off by
|
||||
# more than machine epsilon. It turned out to be due to the fact that getcpen was
|
||||
# modifying fval, among other. It just goes to show that even when you're nearly
|
||||
# perfect, you can still have non trivial bugs.
|
||||
conmat = conmat.copy()
|
||||
cval = cval.copy()
|
||||
fval = fval.copy()
|
||||
sim = sim.copy()
|
||||
simi = simi.copy()
|
||||
|
||||
# Intermediate variables
|
||||
A = np.zeros((np.size(sim, 0), np.size(conmat, 0)))
|
||||
itol = 1
|
||||
|
||||
# Sizes
|
||||
m_lcon = np.size(bvec) if bvec is not None else 0
|
||||
num_constraints = np.size(conmat, 0)
|
||||
num_vars = np.size(sim, 0)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert num_constraints >= 0
|
||||
assert num_vars >= 1
|
||||
assert cpen > 0
|
||||
assert np.size(conmat, 0) == num_constraints and np.size(conmat, 1) == num_vars + 1
|
||||
assert not (np.isnan(conmat) | np.isneginf(conmat)).any()
|
||||
assert np.size(cval) == num_vars + 1 and \
|
||||
not any(cval < 0 | np.isnan(cval) | np.isposinf(cval))
|
||||
assert np.size(fval) == num_vars + 1 and not any(np.isnan(fval) | np.isposinf(fval))
|
||||
assert np.size(sim, 0) == num_vars and np.size(sim, 1) == num_vars + 1
|
||||
assert np.isfinite(sim).all()
|
||||
assert all(np.max(abs(sim[:, :num_vars]), axis=0) > 0)
|
||||
assert np.size(simi, 0) == num_vars and np.size(simi, 1) == num_vars
|
||||
assert np.isfinite(simi).all()
|
||||
assert isinv(sim[:, :num_vars], simi, itol)
|
||||
assert delta >= rho and rho > 0
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
# Initialize INFO which is needed in the postconditions
|
||||
info = INFO_DEFAULT
|
||||
|
||||
# Increase CPEN if necessary to ensure PREREM > 0. Branch back for the next loop
|
||||
# if this change alters the optimal vertex of the current simplex.
|
||||
# Note the following:
|
||||
# 1. In each loop, CPEN is changed only if PREREC > 0 > PREREF, in which case
|
||||
# PREREM is guaranteed positive after the update. Note that PREREC >= 0 and
|
||||
# max(PREREC, PREREF) > 0 in theory. If this holds numerically as well then CPEN
|
||||
# is not changed only if PREREC = 0 or PREREF >= 0, in which case PREREM is
|
||||
# currently positive, explaining why CPEN needs no update.
|
||||
# 2. Even without an upper bound for the loop counter, the loop can occur at most
|
||||
# NUM_VARS+1 times. This is because the update of CPEN does not decrease CPEN,
|
||||
# and hence it can make vertex J (J <= NUM_VARS) become the new optimal vertex
|
||||
# only if CVAL[J] is less than CVAL[NUM_VARS], which can happen at most NUM_VARS
|
||||
# times. See the paragraph below (9) in the COBYLA paper. After the "correct"
|
||||
# optimal vertex is found, one more loop is needed to calculate CPEN, and hence
|
||||
# the loop can occur at most NUM_VARS+1 times.
|
||||
for iter in range(num_vars + 1):
|
||||
# Switch the best vertex of the current simplex to SIM[:, NUM_VARS]
|
||||
conmat, cval, fval, sim, simi, info = updatepole(cpen, conmat, cval, fval, sim,
|
||||
simi)
|
||||
# Check whether to exit due to damaging rounding in UPDATEPOLE
|
||||
if info == DAMAGING_ROUNDING:
|
||||
break
|
||||
|
||||
# Calculate the linear approximations to the objective and constraint functions.
|
||||
g = matprod(fval[:num_vars] - fval[num_vars], simi)
|
||||
A[:, :m_lcon] = amat.T if amat is not None else amat
|
||||
A[:, m_lcon:] = matprod((conmat[m_lcon:, :num_vars] -
|
||||
np.tile(conmat[m_lcon:, num_vars], (num_vars, 1)).T), simi).T
|
||||
|
||||
# Calculate the trust-region trial step D. Note that D does NOT depend on CPEN.
|
||||
d = trstlp(A, -conmat[:, num_vars], delta, g)
|
||||
|
||||
# Predict the change to F (PREREF) and to the constraint violation (PREREC) due
|
||||
# to D.
|
||||
preref = -inprod(d, g) # Can be negative
|
||||
prerec = cval[num_vars] - np.max(np.append(0, conmat[:, num_vars] + matprod(d, A)))
|
||||
|
||||
# PREREC <= 0 or PREREF >=0 or either is NaN
|
||||
if not (prerec > 0 and preref < 0):
|
||||
break
|
||||
|
||||
# Powell's code defines BARMU = -PREREF / PREREC, and CPEN is increased to
|
||||
# 2*BARMU if and only if it is currently less than 1.5*BARMU, a very
|
||||
# "Powellful" scheme. In our implementation, however, we set CPEN directly to
|
||||
# the maximum between its current value and 2*BARMU while handling possible
|
||||
# overflow. The simplifies the scheme without worsening the performance of
|
||||
# COBYLA.
|
||||
cpen = max(cpen, min(-2 * preref / prerec, REALMAX))
|
||||
|
||||
if findpole(cpen, cval, fval) == num_vars:
|
||||
break
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert cpen >= cpen and cpen > 0
|
||||
assert preref + cpen * prerec > 0 or info == DAMAGING_ROUNDING or \
|
||||
not (prerec >= 0 and np.maximum(prerec, preref) > 0) or not np.isfinite(preref)
|
||||
|
||||
return cpen
|
||||
|
||||
|
||||
def fcratio(conmat, fval):
|
||||
'''
|
||||
This function calculates the ratio between the "typical change" of F and that of CONSTR.
|
||||
See equations (12)-(13) in Section 3 of the COBYLA paper for the definition of the ratio.
|
||||
'''
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert np.size(fval) >= 1
|
||||
assert np.size(conmat, 1) == np.size(fval)
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
cmin = np.min(-conmat, axis=1)
|
||||
cmax = np.max(-conmat, axis=1)
|
||||
fmin = min(fval)
|
||||
fmax = max(fval)
|
||||
if any(cmin < 0.5 * cmax) and fmin < fmax:
|
||||
denom = np.min(np.maximum(cmax, 0) - cmin, where=cmin < 0.5 * cmax, initial=np.inf)
|
||||
# Powell mentioned the following alternative in section 4 of his COBYLA paper. According to a test
|
||||
# on 20230610, it does not make much difference to the performance.
|
||||
# denom = np.max(max(*cmax, 0) - cmin, mask=(cmin < 0.5 * cmax))
|
||||
r = (fmax - fmin) / denom
|
||||
else:
|
||||
r = 0
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert r >= 0
|
||||
|
||||
return r
|
||||
@@ -0,0 +1,226 @@
|
||||
'''
|
||||
This module contains subroutines concerning the geometry-improving of the interpolation set.
|
||||
|
||||
Translated from Zaikun Zhang's modern-Fortran reference implementation in PRIMA.
|
||||
|
||||
Dedicated to late Professor M. J. D. Powell FRS (1936--2015).
|
||||
|
||||
Python translation by Nickolai Belakovski.
|
||||
'''
|
||||
|
||||
from ..common.consts import DEBUGGING
|
||||
from ..common.linalg import isinv, matprod, inprod, norm, primasum, primapow2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def setdrop_tr(ximproved, d, delta, rho, sim, simi):
|
||||
'''
|
||||
This function finds (the index) of a current interpolation point to be replaced with
|
||||
the trust-region trial point. See (19)-(22) of the COBYLA paper.
|
||||
N.B.:
|
||||
1. If XIMPROVED == True, then JDROP > 0 so that D is included into XPT. Otherwise,
|
||||
it is a bug.
|
||||
2. COBYLA never sets JDROP = NUM_VARS
|
||||
TODO: Check whether it improves the performance if JDROP = NUM_VARS is allowed when
|
||||
XIMPROVED is True. Note that UPDATEXFC should be revised accordingly.
|
||||
'''
|
||||
|
||||
# Local variables
|
||||
itol = 0.1
|
||||
|
||||
# Sizes
|
||||
num_vars = np.size(sim, 0)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert num_vars >= 1
|
||||
assert np.size(d) == num_vars and all(np.isfinite(d))
|
||||
assert delta >= rho and rho > 0
|
||||
assert np.size(sim, 0) == num_vars and np.size(sim, 1) == num_vars + 1
|
||||
assert np.isfinite(sim).all()
|
||||
assert all(np.max(abs(sim[:, :num_vars]), axis=0) > 0)
|
||||
assert np.size(simi, 0) == num_vars and np.size(simi, 1) == num_vars
|
||||
assert np.isfinite(simi).all()
|
||||
assert isinv(sim[:, :num_vars], simi, itol)
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
# -------------------------------------------------------------------------------------------------- #
|
||||
# The following code is Powell's scheme for defining JDROP.
|
||||
# -------------------------------------------------------------------------------------------------- #
|
||||
# ! JDROP = 0 by default. It cannot be removed, as JDROP may not be set below in some cases (e.g.,
|
||||
# ! when XIMPROVED == FALSE, MAXVAL(ABS(SIMID)) <= 1, and MAXVAL(VETA) <= EDGMAX).
|
||||
# jdrop = 0
|
||||
#
|
||||
# ! SIMID(J) is the value of the J-th Lagrange function at D. It is the counterpart of VLAG in UOBYQA
|
||||
# ! and DEN in NEWUOA/BOBYQA/LINCOA, but it excludes the value of the (N+1)-th Lagrange function.
|
||||
# simid = matprod(simi, d)
|
||||
# if (any(abs(simid) > 1) .or. (ximproved .and. any(.not. is_nan(simid)))) then
|
||||
# jdrop = int(maxloc(abs(simid), mask=(.not. is_nan(simid)), dim=1), kind(jdrop))
|
||||
# !!MATLAB: [~, jdrop] = max(simid, [], 'omitnan');
|
||||
# end if
|
||||
#
|
||||
# ! VETA(J) is the distance from the J-th vertex of the simplex to the best vertex, taking the trial
|
||||
# ! point SIM(:, N+1) + D into account.
|
||||
# if (ximproved) then
|
||||
# veta = sqrt(sum((sim(:, 1:n) - spread(d, dim=2, ncopies=n))**2, dim=1))
|
||||
# !!MATLAB: veta = sqrt(sum((sim(:, 1:n) - d).^2)); % d should be a column! Implicit expansion
|
||||
# else
|
||||
# veta = sqrt(sum(sim(:, 1:n)**2, dim=1))
|
||||
# end if
|
||||
#
|
||||
# ! VSIG(J) (J=1, .., N) is the Euclidean distance from vertex J to the opposite face of the simplex.
|
||||
# vsig = ONE / sqrt(sum(simi**2, dim=2))
|
||||
# sigbar = abs(simid) * vsig
|
||||
#
|
||||
# ! The following JDROP will overwrite the previous one if its premise holds.
|
||||
# mask = (veta > factor_delta * delta .and. (sigbar >= factor_alpha * delta .or. sigbar >= vsig))
|
||||
# if (any(mask)) then
|
||||
# jdrop = int(maxloc(veta, mask=mask, dim=1), kind(jdrop))
|
||||
# !!MATLAB: etamax = max(veta(mask)); jdrop = find(mask & ~(veta < etamax), 1, 'first');
|
||||
# end if
|
||||
#
|
||||
# ! Powell's code does not include the following instructions. With Powell's code, if SIMID consists
|
||||
# ! of only NaN, then JDROP can be 0 even when XIMPROVED == TRUE (i.e., D reduces the merit function).
|
||||
# ! With the following code, JDROP cannot be 0 when XIMPROVED == TRUE, unless VETA is all NaN, which
|
||||
# ! should not happen if X0 does not contain NaN, the trust-region/geometry steps never contain NaN,
|
||||
# ! and we exit once encountering an iterate containing Inf (due to overflow).
|
||||
# if (ximproved .and. jdrop <= 0) then ! Write JDROP <= 0 instead of JDROP == 0 for robustness.
|
||||
# jdrop = int(maxloc(veta, mask=(.not. is_nan(veta)), dim=1), kind(jdrop))
|
||||
# !!MATLAB: [~, jdrop] = max(veta, [], 'omitnan');
|
||||
# end if
|
||||
# -------------------------------------------------------------------------------------------------- #
|
||||
# Powell's scheme ends here.
|
||||
# -------------------------------------------------------------------------------------------------- #
|
||||
|
||||
# The following definition of JDROP is inspired by SETDROP_TR in UOBYQA/NEWUOA/BOBYQA/LINCOA.
|
||||
# It is simpler and works better than Powell's scheme. Note that we allow JDROP to be NUM_VARS+1 if
|
||||
# XIMPROVED is True, whereas Powell's code does not.
|
||||
# See also (4.1) of Scheinberg-Toint-2010: Self-Correcting Geometry in Model-Based Algorithms for
|
||||
# Derivative-Free Unconstrained Optimization, which refers to the strategy here as the "combined
|
||||
# distance/poisedness criteria".
|
||||
|
||||
# DISTSQ[j] is the square of the distance from the jth vertex of the simplex to get "best" point so
|
||||
# far, taking the trial point SIM[:, NUM_VARS] + D into account.
|
||||
distsq = np.zeros(np.size(sim, 1))
|
||||
if ximproved:
|
||||
distsq[:num_vars] = primasum(primapow2(sim[:, :num_vars] - np.tile(d, (num_vars, 1)).T), axis=0)
|
||||
distsq[num_vars] = primasum(d*d)
|
||||
else:
|
||||
distsq[:num_vars] = primasum(primapow2(sim[:, :num_vars]), axis=0)
|
||||
distsq[num_vars] = 0
|
||||
|
||||
weight = np.maximum(1, distsq / primapow2(np.maximum(rho, delta/10))) # Similar to Powell's NEWUOA code.
|
||||
|
||||
# Other possible definitions of weight. They work almost the same as the one above.
|
||||
# weight = distsq # Similar to Powell's LINCOA code, but WRONG. See comments in LINCOA/geometry.f90.
|
||||
# weight = max(1, max(25 * distsq / delta**2)) # Similar to Powell's BOBYQA code, works well.
|
||||
# weight = max(1, max(10 * distsq / delta**2))
|
||||
# weight = max(1, max(1e2 * distsq / delta**2))
|
||||
# weight = max(1, max(distsq / rho**2)) ! Similar to Powell's UOBYQA
|
||||
|
||||
# If 0 <= j < NUM_VARS, SIMID[j] is the value of the jth Lagrange function at D; the value of the
|
||||
# (NUM_VARS+1)th Lagrange function is 1 - sum(SIMID). [SIMID, 1 - sum(SIMID)] is the counterpart of
|
||||
# VLAG in UOBYQA and DEN in NEWUOA/BOBYQA/LINCOA.
|
||||
simid = matprod(simi, d)
|
||||
score = weight * abs(np.array([*simid, 1 - primasum(simid)]))
|
||||
|
||||
# If XIMPROVED = False (D does not render a better X), set SCORE[NUM_VARS] = -1 to avoid JDROP = NUM_VARS.
|
||||
if not ximproved:
|
||||
score[num_vars] = -1
|
||||
|
||||
# score[j] is NaN implies SIMID[j] is NaN, but we want abs(SIMID) to be big. So we
|
||||
# exclude such j.
|
||||
score[np.isnan(score)] = -1
|
||||
|
||||
jdrop = None
|
||||
# The following if statement works a bit better than
|
||||
# `if any(score > 1) or (any(score > 0) and ximproved)` from Powell's UOBYQA and
|
||||
# NEWUOA code.
|
||||
if any(score > 0): # Powell's BOBYQA and LINCOA code.
|
||||
jdrop = np.argmax(score)
|
||||
|
||||
if (ximproved and jdrop is None):
|
||||
jdrop = np.argmax(distsq)
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert jdrop is None or (0 <= jdrop < num_vars + 1)
|
||||
assert jdrop <= num_vars or ximproved
|
||||
assert jdrop >= 0 or not ximproved
|
||||
# JDROP >= 1 when XIMPROVED = TRUE unless NaN occurs in DISTSQ, which should not happen if the
|
||||
# starting point does not contain NaN and the trust-region/geometry steps never contain NaN.
|
||||
|
||||
return jdrop
|
||||
|
||||
|
||||
|
||||
|
||||
def geostep(jdrop, amat, bvec, conmat, cpen, cval, delbar, fval, simi):
|
||||
'''
|
||||
This function calculates a geometry step so that the geometry of the interpolation set is improved
|
||||
when SIM[: JDROP_GEO] is replaced with SIM[:, NUM_VARS] + D. See (15)--(17) of the COBYLA paper.
|
||||
'''
|
||||
|
||||
# Sizes
|
||||
m_lcon = np.size(bvec, 0) if bvec is not None else 0
|
||||
num_constraints = np.size(conmat, 0)
|
||||
num_vars = np.size(simi, 0)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert num_constraints >= m_lcon >= 0
|
||||
assert num_vars >= 1
|
||||
assert delbar > 0
|
||||
assert cpen > 0
|
||||
assert np.size(simi, 0) == num_vars and np.size(simi, 1) == num_vars
|
||||
assert np.isfinite(simi).all()
|
||||
assert np.size(fval) == num_vars + 1 and not any(np.isnan(fval) | np.isposinf(fval))
|
||||
assert np.size(conmat, 0) == num_constraints and np.size(conmat, 1) == num_vars + 1
|
||||
assert not np.any(np.isnan(conmat) | np.isposinf(conmat))
|
||||
assert np.size(cval) == num_vars + 1 and not any(cval < 0 | np.isnan(cval) | np.isposinf(cval))
|
||||
assert 0 <= jdrop < num_vars
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
# SIMI[JDROP, :] is a vector perpendicular to the face of the simplex to the opposite of vertex
|
||||
# JDROP. Set D to the vector in this direction and with length DELBAR.
|
||||
d = simi[jdrop, :]
|
||||
d = delbar * (d / norm(d))
|
||||
|
||||
# The code below chooses the direction of D according to an approximation of the merit function.
|
||||
# See (17) of the COBYLA paper and line 225 of Powell's cobylb.f.
|
||||
|
||||
# Calculate the coefficients of the linear approximations to the objective and constraint functions.
|
||||
# N.B.: CONMAT and SIMI have been updated after the last trust-region step, but G and A have not.
|
||||
# So we cannot pass G and A from outside.
|
||||
g = matprod(fval[:num_vars] - fval[num_vars], simi)
|
||||
A = np.zeros((num_vars, num_constraints))
|
||||
A[:, :m_lcon] = amat.T if amat is not None else amat
|
||||
A[:, m_lcon:] = matprod((conmat[m_lcon:, :num_vars] -
|
||||
np.tile(conmat[m_lcon:, num_vars], (num_vars, 1)).T), simi).T
|
||||
# CVPD and CVND are the predicted constraint violation of D and -D by the linear models.
|
||||
cvpd = np.max(np.append(0, conmat[:, num_vars] + matprod(d, A)))
|
||||
cvnd = np.max(np.append(0, conmat[:, num_vars] - matprod(d, A)))
|
||||
if -inprod(d, g) + cpen * cvnd < inprod(d, g) + cpen * cvpd:
|
||||
d *= -1
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert np.size(d) == num_vars and all(np.isfinite(d))
|
||||
# In theory, ||S|| == DELBAR, which may be false due to rounding, but not too far.
|
||||
# It is crucial to ensure that the geometry step is nonzero, which holds in theory.
|
||||
assert 0.9 * delbar < np.linalg.norm(d) <= 1.1 * delbar
|
||||
return d
|
||||
@@ -0,0 +1,215 @@
|
||||
'''
|
||||
This module contains subroutines for initialization.
|
||||
|
||||
Translated from Zaikun Zhang's modern-Fortran reference implementation in PRIMA.
|
||||
|
||||
Dedicated to late Professor M. J. D. Powell FRS (1936--2015).
|
||||
|
||||
Python translation by Nickolai Belakovski.
|
||||
'''
|
||||
|
||||
from ..common.checkbreak import checkbreak_con
|
||||
from ..common.consts import DEBUGGING, REALMAX
|
||||
from ..common.infos import INFO_DEFAULT
|
||||
from ..common.evaluate import evaluate
|
||||
from ..common.history import savehist
|
||||
from ..common.linalg import inv
|
||||
from ..common.message import fmsg
|
||||
from ..common.selectx import savefilt
|
||||
|
||||
import numpy as np
|
||||
|
||||
def initxfc(calcfc, iprint, maxfun, constr0, amat, bvec, ctol, f0, ftarget, rhobeg, x0,
|
||||
xhist, fhist, chist, conhist, maxhist):
|
||||
'''
|
||||
This subroutine does the initialization concerning X, function values, and
|
||||
constraints.
|
||||
'''
|
||||
|
||||
# Local variables
|
||||
solver = 'COBYLA'
|
||||
srname = "INITIALIZE"
|
||||
|
||||
# Sizes
|
||||
num_constraints = np.size(constr0)
|
||||
m_lcon = np.size(bvec) if bvec is not None else 0
|
||||
m_nlcon = num_constraints - m_lcon
|
||||
num_vars = np.size(x0)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert num_constraints >= 0, f'M >= 0 {srname}'
|
||||
assert num_vars >= 1, f'N >= 1 {srname}'
|
||||
assert abs(iprint) <= 3, f'IPRINT is 0, 1, -1, 2, -2, 3, or -3 {srname}'
|
||||
# assert conmat.shape == (num_constraints , num_vars + 1), f'CONMAT.shape = [M, N+1] {srname}'
|
||||
# assert cval.size == num_vars + 1, f'CVAL.size == N+1 {srname}'
|
||||
# assert maxchist * (maxchist - maxhist) == 0, f'CHIST.shape == 0 or MAXHIST {srname}'
|
||||
# assert conhist.shape[0] == num_constraints and maxconhist * (maxconhist - maxhist) == 0, 'CONHIST.shape[0] == num_constraints, SIZE(CONHIST, 2) == 0 or MAXHIST {srname)}'
|
||||
# assert maxfhist * (maxfhist - maxhist) == 0, f'FHIST.shape == 0 or MAXHIST {srname}'
|
||||
# assert xhist.shape[0] == num_vars and maxxhist * (maxxhist - maxhist) == 0, 'XHIST.shape[0] == N, SIZE(XHIST, 2) == 0 or MAXHIST {srname)}'
|
||||
assert all(np.isfinite(x0)), f'X0 is finite {srname}'
|
||||
assert rhobeg > 0, f'RHOBEG > 0 {srname}'
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
# Initialize info to the default value. At return, a value different from this
|
||||
# value will indicate an abnormal return
|
||||
info = INFO_DEFAULT
|
||||
|
||||
# Initialize the simplex. It will be revised during the initialization.
|
||||
sim = np.eye(num_vars, num_vars+1) * rhobeg
|
||||
sim[:, num_vars] = x0
|
||||
|
||||
# Initialize the matrix simi. In most cases simi is overwritten, but not always.
|
||||
simi = np.eye(num_vars) / rhobeg
|
||||
|
||||
# evaluated[j] = True iff the function/constraint of SIM[:, j] has been evaluated.
|
||||
evaluated = np.zeros(num_vars+1, dtype=bool)
|
||||
|
||||
# Initialize fval
|
||||
fval = np.zeros(num_vars+1) + REALMAX
|
||||
cval = np.zeros(num_vars+1) + REALMAX
|
||||
conmat = np.zeros((num_constraints, num_vars+1)) + REALMAX
|
||||
|
||||
|
||||
for k in range(num_vars + 1):
|
||||
x = sim[:, num_vars].copy()
|
||||
# We will evaluate F corresponding to SIM(:, J).
|
||||
if k == 0:
|
||||
j = num_vars
|
||||
f = f0
|
||||
constr = constr0
|
||||
else:
|
||||
j = k - 1
|
||||
x[j] += rhobeg
|
||||
f, constr = evaluate(calcfc, x, m_nlcon, amat, bvec)
|
||||
cstrv = np.max(np.append(0, constr))
|
||||
|
||||
# Print a message about the function/constraint evaluation according to IPRINT.
|
||||
fmsg(solver, 'Initialization', iprint, k, rhobeg, f, x, cstrv, constr)
|
||||
|
||||
# Save X, F, CONSTR, CSTRV into the history.
|
||||
savehist(maxhist, x, xhist, f, fhist, cstrv, chist, constr, conhist)
|
||||
|
||||
# Save F, CONSTR, and CSTRV to FVAL, CONMAT, and CVAL respectively.
|
||||
evaluated[j] = True
|
||||
fval[j] = f
|
||||
conmat[:, j] = constr
|
||||
cval[j] = cstrv
|
||||
|
||||
# Check whether to exit.
|
||||
subinfo = checkbreak_con(maxfun, k, cstrv, ctol, f, ftarget, x)
|
||||
if subinfo != INFO_DEFAULT:
|
||||
info = subinfo
|
||||
break
|
||||
|
||||
# Exchange the new vertex of the initial simplex with the optimal vertex if necessary.
|
||||
# This is the ONLY part that is essentially non-parallel.
|
||||
if j < num_vars and fval[j] < fval[num_vars]:
|
||||
fval[j], fval[num_vars] = fval[num_vars], fval[j]
|
||||
cval[j], cval[num_vars] = cval[num_vars], cval[j]
|
||||
conmat[:, [j, num_vars]] = conmat[:, [num_vars, j]]
|
||||
sim[:, num_vars] = x
|
||||
sim[j, :j+1] = -rhobeg # SIM[:, :j+1] is lower triangular
|
||||
|
||||
nf = np.count_nonzero(evaluated)
|
||||
|
||||
if evaluated.all():
|
||||
# Initialize SIMI to the inverse of SIM[:, :num_vars]
|
||||
simi = inv(sim[:, :num_vars])
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert nf <= maxfun, f'NF <= MAXFUN {srname}'
|
||||
assert evaluated.size == num_vars + 1, f'EVALUATED.size == Num_vars + 1 {srname}'
|
||||
# assert chist.size == maxchist, f'CHIST.size == MAXCHIST {srname}'
|
||||
# assert conhist.shape== (num_constraints, maxconhist), f'CONHIST.shape == [M, MAXCONHIST] {srname}'
|
||||
assert conmat.shape == (num_constraints, num_vars + 1), f'CONMAT.shape = [M, N+1] {srname}'
|
||||
assert not (np.isnan(conmat).any() or np.isneginf(conmat).any()), f'CONMAT does not contain NaN/-Inf {srname}'
|
||||
assert cval.size == num_vars + 1 and not (any(cval < 0) or any(np.isnan(cval)) or any(np.isposinf(cval))), f'CVAL.shape == Num_vars+1 and CVAL does not contain negative values or NaN/+Inf {srname}'
|
||||
# assert fhist.shape == maxfhist, f'FHIST.shape == MAXFHIST {srname}'
|
||||
# assert maxfhist * (maxfhist - maxhist) == 0, f'FHIST.shape == 0 or MAXHIST {srname}'
|
||||
assert fval.size == num_vars + 1 and not (any(np.isnan(fval)) or any(np.isposinf(fval))), f'FVAL.shape == Num_vars+1 and FVAL is not NaN/+Inf {srname}'
|
||||
# assert xhist.shape == (num_vars, maxxhist), f'XHIST.shape == [N, MAXXHIST] {srname}'
|
||||
assert sim.shape == (num_vars, num_vars + 1), f'SIM.shape == [N, N+1] {srname}'
|
||||
assert np.isfinite(sim).all(), f'SIM is finite {srname}'
|
||||
assert all(np.max(abs(sim[:, :num_vars]), axis=0) > 0), f'SIM(:, 1:N) has no zero column {srname}'
|
||||
assert simi.shape == (num_vars, num_vars), f'SIMI.shape == [N, N] {srname}'
|
||||
assert np.isfinite(simi).all(), f'SIMI is finite {srname}'
|
||||
assert np.allclose(sim[:, :num_vars] @ simi, np.eye(num_vars), rtol=0.1, atol=0.1) or not all(evaluated), f'SIMI = SIM(:, 1:N)^{-1} {srname}'
|
||||
|
||||
return evaluated, conmat, cval, sim, simi, fval, nf, info
|
||||
|
||||
|
||||
def initfilt(conmat, ctol, cweight, cval, fval, sim, evaluated, cfilt, confilt, ffilt, xfilt):
|
||||
'''
|
||||
This function initializes the filter (XFILT, etc) that will be used when selecting
|
||||
x at the end of the solver.
|
||||
N.B.:
|
||||
1. Why not initialize the filters using XHIST, etc? Because the history is empty if
|
||||
the user chooses not to output it.
|
||||
2. We decouple INITXFC and INITFILT so that it is easier to parallelize the former
|
||||
if needed.
|
||||
'''
|
||||
|
||||
# Sizes
|
||||
num_constraints = conmat.shape[0]
|
||||
num_vars = sim.shape[0]
|
||||
maxfilt = len(ffilt)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert num_constraints >= 0
|
||||
assert num_vars >= 1
|
||||
assert maxfilt >= 1
|
||||
assert np.size(confilt, 0) == num_constraints and np.size(confilt, 1) == maxfilt
|
||||
assert np.size(cfilt) == maxfilt
|
||||
assert np.size(xfilt, 0) == num_vars and np.size(xfilt, 1) == maxfilt
|
||||
assert np.size(ffilt) == maxfilt
|
||||
assert np.size(conmat, 0) == num_constraints and np.size(conmat, 1) == num_vars + 1
|
||||
assert not (np.isnan(conmat) | np.isneginf(conmat)).any()
|
||||
assert np.size(cval) == num_vars + 1 and not any(cval < 0 | np.isnan(cval) | np.isposinf(cval))
|
||||
assert np.size(fval) == num_vars + 1 and not any(np.isnan(fval) | np.isposinf(fval))
|
||||
assert np.size(sim, 0) == num_vars and np.size(sim, 1) == num_vars + 1
|
||||
assert np.isfinite(sim).all()
|
||||
assert all(np.max(abs(sim[:, :num_vars]), axis=0) > 0)
|
||||
assert np.size(evaluated) == num_vars + 1
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
|
||||
nfilt = 0
|
||||
for i in range(num_vars+1):
|
||||
if evaluated[i]:
|
||||
if i < num_vars:
|
||||
x = sim[:, i] + sim[:, num_vars]
|
||||
else:
|
||||
x = sim[:, i] # i == num_vars, i.e. the last column
|
||||
nfilt, cfilt, ffilt, xfilt, confilt = savefilt(cval[i], ctol, cweight, fval[i], x, nfilt, cfilt, ffilt, xfilt, conmat[:, i], confilt)
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert nfilt <= maxfilt
|
||||
assert np.size(confilt, 0) == num_constraints and np.size(confilt, 1) == maxfilt
|
||||
assert not (np.isnan(confilt[:, :nfilt]) | np.isneginf(confilt[:, :nfilt])).any()
|
||||
assert np.size(cfilt) == maxfilt
|
||||
assert not any(cfilt[:nfilt] < 0 | np.isnan(cfilt[:nfilt]) | np.isposinf(cfilt[:nfilt]))
|
||||
assert np.size(xfilt, 0) == num_vars and np.size(xfilt, 1) == maxfilt
|
||||
assert not (np.isnan(xfilt[:, :nfilt])).any()
|
||||
# The last calculated X can be Inf (finite + finite can be Inf numerically).
|
||||
assert np.size(ffilt) == maxfilt
|
||||
assert not any(np.isnan(ffilt[:nfilt]) | np.isposinf(ffilt[:nfilt]))
|
||||
|
||||
return nfilt
|
||||
@@ -0,0 +1,492 @@
|
||||
'''
|
||||
This module provides subroutines concerning the trust-region calculations of COBYLA.
|
||||
|
||||
Translated from Zaikun Zhang's modern-Fortran reference implementation in PRIMA.
|
||||
|
||||
Dedicated to late Professor M. J. D. Powell FRS (1936--2015).
|
||||
|
||||
Python translation by Nickolai Belakovski.
|
||||
'''
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
from ..common.consts import DEBUGGING, REALMIN, REALMAX, EPS
|
||||
from ..common.powalg import qradd_Rdiag, qrexc_Rdiag
|
||||
from ..common.linalg import isminor, matprod, inprod, lsqr, primasum
|
||||
|
||||
|
||||
def trstlp(A, b, delta, g):
|
||||
'''
|
||||
This function calculated an n-component vector d by the following two stages. In the first
|
||||
stage, d is set to the shortest vector that minimizes the greatest violation of the constraints
|
||||
A.T @ D <= B, K = 1, 2, 3, ..., M,
|
||||
subject to the Euclidean length of d being at most delta. If its length is strictly less than
|
||||
delta, then the second stage uses the resultant freedom in d to minimize the objective function
|
||||
G.T @ D
|
||||
subject to no increase in any greatest constraint violation.
|
||||
|
||||
It is possible but rare that a degeneracy may prevent d from attaining the target length delta.
|
||||
|
||||
cviol is the largest constraint violation of the current d: max(max(A.T@D - b), 0)
|
||||
icon is the index of a most violated constraint if cviol is positive.
|
||||
|
||||
nact is the number of constraints in the active set and iact[0], ..., iact[nact-1] are their indices,
|
||||
while the remainder of the iact contains a permutation of the remaining constraint indicies.
|
||||
N.B.: nact <= min(num_constraints, num_vars). Obviously nact <= num_constraints. In addition, the constraints
|
||||
in iact[0, ..., nact-1] have linearly independent gradients (see the comments above the instruction
|
||||
that delete a constraint from the active set to make room for the new active constraint with index iact[icon]);
|
||||
it can also be seen from the update of nact: starting from 0, nact is incremented only if nact < n.
|
||||
|
||||
Further, Z is an orthogonal matrix whose first nact columns can be regarded as the result of
|
||||
Gram-Schmidt applied to the active constraint gradients. For j = 0, 1, ..., nact-1, the number
|
||||
zdota[j] is the scalar product of the jth column of Z with the gradient of the jth active
|
||||
constraint. d is the current vector of variables and here the residuals of the active constraints
|
||||
should be zero. Further, the active constraints have nonnegative Lagrange multipliers that are
|
||||
held at the beginning of vmultc. The remainder of this vector holds the residuals of the inactive
|
||||
constraints at d, the ordering of the components of vmultc being in agreement with the permutation
|
||||
of the indices of the constraints that is in iact. All these residuals are nonnegative, which is
|
||||
achieved by the shift cviol that makes the least residual zero.
|
||||
|
||||
N.B.:
|
||||
0. In Powell's implementation, the constraints are A.T @ D >= B. In other words, the A and B in
|
||||
our implementation are the negative of those in Powell's implementation.
|
||||
1. The algorithm was NOT documented in the COBYLA paper. A note should be written to introduce it!
|
||||
2. As a major part of the algorithm (see trstlp_sub), the code maintains and updates the QR
|
||||
factorization of A[iact[:nact]], i.e. the gradients of all the active (linear) constraints. The
|
||||
matrix Z is indeed Q, and the vector zdota is the diagonal of R. The factorization is updated by
|
||||
Givens rotations when an index is added in or removed from iact.
|
||||
3. There are probably better algorithms available for the trust-region linear programming problem.
|
||||
'''
|
||||
|
||||
# Sizes
|
||||
num_constraints = A.shape[1]
|
||||
num_vars = A.shape[0]
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert num_vars >= 1
|
||||
assert num_constraints >= 0
|
||||
assert np.size(g) == num_vars
|
||||
assert np.size(b) == num_constraints
|
||||
assert delta > 0
|
||||
|
||||
|
||||
vmultc = np.zeros(num_constraints + 1)
|
||||
iact = np.zeros(num_constraints + 1, dtype=int)
|
||||
nact = 0
|
||||
d = np.zeros(num_vars)
|
||||
z = np.zeros((num_vars, num_vars))
|
||||
|
||||
# ==================
|
||||
# Calculation starts
|
||||
# ==================
|
||||
|
||||
# Form A_aug and B_aug. This allows the gradient of the objective function to be regarded as the
|
||||
# gradient of a constraint in the second stage.
|
||||
A_aug = np.hstack([A, g.reshape((num_vars, 1))])
|
||||
b_aug = np.hstack([b, 0])
|
||||
|
||||
|
||||
# Scale the problem if A contains large values. Otherwise floating point exceptions may occur.
|
||||
# Note that the trust-region step is scale invariant.
|
||||
for i in range(num_constraints+1): # Note that A_aug.shape[1] == num_constraints+1
|
||||
if (maxval:=max(abs(A_aug[:, i]))) > 1e12:
|
||||
modscal = max(2*REALMIN, 1/maxval)
|
||||
A_aug[:, i] *= modscal
|
||||
b_aug[i] *= modscal
|
||||
|
||||
# Stage 1: minimize the 1+infinity constraint violation of the linearized constraints.
|
||||
iact[:num_constraints], nact, d, vmultc[:num_constraints], z = trstlp_sub(iact[:num_constraints], nact, 1, A_aug[:, :num_constraints], b_aug[:num_constraints], delta, d, vmultc[:num_constraints], z)
|
||||
|
||||
# Stage 2: minimize the linearized objective without increasing the 1_infinity constraint violation.
|
||||
iact, nact, d, vmultc, z = trstlp_sub(iact, nact, 2, A_aug, b_aug, delta, d, vmultc, z)
|
||||
|
||||
# ================
|
||||
# Calculation ends
|
||||
# ================
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert all(np.isfinite(d))
|
||||
# Due to rounding, it may happen that ||D|| > DELTA, but ||D|| > 2*DELTA is highly improbable.
|
||||
assert np.linalg.norm(d) <= 2 * delta
|
||||
|
||||
return d
|
||||
|
||||
def trstlp_sub(iact: npt.NDArray, nact: int, stage, A, b, delta, d, vmultc, z):
|
||||
'''
|
||||
This subroutine does the real calculations for trstlp, both stage 1 and stage 2.
|
||||
Major differences between stage 1 and stage 2:
|
||||
1. Initialization. Stage 2 inherits the values of some variables from stage 1, so they are
|
||||
initialized in stage 1 but not in stage 2.
|
||||
2. cviol. cviol is updated after at iteration in stage 1, while it remains a constant in stage2.
|
||||
3. sdirn. See the definition of sdirn in the code for details.
|
||||
4. optnew. The two stages have different objectives, so optnew is updated differently.
|
||||
5. step. step <= cviol in stage 1.
|
||||
'''
|
||||
zdasav = np.zeros(z.shape[1])
|
||||
vmultd = np.zeros(np.size(vmultc))
|
||||
zdota = np.zeros(np.size(z, 1))
|
||||
|
||||
# Sizes
|
||||
mcon = np.size(A, 1)
|
||||
num_vars = np.size(A, 0)
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert num_vars >= 1
|
||||
assert stage == 1 or stage == 2
|
||||
assert (mcon >= 0 and stage == 1) or (mcon >= 1 and stage == 2)
|
||||
assert np.size(b) == mcon
|
||||
assert np.size(iact) == mcon
|
||||
assert np.size(vmultc) == mcon
|
||||
assert np.size(d) == num_vars
|
||||
assert np.size(z, 0) == num_vars and np.size(z, 1) == num_vars
|
||||
assert delta > 0
|
||||
if stage == 2:
|
||||
assert all(np.isfinite(d)) and np.linalg.norm(d) <= 2 * delta
|
||||
assert nact >= 0 and nact <= np.minimum(mcon, num_vars)
|
||||
assert all(vmultc[:mcon]) >= 0
|
||||
# N.B.: Stage 1 defines only VMULTC(1:M); VMULTC(M+1) is undefined!
|
||||
|
||||
|
||||
# Initialize according to stage
|
||||
if stage == 1:
|
||||
iact = np.linspace(0, mcon-1, mcon, dtype=int)
|
||||
nact = 0
|
||||
d = np.zeros(num_vars)
|
||||
cviol = np.max(np.append(0, -b))
|
||||
vmultc = cviol + b
|
||||
z = np.eye(num_vars)
|
||||
if mcon == 0 or cviol <= 0:
|
||||
# Check whether a quick return is possible. Make sure the in-outputs have been initialized.
|
||||
return iact, nact, d, vmultc, z
|
||||
|
||||
if all(np.isnan(b)):
|
||||
return iact, nact, d, vmultc, z
|
||||
else:
|
||||
icon = np.nanargmax(-b)
|
||||
num_constraints = mcon
|
||||
sdirn = np.zeros(len(d))
|
||||
else:
|
||||
if inprod(d, d) >= delta*delta:
|
||||
# Check whether a quick return is possible.
|
||||
return iact, nact, d, vmultc, z
|
||||
|
||||
iact[mcon-1] = mcon-1
|
||||
vmultc[mcon-1] = 0
|
||||
num_constraints = mcon - 1
|
||||
icon = mcon - 1
|
||||
|
||||
# In Powell's code, stage 2 uses the zdota and cviol calculated by stage1. Here we recalculate
|
||||
# them so that they need not be passed from stage 1 to 2, and hence the coupling is reduced.
|
||||
cviol = np.max(np.append(0, matprod(d, A[:, :num_constraints]) - b[:num_constraints]))
|
||||
zdota[:nact] = [inprod(z[:, k], A[:, iact[k]]) for k in range(nact)]
|
||||
|
||||
# More initialization
|
||||
optold = REALMAX
|
||||
nactold = nact
|
||||
nfail = 0
|
||||
|
||||
# Zaikun 20211011: vmultd is computed from scratch at each iteration, but vmultc is inherited
|
||||
|
||||
# Powell's code can encounter infinite cycling, which did happen when testing the following CUTEst
|
||||
# problems: DANWOODLS, GAUSS1LS, GAUSS2LS, GAUSS3LS, KOEBHELB, TAX13322, TAXR13322. Indeed, in all
|
||||
# these cases, Inf/NaN appear in d due to extremely large values in A (up to 10^219). To resolve
|
||||
# this, we set the maximal number of iterations to maxiter, and terminate if Inf/NaN occurs in d.
|
||||
maxiter = np.minimum(10000, 100*max(num_constraints, num_vars))
|
||||
for iter in range(maxiter):
|
||||
if DEBUGGING:
|
||||
assert all(vmultc >= 0)
|
||||
if stage == 1:
|
||||
optnew = cviol
|
||||
else:
|
||||
optnew = inprod(d, A[:, mcon-1])
|
||||
|
||||
# End the current stage of the calculation if 3 consecutive iterations have either failed to
|
||||
# reduce the best calculated value of the objective function or to increase the number of active
|
||||
# constraints since the best value was calculated. This strategy prevents cycling, but there is
|
||||
# a remote possibility that it will cause premature termination.
|
||||
if optnew < optold or nact > nactold:
|
||||
nactold = nact
|
||||
nfail = 0
|
||||
else:
|
||||
nfail += 1
|
||||
optold = np.minimum(optold, optnew)
|
||||
if nfail == 3:
|
||||
break
|
||||
|
||||
# If icon exceeds nact, then we add the constraint with index iact[icon] to the active set.
|
||||
if icon >= nact: # In Python this needs to be >= since Python is 0-indexed (in Fortran we have 1 > 0, in Python we need 0 >= 0)
|
||||
zdasav[:nact] = zdota[:nact]
|
||||
nactsav = nact
|
||||
z, zdota, nact = qradd_Rdiag(A[:, iact[icon]], z, zdota, nact) # May update nact to nact+1
|
||||
# Indeed it suffices to pass zdota[:min(num_vars, nact+1)] to qradd as follows:
|
||||
# qradd(A[:, iact[icon]], z, zdota[:min(num_vars, nact+1)], nact)
|
||||
|
||||
if nact == nactsav + 1:
|
||||
# N.B.: It is possible to index arrays using [nact, icon] when nact == icon.
|
||||
# Zaikun 20211012: Why should vmultc[nact] = 0?
|
||||
if nact != (icon + 1): # Need to add 1 to Python for 0 indexing
|
||||
vmultc[[icon, nact-1]] = vmultc[nact-1], 0
|
||||
iact[[icon, nact-1]] = iact[[nact-1, icon]]
|
||||
else:
|
||||
vmultc[nact-1] = 0
|
||||
else:
|
||||
# Zaikun 20211011:
|
||||
# 1. VMULTD is calculated from scratch for the first time (out of 2) in one iteration.
|
||||
# 2. Note that IACT has not been updated to replace IACT[NACT] with IACT[ICON]. Thus
|
||||
# A[:, IACT[:NACT]] is the UNUPDATED version before QRADD (note Z[:, :NACT] remains the
|
||||
# same before and after QRADD). Therefore if we supply ZDOTA to LSQR (as Rdiag) as
|
||||
# Powell did, we should use the UNUPDATED version, namely ZDASAV.
|
||||
# vmultd[:nact] = lsqr(A[:, iact[:nact]], A[:, iact[icon]], z[:, :nact], zdasav[:nact])
|
||||
vmultd[:nact] = lsqr(A[:, iact[:nact]], A[:, iact[icon]], z[:, :nact], zdasav[:nact])
|
||||
if not any(np.logical_and(vmultd[:nact] > 0, iact[:nact] <= num_constraints)):
|
||||
# N.B.: This can be triggered by NACT == 0 (among other possibilities)! This is
|
||||
# important, because NACT will be used as an index in the sequel.
|
||||
break
|
||||
# vmultd[NACT+1:mcon] is not used, but we have to initialize it in Fortran, or compilers
|
||||
# complain about the where construct below (another solution: restrict where to 1:NACT).
|
||||
vmultd[nact:mcon] = -1 # len(vmultd) == mcon
|
||||
|
||||
# Revise the Lagrange multipliers. The revision is not applicable to vmultc[nact:num_constraints].
|
||||
fracmult = [vmultc[i]/vmultd[i] if vmultd[i] > 0 and iact[i] <= num_constraints else REALMAX for i in range(nact)]
|
||||
# Only the places with vmultd > 0 and iact <= m is relevant below, if any.
|
||||
frac = min(fracmult[:nact]) # fracmult[nact:mcon] may contain garbage
|
||||
vmultc[:nact] = np.maximum(np.zeros(len(vmultc[:nact])), vmultc[:nact] - frac*vmultd[:nact])
|
||||
|
||||
# Reorder the active constraints so that the one to be replaced is at the end of the list.
|
||||
# Exit if the new value of zdota[nact] is not acceptable. Powell's condition for the
|
||||
# following If: not abs(zdota[nact]) > 0. Note that it is different from
|
||||
# 'abs(zdota[nact]) <=0)' as zdota[nact] can be NaN.
|
||||
# N.B.: We cannot arrive here with nact == 0, which should have triggered a break above
|
||||
if np.isnan(zdota[nact - 1]) or abs(zdota[nact - 1]) <= EPS**2:
|
||||
break
|
||||
vmultc[[icon, nact - 1]] = 0, frac # vmultc[[icon, nact]] is valid as icon > nact
|
||||
iact[[icon, nact - 1]] = iact[[nact - 1, icon]]
|
||||
# end if nact == nactsav + 1
|
||||
|
||||
# In stage 2, ensure that the objective continues to be treated as the last active constraint.
|
||||
# Zaikun 20211011, 20211111: Is it guaranteed for stage 2 that iact[nact-1] = mcon when
|
||||
# iact[nact] != mcon??? If not, then how does the following procedure ensure that mcon is
|
||||
# the last of iact[:nact]?
|
||||
if stage == 2 and iact[nact - 1] != (mcon - 1):
|
||||
if nact <= 1:
|
||||
# We must exit, as nact-2 is used as an index below. Powell's code does not have this.
|
||||
break
|
||||
z, zdota[:nact] = qrexc_Rdiag(A[:, iact[:nact]], z, zdota[:nact], nact - 2) # We pass nact-2 in Python instead of nact-1
|
||||
# Indeed, it suffices to pass Z[:, :nact] to qrexc as follows:
|
||||
# z[:, :nact], zdota[:nact] = qrexc(A[:, iact[:nact]], z[:, :nact], zdota[:nact], nact - 1)
|
||||
iact[[nact-2, nact-1]] = iact[[nact-1, nact-2]]
|
||||
vmultc[[nact-2, nact-1]] = vmultc[[nact-1, nact-2]]
|
||||
# Zaikun 20211117: It turns out that the last few lines do not guarantee iact[nact] == num_vars in
|
||||
# stage 2; the following test cannot be passed. IS THIS A BUG?!
|
||||
# assert iact[nact] == mcon or stage == 1, 'iact[nact] must == mcon in stage 2'
|
||||
|
||||
# Powell's code does not have the following. It avoids subsequent floating points exceptions.
|
||||
if np.isnan(zdota[nact-1]) or abs(zdota[nact-1]) <= EPS**2:
|
||||
break
|
||||
|
||||
# Set sdirn to the direction of the next change to the current vector of variables
|
||||
# Usually during stage 1 the vector sdirn gives a search direction that reduces all the
|
||||
# active constraint violations by one simultaneously.
|
||||
if stage == 1:
|
||||
sdirn -= ((inprod(sdirn, A[:, iact[nact-1]]) + 1)/zdota[nact-1])*z[:, nact-1]
|
||||
else:
|
||||
sdirn = -1/zdota[nact-1]*z[:, nact-1]
|
||||
else: # icon < nact
|
||||
# Delete the constraint with the index iact[icon] from the active set, which is done by
|
||||
# reordering iact[icon:nact] into [iact[icon+1:nact], iact[icon]] and then reduce nact to
|
||||
# nact - 1. In theory, icon > 0.
|
||||
# assert icon > 0, "icon > 0 is required" # For Python I think this is irrelevant
|
||||
z, zdota[:nact] = qrexc_Rdiag(A[:, iact[:nact]], z, zdota[:nact], icon) # qrexc does nothing if icon == nact
|
||||
# Indeed, it suffices to pass Z[:, :nact] to qrexc as follows:
|
||||
# z[:, :nact], zdota[:nact] = qrexc(A[:, iact[:nact]], z[:, :nact], zdota[:nact], icon)
|
||||
iact[icon:nact] = [*iact[icon+1:nact], iact[icon]]
|
||||
vmultc[icon:nact] = [*vmultc[icon+1:nact], vmultc[icon]]
|
||||
nact -= 1
|
||||
|
||||
# Powell's code does not have the following. It avoids subsequent exceptions.
|
||||
# Zaikun 20221212: In theory, nact > 0 in stage 2, as the objective function should always
|
||||
# be considered as an "active constraint" --- more precisely, iact[nact] = mcon. However,
|
||||
# looking at the code, I cannot see why in stage 2 nact must be positive after the reduction
|
||||
# above. It did happen in stage 1 that nact became 0 after the reduction --- this is
|
||||
# extremely rare, and it was never observed until 20221212, after almost one year of
|
||||
# random tests. Maybe nact is theoretically positive even in stage 1?
|
||||
if stage == 2 and nact < 0:
|
||||
break # If this case ever occurs, we have to break, as nact is used as an index below.
|
||||
if nact > 0:
|
||||
if np.isnan(zdota[nact-1]) or abs(zdota[nact-1]) <= EPS**2:
|
||||
break
|
||||
|
||||
# Set sdirn to the direction of the next change to the current vector of variables.
|
||||
if stage == 1:
|
||||
sdirn -= inprod(sdirn, z[:, nact]) * z[:, nact]
|
||||
# sdirn is orthogonal to z[:, nact+1]
|
||||
else:
|
||||
sdirn = -1/zdota[nact-1] * z[:, nact-1]
|
||||
# end if icon > nact
|
||||
|
||||
# Calculate the step to the trust region boundary or take the step that reduces cviol to 0.
|
||||
# ----------------------------------------------------------------------------------------- #
|
||||
# The following calculation of step is adopted from NEWUOA/BOBYQA/LINCOA. It seems to improve
|
||||
# the performance of COBYLA. We also found that removing the precaution about underflows is
|
||||
# beneficial to the overall performance of COBYLA --- the underflows are harmless anyway.
|
||||
dd = delta*delta - inprod(d, d)
|
||||
ss = inprod(sdirn, sdirn)
|
||||
sd = inprod(sdirn, d)
|
||||
if dd <= 0 or ss <= EPS * delta*delta or np.isnan(sd):
|
||||
break
|
||||
# sqrtd: square root of a discriminant. The max avoids sqrtd < abs(sd) due to underflow
|
||||
sqrtd = max(np.sqrt(ss*dd + sd*sd), abs(sd), np.sqrt(ss * dd))
|
||||
if sd > 0:
|
||||
step = dd / (sqrtd + sd)
|
||||
else:
|
||||
step = (sqrtd - sd) / ss
|
||||
# step < 0 should not happen. Step can be 0 or NaN when, e.g., sd or ss becomes inf
|
||||
if step <= 0 or not np.isfinite(step):
|
||||
break
|
||||
|
||||
# Powell's approach and comments are as follows.
|
||||
# -------------------------------------------------- #
|
||||
# The two statements below that include the factor eps prevent
|
||||
# some harmless underflows that occurred in a test calculation
|
||||
# (Zaikun: here, eps is the machine epsilon; Powell's original
|
||||
# code used 1.0e-6, and Powell's code was written in single
|
||||
# precision). Further, we skip the step if it could be 0 within
|
||||
# a reasonable tolerance for computer rounding errors.
|
||||
|
||||
# !dd = delta*delta - sum(d**2, mask=(abs(d) >= EPS * delta))
|
||||
# !ss = inprod(sdirn, sdirn)
|
||||
# !if (dd <= 0) then
|
||||
# ! exit
|
||||
# !end if
|
||||
# !sd = inprod(sdirn, d)
|
||||
# !if (abs(sd) >= EPS * sqrt(ss * dd)) then
|
||||
# ! step = dd / (sqrt(ss * dd + sd*sd) + sd)
|
||||
# !else
|
||||
# ! step = dd / (sqrt(ss * dd) + sd)
|
||||
# !end if
|
||||
# -------------------------------------------------- #
|
||||
|
||||
if stage == 1:
|
||||
if isminor(cviol, step):
|
||||
break
|
||||
step = min(step, cviol)
|
||||
|
||||
# Set dnew to the new variables if step is the steplength, and reduce cviol to the corresponding
|
||||
# maximum residual if stage 1 is being done
|
||||
dnew = d + step * sdirn
|
||||
if stage == 1:
|
||||
cviol = np.max(np.append(0, matprod(dnew, A[:, iact[:nact]]) - b[iact[:nact]]))
|
||||
# N.B.: cviol will be used when calculating vmultd[nact+1:mcon].
|
||||
|
||||
# Zaikun 20211011:
|
||||
# 1. vmultd is computed from scratch for the second (out of 2) time in one iteration.
|
||||
# 2. vmultd[:nact] and vmultd[nact:mcon] are calculated separately with no coupling.
|
||||
# 3. vmultd will be calculated from scratch again in the next iteration.
|
||||
# Set vmultd to the vmultc vector that would occur if d became dnew. A device is included to
|
||||
# force vmultd[k] = 0 if deviations from this value can be attributed to computer rounding
|
||||
# errors. First calculate the new Lagrange multipliers.
|
||||
vmultd[:nact] = -lsqr(A[:, iact[:nact]], dnew, z[:, :nact], zdota[:nact])
|
||||
if stage == 2:
|
||||
vmultd[nact-1] = max(0, vmultd[nact-1]) # This seems never activated.
|
||||
# Complete vmultd by finding the new constraint residuals. (Powell wrote "Complete vmultc ...")
|
||||
cvshift = cviol - (matprod(dnew, A[:, iact]) - b[iact]) # Only cvshift[nact+1:mcon] is needed
|
||||
cvsabs = matprod(abs(dnew), abs(A[:, iact])) + abs(b[iact]) + cviol
|
||||
cvshift[isminor(cvshift, cvsabs)] = 0
|
||||
vmultd[nact:mcon] = cvshift[nact:mcon]
|
||||
|
||||
# Calculate the fraction of the step from d to dnew that will be taken
|
||||
fracmult = [vmultc[i]/(vmultc[i] - vmultd[i]) if vmultd[i] < 0 else REALMAX for i in range(len(vmultd))]
|
||||
# Only the places with vmultd < 0 are relevant below, if any.
|
||||
icon = np.argmin(np.append(1, fracmult)) - 1
|
||||
frac = min(np.append(1, fracmult))
|
||||
|
||||
# Update d, vmultc, and cviol
|
||||
dold = d
|
||||
d = (1 - frac)*d + frac * dnew
|
||||
vmultc = np.maximum(0, (1 - frac)*vmultc + frac*vmultd)
|
||||
# Break in the case of inf/nan in d or vmultc.
|
||||
if not (np.isfinite(primasum(abs(d))) and np.isfinite(primasum(abs(vmultc)))):
|
||||
d = dold # Should we restore also iact, nact, vmultc, and z?
|
||||
break
|
||||
|
||||
if stage == 1:
|
||||
# cviol = (1 - frac) * cvold + frac * cviol # Powell's version
|
||||
# In theory, cviol = np.max(np.append(d@A - b, 0)), yet the
|
||||
# cviol updated as above can be quite different from this value if A has huge entries (e.g., > 1e20)
|
||||
cviol = np.max(np.append(0, matprod(d, A) - b))
|
||||
|
||||
if icon < 0 or icon >= mcon:
|
||||
# In Powell's code, the condition is icon == 0. Indeed, icon < 0 cannot hold unless
|
||||
# fracmult contains only nan, which should not happen; icon >= mcon should never occur.
|
||||
break
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert np.size(iact) == mcon
|
||||
assert np.size(vmultc) == mcon
|
||||
assert all(vmultc >= 0)
|
||||
assert np.size(d) == num_vars
|
||||
assert all(np.isfinite(d))
|
||||
assert np.linalg.norm(d) <= 2 * delta
|
||||
assert np.size(z, 0) == num_vars and np.size(z, 1) == num_vars
|
||||
assert nact >= 0 and nact <= np.minimum(mcon, num_vars)
|
||||
|
||||
return iact, nact, d, vmultc, z
|
||||
|
||||
|
||||
def trrad(delta_in, dnorm, eta1, eta2, gamma1, gamma2, ratio):
|
||||
'''
|
||||
This function updates the trust region radius according to RATIO and DNORM.
|
||||
'''
|
||||
|
||||
# Preconditions
|
||||
if DEBUGGING:
|
||||
assert delta_in >= dnorm > 0
|
||||
assert 0 <= eta1 <= eta2 < 1
|
||||
assert 0 < gamma1 < 1 < gamma2
|
||||
# By the definition of RATIO in ratio.f90, RATIO cannot be NaN unless the
|
||||
# actual reduction is NaN, which should NOT happen due to the moderated extreme
|
||||
# barrier.
|
||||
assert not np.isnan(ratio)
|
||||
|
||||
#====================#
|
||||
# Calculation starts #
|
||||
#====================#
|
||||
|
||||
if ratio <= eta1:
|
||||
delta = gamma1 * dnorm # Powell's UOBYQA/NEWUOA
|
||||
# delta = gamma1 * delta_in # Powell's COBYLA/LINCOA
|
||||
# delta = min(gamma1 * delta_in, dnorm) # Powell's BOBYQA
|
||||
elif ratio <= eta2:
|
||||
delta = max(gamma1 * delta_in, dnorm) # Powell's UOBYQA/NEWUOA/BOBYQA/LINCOA
|
||||
else:
|
||||
delta = max(gamma1 * delta_in, gamma2 * dnorm) # Powell's NEWUOA/BOBYQA
|
||||
# delta = max(delta_in, gamma2 * dnorm) # Modified version. Works well for UOBYQA
|
||||
# For noise-free CUTEst problems of <= 100 variables, Powell's version works slightly better
|
||||
# than the modified one.
|
||||
# delta = max(delta_in, 1.25*dnorm, dnorm + rho) # Powell's UOBYQA
|
||||
# delta = min(max(gamma1 * delta_in, gamma2 * dnorm), gamma3 * delta_in) # Powell's LINCOA, gamma3 = np.sqrt(2)
|
||||
|
||||
# For noisy problems, the following may work better.
|
||||
# if ratio <= eta1:
|
||||
# delta = gamma1 * dnorm
|
||||
# elseif ratio <= eta2: # Ensure DELTA >= DELTA_IN
|
||||
# delta = delta_in
|
||||
# else: # Ensure DELTA > DELTA_IN with a constant factor
|
||||
# delta = max(delta_in * (1 + gamma2) / 2, gamma2 * dnorm)
|
||||
|
||||
#==================#
|
||||
# Calculation ends #
|
||||
#==================#
|
||||
|
||||
# Postconditions
|
||||
if DEBUGGING:
|
||||
assert delta > 0
|
||||
return delta
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user