Skip to content

Commit

Permalink
Simplify some typing and add algebra.apply
Browse files Browse the repository at this point in the history 
Allows a function to be applied to matrix/vector, not necessarily
faster to do than just using list comprehensions, but certainly more
convenient.
  • Loading branch information
@facelessuser
facelessuser committed Jan 28, 2023
1 parent 1934dae commit a566baf
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 112 deletions.
224 changes: 112 additions & 112 deletions coloraide/algebra.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -515,12 +515,12 @@ def cross(a: VectorLike, b: VectorLike) -> Vector:


@overload
def cross(a: MatrixLike, b: VectorLike | MatrixLike) -> Matrix:
def cross(a: MatrixLike, b: Any) -> Matrix:
...


@overload
def cross(a: VectorLike | MatrixLike, b: MatrixLike) -> Matrix:
def cross(a: Any, b: MatrixLike) -> Matrix:
...


Expand Down Expand Up @@ -872,47 +872,57 @@ def _vector_math(op: Callable[..., float], a: VectorLike, b: VectorLike) -> Vect


@overload
def _math(op: Callable[..., float], a: float, b: float, *, dims: Optional[tuple[int, int]] = None) -> float:
...


@overload
def _math(op: Callable[..., float], a: float, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def _math(op: Callable[..., float], a: VectorLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def _math(op: Callable[..., float], a: float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def _math(op: Callable[..., float], a: MatrixLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def _math(
op: Callable[..., float],
a: float,
b: float,
*,
dims: Optional[tuple[int, int]] = None
) -> float:
...


@overload
def _math(op: Callable[..., float], a: VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def _math(
op: Callable[..., float],
a: float | VectorLike,
b: VectorLike,
*,
dims: Optional[tuple[int, int]] = None
) -> Vector:
...


@overload
def _math(op: Callable[..., float], a: VectorLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def _math(
op: Callable[..., float],
a: VectorLike,
b: float | VectorLike,
*,
dims: Optional[tuple[int, int]] = None
) -> Vector:
...


@overload
def _math(op: Callable[..., float], a: MatrixLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def _math(
op: Callable[..., float],
a: MatrixLike,
b: float | ArrayLike,
*,
dims: Optional[tuple[int, int]] = None
) -> Matrix:
...


@overload
def _math(op: Callable[..., float], a: MatrixLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def _math(
op: Callable[..., float],
a: ArrayLike | float,
b: MatrixLike,
*,
dims: Optional[tuple[int, int]] = None
) -> Matrix:
...


Expand Down Expand Up @@ -1008,42 +1018,22 @@ def divide(a: float, b: float, *, dims: Optional[tuple[int, int]] = None) -> flo


@overload
def divide(a: float, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def divide(a: VectorLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def divide(a: float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def divide(a: MatrixLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def divide(a: float | VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def divide(a: VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def divide(a: VectorLike, b: float | VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def divide(a: VectorLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def divide(a: MatrixLike, b: float | ArrayLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def divide(a: MatrixLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def divide(a: MatrixLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def divide(a: ArrayLike | float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


Expand All @@ -1064,42 +1054,22 @@ def multiply(a: float, b: float, *, dims: Optional[tuple[int, int]] = None) -> f


@overload
def multiply(a: float, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def multiply(a: float | VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def multiply(a: VectorLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def multiply(a: VectorLike, b: float | VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def multiply(a: float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def multiply(a: MatrixLike, b: float | ArrayLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def multiply(a: MatrixLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def multiply(a: VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def multiply(a: VectorLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def multiply(a: MatrixLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def multiply(a: MatrixLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def multiply(a: ArrayLike | float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


Expand All @@ -1120,42 +1090,22 @@ def add(a: float, b: float, *, dims: Optional[tuple[int, int]] = None) -> float:


@overload
def add(a: float, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def add(a: VectorLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def add(a: float | VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def add(a: float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def add(a: VectorLike, b: float | VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def add(a: MatrixLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def add(a: MatrixLike, b: float | ArrayLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def add(a: VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def add(a: VectorLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def add(a: MatrixLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def add(a: MatrixLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def add(a: ArrayLike | float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


Expand All @@ -1176,54 +1126,104 @@ def subtract(a: float, b: float, *, dims: Optional[tuple[int, int]] = None) -> f


@overload
def subtract(a: float, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def subtract(a: float | VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def subtract(a: VectorLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def subtract(a: VectorLike, b: float | VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
...


@overload
def subtract(a: float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def subtract(a: MatrixLike, b: float | ArrayLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


@overload
def subtract(a: MatrixLike, b: float, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def subtract(a: ArrayLike | float, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
...


def subtract(
a: float | ArrayLike,
b: float | ArrayLike,
*,
dims: Optional[tuple[int, int]] = None
) -> float | Array:
"""Subtract simple numbers, vectors, and 2D matrices."""

return _math(operator.sub, a, b, dims=dims)


@overload
def subtract(a: VectorLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Vector:
def apply(
fn: Callable[..., float],
a: float,
b: Optional[float] = None,
*,
dims: Optional[tuple[int, int]] = None
) -> float:
...


@overload
def subtract(a: VectorLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def apply(
fn: Callable[..., float],
a: float | VectorLike,
b: VectorLike,
*,
dims: Optional[tuple[int, int]] = None
) -> Vector:
...


@overload
def subtract(a: MatrixLike, b: VectorLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def apply(
fn: Callable[..., float],
a: VectorLike,
b: Optional[float | VectorLike] = None,
*,
dims: Optional[tuple[int, int]] = None
) -> Vector:
...


@overload
def subtract(a: MatrixLike, b: MatrixLike, *, dims: Optional[tuple[int, int]] = None) -> Matrix:
def apply(
fn: Callable[..., float],
a: MatrixLike,
b: Optional[float | ArrayLike] = None,
*,
dims: Optional[tuple[int, int]] = None
) -> Matrix:
...


def subtract(
@overload
def apply(
fn: Callable[..., float],
a: ArrayLike | float,
b: MatrixLike,
*,
dims: Optional[tuple[int, int]] = None
) -> Matrix:
...


def apply(
fn: Callable[..., float],
a: float | ArrayLike,
b: float | ArrayLike,
b: Optional[float | ArrayLike] = None,
*,
dims: Optional[tuple[int, int]] = None
) -> float | Array:
"""Subtract simple numbers, vectors, and 2D matrices."""
"""Apply a given function over each element of the matrix."""

return _math(operator.sub, a, b, dims=dims)
if b is None:
return reshape([fn(f) for f in flatiter(a)], shape(a))

return _math(fn, a, b, dims=dims)


class BroadcastTo:
Expand Down Expand Up @@ -1790,7 +1790,7 @@ def diag(array: VectorLike, k: int = 0) -> Matrix:


@overload
def diag(array: Matrix, k: int = 0) -> Vector:
def diag(array: MatrixLike, k: int = 0) -> Vector:
...


Expand Down
17 changes: 17 additions & 0 deletions tests/test_algebra.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Test Algebra."""
import unittest
from coloraide import algebra as alg
import math


class TestAlgebra(unittest.TestCase):
Expand Down Expand Up @@ -1289,3 +1290,19 @@ def test_interpolate_natural(self):
alg.interpolate([[3, 4], [6, 8], [9, 2]], method='natural').steps(5),
[[3.0, 4.0], [4.5, 6.9375], [6.0, 8.0], [7.5, 5.9375], [9.0, 2.0]]
)

def test_apply_two_inputs(self):
"""Test apply with two inputs."""

self.assertEqual(
alg.apply(alg.npow, [[1, 2, 3], [4, 5, 6]], 2),
[[1, 4, 9], [16, 25, 36]]
)

def test_apply_one_input(self):
"""Test apply with one input."""

self.assertEqual(
alg.apply(math.sqrt, [[1, 4, 9], [16, 25, 36]]),
[[1, 2, 3], [4, 5, 6]]
)

0 comments on commit a566baf

Please sign in to comment.