diff --git a/test_scripts/demo1.py b/test_scripts/demo1.py new file mode 100644 index 0000000..6c40b61 --- /dev/null +++ b/test_scripts/demo1.py @@ -0,0 +1,64 @@ +import numpy as np + + +def one_hot( + indices, + depth, + on_value=None, + off_value=None, + axis=None, + dtype=None, +): + """Convert a vector of indices into a matrix with one-hot encodings as + columns. + + Parameters + ---------- + indices : array_like + Indices to convert to one-hot encodings. + depth : int + Size of one-hot dimension. + on_value : float, optional + Value to fill in output when class is hot, by default 1. + off_value : float, optional + Value to fill in output when class is not hot, by default 0. + axis : int, optional + Axis along which one-hot encodings are added. + dtype : data-type, optional + Data-type of the one-hot matrix. + + Returns + ------- + ret : ndarray + One-hot matrix corresponding to indices. + + Examples + -------- + >>> indices = [0, 1, 2] + >>> one_hot(indices, depth=3) + array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + """ + on_none = on_value is None + off_none = off_value is None + + if dtype is None: + if on_none and off_none: + dtype = np.float32 + else: + if not on_none: + dtype = np.array(on_value).dtype + elif not off_none: + dtype = np.array(off_value).dtype + + res = np.eye(depth, dtype=dtype)[np.array(indices, dtype="int64").reshape(-1)] + res = res.reshape(list(indices.shape) + [depth]) + + if not on_none and not off_none: + res = np.where(res == 1, on_value, off_value) + + if axis is not None: + res = np.moveaxis(res, -1, axis) + + return res