def sum(P, axis=None): if not P.A: if axis is None: return Poly({}, P.dim, (), P.dtype) shape = P.shape[:axis]+P.shape[axis+1:] return Poly({}, P.dim, shape, P.dtype) if isinstance(axis, int): l = len(P.shape) if axis<0: axis += l shape = [0]*(l-1) for i in range(l): if i<axis: shape[i] = P.shape[i] elif i>axis: shape[i-1] = P.shape[i] shape = tuple(shape) else: shape = () A = P.A if P.dtype==f.frac: for key in P.keys: A[key] = f.sum(A[key], axis) else: for key in P.keys: A[key] = np.sum(A[key], axis) return Poly(A, P.dim, shape, P.dtype)
def sum(P, axis=None): if not P.A: if axis is None: return Poly({}, P.dim, (), P.dtype) shape = P.shape[:axis] + P.shape[axis + 1:] return Poly({}, P.dim, shape, P.dtype) if isinstance(axis, int): l = len(P.shape) if axis < 0: axis += l shape = [0] * (l - 1) for i in range(l): if i < axis: shape[i] = P.shape[i] elif i > axis: shape[i - 1] = P.shape[i] shape = tuple(shape) else: shape = () A = P.A if P.dtype == f.frac: for key in P.keys: A[key] = f.sum(A[key], axis) else: for key in P.keys: A[key] = np.sum(A[key], axis) return Poly(A, P.dim, shape, P.dtype)
def sum(A, axis=None): """ Sum the components of a shapeable quantity along a given axis. Parameters ---------- A : Poly, frac, array_like Input data. axis : int, optional Axis over which the sum is taken. By default `axis` is None, and all elements are summed. Returns ------- Q : Poly, frac, array_like Polynomial array with same shape as `P`, with the specified axis removed. If `P` is an 0-d array, or `axis` is None, a (non-iterable) component is returned. Examples -------- >>> P = cp.prange(3) >>> print P [1, q0, q0^2] >>> print cp.sum(P) q0^2+q0+1 """ if isinstance(A, (np.ndarray, float, int, long)): return np.sum(A, axis) elif isinstance(A, f.frac): return f.sum(A, axis) elif isinstance(A, p.Poly): return p.sum(A, axis) raise NotImplementedError