dask-array 0.1.0__py3-none-any.whl

dask_array/_stack.py

@@ -0,0 +1,264 @@
+"""Stack operation - expression and collection function."""
+
+from __future__ import annotations
+
+import functools
+from itertools import product
+
+import numpy as np
+from toolz import concat, first
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Task, TaskRef
+from dask_array._expr import ArrayExpr, unify_chunks_expr
+from dask_array._chunk import getitem
+from dask_array._utils import meta_from_array
+
+
+class Stack(ArrayExpr):
+    _parameters = ["array", "axis", "meta"]
+
+    @functools.cached_property
+    def args(self):
+        return [self.array] + self.operands[len(self._parameters) :]
+
+    @functools.cached_property
+    def _meta(self):
+        return self.operand("meta")
+
+    @functools.cached_property
+    def chunks(self):
+        n = len(self.args)
+        return self.array.chunks[: self.axis] + ((1,) * n,) + self.array.chunks[self.axis :]
+
+    @functools.cached_property
+    def _name(self):
+        return "stack-" + self.deterministic_token
+
+    def _layer(self) -> dict:
+        keys = list(product([self._name], *[range(len(bd)) for bd in self.chunks]))
+        names = [a.name for a in self.args]
+        axis = self.axis
+        ndim = self._meta.ndim - 1
+
+        inputs = [(names[key[axis + 1]],) + key[1 : axis + 1] + key[axis + 2 :] for key in keys]
+        values = [
+            Task(
+                key,
+                getitem,
+                TaskRef(inp),
+                (slice(None, None, None),) * axis + (None,) + (slice(None, None, None),) * (ndim - axis),
+            )
+            for key, inp in zip(keys, inputs)
+        ]
+        return dict(zip(keys, values))
+
+    def _simplify_up(self, parent, dependents):
+        """Allow slice and shuffle operations to push through Stack."""
+        from dask_array._shuffle import Shuffle
+        from dask_array.slicing import SliceSlicesIntegers
+
+        if isinstance(parent, SliceSlicesIntegers):
+            return self._accept_slice(parent)
+        if isinstance(parent, Shuffle):
+            return self._accept_shuffle(parent)
+        return None
+
+    def _accept_shuffle(self, shuffle_expr):
+        """Accept a shuffle being pushed through Stack.
+
+        Stack adds a new dimension at axis. Can't shuffle on the stacked axis.
+        For other axes, adjust axis index for inputs (they have one fewer dim).
+        """
+        from dask_array._shuffle import Shuffle
+
+        stack_axis = self.axis
+        shuffle_axis = shuffle_expr.axis
+
+        # Can't shuffle on the stacked axis itself
+        if shuffle_axis == stack_axis:
+            return None
+
+        # Adjust axis for inputs (they have one fewer dimension)
+        if shuffle_axis > stack_axis:
+            input_shuffle_axis = shuffle_axis - 1
+        else:
+            input_shuffle_axis = shuffle_axis
+
+        # Shuffle each input
+        arrays = self.args
+        shuffled_arrays = [
+            Shuffle(
+                a,
+                shuffle_expr.indexer,
+                input_shuffle_axis,
+                shuffle_expr.operand("name"),
+            )
+            for a in arrays
+        ]
+
+        return type(self)(
+            shuffled_arrays[0],
+            stack_axis,
+            self._meta,
+            *shuffled_arrays[1:],
+        )
+
+    def _accept_slice(self, slice_expr):
+        """Accept a slice being pushed through Stack.
+
+        Stack adds a new dimension at axis. Cases:
+        1. Slice on stacked axis: select subset of inputs
+        2. Slice on other axes: push to all inputs (adjusting for added dim)
+        """
+        from numbers import Integral
+
+        from dask_array._new_collection import new_collection
+        from dask_array._utils import meta_from_array
+
+        axis = self.axis
+        arrays = self.args
+        index = slice_expr.index
+
+        # Pad index to full length (output has one more dim than inputs)
+        full_index = index + (slice(None),) * (self.ndim - len(index))
+
+        # For now, only handle simple slices (no integers that reduce dims)
+        if any(isinstance(idx, Integral) for idx in full_index):
+            return None
+        if any(idx is None for idx in full_index):
+            return None
+
+        # Handle the stacked axis slice
+        stacked_axis_slice = full_index[axis]
+        n_arrays = len(arrays)
+
+        if isinstance(stacked_axis_slice, slice):
+            start, stop, step = stacked_axis_slice.indices(n_arrays)
+            if step != 1:
+                return None
+            selected_arrays = arrays[start:stop]
+        else:
+            return None
+
+        if not selected_arrays:
+            return None
+
+        # Build slice for the other axes (remove the stacked axis from index)
+        other_slices = full_index[:axis] + full_index[axis + 1 :]
+
+        # Check if we need to slice the other axes
+        needs_other_slice = any(s != slice(None) for s in other_slices)
+
+        # Slice each selected array on the other axes
+        sliced_arrays = []
+        for arr in selected_arrays:
+            if needs_other_slice:
+                sliced_arr = new_collection(arr)[other_slices]
+                sliced_arrays.append(sliced_arr.expr)
+            else:
+                sliced_arrays.append(arr)
+
+        # Compute new meta from dtype/ndim. Input metas may use different dummy
+        # shapes even when the real arrays have compatible shapes.
+        dtype = np.result_type(*[a.dtype for a in sliced_arrays])
+        new_meta = meta_from_array(None, ndim=self.ndim, dtype=dtype)
+
+        # Create new Stack with selected/sliced arrays
+        return type(self)(
+            sliced_arrays[0],
+            axis,
+            new_meta,
+            *sliced_arrays[1:],
+        )
+
+
+def stack(seq, axis=0, allow_unknown_chunksizes=False):
+    """
+    Stack arrays along a new axis
+
+    Given a sequence of dask arrays, form a new dask array by stacking them
+    along a new dimension (axis=0 by default)
+
+    Parameters
+    ----------
+    seq: list of dask.arrays
+    axis: int
+        Dimension along which to align all of the arrays
+    allow_unknown_chunksizes: bool
+        Allow unknown chunksizes, such as come from converting from dask
+        dataframes.  Dask.array is unable to verify that chunks line up.  If
+        data comes from differently aligned sources then this can cause
+        unexpected results.
+
+    Examples
+    --------
+
+    Create slices
+
+    >>> import dask_array as da
+    >>> import numpy as np
+
+    >>> data = [da.from_array(np.ones((4, 4)), chunks=(2, 2))
+    ...         for i in range(3)]
+
+    >>> x = da.stack(data, axis=0)
+    >>> x.shape
+    (3, 4, 4)
+
+    >>> da.stack(data, axis=1).shape
+    (4, 3, 4)
+
+    >>> da.stack(data, axis=-1).shape
+    (4, 4, 3)
+
+    Result is a new dask Array
+
+    See Also
+    --------
+    concatenate
+    """
+    from dask_array.creation import empty, empty_like
+
+    # Lazy import to avoid circular dependency
+    from dask_array.core import asarray
+
+    seq = [asarray(a, allow_unknown_chunksizes=allow_unknown_chunksizes) for a in seq]
+
+    if not seq:
+        raise ValueError("Need array(s) to stack")
+    if not allow_unknown_chunksizes and not all(x.shape == seq[0].shape for x in seq):
+        idx = first(i for i in enumerate(seq) if i[1].shape != seq[0].shape)
+        raise ValueError(
+            "Stacked arrays must have the same shape. The first array had shape "
+            f"{seq[0].shape}, while array {idx[0] + 1} has shape {idx[1].shape}."
+        )
+
+    ndim = seq[0].ndim
+    if axis < 0:
+        axis = ndim + axis + 1
+    dtype = np.result_type(*[x.dtype for x in seq])
+    meta = meta_from_array(None, ndim=ndim + 1, dtype=dtype)
+    seq = [x.astype(meta.dtype) for x in seq]
+    shape = tuple(
+        (len(seq) if i == axis else (seq[0].shape[i] if i < axis else seq[0].shape[i - 1])) for i in range(meta.ndim)
+    )
+
+    seq2 = [a for a in seq if a.size]
+    if not seq2:
+        seq2 = seq
+
+    n = len(seq2)
+    if n == 0:
+        try:
+            return empty_like(meta, shape=shape, chunks=shape, dtype=meta.dtype)
+        except TypeError:
+            return empty(shape, chunks=shape, dtype=meta.dtype)
+
+    ind = list(range(ndim))
+    uc_args = list(concat((x.expr, ind) for x in seq2))
+    _, seq2, _ = unify_chunks_expr(*uc_args)
+
+    assert len({a.chunks for a in seq2}) == 1  # same chunks
+
+    return new_collection(Stack(seq2[0], axis, meta, *seq2[1:]))

dask_array/_svg.py

@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+import math
+import re
+from functools import lru_cache
+
+import numpy as np
+
+from dask.utils import cached_cumsum
+
+
+@lru_cache(maxsize=512)
+def svg(chunks, size=200, labels=True, **kwargs):
+    """Convert chunks from Dask Array into an SVG Image
+
+    Parameters
+    ----------
+    chunks: tuple
+    size: int
+        Rough size of the image
+    labels: bool
+        Whether to include dimension labels (default True)
+
+    Returns
+    -------
+    text: An svg string depicting the array as a grid of chunks
+    """
+    shape = tuple(map(sum, chunks))
+    if np.isnan(shape).any():  # don't support unknown sizes
+        raise NotImplementedError(
+            "Can't generate SVG with unknown chunk sizes.\n\n A possible solution is with x.compute_chunk_sizes()"
+        )
+    if not all(shape):
+        raise NotImplementedError("Can't generate SVG with 0-length dimensions")
+    if len(chunks) == 0:
+        raise NotImplementedError("Can't generate SVG with 0 dimensions")
+    if len(chunks) == 1:
+        return svg_1d(chunks, size=size, labels=labels, **kwargs)
+    elif len(chunks) == 2:
+        return svg_2d(chunks, size=size, labels=labels, **kwargs)
+    elif len(chunks) == 3:
+        return svg_3d(chunks, size=size, labels=labels, **kwargs)
+    else:
+        return svg_nd(chunks, size=size, labels=labels, **kwargs)
+
+
+# Modern styling constants - orange from Dask logo
+STROKE_COLOR = "#78716c"  # stone-500 - grid lines
+FILL_COLOR = "#fb923c"  # orange-400 - front face
+FILL_COLOR_TOP = "#fdba74"  # orange-300 - top face (lighter)
+FILL_COLOR_DARK = "#ea580c"  # orange-600 - side face / dense arrays
+TEXT_STYLE = 'font-size="1.0rem" font-weight="400" text-anchor="middle" fill="currentColor"'
+
+# SVG definitions for shadow filter
+SVG_DEFS = """  <defs>
+    <filter id="shadow" x="-20%" y="-20%" width="140%" height="140%">
+      <feDropShadow dx="2" dy="2" stdDeviation="2" flood-opacity="0.15"/>
+    </filter>
+  </defs>
+"""
+
+
+def svg_2d(chunks, offset=(0, 0), skew=(0, 0), size=200, sizes=None, face=None, labels=True):
+    shape = tuple(map(sum, chunks))
+    sizes = sizes or draw_sizes(shape, size=size)
+    y, x = grid_points(chunks, sizes)
+
+    lines, (min_x, max_x, min_y, max_y) = svg_grid(x, y, offset=offset, skew=skew, size=size, face=face)
+
+    # Adjust dimensions based on whether labels are shown
+    margin = 50 if labels else 10
+    header = f'<svg width="{int(max_x + margin)}" height="{int(max_y + margin)}" style="stroke:{STROKE_COLOR};stroke-width:1">\n'
+    header += SVG_DEFS
+    footer = "\n</svg>"
+
+    if labels:
+        if shape[0] >= 100:
+            rotate = -90
+        else:
+            rotate = 0
+
+        text = [
+            "",
+            "  <!-- Text -->",
+            f'  <text x="{max_x / 2}" y="{max_y + 20}" {TEXT_STYLE}>{shape[1]}</text>',
+            f'  <text x="{max_x + 20}" y="{max_y / 2}" {TEXT_STYLE} transform="rotate({rotate},{max_x + 20},{max_y / 2})">{shape[0]}</text>',
+        ]
+    else:
+        text = []
+
+    return header + "\n".join(lines + text) + footer
+
+
+def svg_3d(chunks, size=200, sizes=None, offset=(0, 0), labels=True):
+    shape = tuple(map(sum, chunks))
+    sizes = sizes or draw_sizes(shape, size=size)
+    x, y, z = grid_points(chunks, sizes)
+    ox, oy = offset
+
+    # Left face (side) - darker
+    xy, (mnx, mxx, mny, mxy) = svg_grid(x / 1.7, y, offset=(ox + 10, oy + 0), skew=(1, 0), size=size, face="side")
+
+    # Top face - lighter
+    zx, (_, _, _, max_x) = svg_grid(z, x / 1.7, offset=(ox + 10, oy + 0), skew=(0, 1), size=size, face="top")
+
+    # Front face - normal
+    zy, (min_z, max_z, min_y, max_y) = svg_grid(
+        z, y, offset=(ox + max_x + 10, oy + max_x), skew=(0, 0), size=size, face="front"
+    )
+
+    margin = 50 if labels else 10
+    header = f'<svg width="{int(max_z + margin)}" height="{int(max_y + margin)}" style="stroke:{STROKE_COLOR};stroke-width:1">\n'
+    header += SVG_DEFS
+    footer = "\n</svg>"
+
+    if labels:
+        if shape[1] >= 100:
+            rotate = -90
+        else:
+            rotate = 0
+
+        text = [
+            "",
+            "  <!-- Text -->",
+            f'  <text x="{(min_z + max_z) / 2}" y="{max_y + 20}" {TEXT_STYLE}>{shape[2]}</text>',
+            f'  <text x="{max_z + 20}" y="{(min_y + max_y) / 2}" {TEXT_STYLE} transform="rotate({rotate},{max_z + 20},{(min_y + max_y) / 2})">{shape[1]}</text>',
+            f'  <text x="{(mnx + mxx) / 2 - 10}" y="{mxy - (mxx - mnx) / 2 + 20}" {TEXT_STYLE} transform="rotate(45,{(mnx + mxx) / 2 - 10},{mxy - (mxx - mnx) / 2 + 20})">{shape[0]}</text>',
+        ]
+    else:
+        text = []
+
+    return header + "\n".join(xy + zx + zy + text) + footer
+
+
+def svg_nd(chunks, size=200, labels=True):
+    if len(chunks) % 3 == 1:
+        chunks = ((1,),) + chunks
+    shape = tuple(map(sum, chunks))
+    sizes = draw_sizes(shape, size=size)
+
+    chunks2 = chunks
+    sizes2 = sizes
+    out = []
+    left = 0
+    total_height = 0
+    while chunks2:
+        n = len(chunks2) % 3 or 3
+        o = svg(chunks2[:n], sizes=sizes2[:n], offset=(left, 0), labels=labels)
+        chunks2 = chunks2[n:]
+        sizes2 = sizes2[n:]
+
+        lines = o.split("\n")
+        header = lines[0]
+        height = float(re.search(r'height="(\d*\.?\d*)"', header).groups()[0])
+        total_height = max(total_height, height)
+        width = float(re.search(r'width="(\d*\.?\d*)"', header).groups()[0])
+        left += width + 10
+        o = "\n".join(lines[1:-1])  # remove header and footer
+
+        out.append(o)
+
+    header = f'<svg width="{int(left)}" height="{int(total_height)}" style="stroke:{STROKE_COLOR};stroke-width:1">\n'
+    header += SVG_DEFS
+    footer = "\n</svg>"
+    return header + "\n\n".join(out) + footer
+
+
+def svg_lines(x1, y1, x2, y2, max_n=20):
+    """Convert points into lines of text for an SVG plot
+
+    Examples
+    --------
+    >>> svg_lines([0, 1], [0, 0], [10, 11], [1, 1])  # doctest: +NORMALIZE_WHITESPACE
+    ['  <line x1="0" y1="0" x2="10" y2="1" style="stroke-width:2" />',
+     '  <line x1="1" y1="0" x2="11" y2="1" style="stroke-width:2" />']
+    """
+    n = len(x1)
+
+    if n > max_n:
+        indices = np.linspace(0, n - 1, max_n, dtype="int")
+    else:
+        indices = range(n)
+
+    lines = [
+        '  <line x1="{}" y1="{}" x2="{}" y2="{}" />'.format(int(x1[i]), int(y1[i]), int(x2[i]), int(y2[i]))
+        for i in indices
+    ]
+
+    lines[0] = lines[0].replace(" /", ' style="stroke-width:2" /')
+    lines[-1] = lines[-1].replace(" /", ' style="stroke-width:2" /')
+    return lines
+
+
+def svg_grid(x, y, offset=(0, 0), skew=(0, 0), size=200, face=None):
+    """Create lines of SVG text that show a grid
+
+    Parameters
+    ----------
+    x: numpy.ndarray
+    y: numpy.ndarray
+    offset: tuple
+        translational displacement of the grid in SVG coordinates
+    skew: tuple
+    face: str, optional
+        For 3D: "top", "front", or "side" to determine shading
+    """
+    # Horizontal lines
+    x1 = np.zeros_like(y) + offset[0]
+    y1 = y + offset[1]
+    x2 = np.full_like(y, x[-1]) + offset[0]
+    y2 = y + offset[1]
+
+    if skew[0]:
+        y2 += x.max() * skew[0]
+    if skew[1]:
+        x1 += skew[1] * y
+        x2 += skew[1] * y
+
+    min_x = min(x1.min(), x2.min())
+    min_y = min(y1.min(), y2.min())
+    max_x = max(x1.max(), x2.max())
+    max_y = max(y1.max(), y2.max())
+    max_n = size // 6
+
+    h_lines = ["", "  <!-- Horizontal lines -->"] + svg_lines(x1, y1, x2, y2, max_n)
+
+    # Vertical lines
+    x1 = x + offset[0]
+    y1 = np.zeros_like(x) + offset[1]
+    x2 = x + offset[0]
+    y2 = np.full_like(x, y[-1]) + offset[1]
+
+    if skew[0]:
+        y1 += skew[0] * x
+        y2 += skew[0] * x
+    if skew[1]:
+        x2 += skew[1] * y.max()
+
+    v_lines = ["", "  <!-- Vertical lines -->"] + svg_lines(x1, y1, x2, y2, max_n)
+
+    # Determine fill color based on face
+    if face == "top":
+        color = FILL_COLOR_TOP
+    elif face == "side":
+        color = FILL_COLOR_DARK
+    else:
+        color = FILL_COLOR
+
+    corners = f"{x1[0]},{y1[0]} {x1[-1]},{y1[-1]} {x2[-1]},{y2[-1]} {x2[0]},{y2[0]}"
+    rect = [
+        "",
+        "  <!-- Colored Rectangle -->",
+        f'  <polygon points="{corners}" style="fill:{color};fill-opacity:0.7;stroke-width:0" filter="url(#shadow)"/>',
+    ]
+
+    return h_lines + v_lines + rect, (min_x, max_x, min_y, max_y)
+
+
+def svg_1d(chunks, sizes=None, labels=True, **kwargs):
+    return svg_2d(((1,),) + chunks, labels=labels, **kwargs)
+
+
+def grid_points(chunks, sizes):
+    cumchunks = [np.array(cached_cumsum(c, initial_zero=True)) for c in chunks]
+    points = [x * size / x[-1] for x, size in zip(cumchunks, sizes)]
+    return points
+
+
+def draw_sizes(shape, size=200):
+    """Get size in pixels for all dimensions"""
+    mx = max(shape)
+    ratios = [mx / max(0.1, d) for d in shape]
+    ratios = [ratio_response(r) for r in ratios]
+    return tuple(size / r for r in ratios)
+
+
+def ratio_response(x):
+    """How we display actual size ratios
+
+    Common ratios in sizes span several orders of magnitude,
+    which is hard for us to perceive.
+
+    We keep ratios in the 1-3 range accurate, and then apply a logarithm to
+    values up until about 100 or so, at which point we stop scaling.
+    """
+    if x < math.e:
+        return x
+    elif x <= 100:
+        return math.log(x + 12.4)  # f(e) == e
+    else:
+        return math.log(100 + 12.4)

dask_array/_templates.py

@@ -0,0 +1,29 @@
+"""Template loading for array HTML representations."""
+
+from __future__ import annotations
+
+import os.path
+
+from jinja2 import Environment, FileSystemLoader, Template
+from jinja2.exceptions import TemplateNotFound
+
+from dask.utils import typename
+
+FILTERS = {
+    "type": type,
+    "typename": typename,
+}
+
+TEMPLATE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates")
+
+
+def get_template(name: str) -> Template:
+    """Load a Jinja2 template from the templates directory."""
+    loader = FileSystemLoader([TEMPLATE_PATH])
+    environment = Environment(loader=loader)
+    environment.filters.update(FILTERS)
+
+    try:
+        return environment.get_template(name)
+    except TemplateNotFound as e:
+        raise TemplateNotFound(f"Unable to find {name} in {TEMPLATE_PATH}") from e