def test_ndarray_matmul(): np_v = np.array([1, 2]) np_m = np.array([[1, 2], [3, 4]]) np_r = np.array([[1, 2, 3], [4, 5, 6]]) np_cube = np.arange(8).reshape((2, 2, 2)) np_rect_prism = np.arange(12).reshape((3, 2, 2)) np_broadcasted_mat = np.arange(4).reshape((1, 2, 2)) np_six_dim_tensor = np.arange(3 * 7 * 1 * 9 * 4 * 5).reshape((3, 7, 1, 9, 4, 5)) np_five_dim_tensor = np.arange(7 * 5 * 1 * 5 * 3).reshape((7, 5, 1, 5, 3)) v = hl._nd.array(np_v) m = hl._nd.array(np_m) r = hl._nd.array(np_r) cube = hl._nd.array(np_cube) rect_prism = hl._nd.array(np_rect_prism) broadcasted_mat = hl._nd.array(np_broadcasted_mat) six_dim_tensor = hl._nd.array(np_six_dim_tensor) five_dim_tensor = hl._nd.array(np_five_dim_tensor) assert_ndarrays_eq( (v @ v, np_v @ np_v), (m @ m, np_m @ np_m), (m @ m.T, np_m @ np_m.T), (r @ r.T, np_r @ np_r.T), (v @ m, np_v @ np_m), (m @ v, np_m @ np_v), (cube @ cube, np_cube @ np_cube), (cube @ v, np_cube @ np_v), (v @ cube, np_v @ np_cube), (cube @ m, np_cube @ np_m), (m @ cube, np_m @ np_cube), (rect_prism @ m, np_rect_prism @ np_m), (m @ rect_prism, np_m @ np_rect_prism), (m @ rect_prism.T, np_m @ np_rect_prism.T), (broadcasted_mat @ rect_prism, np_broadcasted_mat @ np_rect_prism), (six_dim_tensor @ five_dim_tensor, np_six_dim_tensor @ np_five_dim_tensor) ) assert hl.eval(hl.null(hl.tndarray(hl.tfloat64, 2)) @ hl.null(hl.tndarray(hl.tfloat64, 2))) is None assert hl.eval(hl.null(hl.tndarray(hl.tint64, 2)) @ hl._nd.array(np.arange(10).reshape(5, 2))) is None assert hl.eval(hl._nd.array(np.arange(10).reshape(5, 2)) @ hl.null(hl.tndarray(hl.tint64, 2))) is None with pytest.raises(ValueError): m @ 5 with pytest.raises(ValueError): m @ hl._nd.array(5) with pytest.raises(ValueError): cube @ hl._nd.array(5) with pytest.raises(FatalError) as exc: hl.eval(r @ r) assert "Matrix dimensions incompatible: 3 2" in str(exc) with pytest.raises(FatalError) as exc: hl.eval(hl._nd.array([1, 2]) @ hl._nd.array([1, 2, 3])) assert "Matrix dimensions incompatible" in str(exc)
def solve_triangular(nd_coef, nd_dep, lower=False): """Solve a triangular linear system. Parameters ---------- nd_coef : :class:`.NDArrayNumericExpression`, (N, N) Triangular coefficient matrix. nd_dep : :class:`.NDArrayNumericExpression`, (N,) or (N, K) Dependent variables. lower : `bool`: If true, nd_coef is interpreted as a lower triangular matrix If false, nd_coef is interpreted as a upper triangular matrix Returns ------- :class:`.NDArrayNumericExpression`, (N,) or (N, K) Solution to the triangular system Ax = B. Shape is same as shape of B. """ nd_dep_ndim_orig = nd_dep.ndim nd_coef, nd_dep = solve_helper(nd_coef, nd_dep, nd_dep_ndim_orig) return_type = hl.tndarray(hl.tfloat64, 2) ir = Apply("linear_triangular_solve", return_type, nd_coef._ir, nd_dep._ir, lower._ir) result = construct_expr(ir, return_type, nd_coef._indices, nd_coef._aggregations) if nd_dep_ndim_orig == 1: result = result.reshape((-1)) return result
def test_ndarray_transpose(): np_v = np.array([1, 2, 3]) np_m = np.array([[1, 2, 3], [4, 5, 6]]) np_cube = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) v = hl.nd.array(np_v) m = hl.nd.array(np_m) cube = hl.nd.array(np_cube) assert_ndarrays_eq((v.T, np_v.T), (v.T, np_v), (m.T, np_m.T), (cube.transpose((0, 2, 1)), np_cube.transpose( (0, 2, 1))), (cube.T, np_cube.T)) assert hl.eval(hl.null(hl.tndarray(hl.tfloat, 1)).T) is None with pytest.raises(ValueError) as exc: v.transpose((1, )) assert "Invalid axis: 1" in str(exc.value) with pytest.raises(ValueError) as exc: cube.transpose((1, 1)) assert "Expected 3 axes, got 2" in str(exc.value) with pytest.raises(ValueError) as exc: cube.transpose((1, 1, 1)) assert "Axes cannot contain duplicates" in str(exc.value)
def test_loop_errors(self): with pytest.raises( TypeError, match= "requested type ndarray<int32, 2> does not match inferred type ndarray<float64, 2>" ): result = hl.experimental.loop( lambda f, my_nd: hl.if_else(my_nd[0, 0] == 1000, my_nd, f(my_nd + 1)), hl.tndarray(hl.tint32, 2), hl.nd.zeros((20, 10), hl.tfloat64))
def test_ndarray_map(): a = hl.nd.array([[2, 3, 4], [5, 6, 7]]) b = hl.map(lambda x: -x, a) c = hl.map(lambda x: True, a) assert_ndarrays_eq((b, [[-2, -3, -4], [-5, -6, -7]]), (c, [[True, True, True], [True, True, True]])) assert hl.eval(hl.null(hl.tndarray(hl.tfloat, 1)).map(lambda x: x * 2)) is None
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, hl.ttuple(hl.tndarray(hl.tfloat64, 2), hl.tndarray(hl.tfloat64, 1))) elif mode == "r": return construct_expr(ir, hl.tndarray(hl.tfloat64, 2)) elif mode in ["complete", "reduced"]: return construct_expr( ir, hl.ttuple(hl.tndarray(hl.tfloat64, 2), hl.tndarray(hl.tfloat64, 2)))
def ndarray_floating_point_divide(arg_type, ret_type): register_function("/", ( arg_type, hl.tndarray(arg_type, NatVariable()), ), hl.tndarray(ret_type, NatVariable())) register_function("/", (hl.tndarray(arg_type, NatVariable()), arg_type), hl.tndarray(ret_type, NatVariable())) register_function("/", (hl.tndarray( arg_type, NatVariable()), hl.tndarray(arg_type, NatVariable())), hl.tndarray(ret_type, NatVariable()))
def test_ndarray_mixed(): assert hl.eval( hl.null(hl.tndarray(hl.tint64, 2)).map(lambda x: x * x).reshape( (4, 5)).T) is None assert hl.eval((hl.nd.zeros((5, 10)).map(lambda x: x - 2) + hl.nd.ones( (5, 10)).map(lambda x: x + 5)).reshape( hl.null(hl.ttuple(hl.tint64, hl.tint64))).T.reshape( (10, 5))) is None assert hl.eval( hl.or_missing( False, hl.nd.array(np.arange(10)).reshape( (5, 2)).map(lambda x: x * 2)).map(lambda y: y * 2)) is None
def test_parses(self): env = {'c': hl.tbool, 'a': hl.tarray(hl.tint32), 'aa': hl.tarray(hl.tarray(hl.tint32)), 'da': hl.tarray(hl.ttuple(hl.tint32, hl.tstr)), 'nd': hl.tndarray(hl.tfloat64, 1), 'v': hl.tint32, 's': hl.tstruct(x=hl.tint32, y=hl.tint64, z=hl.tfloat64), 't': hl.ttuple(hl.tint32, hl.tint64, hl.tfloat64), 'call': hl.tcall, 'x': hl.tint32} env = {name: t._parsable_string() for name, t in env.items()} for x in self.value_irs(): Env.hail().expr.ir.IRParser.parse_value_ir(str(x), env, {})
def solve(a, b): """Solve a linear system. Parameters ---------- a : :class:`.NDArrayNumericExpression`, (N, N) Coefficient matrix. b : :class:`.NDArrayNumericExpression`, (N,) or (N, K) Dependent variables. Returns ------- :class:`.NDArrayNumericExpression`, (N,) or (N, K) Solution to the system Ax = B. Shape is same as shape of B. """ assert a.ndim == 2 assert b.ndim == 1 or b.ndim == 2 b_ndim_orig = b.ndim if b_ndim_orig == 1: b = b.reshape((-1, 1)) if a.dtype.element_type != hl.tfloat64: a = a.map(lambda e: hl.float64(e)) if b.dtype.element_type != hl.tfloat64: b = b.map(lambda e: hl.float64(e)) ir = Apply("linear_solve", hl.tndarray(hl.tfloat64, 2), a._ir, b._ir) result = construct_expr(ir, hl.tndarray(hl.tfloat64, 2), a._indices, a._aggregations) if b_ndim_orig == 1: result = result.reshape((-1)) return result
def solve(a, b, no_crash=False): """Solve a linear system. Parameters ---------- a : :class:`.NDArrayNumericExpression`, (N, N) Coefficient matrix. b : :class:`.NDArrayNumericExpression`, (N,) or (N, K) Dependent variables. Returns ------- :class:`.NDArrayNumericExpression`, (N,) or (N, K) Solution to the system Ax = B. Shape is same as shape of B. """ b_ndim_orig = b.ndim a, b = solve_helper(a, b, b_ndim_orig) if no_crash: name = "linear_solve_no_crash" return_type = hl.tstruct(solution=hl.tndarray(hl.tfloat64, 2), failed=hl.tbool) else: name = "linear_solve" return_type = hl.tndarray(hl.tfloat64, 2) ir = Apply(name, return_type, a._ir, b._ir) result = construct_expr(ir, return_type, a._indices, a._aggregations) if b_ndim_orig == 1: if no_crash: result = hl.struct(solution=result.solution.reshape((-1)), failed=result.failed) else: result = result.reshape((-1)) return result
def test_parses(self): env = {'c': hl.tbool, 'a': hl.tarray(hl.tint32), 'st': hl.tstream(hl.tint32), 'aa': hl.tarray(hl.tarray(hl.tint32)), 'sta': hl.tstream(hl.tarray(hl.tint32)), 'da': hl.tarray(hl.ttuple(hl.tint32, hl.tstr)), 'nd': hl.tndarray(hl.tfloat64, 1), 'v': hl.tint32, 's': hl.tstruct(x=hl.tint32, y=hl.tint64, z=hl.tfloat64), 't': hl.ttuple(hl.tint32, hl.tint64, hl.tfloat64), 'call': hl.tcall, 'x': hl.tint32} for x in self.value_irs(): Env.spark_backend('ValueIRTests.test_parses')._parse_value_ir(str(x), env)
def test_ndarray_matmul(): np_v = np.array([1, 2]) np_y = np.array([1, 1, 1]) np_m = np.array([[1, 2], [3, 4]]) np_m_f32 = np_m.astype(np.float32) np_m_f64 = np_m.astype(np.float64) np_r = np.array([[1, 2, 3], [4, 5, 6]]) np_r_f32 = np_r.astype(np.float32) np_r_f64 = np_r.astype(np.float64) np_cube = np.arange(8).reshape((2, 2, 2)) np_rect_prism = np.arange(12).reshape((3, 2, 2)) np_broadcasted_mat = np.arange(4).reshape((1, 2, 2)) np_six_dim_tensor = np.arange(3 * 7 * 1 * 9 * 4 * 5).reshape( (3, 7, 1, 9, 4, 5)) np_five_dim_tensor = np.arange(7 * 5 * 1 * 5 * 3).reshape((7, 5, 1, 5, 3)) np_ones_int32 = np.ones((4, 4), dtype=np.int32) np_ones_float64 = np.ones((4, 4), dtype=np.float64) np_zero_by_four = np.array([], dtype=np.float64).reshape((0, 4)) v = hl.nd.array(np_v) y = hl.nd.array(np_y) m = hl.nd.array(np_m) m_f32 = hl.nd.array(np_m_f32) m_f64 = hl.nd.array(np_m_f64) r = hl.nd.array(np_r) r_f32 = hl.nd.array(np_r_f32) r_f64 = hl.nd.array(np_r_f64) cube = hl.nd.array(np_cube) rect_prism = hl.nd.array(np_rect_prism) broadcasted_mat = hl.nd.array(np_broadcasted_mat) six_dim_tensor = hl.nd.array(np_six_dim_tensor) five_dim_tensor = hl.nd.array(np_five_dim_tensor) ones_int32 = hl.nd.array(np_ones_int32) ones_float64 = hl.nd.array(np_ones_float64) zero_by_four = hl.nd.array(np_zero_by_four) assert_ndarrays_eq( (v @ v, np_v @ np_v), (m @ m, np_m @ np_m), (m_f32 @ m_f32, np_m_f32 @ np_m_f32), (m_f64 @ m_f64, np_m_f64 @ np_m_f64), (m @ m.T, np_m @ np_m.T), (m_f64 @ m_f64.T, np_m_f64 @ np_m_f64.T), (r @ r.T, np_r @ np_r.T), (r_f32 @ r_f32.T, np_r_f32 @ np_r_f32.T), (r_f64 @ r_f64.T, np_r_f64 @ np_r_f64.T), (v @ m, np_v @ np_m), (m @ v, np_m @ np_v), (v @ r, np_v @ np_r), (r @ y, np_r @ np_y), (cube @ cube, np_cube @ np_cube), (cube @ v, np_cube @ np_v), (v @ cube, np_v @ np_cube), (cube @ m, np_cube @ np_m), (m @ cube, np_m @ np_cube), (rect_prism @ m, np_rect_prism @ np_m), (m @ rect_prism, np_m @ np_rect_prism), (m @ rect_prism.T, np_m @ np_rect_prism.T), (broadcasted_mat @ rect_prism, np_broadcasted_mat @ np_rect_prism), (six_dim_tensor @ five_dim_tensor, np_six_dim_tensor @ np_five_dim_tensor), (zero_by_four @ ones_float64, np_zero_by_four, np_ones_float64), (zero_by_four.transpose() @ zero_by_four, np_zero_by_four.transpose() @ np_zero_by_four)) assert hl.eval( hl.null(hl.tndarray(hl.tfloat64, 2)) @ hl.null( hl.tndarray(hl.tfloat64, 2))) is None assert hl.eval( hl.null(hl.tndarray(hl.tint64, 2)) @ hl.nd.array( np.arange(10).reshape(5, 2))) is None assert hl.eval( hl.nd.array(np.arange(10).reshape(5, 2)) @ hl.null( hl.tndarray(hl.tint64, 2))) is None assert np.array_equal(hl.eval(ones_int32 @ ones_float64), np_ones_int32 @ np_ones_float64) with pytest.raises(ValueError): m @ 5 with pytest.raises(ValueError): m @ hl.nd.array(5) with pytest.raises(ValueError): cube @ hl.nd.array(5) with pytest.raises(FatalError) as exc: hl.eval(r @ r) assert "Matrix dimensions incompatible: 3 2" in str(exc) with pytest.raises(FatalError) as exc: hl.eval(hl.nd.array([1, 2]) @ hl.nd.array([1, 2, 3])) assert "Matrix dimensions incompatible" in str(exc)
def test_ndarray_map2(): a = 2.0 b = 3.0 x = np.array([a, b]) y = np.array([b, a]) row_vec = np.array([[1, 2]]) cube1 = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) cube2 = np.array([[[9, 10], [11, 12]], [[13, 14], [15, 16]]]) na = hl.nd.array(a) nx = hl.nd.array(x) ny = hl.nd.array(y) nrow_vec = hl.nd.array(row_vec) ncube1 = hl.nd.array(cube1) ncube2 = hl.nd.array(cube2) assert_ndarrays_eq( # with lists/numerics (na + b, np.array(a + b)), (b + na, np.array(a + b)), (nx + y, x + y), (ncube1 + cube2, cube1 + cube2), # Addition (na + na, np.array(a + a)), (nx + ny, x + y), (ncube1 + ncube2, cube1 + cube2), # Broadcasting (ncube1 + na, cube1 + a), (na + ncube1, a + cube1), (ncube1 + ny, cube1 + y), (ny + ncube1, y + cube1), (nrow_vec + ncube1, row_vec + cube1), (ncube1 + nrow_vec, cube1 + row_vec), # Subtraction (na - na, np.array(a - a)), (nx - nx, x - x), (ncube1 - ncube2, cube1 - cube2), # Broadcasting (ncube1 - na, cube1 - a), (na - ncube1, a - cube1), (ncube1 - ny, cube1 - y), (ny - ncube1, y - cube1), (ncube1 - nrow_vec, cube1 - row_vec), (nrow_vec - ncube1, row_vec - cube1), # Multiplication (na * na, np.array(a * a)), (nx * nx, x * x), (nx * na, x * a), (na * nx, a * x), (ncube1 * ncube2, cube1 * cube2), # Broadcasting (ncube1 * na, cube1 * a), (na * ncube1, a * cube1), (ncube1 * ny, cube1 * y), (ny * ncube1, y * cube1), (ncube1 * nrow_vec, cube1 * row_vec), (nrow_vec * ncube1, row_vec * cube1), # Floor div (na // na, np.array(a // a)), (nx // nx, x // x), (nx // na, x // a), (na // nx, a // x), (ncube1 // ncube2, cube1 // cube2), # Broadcasting (ncube1 // na, cube1 // a), (na // ncube1, a // cube1), (ncube1 // ny, cube1 // y), (ny // ncube1, y // cube1), (ncube1 // nrow_vec, cube1 // row_vec), (nrow_vec // ncube1, row_vec // cube1)) # Division assert_ndarrays_almost_eq( (na / na, np.array(a / a)), (nx / nx, x / x), (nx / na, x / a), (na / nx, a / x), (ncube1 / ncube2, cube1 / cube2), # Broadcasting (ncube1 / na, cube1 / a), (na / ncube1, a / cube1), (ncube1 / ny, cube1 / y), (ny / ncube1, y / cube1), (ncube1 / nrow_vec, cube1 / row_vec), (nrow_vec / ncube1, row_vec / cube1)) # Missingness tests missing = hl.null(hl.tndarray(hl.tfloat64, 2)) present = hl.nd.array(np.arange(10).reshape(5, 2)) assert hl.eval(missing + missing) is None assert hl.eval(missing + present) is None assert hl.eval(present + missing) is None
def ndarray_floating_point_divide(arg_type, ret_type): register_function("/", (arg_type, hl.tndarray(arg_type, NatVariable()),), hl.tndarray(ret_type, NatVariable())) register_function("/", (hl.tndarray(arg_type, NatVariable()), arg_type), hl.tndarray(ret_type, NatVariable())) register_function("/", (hl.tndarray(arg_type, NatVariable()), hl.tndarray(arg_type, NatVariable())), hl.tndarray(ret_type, NatVariable()))
def test_ndarray_reshape(): np_single = np.array([8]) single = hl.nd.array([8]) np_zero_dim = np.array(4) zero_dim = hl.nd.array(4) np_a = np.array([1, 2, 3, 4, 5, 6]) a = hl.nd.array(np_a) np_cube = np.array([0, 1, 2, 3, 4, 5, 6, 7]).reshape((2, 2, 2)) cube = hl.nd.array([0, 1, 2, 3, 4, 5, 6, 7]).reshape((2, 2, 2)) cube_to_rect = cube.reshape((2, 4)) np_cube_to_rect = np_cube.reshape((2, 4)) cube_t_to_rect = cube.transpose((1, 0, 2)).reshape((2, 4)) np_cube_t_to_rect = np_cube.transpose((1, 0, 2)).reshape((2, 4)) np_hypercube = np.arange(3 * 5 * 7 * 9).reshape((3, 5, 7, 9)) hypercube = hl.nd.array(np_hypercube) np_shape_zero = np.array([]) shape_zero = hl.nd.array(np_shape_zero) assert_ndarrays_eq((single.reshape(()), np_single.reshape( ())), (zero_dim.reshape(()), np_zero_dim.reshape( ())), (zero_dim.reshape((1, )), np_zero_dim.reshape( (1, ))), (a.reshape((6, )), np_a.reshape((6, ))), (a.reshape( (2, 3)), np_a.reshape((2, 3))), (a.reshape( (3, 2)), np_a.reshape((3, 2))), (a.reshape( (3, -1)), np_a.reshape((3, -1))), (a.reshape( (-1, 2)), np_a.reshape( (-1, 2))), (cube_to_rect, np_cube_to_rect), (cube_t_to_rect, np_cube_t_to_rect), (hypercube.reshape( (5, 7, 9, 3)).reshape( (7, 9, 3, 5)), np_hypercube.reshape( (7, 9, 3, 5))), (hypercube.reshape(hl.tuple( [5, 7, 9, 3])), np_hypercube.reshape( (5, 7, 9, 3))), (shape_zero.reshape( (0, 5)), np_shape_zero.reshape((0, 5))), (shape_zero.reshape( (-1, 5)), np_shape_zero.reshape((-1, 5)))) assert hl.eval(hl.null(hl.tndarray(hl.tfloat, 2)).reshape((4, 5))) is None assert hl.eval( hl.nd.array(hl.range(20)).reshape( hl.null(hl.ttuple(hl.tint64, hl.tint64)))) is None with pytest.raises(FatalError) as exc: hl.eval(hl.literal(np_cube).reshape((-1, -1))) assert "more than one -1" in str(exc) with pytest.raises(FatalError) as exc: hl.eval(hl.literal(np_cube).reshape((20, ))) assert "requested shape is incompatible with number of elements" in str( exc) with pytest.raises(FatalError) as exc: hl.eval(a.reshape((3, ))) assert "requested shape is incompatible with number of elements" in str( exc) with pytest.raises(FatalError) as exc: hl.eval(a.reshape(())) assert "requested shape is incompatible with number of elements" in str( exc) with pytest.raises(FatalError) as exc: hl.eval(hl.literal(np_cube).reshape((0, 2, 2))) assert "requested shape is incompatible with number of elements" in str( exc) with pytest.raises(FatalError) as exc: hl.eval(hl.literal(np_cube).reshape((2, 2, -2))) assert "must contain only nonnegative numbers or -1" in str(exc) with pytest.raises(FatalError) as exc: hl.eval(shape_zero.reshape((0, -1))) assert "Can't reshape" in str(exc)
def visit_ndarray(self, node, visited_children): tndarray, _, angle_bracket, elem_t, comma, ndim, angle_bracket = visited_children return hl.tndarray(elem_t, ndim)
def visit_ndarray(self, node, visited_children): tndarray, _, angle_bracket, t, angle_bracket = visited_children return hl.tndarray(t)