def __init__(self, a, N=2): arrays = a if all(isinstance(x, _hfarray) for x in a): arrays = make_same_dims_list(arrays) else: raise Exception("Can only broadcast hfarrays") try: N[0] except TypeError: N = [N] * len(a) self.Nlist = Nlist = N _matrixshapes = [get_shape_helper(x.shape, Nelem) for x, Nelem in zip(arrays, Nlist)] _matrixstrides = [get_shape_helper(x.strides, Nelem) for x, Nelem in zip(arrays, Nlist)] self._matrixshapes = _matrixshapes self._matrixstrides = _matrixstrides dimslist = [x.dims for x, Nelem in zip(arrays, Nlist)] firstelems = broadcast_arrays(*[firstpos(x, Nelem) for x, Nelem in zip(arrays, Nlist)]) self._broadcasted = broadcasted = [] for o, endshapes, endstrides, dims in zip(firstelems, _matrixshapes, _matrixstrides, dimslist): x = as_strided(o, o.shape + endshapes, o.strides + endstrides) broadcasted.append(hfarray(x, dims=dims, copy=False)) self.outershape = broadcasted[0].shape[:-Nlist[0]]
def make_matrix(a, b, c, d): a, b, c, d = make_same_dims_list([a, b, c, d]) abcdshape = zip(a.shape, b.shape, c.shape, d.shape) maxshape = (tuple(max(x) for x in abcdshape) + (2, 2)) res = zeros(maxshape, a.dtype) res[..., 0, 0] = a res[..., 0, 1] = b res[..., 1, 0] = c res[..., 1, 1] = d dims = a.dims + (DimMatrix_i("i", 2), DimMatrix_j("j", 2),) out = hfarray(res, dims=dims) return out