├── .codecov.yml ├── src └── npx │ ├── py.typed │ ├── __init__.py │ ├── _isin.py │ ├── _mean.py │ ├── _unique.py │ └── _main.py ├── MANIFEST.in ├── .flake8 ├── .gitignore ├── tox.ini ├── tests ├── test_isin.py ├── test_unique.py ├── test_dot_solve.py ├── speedtest.py ├── test_at.py ├── test_unique_rows.py └── test_mean.py ├── justfile ├── .pre-commit-config.yaml ├── CONTRIBUTING.md ├── .github └── workflows │ ├── release.yml │ └── tests.yml ├── LICENSE ├── pyproject.toml ├── CODE_OF_CONDUCT.md └── README.md /.codecov.yml: -------------------------------------------------------------------------------- 1 | comment: no 2 | -------------------------------------------------------------------------------- /src/npx/py.typed: -------------------------------------------------------------------------------- 1 | # Marker file for PEP 561. 2 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include tox.ini 2 | include tests/* 3 | include src/npx/py.typed 4 | -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | ignore = E203, E266, E501, W503 3 | max-line-length = 80 4 | max-complexity = 18 5 | select = B,C,E,F,W,T4,B9 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.swp 3 | *.prof 4 | MANIFEST 5 | dist/ 6 | build/ 7 | .coverage 8 | .cache/ 9 | *.egg-info/ 10 | .pytest_cache/ 11 | .tox/ 12 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = py3 3 | isolated_build = True 4 | 5 | [testenv] 6 | deps = 7 | pytest 8 | pytest-codeblocks 9 | pytest-cov 10 | extras = all 11 | commands = 12 | pytest {posargs} --codeblocks 13 | -------------------------------------------------------------------------------- /tests/test_isin.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import npx 4 | 5 | 6 | def test_isin(): 7 | a = [[0, 3], [1, 0]] 8 | b = [[1, 0], [7, 12], [-1, 5]] 9 | 10 | out = npx.isin_rows(a, b) 11 | assert np.all(out == [False, True]) 12 | 13 | 14 | def test_scalar(): 15 | a = [0, 3, 5] 16 | b = [-1, 6, 5, 0, 0, 0] 17 | 18 | out = npx.isin_rows(a, b) 19 | assert np.all(out == [True, False, True]) 20 | -------------------------------------------------------------------------------- /src/npx/__init__.py: -------------------------------------------------------------------------------- 1 | from ._isin import isin_rows 2 | from ._main import add_at, dot, outer, solve, subtract_at, sum_at 3 | from ._mean import mean 4 | from ._unique import unique, unique_rows 5 | 6 | __all__ = [ 7 | "dot", 8 | "outer", 9 | "solve", 10 | "sum_at", 11 | "add_at", 12 | "subtract_at", 13 | "unique_rows", 14 | "isin_rows", 15 | "mean", 16 | "unique", 17 | "unique_rows", 18 | ] 19 | -------------------------------------------------------------------------------- /tests/test_unique.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import npx 4 | 5 | 6 | def test_unique_tol(): 7 | a = [0.1, 0.15, 0.7] 8 | 9 | a_unique = npx.unique(a, 2.0e-1) 10 | print(a_unique) 11 | assert np.all(a_unique == [0.1, 0.7]) 12 | 13 | a_unique, inv = npx.unique(a, 2.0e-1, return_inverse=True) 14 | assert np.all(a_unique == [0.1, 0.7]) 15 | assert np.all(inv == [0, 0, 1]) 16 | 17 | 18 | def test_unique_edge_case(): 19 | # 1.1 + 2.2 = 3.3000000000000003 20 | out = npx.unique([1.1 + 2.2, 3.3], tol=0.1) 21 | assert len(out) == 1 22 | -------------------------------------------------------------------------------- /justfile: -------------------------------------------------------------------------------- 1 | version := `python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])"` 2 | 3 | default: 4 | @echo "\"just publish\"?" 5 | 6 | publish: 7 | @if [ "$(git rev-parse --abbrev-ref HEAD)" != "main" ]; then exit 1; fi 8 | gh release create "v{{version}}" 9 | 10 | clean: 11 | @find . | grep -E "(__pycache__|\.pyc|\.pyo$)" | xargs rm -rf 12 | @rm -rf src/*.egg-info/ build/ dist/ .tox/ .mypy_cache/ 13 | 14 | format: 15 | ruff --fix src/ tests/ 16 | black src/ tests/ 17 | blacken-docs README.md 18 | 19 | lint: 20 | pre-commit run --all 21 | -------------------------------------------------------------------------------- /tests/test_dot_solve.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import npx 4 | 5 | rng = np.random.default_rng(0) 6 | 7 | 8 | def test_dot(): 9 | a = rng.random((1, 2, 3)) 10 | b = rng.random((3, 4, 5)) 11 | c = npx.dot(a, b) 12 | assert c.shape == (1, 2, 4, 5) 13 | 14 | 15 | def test_solve(): 16 | a = rng.random((3, 3)) 17 | b = rng.random((3, 4, 5)) 18 | c = npx.solve(a, b) 19 | assert c.shape == b.shape 20 | 21 | 22 | def test_outer(): 23 | a = rng.random((1, 2)) 24 | b = rng.random((3, 4)) 25 | c = npx.outer(a, b) 26 | assert c.shape == (1, 2, 3, 4) 27 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/charliermarsh/ruff-pre-commit 3 | rev: v0.4.10 4 | hooks: 5 | - id: ruff 6 | 7 | - repo: https://github.com/pre-commit/mirrors-prettier 8 | rev: v3.1.0 9 | hooks: 10 | - id: prettier 11 | 12 | - repo: https://github.com/codespell-project/codespell 13 | rev: v2.3.0 14 | hooks: 15 | - id: codespell 16 | # args: ["-L", "sur,nd"] 17 | 18 | - repo: https://github.com/pre-commit/mirrors-mypy 19 | rev: v1.10.0 20 | hooks: 21 | - id: mypy 22 | files: src 23 | args: ["--install-types", "--non-interactive"] 24 | -------------------------------------------------------------------------------- /tests/speedtest.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import perfplot 3 | 4 | import npx 5 | 6 | rng = np.random.default_rng(0) 7 | 8 | m = 100 9 | 10 | 11 | def setup(n): 12 | idx = rng.randomint(0, m, n) 13 | b = rng.random(n) 14 | return idx, b 15 | 16 | 17 | def np_add_at(data): 18 | a = np.zeros(m) 19 | idx, b = data 20 | np.add.at(a, idx, b) 21 | return a 22 | 23 | 24 | def npx_add_at(data): 25 | a = np.zeros(m) 26 | idx, b = data 27 | npx.add_at(a, idx, b) 28 | return a 29 | 30 | 31 | def npx_sum_at(data): 32 | idx, b = data 33 | return npx.sum_at(b, idx, minlength=m) 34 | 35 | 36 | b = perfplot.bench( 37 | setup=setup, 38 | kernels=[np_add_at, npx_add_at, npx_sum_at], 39 | n_range=[2**k for k in range(23)], 40 | ) 41 | b.save("perf-add-at.svg") 42 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # npx contributing guidelines 2 | 3 | The npx community appreciates your contributions via issues and 4 | pull requests. Note that the [code of conduct](CODE_OF_CONDUCT.md) 5 | applies to all interactions with the npx project, including 6 | issues and pull requests. 7 | 8 | When submitting pull requests, please follow the style guidelines of 9 | the project, ensure that your code is tested and documented, and write 10 | good commit messages, e.g., following [these 11 | guidelines](https://chris.beams.io/posts/git-commit/). 12 | 13 | By submitting a pull request, you are licensing your code under the 14 | project [license](LICENSE) and affirming that you either own copyright 15 | (automatic for most individuals) or are authorized to distribute under 16 | the project license (e.g., in case your employer retains copyright on 17 | your work). 18 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release 2 | 3 | on: 4 | release: 5 | types: [released] 6 | 7 | jobs: 8 | build-upload: 9 | name: Build and upload to PyPI 10 | runs-on: ubuntu-latest 11 | permissions: 12 | id-token: write # required for pypi upload 13 | contents: read # required for checkout 14 | 15 | steps: 16 | - uses: actions/checkout@v4 17 | 18 | - uses: actions/setup-python@v5 19 | with: 20 | python-version: "3.11" 21 | 22 | - name: Build stubs 23 | run: | 24 | pip install mypy 25 | stubgen --include-docstrings src/ -o src/ 26 | 27 | - name: Build wheels 28 | run: | 29 | pip install build 30 | python3 -m build --wheel 31 | 32 | - name: Upload to PyPI 33 | uses: pypa/gh-action-pypi-publish@release/v1 34 | -------------------------------------------------------------------------------- /tests/test_at.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import npx 4 | 5 | 6 | def test_sum_at(): 7 | a = [1.0, 2.0, 3.0] 8 | idx = [0, 1, 0] 9 | out = npx.sum_at(a, idx, minlength=4) 10 | 11 | tol = 1.0e-13 12 | ref = np.array([4.0, 2.0, 0.0, 0.0]) 13 | assert np.all(np.abs(out - ref) < (1 + np.abs(ref)) * tol) 14 | 15 | 16 | def test_add_at(): 17 | a = [1.0, 2.0, 3.0] 18 | idx = [0, 1, 0] 19 | out = np.zeros(2) 20 | npx.add_at(out, idx, a) 21 | 22 | tol = 1.0e-13 23 | ref = np.array([4.0, 2.0]) 24 | assert np.all(np.abs(out - ref) < (1 + np.abs(ref)) * tol) 25 | 26 | 27 | def test_subtract_at(): 28 | a = [1.0, 2.0, 3.0] 29 | idx = [0, 1, 0] 30 | out = np.ones(2) 31 | npx.subtract_at(out, idx, a) 32 | 33 | tol = 1.0e-13 34 | ref = np.array([-3.0, -1.0]) 35 | assert np.all(np.abs(out - ref) < (1 + np.abs(ref)) * tol) 36 | -------------------------------------------------------------------------------- /src/npx/_isin.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy.typing import ArrayLike 3 | 4 | 5 | def isin_rows(a: ArrayLike, b: ArrayLike) -> np.ndarray: 6 | a = np.asarray(a) 7 | b = np.asarray(b) 8 | if not np.issubdtype(a.dtype, np.integer): 9 | msg = f"Input array must be integer type, got {a.dtype}." 10 | raise ValueError(msg) 11 | if not np.issubdtype(b.dtype, np.integer): 12 | msg = f"Input array must be integer type, got {b.dtype}." 13 | raise ValueError(msg) 14 | 15 | a = a.reshape(a.shape[0], np.prod(a.shape[1:], dtype=int)) 16 | b = b.reshape(b.shape[0], np.prod(b.shape[1:], dtype=int)) 17 | 18 | a_view = np.ascontiguousarray(a).view( 19 | np.dtype((np.void, a.dtype.itemsize * a.shape[1])), 20 | ) 21 | b_view = np.ascontiguousarray(b).view( 22 | np.dtype((np.void, b.dtype.itemsize * b.shape[1])), 23 | ) 24 | 25 | out = np.isin(a_view, b_view) 26 | 27 | return out.reshape(a.shape[0]) 28 | -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: ci 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | branches: 9 | - main 10 | 11 | jobs: 12 | lint: 13 | runs-on: ubuntu-latest 14 | steps: 15 | - name: Check out repo 16 | uses: actions/checkout@v4 17 | - name: Run pre-commit 18 | uses: pre-commit/action@v3.0.1 19 | 20 | build: 21 | runs-on: ubuntu-latest 22 | strategy: 23 | matrix: 24 | python-version: ["3.7", "3.12"] 25 | steps: 26 | - uses: actions/setup-python@v5 27 | with: 28 | python-version: ${{ matrix.python-version }} 29 | allow-prereleases: true 30 | 31 | - uses: actions/checkout@v4 32 | 33 | - name: Test with tox 34 | run: | 35 | pip install tox 36 | tox -- --cov npx --cov-report xml --cov-report term 37 | 38 | - name: Submit to codecov 39 | uses: codecov/codecov-action@v4 40 | if: ${{ matrix.python-version == '3.9' }} 41 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2021-present Nico Schlömer 2 | 3 | Redistribution and use in source and binary forms, with or without modification, are 4 | permitted provided that the following conditions are met: 5 | 6 | 1. Redistributions of source code must retain the above copyright notice, this list of 7 | conditions and the following disclaimer. 8 | 9 | 2. Redistributions in binary form must reproduce the above copyright notice, this list 10 | of conditions and the following disclaimer in the documentation and/or other 11 | materials provided with the distribution. 12 | 13 | 3. Neither the name of the copyright holder nor the names of its contributors may be 14 | used to endorse or promote products derived from this software without specific prior 15 | written permission. 16 | 17 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY 18 | EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF 19 | MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL 20 | THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 21 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT 22 | OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 23 | INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 24 | LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /src/npx/_mean.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy.typing import ArrayLike 3 | 4 | 5 | # There also is 6 | # , 7 | # but implementation is easy enough 8 | def _logsumexp(x: ArrayLike): 9 | c = np.max(x) 10 | return c + np.log(np.sum(np.exp(x - c))) 11 | 12 | 13 | def mean(x: ArrayLike, p: float = 1) -> np.ndarray: 14 | """Generalized mean. 15 | 16 | See for the numpy issue. 17 | """ 18 | x = np.asarray(x) 19 | 20 | n = len(x) 21 | if p == 1: 22 | return np.mean(x) 23 | 24 | if p == -np.inf: 25 | return np.min(np.abs(x)) 26 | 27 | if p == 0: 28 | # first compute the root, then the product, to avoid numerical 29 | # difficulties with too small values of prod(x) 30 | if np.any(x < 0.0): 31 | msg = "p=0 only works with nonnegative x." 32 | raise ValueError(msg) 33 | return np.prod(np.power(x, 1 / n)) 34 | # alternative: 35 | # return np.exp(np.mean(np.log(x))) 36 | 37 | if p == np.inf: 38 | return np.max(np.abs(x)) 39 | 40 | if np.all(x > 0.0): 41 | # logsumexp trick to avoid overflow for large p 42 | # only works for positive x though 43 | return np.exp((_logsumexp(p * np.log(x)) - np.log(n)) / p) 44 | 45 | if not isinstance(p, (int, np.integer)): 46 | msg = f"Non-integer p (={p}) only work with nonnegative x." 47 | raise TypeError(msg) 48 | 49 | return (np.sum(x**p) / n) ** (1.0 / p) 50 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=61"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "npx" 7 | version = "0.1.6" 8 | authors = [{name = "Nico Schlömer", email = "nico.schloemer@gmail.com"}] 9 | description = "Some useful extensions for NumPy" 10 | readme = "README.md" 11 | license = {file = "LICENSE"} 12 | classifiers = [ 13 | "Development Status :: 4 - Beta", 14 | "Intended Audience :: Science/Research", 15 | "License :: OSI Approved :: BSD License", 16 | "Operating System :: OS Independent", 17 | "Programming Language :: Python", 18 | "Programming Language :: Python :: 3", 19 | "Programming Language :: Python :: 3.7", 20 | "Programming Language :: Python :: 3.8", 21 | "Programming Language :: Python :: 3.9", 22 | "Programming Language :: Python :: 3.10", 23 | "Programming Language :: Python :: 3.11", 24 | "Programming Language :: Python :: 3.12", 25 | "Topic :: Scientific/Engineering", 26 | "Topic :: Utilities", 27 | ] 28 | requires-python = ">=3.7" 29 | dependencies = [ 30 | "numpy >= 1.20.0", 31 | # "scipy >= 1.8", 32 | ] 33 | 34 | [project.urls] 35 | Code = "https://github.com/sigma-py/npx" 36 | Issues = "https://github.com/sigma-py/npx/issues" 37 | 38 | [tool.ruff] 39 | src = ["src", "tests"] 40 | line-length = 88 41 | select = ["ALL"] 42 | ignore = [ 43 | "ANN", "S101", "D", "T201", "ERA", "N803", "PLR2004" 44 | # "ANN", "C901", "D", "E741", "ERA", "FBT", "INP001", 45 | # "N", "PLR", "S101", "T201", "TID252", "TD", "FIX002" 46 | ] 47 | target-version = "py38" 48 | 49 | [tool.mypy] 50 | ignore_missing_imports = true 51 | -------------------------------------------------------------------------------- /tests/test_unique_rows.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import npx 4 | 5 | 6 | def test_1d(): 7 | a = [1, 2, 1] 8 | a_unique = npx.unique_rows(a) 9 | assert np.all(a_unique == [1, 2]) 10 | 11 | 12 | def test_2d(): 13 | a = [[1, 2], [1, 4], [1, 2]] 14 | a_unique = npx.unique_rows(a) 15 | assert np.all(a_unique == [[1, 2], [1, 4]]) 16 | 17 | 18 | def test_3d(): 19 | # entries are matrices 20 | # fails for some reason. keep an eye on 21 | # 22 | a = [[[3, 4], [-1, 2]], [[3, 4], [-1, 2]]] 23 | a_unique = npx.unique_rows(a) 24 | assert np.all(a_unique == [[[3, 4], [-1, 2]]]) 25 | 26 | 27 | def test_return_all(): 28 | a = [[1, 2], [1, 4], [1, 2]] 29 | a_unique, inv, count = npx.unique_rows(a, return_inverse=True, return_counts=True) 30 | assert np.all(a_unique == [[1, 2], [1, 4]]) 31 | assert np.all(inv == [0, 1, 0]) 32 | assert np.all(count == [2, 1]) 33 | 34 | 35 | def test_empty(): 36 | # empty mesh 37 | a = np.empty((1, 0), dtype=int) 38 | a_unique = npx.unique_rows(a) 39 | assert np.all(a_unique == [[]]) 40 | 41 | a = np.empty((0, 2), dtype=int) 42 | a_unique = npx.unique_rows(a) 43 | assert np.all(a_unique == a) 44 | 45 | 46 | def test_float(): 47 | a = [1.1, 1.2, 1.1] 48 | out = npx.unique_rows(a) 49 | ref = np.array([1.2, 1.1]) 50 | assert np.all(np.abs(out - ref) < 1.0e-14) 51 | 52 | 53 | def test_float_rows(): 54 | a = [[1.1, 0.7], [1.0, 1.2], [1.1, 0.7]] 55 | out = npx.unique_rows(a) 56 | ref = [[1.0, 1.2], [1.1, 0.7]] 57 | assert np.all(np.abs(out - ref) < 1.0e-14) 58 | -------------------------------------------------------------------------------- /tests/test_mean.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | 4 | import npx 5 | 6 | 7 | @pytest.mark.parametrize( 8 | ("p", "ref"), 9 | [ 10 | (-np.inf, 1.0), # min 11 | (-20000, 1.0000693171203765), 12 | (-1, 1.9672131147540985), # harmonic mean 13 | (-0.1, 2.3000150292740735), 14 | (0, 2.3403473193207156), # geometric mean 15 | (0.1, 2.3810581190184337), 16 | (1, 2.75), # arithmetic mean 17 | (2, np.sqrt(9.75)), # root mean square 18 | (10000, 4.999306900862521), 19 | (np.inf, 5.0), # max 20 | ], 21 | ) 22 | def test_mean_pos(p, ref): 23 | a = [1.0, 2.0, 3.0, 5.0] 24 | val = npx.mean(a, p) 25 | # val = pmean(a, p) 26 | # print(p, val) 27 | assert abs(val - ref) < 1.0e-13 * abs(ref) 28 | 29 | 30 | @pytest.mark.parametrize( 31 | ("p", "ref"), 32 | [ 33 | (-np.inf, 1.0), # absmin 34 | (-1, -1.9672131147540985), # harmonic mean 35 | # (0, 2.3403473193207156), # geometric mean 36 | (1, -2.75), # arithmetic mean 37 | (2, np.sqrt(9.75)), # root mean square 38 | (np.inf, 5.0), # absmax 39 | ], 40 | ) 41 | def test_mean_neg(p, ref): 42 | a = [-1.0, -2.0, -3.0, -5.0] 43 | val = npx.mean(a, p) 44 | # val = pmean(a, p) 45 | # print(p, val) 46 | assert abs(val - ref) < 1.0e-13 * abs(ref) 47 | 48 | 49 | def test_errors(): 50 | a = [-1.0, -2.0, -3.0, -5.0] 51 | with pytest.raises(TypeError, match="Non-integer p.*"): 52 | npx.mean(a, 0.5) 53 | 54 | with pytest.raises(ValueError, match="p=0 only works with nonnegative x."): 55 | npx.mean(a, 0) 56 | -------------------------------------------------------------------------------- /src/npx/_unique.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from typing import TYPE_CHECKING, Callable 4 | 5 | import numpy as np 6 | 7 | if TYPE_CHECKING: 8 | from numpy.typing import ArrayLike 9 | 10 | 11 | def _unique_tol( 12 | unique_fun: Callable, 13 | a: ArrayLike, 14 | tol: float, 15 | **kwargs, 16 | ) -> np.ndarray | tuple[np.ndarray, ...]: 17 | a = np.asarray(a) 18 | # compute 1/tol first. Difference: 19 | # 20 | # int(3.3 / 0.1) = int(32.99999999999999) = 32 21 | # int(3.3 * (1.0 / 0.1)) = int(33.0) = 33 22 | # 23 | aint = (a * (1.0 / tol)).astype(int) 24 | 25 | return_index = kwargs.pop("return_index", False) 26 | 27 | _, idx, *out = unique_fun(aint, return_index=True, **kwargs) 28 | unique_a = a[idx] 29 | 30 | if return_index: 31 | out = [idx, *out] 32 | 33 | if len(out) == 0: 34 | return unique_a 35 | 36 | return (unique_a, *out) 37 | 38 | 39 | def unique( 40 | a: ArrayLike, 41 | tol: float = 0.0, 42 | **kwargs, 43 | ) -> np.ndarray | tuple[np.ndarray, ...]: 44 | assert tol >= 0.0 45 | if tol > 0.0: 46 | return _unique_tol(np.unique, a, tol, **kwargs) 47 | 48 | return np.unique(a, **kwargs) 49 | 50 | 51 | def unique_rows(a: ArrayLike, **kwargs) -> np.ndarray | tuple[np.ndarray, ...]: 52 | # The numpy alternative `np.unique(a, axis=0)` is slow; cf. 53 | # . 54 | a = np.asarray(a) 55 | 56 | a_shape = a.shape 57 | a = a.reshape(a.shape[0], np.prod(a.shape[1:], dtype=int)) 58 | 59 | b = np.ascontiguousarray(a).view(np.dtype((np.void, a.dtype.itemsize * a.shape[1]))) 60 | out = np.unique(b, **kwargs) 61 | 62 | # out[0] are the sorted, unique rows 63 | if isinstance(out, tuple): 64 | out = list(out) 65 | out[0] = out[0].view(a.dtype).reshape(out[0].shape[0], *a_shape[1:]) 66 | if len(out) > 1: 67 | out[1] = out[1].flatten() 68 | if len(out) > 2: 69 | out[2] = out[2].flatten() 70 | out = tuple(out) 71 | else: 72 | out = out.view(a.dtype).reshape(out.shape[0], *a_shape[1:]) 73 | 74 | return out 75 | -------------------------------------------------------------------------------- /src/npx/_main.py: -------------------------------------------------------------------------------- 1 | from functools import reduce 2 | from operator import mul 3 | 4 | import numpy as np 5 | from numpy.typing import ArrayLike 6 | 7 | 8 | # math.prod in 3.8 9 | # https://docs.python.org/3/library/math.html#math.prod 10 | def _prod(a): 11 | return reduce(mul, a, 1) 12 | 13 | 14 | def dot(a: ArrayLike, b: np.ndarray) -> np.ndarray: 15 | """Take arrays `a` and `b` and form the dot product between the last axis 16 | of `a` and the first of `b`. 17 | """ 18 | return np.tensordot(a, b, 1) 19 | 20 | 21 | def outer(a: ArrayLike, b: ArrayLike) -> np.ndarray: 22 | """Compute the outer product of two arrays `a` and `b` such that the shape 23 | of the resulting array is `(*a.shape, *b.shape)`. 24 | """ 25 | a = np.asarray(a) 26 | b = np.asarray(b) 27 | return np.outer(a, b).reshape(*a.shape, *b.shape) 28 | 29 | 30 | def solve(A: np.ndarray, x: np.ndarray) -> np.ndarray: 31 | """Solves a linear equation system with a matrix of shape (n, n) and an array of 32 | shape (n, ...). The output has the same shape as the second argument. 33 | """ 34 | # https://stackoverflow.com/a/48387507/353337 35 | x = np.asarray(x) 36 | return np.linalg.solve(A, x.reshape(x.shape[0], -1)).reshape(x.shape) 37 | 38 | 39 | def sum_at(a: ArrayLike, indices: ArrayLike, minlength: int): 40 | """Sums up values `a` with `indices` into an output array of at least 41 | length `minlength` while treating dimensionality correctly. It's a lot 42 | faster than numpy's own np.add.at (see 43 | https://github.com/numpy/numpy/issues/5922#issuecomment-511477435). 44 | 45 | Typically, `indices` will be a one-dimensional array; `a` can have any 46 | dimensionality. In this case, the output array will have shape (minlength, 47 | a.shape[1:]). 48 | 49 | `indices` may have arbitrary shape, too, but then `a` has to start out the 50 | same. (Those dimensions are flattened out in the computation.) 51 | """ 52 | a = np.asarray(a) 53 | indices = np.asarray(indices) 54 | 55 | if len(a.shape) < len(indices.shape): 56 | msg = ( 57 | f"a.shape = {a.shape}, indices.shape = {indices.shape}, " 58 | "but len(a.shape) >= len(indices.shape) is required." 59 | ) 60 | raise RuntimeError(msg) 61 | 62 | m = len(indices.shape) 63 | assert indices.shape == a.shape[:m] 64 | 65 | out_shape = (minlength, *a.shape[m:]) 66 | 67 | indices = indices.reshape(-1) 68 | a = a.reshape(_prod(a.shape[:m]), _prod(a.shape[m:])) 69 | 70 | # Cast to int; bincount doesn't work for uint64 yet 71 | # https://github.com/numpy/numpy/issues/17760 72 | indices = indices.astype(int) 73 | 74 | return np.array( 75 | [ 76 | np.bincount(indices, weights=a[:, k], minlength=minlength) 77 | for k in range(a.shape[1]) 78 | ], 79 | ).T.reshape(out_shape) 80 | 81 | 82 | def add_at(a: ArrayLike, indices: ArrayLike, b: ArrayLike): 83 | a = np.asarray(a) 84 | indices = np.asarray(indices) 85 | b = np.asarray(b) 86 | 87 | m = len(indices.shape) 88 | assert a.shape[1:] == b.shape[m:] 89 | a += sum_at(b, indices, a.shape[0]) 90 | 91 | 92 | def subtract_at(a: ArrayLike, indices: ArrayLike, b: ArrayLike): 93 | b = np.asarray(b) 94 | add_at(a, indices, -b) 95 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # npx Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to making participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socioeconomic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | - Using welcoming and inclusive language 18 | - Being respectful of differing viewpoints and experiences 19 | - Gracefully accepting constructive criticism 20 | - Focusing on what is best for the community 21 | - Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | - The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | - Trolling, insulting/derogatory comments, and personal or political attacks 28 | - Public or private harassment 29 | - Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | - Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies within all project spaces, and it also applies when 49 | an individual is representing the project or its community in public spaces. 50 | Examples of representing a project or community include using an official 51 | project e-mail address, posting via an official social media account, or acting 52 | as an appointed representative at an online or offline event. Representation of 53 | a project may be further defined and clarified by project maintainers. 54 | 55 | ## Enforcement 56 | 57 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 58 | reported by contacting the project team at [INSERT EMAIL ADDRESS]. All 59 | complaints will be reviewed and investigated and will result in a response that 60 | is deemed necessary and appropriate to the circumstances. The project team is 61 | obligated to maintain confidentiality with regard to the reporter of an incident. 62 | Further details of specific enforcement policies may be posted separately. 63 | 64 | Project maintainers who do not follow or enforce the Code of Conduct in good 65 | faith may face temporary or permanent repercussions as determined by other 66 | members of the project's leadership. 67 | 68 | ## Attribution 69 | 70 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 71 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct/ 72 | 73 | [homepage]: https://www.contributor-covenant.org/ 74 | 75 | For answers to common questions about this code of conduct, see 76 | https://www.contributor-covenant.org/faq/ 77 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # npx 2 | 3 | [![PyPi Version](https://img.shields.io/pypi/v/npx.svg?style=flat-square)](https://pypi.org/project/npx/) 4 | [![PyPI pyversions](https://img.shields.io/pypi/pyversions/npx.svg?style=flat-square)](https://pypi.org/project/npx/) 5 | [![GitHub stars](https://img.shields.io/github/stars/nschloe/npx.svg?style=flat-square&logo=github&label=Stars&logoColor=white)](https://github.com/nschloe/npx) 6 | [![Downloads](https://pepy.tech/badge/npx/month?style=flat-square)](https://pepy.tech/project/npx) 7 | 8 | 9 | 10 | [![gh-actions](https://img.shields.io/github/workflow/status/nschloe/npx/ci?style=flat-square)](https://github.com/nschloe/npx/actions?query=workflow%3Aci) 11 | [![codecov](https://img.shields.io/codecov/c/github/nschloe/npx.svg?style=flat-square)](https://app.codecov.io/gh/nschloe/npx) 12 | [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg?style=flat-square)](https://github.com/psf/black) 13 | 14 | [NumPy](https://numpy.org/) is a large library used everywhere in scientific computing. 15 | That's why breaking backwards-compatibility comes at a significant cost and is almost 16 | always avoided, even if the API of some methods is arguably lacking. This package 17 | provides drop-in wrappers "fixing" those. 18 | 19 | [scipyx](https://github.com/nschloe/scipyx) does the same for 20 | [SciPy](https://www.scipy.org/). 21 | 22 | If you have a fix for a NumPy method that can't go upstream for some reason, feel free 23 | to PR here. 24 | 25 | #### `dot` 26 | 27 | ```python 28 | import npx 29 | import numpy as np 30 | 31 | a = np.random.rand(3, 4, 5) 32 | b = np.random.rand(5, 2, 2) 33 | 34 | out = npx.dot(a, b) 35 | # out.shape == (3, 4, 2, 2) 36 | ``` 37 | 38 | Forms the dot product between the last axis of `a` and the _first_ axis of `b`. 39 | 40 | (Not the second-last axis of `b` as `numpy.dot(a, b)`.) 41 | 42 | #### `np.solve` 43 | 44 | ```python 45 | import npx 46 | import numpy as np 47 | 48 | A = np.random.rand(3, 3) 49 | b = np.random.rand(3, 10, 4) 50 | 51 | out = npx.solve(A, b) 52 | # out.shape == (3, 10, 4) 53 | ``` 54 | 55 | Solves a linear equation system with a matrix of shape `(n, n)` and an array of shape 56 | `(n, ...)`. The output has the same shape as the second argument. 57 | 58 | #### `sum_at`/`add_at` 59 | 60 | 61 | 62 | ```python 63 | npx.sum_at(a, idx, minlength=0) 64 | npx.add_at(out, idx, a) 65 | ``` 66 | 67 | Returns an array with entries of `a` summed up at indices `idx` with a minimum length of 68 | `minlength`. `idx` can have any shape as long as it's matching `a`. The output shape is 69 | `(minlength,...)`. 70 | 71 | The numpy equivalent `numpy.add.at` is _much_ 72 | slower: 73 | 74 | memory usage 75 | 76 | Relevant issue reports: 77 | 78 | - [ufunc.at (and possibly other methods) 79 | slow](https://github.com/numpy/numpy/issues/11156) 80 | 81 | #### `unique` 82 | 83 | ```python 84 | import npx 85 | 86 | a = [0.1, 0.15, 0.7] 87 | a_unique = npx.unique(a, tol=2.0e-1) 88 | 89 | assert all(a_unique == [0.1, 0.7]) 90 | ``` 91 | 92 | npx's `unique()` works just like NumPy's, except that it provides a parameter 93 | `tol` (default `0.0`) which allows the user to set a tolerance. The real line 94 | is essentially partitioned into bins of size `tol` and at most one 95 | representative of each bin is returned. 96 | 97 | #### `unique_rows` 98 | 99 | ```python 100 | import npx 101 | import numpy as np 102 | 103 | a = np.random.randint(0, 5, size=(100, 2)) 104 | 105 | npx.unique_rows(a, return_inverse=False, return_counts=False) 106 | ``` 107 | 108 | Returns the unique rows of the integer array `a`. The numpy alternative `np.unique(a, axis=0)` is slow. 109 | 110 | Relevant issue reports: 111 | 112 | - [unique() needlessly slow](https://github.com/numpy/numpy/issues/11136) 113 | 114 | #### `isin_rows` 115 | 116 | ```python 117 | import npx 118 | import numpy as np 119 | 120 | a = [[0, 1], [0, 2]] 121 | b = np.random.randint(0, 5, size=(100, 2)) 122 | 123 | npx.isin_rows(a, b) 124 | ``` 125 | 126 | Returns a boolean array of length `len(a)` specifying if the rows `a[k]` appear in `b`. 127 | Similar to NumPy's own `np.isin` which only works for scalars. 128 | 129 | #### `mean` 130 | 131 | ```python 132 | import npx 133 | 134 | a = [1.0, 2.0, 5.0] 135 | npx.mean(a, p=3) 136 | ``` 137 | 138 | Returns the [generalized mean](https://en.wikipedia.org/wiki/Generalized_mean) of a 139 | given list. Handles the cases `+-np.inf` (max/min) and`0` (geometric mean) correctly. 140 | Also does well for large `p`. 141 | 142 | Relevant NumPy issues: 143 | 144 | - [generalized mean](https://github.com/numpy/numpy/issues/19341) 145 | 146 | ### License 147 | 148 | This software is published under the [BSD-3-Clause 149 | license](https://spdx.org/licenses/BSD-3-Clause.html). 150 | --------------------------------------------------------------------------------