Play around with NumPy

datanerd.py

import numpy as np


# https://numpy.org/doc/stable/user/quickstart.html
def main():
    # a_reshape (3 x 5)
    a_reshape: np.ndarray = np.arange(15).reshape(3, 5)
    assert a_reshape.shape == (3, 5)
    assert a_reshape.ndim == 2
    assert a_reshape.dtype.name == "int64"
    assert a_reshape.size == 15

    # a_i64 (3)
    a_i64 = np.array([2, 3, 4])
    assert a_i64.shape == (3,)
    assert a_i64.ndim == 1
    assert a_i64.dtype.name == "int64"

    # a_f64 (3)
    a_f64 = np.array([1.5, 3.5, 5.5])
    assert a_f64.shape == a_i64.shape
    assert a_f64.ndim == a_i64.ndim
    assert a_f64.dtype.name == "float64"

    # a_nested (1 x 2 x 3)
    a_nested = np.array([[[1, 2, 3], [4, 5, 6]]])
    assert a_nested.shape == (1, 2, 3)
    assert a_nested.ndim == 3
    assert a_nested.dtype.name == "int64"

    # a_complex (2 x 2)
    a_complex = np.array([[1, 2], [3, 4]], dtype=complex)
    assert a_complex.dtype.name == "complex128"

    # a_zeros (50 x 1)
    a_zeros = np.zeros((50, 1), dtype=int)
    assert a_zeros.shape == (50, 1)
    assert a_zeros.ndim == 2
    assert not a_zeros.all()

    # a_ones (50 x 1)
    a_ones = np.ones((50, 1), dtype=int)
    assert a_zeros.shape == (50, 1)
    assert a_zeros.ndim == 2
    assert a_ones.all()

    # element-wise addition (50 x 1)
    assert (a_ones + a_ones == 2).all()

    # element-wise subtraction (50 x 1)
    assert (a_ones - a_ones == 0).all()

    # element-wise multiplication (50 x 1)
    assert (a_ones * a_ones == 1).all()

    # a_twos_t (1 x 50)
    a_twos_t = 2 * np.transpose(a_ones)

    # matrix-wise multiplication (1 x 1)
    assert (a_twos_t @ a_ones == 2 * a_ones.size).all()

    # a_random (3 x 3)
    rg = np.random.default_rng(1)
    a_random = rg.random((3, 3))
    assert (a_random > 0).all() and (a_random < 1).all()

    # a_random can be flattened from a matrix to a set
    assert {int(num) for num in a_random.flat} == {0}

    # a_one_stack (50 x 4)
    a_ones_stack = np.column_stack((a_ones, a_ones))
    a_ones_stack = np.column_stack((a_ones_stack, a_ones_stack))
    assert a_ones_stack.shape == (50, 4)

    # a_twos_t_stack (4 x 50)
    a_twos_t_stack = np.row_stack((a_twos_t, a_twos_t))
    a_twos_t_stack = np.row_stack((a_twos_t_stack, a_twos_t_stack))
    assert a_twos_t_stack.shape == (4, 50)


if __name__ == "__main__":
    main()