dask-array 0.1.0__py3-none-any.whl

dask_array/linalg/_cholesky.py

@@ -0,0 +1,234 @@
+"""Cholesky decomposition for array-expr."""
+
+from __future__ import annotations
+
+import functools
+import operator
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Alias, List, Task, TaskRef
+from dask_array._expr import ArrayExpr
+from dask_array.linalg._utils import _solve_triangular_lower
+from dask_array._utils import meta_from_array
+
+
+def _cholesky_lower(a):
+    """Compute Cholesky decomposition (lower triangular)."""
+    return np.linalg.cholesky(a)
+
+
+def _zeros_like_shape(arr_meta, shape):
+    """Create zeros with specific shape, dtype from arr_meta."""
+    return np.zeros(shape, dtype=arr_meta.dtype)
+
+
+class Cholesky(ArrayExpr):
+    """Block Cholesky decomposition.
+
+    Computes both lower and upper triangular factors.
+    """
+
+    _parameters = ["array"]
+
+    @functools.cached_property
+    def _meta(self):
+        from dask_array._utils import array_safe
+
+        arr_meta = meta_from_array(self.array._meta)
+        cho = np.linalg.cholesky(array_safe([[1, 2], [2, 5]], dtype=self.array.dtype, like=arr_meta))
+        return meta_from_array(self.array._meta, dtype=cho.dtype)
+
+    @functools.cached_property
+    def chunks(self):
+        return self.array.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return f"cholesky-{self.deterministic_token}"
+
+    def _layer(self):
+        vdim = len(self.array.chunks[0])
+        hdim = len(self.array.chunks[1])
+
+        name = self._name
+        name_upper = f"cholesky-upper-{self.deterministic_token}"
+        name_lt_dot = f"cholesky-lt-dot-{self.deterministic_token}"
+
+        dsk = {}
+
+        for j in range(hdim):
+            for i in range(vdim):
+                if i < j:
+                    chunk_shape = (
+                        self.array.chunks[0][i],
+                        self.array.chunks[1][j],
+                    )
+                    dsk[(name, i, j)] = Task(
+                        (name, i, j),
+                        _zeros_like_shape,
+                        self.array._meta,
+                        chunk_shape,
+                    )
+                    dsk[(name_upper, j, i)] = TaskRef((name, i, j))
+                elif i == j:
+                    target = TaskRef((self.array._name, i, j))
+                    if i > 0:
+                        prevs = []
+                        for p in range(i):
+                            prev = (name_lt_dot, i, p, i, p)
+                            dsk[prev] = Task(
+                                prev,
+                                np.dot,
+                                TaskRef((name, i, p)),
+                                TaskRef((name_upper, p, i)),
+                            )
+                            prevs.append(TaskRef(prev))
+                        target = Task(None, operator.sub, target, Task(None, sum, List(*prevs)))
+                    dsk[(name, i, i)] = Task((name, i, i), _cholesky_lower, target)
+                    dsk[(name_upper, i, i)] = Task((name_upper, i, i), np.transpose, TaskRef((name, i, i)))
+                else:
+                    target = TaskRef((self.array._name, j, i))
+                    if j > 0:
+                        prevs = []
+                        for p in range(j):
+                            prev = (name_lt_dot, j, p, i, p)
+                            dsk[prev] = Task(
+                                prev,
+                                np.dot,
+                                TaskRef((name, j, p)),
+                                TaskRef((name_upper, p, i)),
+                            )
+                            prevs.append(TaskRef(prev))
+                        target = Task(None, operator.sub, target, Task(None, sum, List(*prevs)))
+                    dsk[(name_upper, j, i)] = Task(
+                        (name_upper, j, i),
+                        _solve_triangular_lower,
+                        TaskRef((name, j, j)),
+                        target,
+                    )
+                    dsk[(name, i, j)] = Task((name, i, j), np.transpose, TaskRef((name_upper, j, i)))
+
+        return dsk
+
+
+class CholeskyLower(ArrayExpr):
+    """Extract lower triangular from Cholesky.
+
+    This is a view into the lower triangular portion of the Cholesky
+    factorization. It uses Alias tasks to reference the parent Cholesky's
+    tasks.
+    """
+
+    _parameters = ["chol"]
+
+    @functools.cached_property
+    def _meta(self):
+        return self.chol._meta
+
+    @functools.cached_property
+    def chunks(self):
+        return self.chol.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return f"cholesky-lower-{self.chol.deterministic_token}"
+
+    def _layer(self):
+        vdim = len(self.chol.chunks[0])
+        hdim = len(self.chol.chunks[1])
+        parent_name = self.chol._name
+
+        dsk = {}
+        for i in range(vdim):
+            for j in range(hdim):
+                out_key = (self._name, i, j)
+                in_key = (parent_name, i, j)
+                dsk[out_key] = Alias(out_key, in_key)
+        return dsk
+
+
+class CholeskyUpper(ArrayExpr):
+    """Extract upper triangular from Cholesky.
+
+    This is a view into the upper triangular portion of the Cholesky
+    factorization. It uses Alias tasks to reference the parent Cholesky's
+    tasks.
+    """
+
+    _parameters = ["chol"]
+
+    @functools.cached_property
+    def _meta(self):
+        return self.chol._meta
+
+    @functools.cached_property
+    def chunks(self):
+        return self.chol.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return f"cholesky-upper-view-{self.chol.deterministic_token}"
+
+    def _layer(self):
+        vdim = len(self.chol.chunks[0])
+        hdim = len(self.chol.chunks[1])
+        parent_name = f"cholesky-upper-{self.chol.deterministic_token}"
+
+        dsk = {}
+        for i in range(vdim):
+            for j in range(hdim):
+                out_key = (self._name, i, j)
+                in_key = (parent_name, i, j)
+                dsk[out_key] = Alias(out_key, in_key)
+        return dsk
+
+
+def _cholesky(a):
+    """Private function to compute both L and U Cholesky factors."""
+    from dask_array.core import asanyarray
+
+    a = asanyarray(a)
+
+    if a.ndim != 2:
+        raise ValueError("Dimension must be 2 to perform cholesky decomposition")
+
+    xdim, ydim = a.shape
+    if xdim != ydim:
+        raise ValueError("Input must be a square matrix to perform cholesky decomposition")
+    if len(set(a.chunks[0] + a.chunks[1])) != 1:
+        msg = (
+            "All chunks must be a square matrix to perform cholesky decomposition. "
+            "Use .rechunk method to change the size of chunks."
+        )
+        raise ValueError(msg)
+
+    chol_expr = Cholesky(a.expr)
+    lower_expr = CholeskyLower(chol_expr)
+    upper_expr = CholeskyUpper(chol_expr)
+
+    return new_collection(lower_expr), new_collection(upper_expr)
+
+
+def cholesky(a, lower=False):
+    """Returns the Cholesky decomposition of a Hermitian positive-definite matrix.
+
+    Parameters
+    ----------
+    a : (M, M) array_like
+        Matrix to be decomposed
+    lower : bool, optional
+        Whether to compute the upper or lower triangular Cholesky
+        factorization. Default is upper-triangular.
+
+    Returns
+    -------
+    c : (M, M) Array
+        Upper- or lower-triangular Cholesky factor of `a`.
+    """
+    l, u = _cholesky(a)
+    if lower:
+        return l
+    else:
+        return u

dask_array/linalg/_lu.py

@@ -0,0 +1,300 @@
+"""LU decomposition for array-expr."""
+
+from __future__ import annotations
+
+import functools
+import operator
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import List, Task, TaskRef
+from dask_array._expr import ArrayExpr
+from dask_array.linalg._utils import (
+    _solve_triangular_lower,
+    _transpose,
+)
+from dask_array._utils import meta_from_array
+
+
+def _scipy_lu(a):
+    """Compute LU decomposition using scipy."""
+    import scipy.linalg
+
+    return scipy.linalg.lu(a)
+
+
+class LU(ArrayExpr):
+    """Block LU decomposition returning (P, L, U) tuple.
+
+    Uses scipy.linalg.lu for diagonal blocks and propagates through
+    off-diagonal blocks with forward/backward substitution.
+    """
+
+    _parameters = ["array"]
+
+    @functools.cached_property
+    def _meta(self):
+        import scipy.linalg
+
+        pp, ll, uu = scipy.linalg.lu(np.ones(shape=(1, 1), dtype=self.array.dtype))
+        arr_meta = self.array._meta
+        return (
+            meta_from_array(arr_meta, ndim=2, dtype=pp.dtype),
+            meta_from_array(arr_meta, ndim=2, dtype=ll.dtype),
+            meta_from_array(arr_meta, ndim=2, dtype=uu.dtype),
+        )
+
+    @functools.cached_property
+    def chunks(self):
+        return self.array.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return f"lu-{self.deterministic_token}"
+
+    def _layer(self):
+        vdim = len(self.array.chunks[0])
+        hdim = len(self.array.chunks[1])
+
+        name_lu = f"lu-lu-{self.deterministic_token}"
+        name_p = f"lu-p-{self.deterministic_token}"
+        name_l = f"lu-l-{self.deterministic_token}"
+        name_u = f"lu-u-{self.deterministic_token}"
+        name_transpose = f"lu-p-inv-{self.deterministic_token}"
+        name_l_permuted = f"lu-l-permute-{self.deterministic_token}"
+        name_transposed = f"lu-u-transpose-{self.deterministic_token}"
+        name_plu_dot = f"lu-plu-dot-{self.deterministic_token}"
+        name_lu_dot = f"lu-lu-dot-{self.deterministic_token}"
+
+        dsk = {}
+
+        for i in range(min(vdim, hdim)):
+            target = TaskRef((self.array._name, i, i))
+            if i > 0:
+                prevs = []
+                for p in range(i):
+                    prev = (name_plu_dot, i, p, p, i)
+                    dsk[prev] = Task(
+                        prev,
+                        np.dot,
+                        TaskRef((name_l_permuted, i, p)),
+                        TaskRef((name_u, p, i)),
+                    )
+                    prevs.append(TaskRef(prev))
+                target = Task(None, operator.sub, target, Task(None, sum, List(*prevs)))
+            dsk[(name_lu, i, i)] = Task((name_lu, i, i), _scipy_lu, target)
+
+            for j in range(i + 1, hdim):
+                target_h = Task(
+                    None,
+                    np.dot,
+                    TaskRef((name_transpose, i, i)),
+                    TaskRef((self.array._name, i, j)),
+                )
+                if i > 0:
+                    prevs = []
+                    for p in range(i):
+                        prev = (name_lu_dot, i, p, p, j)
+                        dsk[prev] = Task(
+                            prev,
+                            np.dot,
+                            TaskRef((name_l, i, p)),
+                            TaskRef((name_u, p, j)),
+                        )
+                        prevs.append(TaskRef(prev))
+                    target_h = Task(None, operator.sub, target_h, Task(None, sum, List(*prevs)))
+                dsk[(name_lu, i, j)] = Task(
+                    (name_lu, i, j),
+                    _solve_triangular_lower,
+                    TaskRef((name_l, i, i)),
+                    target_h,
+                )
+
+            for k in range(i + 1, vdim):
+                target_v = TaskRef((self.array._name, k, i))
+                if i > 0:
+                    prevs = []
+                    for p in range(i):
+                        prev = (name_plu_dot, k, p, p, i)
+                        dsk[prev] = Task(
+                            prev,
+                            np.dot,
+                            TaskRef((name_l_permuted, k, p)),
+                            TaskRef((name_u, p, i)),
+                        )
+                        prevs.append(TaskRef(prev))
+                    target_v = Task(None, operator.sub, target_v, Task(None, sum, List(*prevs)))
+                dsk[(name_lu, k, i)] = Task(
+                    (name_lu, k, i),
+                    np.transpose,
+                    Task(
+                        None,
+                        _solve_triangular_lower,
+                        TaskRef((name_transposed, i, i)),
+                        Task(None, np.transpose, target_v),
+                    ),
+                )
+
+        for i in range(min(vdim, hdim)):
+            for j in range(min(vdim, hdim)):
+                if i == j:
+                    dsk[(name_p, i, j)] = Task(
+                        (name_p, i, j),
+                        operator.getitem,
+                        TaskRef((name_lu, i, j)),
+                        0,
+                    )
+                    dsk[(name_l, i, j)] = Task(
+                        (name_l, i, j),
+                        operator.getitem,
+                        TaskRef((name_lu, i, j)),
+                        1,
+                    )
+                    dsk[(name_u, i, j)] = Task(
+                        (name_u, i, j),
+                        operator.getitem,
+                        TaskRef((name_lu, i, j)),
+                        2,
+                    )
+                    dsk[(name_l_permuted, i, j)] = Task(
+                        (name_l_permuted, i, j),
+                        np.dot,
+                        TaskRef((name_p, i, j)),
+                        TaskRef((name_l, i, j)),
+                    )
+                    dsk[(name_transposed, i, j)] = Task(
+                        (name_transposed, i, j),
+                        _transpose,
+                        TaskRef((name_u, i, j)),
+                    )
+                    dsk[(name_transpose, i, j)] = Task(
+                        (name_transpose, i, j),
+                        _transpose,
+                        TaskRef((name_p, i, j)),
+                    )
+                elif i > j:
+                    chunk_shape = (
+                        self.array.chunks[0][i],
+                        self.array.chunks[1][j],
+                    )
+                    dsk[(name_p, i, j)] = Task((name_p, i, j), np.zeros, chunk_shape)
+                    dsk[(name_l, i, j)] = Task(
+                        (name_l, i, j),
+                        np.dot,
+                        TaskRef((name_transpose, i, i)),
+                        TaskRef((name_lu, i, j)),
+                    )
+                    dsk[(name_u, i, j)] = Task((name_u, i, j), np.zeros, chunk_shape)
+                    dsk[(name_l_permuted, i, j)] = TaskRef((name_lu, i, j))
+                else:
+                    chunk_shape = (
+                        self.array.chunks[0][i],
+                        self.array.chunks[1][j],
+                    )
+                    dsk[(name_p, i, j)] = Task((name_p, i, j), np.zeros, chunk_shape)
+                    dsk[(name_l, i, j)] = Task((name_l, i, j), np.zeros, chunk_shape)
+                    dsk[(name_u, i, j)] = TaskRef((name_lu, i, j))
+
+        return dsk
+
+
+class LUGetP(ArrayExpr):
+    """Extract P from LU decomposition."""
+
+    _parameters = ["lu"]
+
+    @functools.cached_property
+    def _meta(self):
+        return self.lu._meta[0]
+
+    @functools.cached_property
+    def chunks(self):
+        return self.lu.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return f"lu-p-{self.lu.deterministic_token}"
+
+    def _layer(self):
+        return {}
+
+
+class LUGetL(ArrayExpr):
+    """Extract L from LU decomposition."""
+
+    _parameters = ["lu"]
+
+    @functools.cached_property
+    def _meta(self):
+        return self.lu._meta[1]
+
+    @functools.cached_property
+    def chunks(self):
+        return self.lu.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return f"lu-l-{self.lu.deterministic_token}"
+
+    def _layer(self):
+        return {}
+
+
+class LUGetU(ArrayExpr):
+    """Extract U from LU decomposition."""
+
+    _parameters = ["lu"]
+
+    @functools.cached_property
+    def _meta(self):
+        return self.lu._meta[2]
+
+    @functools.cached_property
+    def chunks(self):
+        return self.lu.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return f"lu-u-{self.lu.deterministic_token}"
+
+    def _layer(self):
+        return {}
+
+
+def lu(a):
+    """Compute the LU decomposition of a matrix.
+
+    Examples
+    --------
+    >>> p, l, u = da.linalg.lu(x)  # doctest: +SKIP
+
+    Returns
+    -------
+    p : Array, permutation matrix
+    l : Array, lower triangular matrix with unit diagonal.
+    u : Array, upper triangular matrix
+    """
+    from dask_array.core import asanyarray
+
+    a = asanyarray(a)
+
+    if a.ndim != 2:
+        raise ValueError("Dimension must be 2 to perform lu decomposition")
+
+    xdim, ydim = a.shape
+    if xdim != ydim:
+        raise ValueError("Input must be a square matrix to perform lu decomposition")
+    if len(set(a.chunks[0] + a.chunks[1])) != 1:
+        msg = (
+            "All chunks must be a square matrix to perform lu decomposition. "
+            "Use .rechunk method to change the size of chunks."
+        )
+        raise ValueError(msg)
+
+    lu_expr = LU(a.expr)
+    p_expr = LUGetP(lu_expr)
+    l_expr = LUGetL(lu_expr)
+    u_expr = LUGetU(lu_expr)
+
+    return new_collection(p_expr), new_collection(l_expr), new_collection(u_expr)

dask_array/linalg/_norm.py

@@ -0,0 +1,94 @@
+"""Matrix and vector norms for array-expr."""
+
+from __future__ import annotations
+
+from numbers import Number
+
+import numpy as np
+
+from dask.utils import derived_from
+
+
+@derived_from(np.linalg)
+def norm(x, ord=None, axis=None, keepdims=False):
+    """Matrix or vector norm.
+
+    This function uses array operations (abs, sum, max, min) which are
+    already implemented in array-expr.
+    """
+    from dask_array.core import asanyarray
+    from dask_array.linalg._svd import svd
+
+    x = asanyarray(x)
+
+    if axis is None:
+        axis = tuple(range(x.ndim))
+    elif isinstance(axis, Number):
+        axis = (int(axis),)
+    else:
+        axis = tuple(axis)
+
+    if len(axis) > 2:
+        raise ValueError("Improper number of dimensions to norm.")
+
+    if ord == "fro":
+        ord = None
+        if len(axis) == 1:
+            raise ValueError("Invalid norm order for vectors.")
+
+    r = abs(x)
+
+    if ord is None:
+        r = (r**2).sum(axis=axis, keepdims=keepdims) ** 0.5
+    elif ord == "nuc":
+        if len(axis) == 1:
+            raise ValueError("Invalid norm order for vectors.")
+        if x.ndim > 2:
+            raise NotImplementedError("SVD based norm not implemented for ndim > 2")
+
+        r = svd(x)[1][None].sum(keepdims=keepdims)
+    elif ord == np.inf:
+        if len(axis) == 1:
+            r = r.max(axis=axis, keepdims=keepdims)
+        else:
+            r = r.sum(axis=axis[1], keepdims=True).max(axis=axis[0], keepdims=True)
+            if keepdims is False:
+                r = r.squeeze(axis=axis)
+    elif ord == -np.inf:
+        if len(axis) == 1:
+            r = r.min(axis=axis, keepdims=keepdims)
+        else:
+            r = r.sum(axis=axis[1], keepdims=True).min(axis=axis[0], keepdims=True)
+            if keepdims is False:
+                r = r.squeeze(axis=axis)
+    elif ord == 0:
+        if len(axis) == 2:
+            raise ValueError("Invalid norm order for matrices.")
+
+        r = (r != 0).astype(r.dtype).sum(axis=axis, keepdims=keepdims)
+    elif ord == 1:
+        if len(axis) == 1:
+            r = r.sum(axis=axis, keepdims=keepdims)
+        else:
+            r = r.sum(axis=axis[0], keepdims=True).max(axis=axis[1], keepdims=True)
+            if keepdims is False:
+                r = r.squeeze(axis=axis)
+    elif len(axis) == 2 and ord == -1:
+        r = r.sum(axis=axis[0], keepdims=True).min(axis=axis[1], keepdims=True)
+        if keepdims is False:
+            r = r.squeeze(axis=axis)
+    elif len(axis) == 2 and ord == 2:
+        if x.ndim > 2:
+            raise NotImplementedError("SVD based norm not implemented for ndim > 2")
+        r = svd(x)[1][None].max(keepdims=keepdims)
+    elif len(axis) == 2 and ord == -2:
+        if x.ndim > 2:
+            raise NotImplementedError("SVD based norm not implemented for ndim > 2")
+        r = svd(x)[1][None].min(keepdims=keepdims)
+    else:
+        if len(axis) == 2:
+            raise ValueError("Invalid norm order for matrices.")
+
+        r = (r**ord).sum(axis=axis, keepdims=keepdims) ** (1.0 / ord)
+
+    return r