def variable(dims=1): """ Simple constructor to create single variables to create polynomials. Args: dims (int): Number of dimensions in the array. Returns: (Poly): Polynomial array with unit components in each dimension. Examples: >>> print(variable()) q0 >>> print(variable(3)) [q0, q1, q2] """ if dims == 1: return Poly({(1, ): 1}, dim=1, shape=()) return Poly( {tuple(indices): indices for indices in numpy.eye(dims, dtype=int)}, dim=dims, shape=(dims, ))
def dot(poly1, poly2): """ Dot product of polynomial vectors. Args: poly1 (Poly) : left part of product. poly2 (Poly) : right part of product. Returns: (Poly) : product of poly1 and poly2. Examples: >>> poly = cp.prange(3, 1) >>> print(poly) [1, q0, q0^2] >>> print(cp.dot(poly, numpy.arange(3))) 2q0^2+q0 >>> print(cp.dot(poly, poly)) q0^4+q0^2+1 """ if not isinstance(poly1, Poly) and not isinstance(poly2, Poly): return numpy.dot(poly1, poly2) poly1 = Poly(poly1) poly2 = Poly(poly2) poly = poly1 * poly2 if numpy.prod(poly1.shape) <= 1 or numpy.prod(poly2.shape) <= 1: return poly return chaospy.poly.sum(poly, 0)
def variable(dims=1): """ Simple constructor to create single variables to create polynomials. Args: dims (int) : Number of dimensions in the array. Returns: (Poly) : Polynomial array with unit components in each dimension. Examples: >>> print(variable()) q0 >>> print(variable(3)) [q0, q1, q2] """ if dims == 1: return Poly({(1, ): np.array(1)}, dim=1, shape=()) r = np.arange(dims, dtype=int) A = {} for i in range(dims): A[tuple(1 * (r == i))] = 1 * (r == i) return Poly(A, dim=dims, shape=(dims, ))
def decompose(P): """ Decompose a polynomial to component form. In array missing values are padded with 0 to make decomposition compatible with `cp.sum(Q, 0)`. Args: P (Poly) : Input data. Returns: (Poly) : Decomposed polynomial with `P.shape==(M,)+Q.shape` where `M` is the number of components in `P`. Examples: >>> q = cp.variable() >>> P = cp.Poly([q**2-1, 2]) >>> print(P) [q0^2-1, 2] >>> print(cp.decompose(P)) [[-1, 2], [q0^2, 0]] >>> print(cp.sum(cp.decompose(P), 0)) [q0^2-1, 2] """ P = P.copy() if not P: return P out = [Poly({key:P.A[key]}) for key in P.keys] return Poly(out, None, None, None)
def dot(P, Q): P = Poly(P) Q = Poly(Q) if numpy.prod(P.shape) <= 1 or numpy.prod(Q.shape) <= 1: return P * Q return sum(P * Q, -1)
def mul(*args): """Polynomial multiplication.""" if len(args) > 2: return add(args[0], add(args[1], args[1:])) if len(args) == 1: return args[0] part1, part2 = args if not isinstance(part2, Poly): if isinstance(part2, (float, int)): part2 = np.asarray(part2) if not part2.shape: core = part1.A.copy() dtype = chaospy.poly.typing.dtyping(part1.dtype, part2.dtype) for key in part1.keys: core[key] = np.asarray(core[key] * part2, dtype) return Poly(core, part1.dim, part1.shape, dtype) part2 = Poly(part2) if part2.dim > part1.dim: part1 = chaospy.dimension.setdim(part1, part2.dim) elif part2.dim < part1.dim: part2 = chaospy.dimension.setdim(part2, part1.dim) if np.prod(part1.shape) >= np.prod(part2.shape): shape = part1.shape else: shape = part2.shape dtype = chaospy.poly.typing.dtyping(part1.dtype, part2.dtype) if part1.dtype != part2.dtype: if part1.dtype == dtype: part2 = chaospy.poly.typing.asfloat(part2) else: part1 = chaospy.poly.typing.asfloat(part1) core = {} for idx1 in part2.A: for idx2 in part1.A: key = tuple(np.array(idx1) + np.array(idx2)) core[key] = np.asarray( core.get(key, 0) + part2.A[idx1] * part1.A[idx2]) core = {key: value for key, value in core.items() if np.any(value)} out = Poly(core, part1.dim, shape, dtype) return out
def transpose(vari): """ Transpose a shapeable quantety. Args: vari (chaospy.poly.base.Poly, numpy.ndarray): Quantety of interest. Returns: (chaospy.poly.base.Poly, numpy.ndarray): Same type as ``vari``. Examples: >>> P = chaospy.reshape(chaospy.prange(4), (2,2)) >>> print(P) [[1, q0], [q0^2, q0^3]] >>> print(chaospy.transpose(P)) [[1, q0^2], [q0, q0^3]] """ if isinstance(vari, Poly): core = vari.A.copy() for key in vari.keys: core[key] = transpose(core[key]) return Poly(core, vari.dim, vari.shape[::-1], vari.dtype) return numpy.transpose(vari)
def around(A, decimals=0): """ Evenly round to the given number of decimals. Args: A (Poly, numpy.ndarray): Input data. decimals (int): Number of decimal places to round to (default: 0). If decimals is negative, it specifies the number of positions to the left of the decimal point. Returns: (Poly, numpy.ndarray): Same type as A. Examples: >>> P = chaospy.prange(3)*2**-numpy.arange(0, 6, 2, float) >>> print(P) [1.0, 0.25q0, 0.0625q0^2] >>> print(chaospy.around(P)) [1.0, 0.0, 0.0] >>> print(chaospy.around(P, 2)) [1.0, 0.25q0, 0.06q0^2] """ if isinstance(A, Poly): B = A.A.copy() for key in A.keys: B[key] = around(B[key], decimals) return Poly(B, A.dim, A.shape, A.dtype) return numpy.around(A, decimals)
def cutoff(poly, *args): """ Remove polynomial components with order outside a given interval. Args: poly (Poly) : Input data. low (int, optional) : The lowest order that is allowed to be included. Defaults to 0. high (int) : The upper threshold for the cutoff range. Returns: (Poly) : The same as `P`, except that all terms that have a order not within the bound `low<=order<high` are removed. Examples: >>> poly = cp.prange(4, 1) + cp.prange(4, 2)[::-1] >>> print(poly) [q1^3+1, q1^2+q0, q0^2+q1, q0^3+1] >>> print(cp.cutoff(poly, 3)) [1, q1^2+q0, q0^2+q1, 1] >>> print(cp.cutoff(poly, 1, 3)) [0, q1^2+q0, q0^2+q1, 0] """ if len(args) == 1: low, high = 0, args[0] else: low, high = args[:2] core_old = poly.A core_new = {} for key in poly.keys: if low <= np.sum(key) < high: core_new[key] = core_old[key] return Poly(core_new, poly.dim, poly.shape, poly.dtype)
def cumsum(vari, axis=None): """ Cumulative sum the components of a shapeable quantity along a given axis. Args: vari (chaospy.poly.base.Poly, numpy.ndarray): Input data. axis (int): Axis over which the sum is taken. By default ``axis`` is None, and all elements are summed. Returns: (chaospy.poly.base.Poly, numpy.ndarray): Polynomial array with same shape as ``vari``. Examples: >>> poly = cp.prange(3) >>> print(poly) [1, q0, q0^2] >>> print(cp.cumsum(poly)) [1, q0+1, q0^2+q0+1] """ if isinstance(vari, Poly): core = vari.A.copy() for key, val in core.items(): core[key] = cumsum(val, axis) return Poly(core, vari.dim, None, vari.dtype) return np.cumsum(vari, axis)
def sum(vari, axis=None): # pylint: disable=redefined-builtin """ Sum the components of a shapeable quantity along a given axis. Args: vari (chaospy.poly.base.Poly, numpy.ndarray): Input data. axis (int): Axis over which the sum is taken. By default ``axis`` is None, and all elements are summed. Returns: (chaospy.poly.base.Poly, numpy.ndarray): Polynomial array with same shape as ``vari``, with the specified axis removed. If ``vari`` is an 0-d array, or ``axis`` is None, a (non-iterable) component is returned. Examples: >>> vari = cp.prange(3) >>> print(vari) [1, q0, q0^2] >>> print(cp.sum(vari)) q0^2+q0+1 """ if isinstance(vari, Poly): core = vari.A.copy() for key in vari.keys: core[key] = sum(core[key], axis) return Poly(core, vari.dim, None, vari.dtype) return np.sum(vari, axis)
def cumprod(vari, axis=None): """ Perform the cumulative product of a shapeable quantity over a given axis. Args: vari (Poly, array_like) : Input data. axis (int, optional) : Axis over which the product is taken. By default, the product of all elements is calculated. Returns: (Poly) : An array shaped as `vari` but with the specified axis removed. Examples: >>> vari = cp.prange(4) >>> print(vari) [1, q0, q0^2, q0^3] >>> print(cp.cumprod(vari)) [1, q0, q0^3, q0^6] """ if isinstance(vari, Poly): if np.prod(vari.shape) == 1: return vari.copy() if axis is None: vari = chaospy.poly.shaping.flatten(vari) axis = 0 vari = chaospy.poly.shaping.rollaxis(vari, axis) out = [vari[0]] for poly in vari[1:]: out.append(out[-1] * poly) return Poly(out, vari.dim, vari.shape, vari.dtype) return np.cumprod(vari, axis)
def reshape(vari, shape): """ Reshape the shape of a shapeable quantity. Args: vari (chaospy.poly.base.Poly, numpy.ndarray): Shapeable input quantity. shape (tuple): The polynomials new shape. Must be compatible with the number of elements in ``vari``. Returns: (chaospy.poly.base.Poly, numpy.ndarray): Same type as ``vari``. Examples: >>> poly = chaospy.prange(6) >>> print(poly) [1, q0, q0^2, q0^3, q0^4, q0^5] >>> print(chaospy.reshape(poly, (2,3))) [[1, q0, q0^2], [q0^3, q0^4, q0^5]] """ if isinstance(vari, Poly): core = vari.A.copy() for key in vari.keys: core[key] = reshape(core[key], shape) out = Poly(core, vari.dim, shape, vari.dtype) return out return numpy.asarray(vari).reshape(shape)
def prange(N=1, dim=1): """ Constructor to create a range of polynomials where the exponent vary. Args: N (int): Number of polynomials in the array. dim (int): The dimension the polynomial should span. Returns: (Poly): A polynomial array of length N containing simple polynomials with increasing exponent. Examples: >>> print(prange(4)) [1, q0, q0^2, q0^3] >>> print(prange(4, dim=3)) [1, q2, q2^2, q2^3] """ A = {} r = numpy.arange(N, dtype=int) key = numpy.zeros(dim, dtype=int) for i in range(N): key[-1] = i A[tuple(key)] = 1 * (r == i) return Poly(A, dim, (N, ), int)
def repeat(A, repeats, axis=None): if isinstance(A, Poly): core = A.A.copy() for key in A.keys: core[key] = repeat(core[key], repeats, axis) return Poly(core, A.dim, None, A.dtype) return numpy.repeat(A, repeats, axis)
def trace(A, offset=0, ax1=0, ax2=1): if isinstance(A, Poly): core = A.A.copy() for key in A.keys: core[key] = trace(core[key], ax1, ax2) return Poly(core, A.dim, None, A.dtype) return numpy.trace(A, offset, ax1, ax2)
def roll(vari, shift, axis=None): """Roll array elements along a given axis.""" if isinstance(vari, Poly): core = vari.A.copy() for key in vari.keys: core[key] = roll(core[key], shift, axis) return Poly(core, vari.dim, None, vari.dtype) return numpy.roll(vari, shift, axis)
def diag(A, k=0): """Extract or construct a diagonal polynomial array.""" if isinstance(A, Poly): core, core_new = A.A, {} for key in A.keys: core_new[key] = numpy.diag(core[key], k) return Poly(core_new, A.dim, None, A.dtype) return numpy.diag(A, k)
def swapaxes(vari, ax1, ax2): """Interchange two axes of a polynomial.""" if isinstance(vari, Poly): core = vari.A.copy() for key in vari.keys: core[key] = swapaxes(core[key], ax1, ax2) return Poly(core, vari.dim, None, vari.dtype) return numpy.swapaxes(vari, ax1, ax2)
def lagrange(X): X = numpy.array(X) if len(X.shape) < 2: X = X.reshape(1, *X.shape) if len(X.shape) < 2: X = X.reshape(1, *X.shape) dim, K = X.shape return Poly(poly, dim)
def differential(P, Q): """ Polynomial differential operator. Args: P (Poly): Polynomial to be differentiated. Q (Poly): Polynomial to differentiate by. Must be decomposed. If polynomial array, the output is the Jacobian matrix. """ P, Q = Poly(P), Poly(Q) if not chaospy.poly.is_decomposed(Q): differential(chaospy.poly.decompose(Q)).sum(0) if Q.shape: return Poly([differential(P, q) for q in Q]) if Q.dim > P.dim: P = chaospy.poly.setdim(P, Q.dim) else: Q = chaospy.poly.setdim(Q, P.dim) qkey = Q.keys[0] A = {} for key in P.keys: newkey = numpy.array(key) - numpy.array(qkey) if numpy.any(newkey < 0): continue A[tuple(newkey)] = P.A[key]*numpy.prod([factorial(key[i], \ exact=True)/factorial(newkey[i], exact=True) \ for i in range(P.dim)]) return Poly(B, P.dim, P.shape, P.dtype)
def tricu(P, k=0): """Cross-diagonal upper triangle.""" tri = numpy.sum(numpy.mgrid[[slice(0, _, 1) for _ in P.shape]], 0) tri = tri < len(tri) + k if isinstance(P, Poly): A = P.A.copy() B = {} for key in P.keys: B[key] = A[key] * tri return Poly(B, shape=P.shape, dim=P.dim, dtype=P.dtype) out = P * tri return out
def swapdim(P, dim1=1, dim2=0): """ Swap the dim between two variables. Args: P (Poly): Input polynomial. dim1 (int): First dim dim2 (int): Second dim. Returns: (Poly): Polynomial with swapped dimensions. Examples: >>> x,y = variable(2) >>> P = x**4-y >>> print(P) q0^4-q1 >>> print(swapdim(P)) q1^4-q0 """ if not isinstance(P, Poly): return numpy.swapaxes(P, dim1, dim2) dim = P.dim shape = P.shape dtype = P.dtype if dim1 == dim2: return P m = max(dim1, dim2) if P.dim <= m: P = chaospy.poly.dimension.setdim(P, m + 1) dim = m + 1 A = {} for key in P.keys: val = P.A[key] key = list(key) key[dim1], key[dim2] = key[dim2], key[dim1] A[tuple(key)] = val return Poly(A, dim, shape, dtype)
def asfloat(vari, limit=10**300): """ Convert dtype of polynomial coefficients to float. Example: >>> poly = 2*cp.variable()+1 >>> print(poly) 2q0+1 >>> print(cp.asfloat(poly)) 2.0q0+1.0 """ if isinstance(vari, Poly): core = vari.A.copy() for key in vari.keys: core[key] = core[key] * 1. return Poly(core, vari.dim, vari.shape, float) return numpy.asfarray(vari)
def toarray(vari): """ Convert polynomial array into a numpy.asarray of polynomials. Args: vari (Poly, numpy.ndarray): Input data. Returns: (numpy.ndarray): A numpy array with ``Q.shape==A.shape``. Examples: >>> poly = cp.prange(3) >>> print(poly) [1, q0, q0^2] >>> array = cp.toarray(poly) >>> print(isinstance(array, numpy.ndarray)) True >>> print(array[1]) q0 """ if isinstance(vari, Poly): shape = vari.shape out = numpy.asarray([{} for _ in range(numpy.prod(shape))], dtype=object) core = vari.A.copy() for key in core.keys(): core[key] = core[key].flatten() for i in range(numpy.prod(shape)): if not numpy.all(core[key][i] == 0): out[i][key] = core[key][i] for i in range(numpy.prod(shape)): out[i] = Poly(out[i], vari.dim, (), vari.dtype) out = out.reshape(shape) return out return numpy.asarray(vari)
def asint(vari): """ Convert dtype of polynomial coefficients to float. Example: >>> poly = 1.5*cp.variable()+2.25 >>> print(poly) 1.5q0+2.25 >>> print(cp.asint(poly)) q0+2 """ if isinstance(vari, Poly): core = vari.A.copy() for key in vari.keys: core[key] = numpy.asarray(core[key], dtype=int) return Poly(core, vari.dim, vari.shape, int) return numpy.asarray(vari, dtype=int)
def rollaxis(vari, axis, start=0): """ Roll the specified axis backwards, until it lies in a given position. Args: vari (chaospy.poly.base.Poly, numpy.ndarray): Input array or polynomial. axis (int): The axis to roll backwards. The positions of the other axes do not change relative to one another. start (int): The axis is rolled until it lies before thes position. """ if isinstance(vari, Poly): core_old = vari.A.copy() core_new = {} for key in vari.keys: core_new[key] = rollaxis(core_old[key], axis, start) return Poly(core_new, vari.dim, None, vari.dtype) return numpy.rollaxis(vari, axis, start)
def prune(A, threshold): """ Remove coefficients that is not larger than a given threshold. Args: A (Poly): Input data. threshold (float): Threshold for which values to cut. Returns: (Poly): Same type as A. Examples: >>> P = chaospy.sum(chaospy.prange(3)*2**-numpy.arange(0, 6, 2, float)) >>> print(P) 0.0625q0^2+0.25q0+1.0 >>> print(chaospy.prune(P, 0.1)) 0.25q0+1.0 >>> print(chaospy.prune(P, 0.5)) 1.0 >>> print(chaospy.prune(P, 1.5)) 0.0 """ if isinstance(A, Poly): B = A.A.copy() for key in A.keys: values = B[key].copy() values[numpy.abs(values) < threshold] = 0. B[key] = values return Poly(B, A.dim, A.shape, A.dtype) A = A.copy() A[numpy.abs(A) < threshold] = 0. return A
def rolldim(P, n=1): """ Roll the axes. Args: P (Poly) : Input polynomial. n (int) : The axis that after rolling becomes the 0th axis. Returns: (Poly) : Polynomial with new axis configuration. Examples: >>> x,y,z = variable(3) >>> P = x*x*x + y*y + z >>> print(P) q0^3+q1^2+q2 >>> print(rolldim(P)) q0^2+q2^3+q1 """ dim = P.dim shape = P.shape dtype = P.dtype A = dict(((key[n:] + key[:n], P.A[key]) for key in P.keys)) return Poly(A, dim, shape, dtype)
def substitute(P, x0, x1, V=0): """ Substitute a variable in a polynomial array. Args: P (Poly) : Input data. x0 (Poly, int) : The variable to substitute. Indicated with either unit variable, e.g. `x`, `y`, `z`, etc. or through an integer matching the unit variables dimension, e.g. `x==0`, `y==1`, `z==2`, etc. x1 (Poly) : Simple polynomial to substitute `x0` in `P`. If `x1` is an polynomial array, an error will be raised. Returns: (Poly) : The resulting polynomial (array) where `x0` is replaced with `x1`. Examples: >>> x,y = cp.variable(2) >>> P = cp.Poly([y*y-1, y*x]) >>> print(cp.substitute(P, y, x+1)) [q0^2+2q0, q0^2+q0] With multiple substitutions: >>> print(cp.substitute(P, [x,y], [y,x])) [q0^2-1, q0q1] """ x0,x1 = map(Poly, [x0,x1]) dim = np.max([p.dim for p in [P,x0,x1]]) dtype = chaospy.poly.typing.dtyping(P.dtype, x0.dtype, x1.dtype) P, x0, x1 = [chaospy.poly.dimension.setdim(p, dim) for p in [P,x0,x1]] if x0.shape: x0 = [x for x in x0] else: x0 = [x0] if x1.shape: x1 = [x for x in x1] else: x1 = [x1] # Check if substitution is needed. valid = False C = [x.keys[0].index(1) for x in x0] for key in P.keys: if np.any([key[c] for c in C]): valid = True break if not valid: return P dims = [tuple(np.array(x.keys[0])!=0).index(True) for x in x0] dec = is_decomposed(P) if not dec: P = decompose(P) P = chaospy.poly.dimension.dimsplit(P) shape = P.shape P = [p for p in chaospy.poly.shaping.flatten(P)] for i in range(len(P)): for j in range(len(dims)): if P[i].keys and P[i].keys[0][dims[j]]: P[i] = x1[j].__pow__(P[i].keys[0][dims[j]]) break P = Poly(P, dim, None, dtype) P = chaospy.poly.shaping.reshape(P, shape) P = chaospy.poly.collection.prod(P, 0) if not dec: P = chaospy.poly.collection.sum(P, 0) return P