Source code for hail.nd.nd

from functools import reduce

import hail as hl
from hail.expr.functions import _ndarray
from hail.expr.functions import array as aarray
from hail.expr.types import HailType, tfloat64, ttuple, tndarray
from hail.typecheck import typecheck, nullable, oneof, tupleof, sequenceof
from hail.expr.expressions import (
    expr_int32, expr_int64, expr_tuple, expr_any, expr_array, expr_ndarray,
    expr_numeric, Int64Expression, cast_expr, construct_expr)
from hail.expr.expressions.typed_expressions import NDArrayNumericExpression
from hail.ir import NDArrayQR, NDArrayInv, NDArrayConcat

tsequenceof_nd = oneof(sequenceof(expr_ndarray()), expr_array(expr_ndarray()))
shape_type = oneof(expr_int64, tupleof(expr_int64), expr_tuple())


[docs]def array(input_array, dtype=None): """Construct an :class:`.NDArrayExpression` Examples -------- >>> hl.eval(hl.nd.array([1, 2, 3, 4])) array([1, 2, 3, 4], dtype=int32) >>> hl.eval(hl.nd.array([[1, 2, 3], [4, 5, 6]])) array([[1, 2, 3], [4, 5, 6]], dtype=int32) >>> hl.eval(hl.nd.array(np.identity(3))) array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]) >>> hl.eval(hl.nd.array(hl.range(10, 20))) array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=int32) Parameters ---------- input_array : :class:`.ArrayExpression`, numpy ndarray, or nested python lists/tuples dtype : :class:`.HailType` Desired hail type. Default: `float64`. Returns ------- :class:`.NDArrayExpression` An ndarray based on the input array. """ return _ndarray(input_array, dtype=dtype)
@typecheck(a=expr_array(), shape=shape_type) def from_column_major(a, shape): assert len(shape) == 2 return array(a).reshape(tuple(reversed(shape))).T
[docs]@typecheck(start=expr_int32, stop=nullable(expr_int32), step=expr_int32) def arange(start, stop=None, step=1) -> NDArrayNumericExpression: """Returns a 1-dimensions ndarray of integers from `start` to `stop` by `step`. Examples -------- >>> hl.eval(hl.nd.arange(10)) array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32) >>> hl.eval(hl.nd.arange(3, 10)) array([3, 4, 5, 6, 7, 8, 9], dtype=int32) >>> hl.eval(hl.nd.arange(0, 10, step=3)) array([0, 3, 6, 9], dtype=int32) Notes ----- The range includes `start`, but excludes `stop`. If provided exactly one argument, the argument is interpreted as `stop` and `start` is set to zero. This matches the behavior of Python's ``range``. Parameters ---------- start : int or :class:`.Expression` of type :py:data:`.tint32` Start of range. stop : int or :class:`.Expression` of type :py:data:`.tint32` End of range. step : int or :class:`.Expression` of type :py:data:`.tint32` Step of range. Returns ------- :class:`.NDArrayNumericExpression` A 1-dimensional ndarray from `start` to `stop` by `step`. """ return array(hl.range(start, stop, step))
[docs]@typecheck(shape=shape_type, value=expr_any, dtype=nullable(HailType)) def full(shape, value, dtype=None): """Creates a hail :class:`.NDArrayNumericExpression` full of the specified value. Examples -------- Create a 5 by 7 NDArray of type :py:data:`.tfloat64` 9s. >>> hl.nd.full((5, 7), 9) It is possible to specify a type other than :py:data:`.tfloat64` with the `dtype` argument. >>> hl.nd.full((5, 7), 9, dtype=hl.tint32) Parameters ---------- shape : `tuple` or :class:`.TupleExpression` Desired shape. value : :class:`.Expression` or python value Value to fill ndarray with. dtype : :class:`.HailType` Desired hail type. Returns ------- :class:`.NDArrayNumericExpression` An ndarray of the specified shape filled with the specified value. """ if isinstance(shape, Int64Expression): shape_product = shape else: shape_product = reduce(lambda a, b: a * b, shape) return arange(hl.int32(shape_product)).map(lambda x: cast_expr(value, dtype)).reshape(shape)
[docs]@typecheck(shape=shape_type, dtype=HailType) def zeros(shape, dtype=tfloat64): """Creates a hail :class:`.NDArrayNumericExpression` full of zeros. Examples -------- Create a 5 by 7 NDArray of type :py:data:`.tfloat64` zeros. >>> hl.nd.zeros((5, 7)) It is possible to specify a type other than :py:data:`.tfloat64` with the `dtype` argument. >>> hl.nd.zeros((5, 7), dtype=hl.tfloat32) Parameters ---------- shape : `tuple` or :class:`.TupleExpression` Desired shape. dtype : :class:`.HailType` Desired hail type. Default: `float64`. See Also -------- :func:`.full` Returns ------- :class:`.NDArrayNumericExpression` ndarray of the specified size full of zeros. """ return full(shape, 0, dtype)
[docs]@typecheck(shape=shape_type, dtype=HailType) def ones(shape, dtype=tfloat64): """Creates a hail :class:`.NDArrayNumericExpression` full of ones. Examples -------- Create a 5 by 7 NDArray of type :py:data:`.tfloat64` ones. >>> hl.nd.ones((5, 7)) It is possible to specify a type other than :py:data:`.tfloat64` with the `dtype` argument. >>> hl.nd.ones((5, 7), dtype=hl.tfloat32) Parameters ---------- shape : `tuple` or :class:`.TupleExpression` Desired shape. dtype : :class:`.HailType` Desired hail type. Default: `float64`. See Also -------- :func:`.full` Returns ------- :class:`.NDArrayNumericExpression` ndarray of the specified size full of ones. """ return full(shape, 1, dtype)
@typecheck(nd=expr_ndarray()) def diagonal(nd): """Gets the diagonal of a 2 dimensional NDArray. Examples -------- >>> hl.eval(hl.nd.diagonal(hl.nd.array([[1, 2], [3, 4]]))) array([1, 4], dtype=int32) :param nd: A 2 dimensional NDArray, shape(M, N). :return: A 1 dimension NDArray of length min (M, N), containing the diagonal of `nd`. """ assert nd.ndim == 2, "diagonal requires 2 dimensional ndarray" shape_min = hl.min(nd.shape[0], nd.shape[1]) return hl.nd.array(hl.range(hl.int32(shape_min)).map(lambda i: nd[i, i]))
[docs]@typecheck(nd=expr_ndarray(), mode=str) def qr(nd, mode="reduced"): """Performs a QR decomposition. :param nd: A 2 dimensional ndarray, shape(M, N) :param mode: One of "reduced", "complete", "r", or "raw". If K = min(M, N), then: - `reduced`: returns q and r with dimensions (M, K), (K, N) - `complete`: returns q and r with dimensions (M, M), (M, N) - `r`: returns only r with dimensions (K, N) - `raw`: returns h, tau with dimensions (N, M), (K,) Returns ------- - q: ndarray of float64 A matrix with orthonormal columns. - r: ndarray of float64 The upper-triangular matrix R. - (h, tau): ndarrays of float64 The array h contains the Householder reflectors that generate q along with r. The tau array contains scaling factors for the reflectors """ assert nd.ndim == 2, "QR decomposition requires 2 dimensional ndarray" if mode not in ["reduced", "r", "raw", "complete"]: raise ValueError(f"Unrecognized mode '{mode}' for QR decomposition") float_nd = nd.map(lambda x: hl.float64(x)) ir = NDArrayQR(float_nd._ir, mode) if mode == "raw": return construct_expr(ir, ttuple(tndarray(tfloat64, 2), tndarray(tfloat64, 1))) elif mode == "r": return construct_expr(ir, tndarray(tfloat64, 2)) elif mode in ["complete", "reduced"]: return construct_expr(ir, ttuple(tndarray(tfloat64, 2), tndarray(tfloat64, 2)))
[docs]@typecheck(nd=expr_ndarray()) def inv(nd): """Performs a matrix inversion. :param nd: A 2 dimensional ndarray, shape(M, N) Returns ------- - a: ndarray of float64 The inverted matrix """ assert nd.ndim == 2, "Matrix inversion requires 2 dimensional ndarray" float_nd = nd.map(lambda x: hl.float64(x)) ir = NDArrayInv(float_nd._ir) return construct_expr(ir, tndarray(tfloat64, 2))
[docs]@typecheck(nds=tsequenceof_nd, axis=int) def concatenate(nds, axis=0): """Join a sequence of arrays along an existing axis. Examples -------- >>> x = hl.nd.array([[1., 2.], [3., 4.]]) >>> y = hl.nd.array([[5.], [6.]]) >>> hl.eval(hl.nd.concatenate([x, y], axis=1)) array([[1., 2., 5.], [3., 4., 6.]]) >>> x = hl.nd.array([1., 2.]) >>> y = hl.nd.array([3., 4.]) >>> hl.eval(hl.nd.concatenate((x, y), axis=0)) array([1., 2., 3., 4.]) Parameters ---------- :param nds: a1, a2, …sequence of array_like The arrays must have the same shape, except in the dimension corresponding to axis (the first, by default). Note: unlike Numpy, the numerical element type of each array_like must match. :param axis: int, optional The axis along which the arrays will be joined. Default is 0. Note: unlike Numpy, if provided, axis cannot be None. Returns ------- - res: ndarray The concatenated array """ head_nd = nds[0] head_ndim = head_nd.ndim hl.case().when(hl.all(lambda a: a.ndim == head_ndim, nds), True).or_error("Mismatched ndim") makearr = aarray(nds) concat_ir = NDArrayConcat(makearr._ir, axis) return construct_expr(concat_ir, tndarray(head_nd._type.element_type, head_ndim))
[docs]@typecheck(N=expr_numeric, M=nullable(expr_numeric), dtype=HailType) def eye(N, M=None, dtype=hl.tfloat64): """ Construct a 2-D :class:`.NDArrayExpression` with ones on the *main* diagonal and zeros elsewhere. Parameters ---------- N : :class:`.NumericExpression` or Python number Number of rows in the output. M : :class:`.NumericExpression` or Python number, optional Number of columns in the output. If None, defaults to `N`. dtype : numeric :class:`.HailType`, optional Element type of the returned array. Defaults to :py:data:`.tfloat64` Returns ------- I : :class:`.NDArrayExpression` representing a Hail ndarray of shape (N,M) An ndarray whose elements are equal to one on the main diagonal, zeroes elsewhere. See Also -------- :func:`.identity` :func:`.diagonal` Examples -------- >>> hl.eval(hl.nd.eye(3)) array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]) >>> hl.eval(hl.nd.eye(2, 5, dtype=hl.tint32)) array([[1, 0, 0, 0, 0], [0, 1, 0, 0, 0]], dtype=int32) """ n_row = hl.int32(N) if M is None: n_col = n_row else: n_col = hl.int32(M) return hl.nd.array(hl.range(0, n_row * n_col).map( lambda i: hl.cond((i // n_col) == (i % n_col), hl.literal(1, dtype), hl.literal(0, dtype)) )).reshape((n_row, n_col))
[docs]@typecheck(N=expr_numeric, dtype=HailType) def identity(N, dtype=hl.tfloat64): """ Constructs a 2-D :class:`.NDArrayExpression` representing the identity array. The identity array is a square array with ones on the main diagonal. Parameters ---------- n : :class:`.NumericExpression` or Python number Number of rows and columns in the output. dtype : numeric :class:`.HailType`, optional Element type of the returned array. Defaults to :py:data:`.tfloat64` Returns ------- out : :class:`.NDArrayExpression` `n` x `n` ndarray with its main diagonal set to one, and all other elements 0. See Also -------- :func:`.eye` Examples -------- >>> hl.eval(hl.nd.identity(3)) array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]) """ return eye(N, dtype=dtype)
[docs]@typecheck(arrs=tsequenceof_nd) def vstack(arrs): """ Stack arrays in sequence vertically (row wise). 1-D arrays of shape `(N,)`, will reshaped to `(1,N)` before concatenation. For all other arrays, equivalent to :func:`.concatenate` with axis=0. Parameters ---------- arrs : sequence of :class:`.NDArrayExpression` The arrays must have the same shape along all but the first axis. 1-D arrays must have the same length. Returns ------- stacked : :class:`.NDArrayExpression` The array formed by stacking the given arrays, will be at least 2-D. See Also -------- :func:`.concatenate` : Join a sequence of arrays along an existing axis. Examples -------- >>> a = hl.nd.array([1, 2, 3]) >>> b = hl.nd.array([2, 3, 4]) >>> hl.eval(hl.nd.vstack((a,b))) array([[1, 2, 3], [2, 3, 4]], dtype=int32) >>> a = hl.nd.array([[1], [2], [3]]) >>> b = hl.nd.array([[2], [3], [4]]) >>> hl.eval(hl.nd.vstack((a,b))) array([[1], [2], [3], [2], [3], [4]], dtype=int32) """ head_ndim = arrs[0].ndim if head_ndim == 1: return concatenate(hl.map(lambda a: a._broadcast(2), arrs), 0) return concatenate(arrs, 0)
@typecheck(arrs=tsequenceof_nd) def hstack(arrs): """ Stack arrays in sequence horizontally (column wise). Equivalent to concatenation along the second axis, except for 1-D arrays where it concatenates along the first axis. This function makes most sense for arrays with up to 3 dimensions. :func:`.concatenate` provides more general stacking and concatenation operations. Parameters ---------- tup : sequence of :class:`.NDArrayExpression` The arrays must have the same shape along all but the second axis, except 1-D arrays which can be any length. Returns ------- stacked : :class:`.NDArrayExpression` The array formed by stacking the given arrays. See Also -------- :func:`.concatenate` :func:`.vstack` Examples -------- >>> a = hl.nd.array([1,2,3]) >>> b = hl.nd.array([2,3,4]) >>> hl.eval(hl.nd.hstack((a,b))) array([1, 2, 3, 2, 3, 4], dtype=int32) >>> a = hl.nd.array([[1],[2],[3]]) >>> b = hl.nd.array([[2],[3],[4]]) >>> hl.eval(hl.nd.hstack((a,b))) array([[1, 2], [2, 3], [3, 4]], dtype=int32) """ head_ndim = arrs[0].ndim if head_ndim == 1: axis = 0 else: axis = 1 return concatenate(arrs, axis)