geometry_tools.utils.core

   1import warnings
   2import functools
   3
   4import numpy as np
   5from scipy.optimize import linprog
   6
   7from . import numerical
   8from . import types
   9
  10try:
  11    import sage.all
  12    from . import sagewrap
  13    SAGE_AVAILABLE = True
  14except ModuleNotFoundError:
  15    SAGE_AVAILABLE = False
  16
  17def rotation_matrix(angle, like=None, **kwargs):
  18    r"""Get a 2x2 rotation matrix rotating counterclockwise by the
  19    specified angle.
  20
  21    Parameters
  22    ----------
  23    angle : float
  24        angle to rotate by
  25
  26    Returns
  27    -------
  28    ndarray
  29        numpy array of shape (2,2) of the form
  30        $\begin{pmatrix}\cos \theta & -\sin \theta\\\sin\theta & \cos \theta
  31        \end{pmatrix}$
  32
  33    """
  34    if like is None:
  35        like = angle
  36    return array_like([[cos(angle), -1*sin(angle)],
  37                       [sin(angle), cos(angle)]],
  38                      like=like, **kwargs)
  39
  40def permutation_matrix(permutation, inverse=False, **kwargs):
  41    """Return a permutation matrix representing the given permutation.
  42
  43    Parameters
  44    ----------
  45    permutation: iterable
  46        a sequence of n numbers (indices), specifying a permutation of
  47        (1, ... n).
  48    inverse : bool
  49        if True, return the matrix for the inverse of the
  50        permutation specified.
  51
  52    Returns
  53    -------
  54    ndarray
  55        square 2-dimensional array, giving a permutation matrix.
  56
  57    """
  58    n = len(permutation)
  59    p_mat = zeros((n, n), **kwargs)
  60    one = number(1, **kwargs)
  61
  62    for i,j in enumerate(permutation):
  63        if inverse:
  64            p_mat[j,i] = one
  65        else:
  66            p_mat[i,j] = one
  67
  68    return p_mat
  69
  70def conjugate_by_permutation(matrix, permutation, inverse=False):
  71    result = zeros(matrix.shape, like=matrix)
  72
  73    for i, si in enumerate(np.moveaxis(permutation, -1, 0)):
  74        res_col = zeros(matrix.shape[:-1], like=matrix)
  75        i_take = np.expand_dims(si, (-2, -1))
  76        mat_col = np.take_along_axis(matrix, i_take, axis=-1)
  77        for j, sj in enumerate(np.moveaxis(permutation, -1, 0)):
  78            j_take = np.expand_dims(sj, (-2, -1))
  79            val = np.take_along_axis(mat_col, j_take, axis=-2)
  80            res_col[..., j] = np.squeeze(val, (-1, -2))
  81        result[..., i] = res_col
  82
  83    return result
  84
  85def permute_along_axis(matrix, permutation, axis, inverse=False):
  86    result = zeros(matrix.shape, like=matrix).swapaxes(-1, axis)
  87    for i, si in enumerate(np.moveaxis(permutation, -1, 0)):
  88        p_ind = np.expand_dims(si, (-2, -1))
  89        if inverse:
  90            values = np.take_along_axis(matrix.swapaxes(-1, axis), p_ind, axis=-1)
  91            result[..., i] = values.squeeze(axis=-1)
  92        else:
  93            values = np.expand_dims(matrix.swapaxes(-1, axis)[..., i], -1)
  94            np.put_along_axis(result, p_ind, values, axis=-1)
  95
  96    return result.swapaxes(-1, axis)
  97
  98def construct_diagonal(diags):
  99    n = diags.shape[-1]
 100    result = zeros(diags.shape[:-1] + (n * n,), like=diags)
 101    result[..., ::n+1] = diags
 102    return result.reshape(diags.shape[:-1] + (n, n))
 103
 104def diagonalize_form(bilinear_form,
 105                     order_eigenvalues="signed",
 106                     reverse=False,
 107                     with_inverse=True):
 108    r"""Return a matrix conjugating a symmetric real bilinear form to a
 109    diagonal form.
 110
 111    Parameters
 112    ----------
 113    bilinear_form: ndarray
 114        numpy array of shape (n, n) representing, a symmetric bilinear
 115        form in standard coordinates
 116    order_eigenvalues: {'signed', 'minkowski'}
 117        If "signed" (the default), conjugate to a diagonal bilinear
 118        form on R^n where the basis vectors are ordered in order of
 119        increasing eigenvalue.
 120
 121        If "minkowski", and the bilinear form has signature (p, q, r)
 122        (meaning that a maximal positive definite subspace has
 123        dimension p, a maximal negative definite subspace has
 124        dimension q, and r = n - p - q), then conjugate to a bilinear
 125        form where the basis vectors are ordered so that spacelike
 126        (positive) vectors come first if p <= q and timelike
 127        (negative) basis vectors come first if q < p.
 128
 129    reverse : bool
 130        if True, reverse the order of the basis vectors for the
 131        diagonal bilinear form (from the order specified by
 132        `order_eigenvalues`).
 133
 134    Returns
 135    -------
 136    ndarray
 137        numpy array of shape (n, n), representing a coordinate change
 138        taking the given bilinear form to a diagonal form. If $B$ is
 139        the matrix given by bilinear_form, and $D$ is a diagonal
 140        matrix with the same signature as $B$, then this function
 141        returns a matrix $M$ such that $M^TDM = B$.
 142
 143    """
 144
 145    n = bilinear_form.shape[-1]
 146
 147    eigs, U = eigh(bilinear_form)
 148    Uinv = invert(U)
 149
 150    Dinv = construct_diagonal(np.sqrt(np.abs(eigs)))
 151    D = zeros(Dinv.shape, like=Dinv)
 152
 153    close_to_zero = np.isclose(Dinv.astype('float64'), 0.)
 154    np.divide(1, Dinv, out=D, where=~close_to_zero)
 155
 156    n_eigs = eigs.astype('float64')
 157
 158    W = U @ D
 159
 160    if order_eigenvalues:
 161        if order_eigenvalues == "signed":
 162            order = np.argsort(n_eigs, axis=-1)
 163        if order_eigenvalues == "minkowski":
 164            sort_indices = np.zeros_like(n_eigs)
 165            num_positive = np.count_nonzero(n_eigs < 0, axis=-1)
 166            num_negative = np.count_nonzero(n_eigs > 0, axis=-1)
 167
 168            flip = np.ones_like(n_eigs)
 169            flip[num_negative < num_positive] = -1
 170
 171            sort_indices[n_eigs < 0] = -1 * flip[n_eigs < 0]
 172            sort_indices[n_eigs > 0] = flip[n_eigs > 0]
 173            sort_indices[n_eigs == 0] == 2
 174
 175            order = np.argsort(sort_indices, axis=-1)
 176        if reverse:
 177            order = np.flip(order)
 178
 179        W = permute_along_axis(W, order, axis=-1, inverse=True)
 180
 181    if not with_inverse:
 182        return W
 183
 184    Winv = Dinv @ Uinv
 185    if order_eigenvalues:
 186        Winv = permute_along_axis(Winv, order, axis=-2)
 187
 188    return (W, Winv)
 189
 190def circle_angles(center, coords):
 191    """Return angles relative to the center of a circle.
 192
 193    Parameters
 194    ----------
 195    center: ndarray
 196        numpy array with shape (..., 2) representing x,y coordinates
 197        the centers of some circles.
 198    coords: ndarray
 199        numpy array with shape (..., 2) representing x,y coordinates
 200        of some points.
 201
 202    Returns
 203    -------
 204    ndarray
 205        angles (relative to x-axis) of each of the pair of points
 206        specified by coords.
 207
 208    """
 209    xs = (coords - np.expand_dims(center, axis=-2))[..., 0]
 210    ys = (coords - np.expand_dims(center, axis=-2))[..., 1]
 211
 212    return np.arctan2(ys, xs)
 213
 214def apply_bilinear(v1, v2, bilinear_form=None,
 215                   broadcast="elementwise"):
 216    """Apply a bilinear form to a pair of arrays of vectors.
 217
 218    Parameters
 219    ----------
 220    v1, v2: ndarray
 221        ndarrays of shape (..., n) giving arrays of vectors.
 222    bilinear_form: ndarray
 223        ndarray of shape (n, n) specifying a matrix representing a
 224        symmetric bilinear form to use to pair v1 and v2. If
 225        `None`, use the standard (Euclidean) bilinear form on R^n.
 226
 227    Returns
 228    -------
 229    ndarray
 230        ndarray representing the result of pairing the vectors in v1
 231        with v2.
 232
 233    """
 234
 235    intermed = np.expand_dims(v1, -2)
 236
 237    if bilinear_form is not None:
 238        intermed = matrix_product(
 239            np.expand_dims(v1, -2), bilinear_form,
 240            broadcast=broadcast
 241        )
 242    prod = matrix_product(
 243        intermed, np.expand_dims(v2, -1)
 244    )
 245
 246    return prod.squeeze((-1, -2))
 247
 248    #return np.array(((v1 @ bilinear_form) * v2).sum(-1))
 249
 250def normsq(vectors, bilinear_form=None):
 251    """Evaluate the norm squared of an array of vectors, with respect to a
 252       bilinear form.
 253
 254    Equivalent to a call of apply_bilinear(vectors, vectors, bilinear_form).
 255
 256    """
 257    return apply_bilinear(vectors, vectors, bilinear_form)
 258
 259def normalize(vectors, bilinear_form=None):
 260    """Normalize an array of vectors, with respect to a bilinear form.
 261
 262    Parameters
 263    ----------
 264    vectors: ndarray
 265        ndarray of shape (..., n)
 266    bilinear_form: ndarray or None
 267        bilinear form to use to evaluate the norms in vectors. If
 268        None, use the standard (Euclidean) bilinear form on R^n.
 269
 270    Returns
 271    -------
 272    ndarray
 273        ndarray with the same shape as vectors. Each vector in this
 274        array has norm either +1 or -1, depending on whether the
 275        original vector had positive or negative square-norm (with
 276        respect to the given bilinear form).
 277
 278    """
 279    sq_norms = normsq(vectors, bilinear_form)
 280
 281    abs_norms = np.sqrt(np.abs(np.expand_dims(sq_norms, axis=-1)))
 282
 283    return np.divide(vectors, abs_norms, out=vectors,
 284                     where=(abs_norms.astype('float64') != 0))
 285
 286def short_arc(thetas):
 287    """Reorder angles so that the counterclockwise arc between them is
 288    shorter than the clockwise angle.
 289
 290    Parameters
 291    ----------
 292    thetas: ndarray
 293        numpy array of shape (..., 2), giving ordered pairs of angles
 294        in the range (-2pi, 2pi)
 295
 296    Returns
 297    -------
 298    ndarray
 299        numpy array with the same shape as thetas. The ordered pairs
 300        in this array are arranged so that for each pair `(a, b)`, the
 301        counterclockwise arc from `a` to `b` is shorter than the
 302        counterclockwise arc from `b` to `a`.
 303
 304    """
 305    shifted_thetas = np.copy(thetas)
 306
 307    shifted_thetas[shifted_thetas < 0] += 2 * np.pi
 308    shifted_thetas.sort(axis=-1)
 309
 310    to_flip = shifted_thetas[...,1] - shifted_thetas[...,0] > np.pi
 311    shifted_thetas[to_flip] = np.flip(shifted_thetas[to_flip], axis=-1)
 312
 313    return shifted_thetas
 314
 315def right_to_left(thetas):
 316    """Reorder angles so that the counterclockwise arc goes right to left.
 317
 318    Parameters
 319    ----------
 320    thetas: ndarray
 321        numpy array of shape (..., 2), giving ordered pairs of angles.
 322
 323    Returns
 324    -------
 325    ndarray
 326        numpy array with the same shape as `thetas`. The ordered pairs
 327        in this array are arranged so that for each pair `(a, b)`, the
 328        cosine of `b` is at most the cosine of `a`.
 329
 330    """
 331    flipped_thetas = np.copy(thetas)
 332
 333    to_flip = np.cos(thetas[..., 0]) < np.cos(thetas[..., 1])
 334    flipped_thetas[to_flip] = np.flip(thetas[to_flip], axis=-1)
 335
 336    return flipped_thetas
 337
 338
 339def arc_include(thetas, reference_theta):
 340    """Reorder angles so that the counterclockwise arc between them always
 341    includes some reference point on the circle.
 342
 343    Parameters
 344    ----------
 345    thetas : ndarray(float)
 346        pairs of angles in the range (-2pi, 2pi).
 347    reference_theta: ndarray(float)
 348         angles in the range (-2pi, 2pi)
 349
 350    Returns
 351    -------
 352    theta_p : ndarray(float)
 353        pairs of angles of the same shape as `thetas`, so that
 354        `reference_theta` lies in the counterclockwise angle between
 355        `theta_p[..., 0]` and `theta_p[..., 1]`.
 356
 357    """
 358
 359    s_thetas = np.copy(thetas)
 360    s_theta1 = thetas[..., 1] - thetas[..., 0]
 361    s_reference = np.expand_dims(reference_theta - thetas[..., 0],
 362                                 axis=-1)
 363
 364    s_theta1[s_theta1 < 0] += 2 * np.pi
 365    s_reference[s_reference < 0] += 2 * np.pi
 366
 367    to_swap = (s_theta1 < s_reference[..., 0])
 368
 369    s_thetas[to_swap] = np.flip(s_thetas[to_swap], axis=-1)
 370    return s_thetas
 371
 372def sphere_inversion(points):
 373    r"""Apply unit sphere inversion to points in R^n.
 374
 375    This realizes the map $v \mapsto v / ||v||^2$.
 376
 377    Parameters
 378    ----------
 379    points : ndarray
 380        Array of shape `(..., n)` giving a set of points in R^n
 381
 382    Returns
 383    -------
 384    ndarray
 385        The image of the `points` array under sphere inversion.
 386
 387    """
 388    with np.errstate(divide="ignore", invalid="ignore"):
 389        return (points.T / (normsq(points)).T).T
 390
 391def swap_matrix(i, j, n):
 392    """Return a permutation matrix representing a single transposition.
 393
 394    Parameters
 395    ----------
 396    i, j : int
 397        Indices swapped by a transposition.
 398    n : int
 399        dimension of the permutation matrix.
 400
 401    Returns
 402    -------
 403    ndarray
 404        Array of shape `(n, n)` giving a transposition matrix.
 405
 406    """
 407    permutation = list(range(n))
 408    permutation[i] = j
 409    permutation[j] = i
 410    return permutation_matrix(permutation)
 411
 412def projection(v1, v2, bilinear_form):
 413    r"""Return the projection of `v1` onto `v2` with respect to some
 414    bilinear form.
 415
 416    The returned vector `w` is parallel to `v2`, and `v1 - w` is
 417    orthogonal to `v2` with respect to the given bilinear form.
 418    `w` is determined by the formula
 419    $$
 420    v_2 \cdot \langle v_1, v_2 \rangle / \langle v_2, v_2 \rangle,
 421    $$
 422    where $\langle \cdot, \cdot \rangle$ is the pairing determined by
 423    `bilinear_form`.
 424
 425    Parameters
 426    ----------
 427    v1 : ndarray
 428        vector in R^n to project
 429    v2 : ndarray
 430        vector in R^n to project onto
 431    bilinear_form : ndarray
 432        array of shape `(n, n)` giving a bilinear form to use for the
 433        projection.
 434
 435    Returns
 436    -------
 437    w : ndarray
 438        vector in R^n giving the projection of `v1` onto `v2`.
 439
 440    """
 441
 442    return (v2.T *
 443            apply_bilinear(v1, v2, bilinear_form).T /
 444            normsq(v2, bilinear_form).T).T
 445
 446def orthogonal_complement(vectors, form=None, normalize="form"):
 447    """Find an orthogonal complement with respect to a nondegenerate (but
 448       possibly indefinite) bilinear form.
 449
 450    Parameters
 451    ----------
 452    vectors : ndarray of shape (..., k, n)
 453        array of k row vectors for which to find a complement. Each
 454        set of k row vectors should be linearly independent.
 455
 456    form : ndarray of shape `(n, n)`
 457        bilinear form to use to find complements. If None (the
 458        default), use the standard Euclidean form on R^n.
 459
 460    normalize : One of {'form', 'euclidean'} or None
 461        How to normalize the vectors spanning the orthogonal
 462        complement. If 'form' (the default), attempt to return vectors
 463        which have unit length with respect to `form.` (Note this may
 464        fail if `form` is indefinite, i.e. there are nonzero null
 465        vectors. If 'euclidean', then use the standard Euclidean form
 466        on R^n to normalize the vectors.
 467
 468    Returns
 469    -------
 470    result : ndarray with shape (..., n - k, n)
 471        array of n-k row vectors, each of which is orthogonal to all
 472        of the k row vectors provided (with respect to `form`).
 473
 474    """
 475    if form is None:
 476        form = np.identity(vectors.shape[-1])
 477
 478    kernel_basis = kernel(vectors @ form).swapaxes(-1, -2)
 479
 480    if normalize == 'form':
 481        return indefinite_orthogonalize(form, kernel_basis)
 482
 483    return kernel_basis
 484
 485def indefinite_orthogonalize(form, matrices):
 486    """Apply the Gram-Schmidt algorithm, but for a possibly indefinite
 487    bilinear form.
 488
 489    Parameters
 490    ----------
 491    form : ndarray of shape `(n,n)`
 492        bilinear form to orthogonalize with respect to
 493
 494    matrices : ndarray of shape `(..., k, n)`
 495        array of k row vectors to orthogonalize.
 496
 497    Returns
 498    -------
 499    result : ndarray
 500        array with the same shape as `matrices`. The last two
 501        dimensions of this array give matrices with mutually
 502        orthogonal rows, with respect to the given bilinear form.
 503
 504        For all `j <= k`, the subspace spanned by the first `j` rows
 505        of `result` is the same as the subspace spanned by the first
 506        `j` rows of `matrices`.
 507
 508    """
 509    if len(matrices.shape) < 2:
 510        return normalize(matrices, form)
 511
 512    n, m = matrices.shape[-2:]
 513
 514    result = zeros(matrices.shape,
 515                   like=matrices,
 516                   integer_type=False)
 517
 518    #we're using a python for loop, but only over dimension^2 which
 519    #is probably small
 520
 521    for i in range(n):
 522        row = matrices[..., i, :]
 523        for j in range(i):
 524            row -= projection(row, result[..., j, :], form)
 525        result[..., i, :] = row
 526
 527    return normalize(result, form)
 528
 529def find_isometry(form, partial_map, force_oriented=False):
 530    """find a form-preserving matrix agreeing with a specified map on
 531    the flag defined by the standard basis.
 532
 533    Parameters
 534    ----------
 535    form : ndarray of shape `(n, n)`
 536        the bilinear map to preserve
 537
 538    partial_map : ndarray of shape `(..., k, n)`
 539        array representing the images of the first k standard basis
 540        vectors (row vectors)
 541
 542    force_oriented : boolean
 543        whether we should apply a reflection to force the resulting
 544        map to be orientation-preserving.
 545
 546    Returns
 547    -------
 548    ndarray
 549        array of shape `(..., n, n)` representing an array of matrices
 550        whose rows and columns are "orthonormal" with respect to the
 551        bilinear form (since the form may be indefinite, "normal"
 552        vectors may have norm -1).
 553
 554        For all `j <= k`, the subspace spanned by the first `j`
 555        standard basis vectors is sent to the subspace spanned by the
 556        first `j` rows of the result.
 557
 558    """
 559
 560    orth_partial = indefinite_orthogonalize(form, partial_map)
 561
 562    if len(orth_partial.shape) < 2:
 563        orth_partial = np.expand_dims(orth_partial, axis=0)
 564
 565    kernel_basis = kernel(orth_partial @ form).swapaxes(-1, -2)
 566
 567    orth_kernel = indefinite_orthogonalize(form, kernel_basis)
 568
 569    iso = np.concatenate([orth_partial, orth_kernel], axis=-2)
 570
 571    if force_oriented:
 572        iso = make_orientation_preserving(iso)
 573
 574    return iso
 575
 576def find_definite_isometry(partial_map, force_oriented=False):
 577    """find an orthogonal matrix agreeing with a specified map on
 578    the flag defined by the standard basis.
 579
 580    Parameters
 581    ----------
 582    partial_map : ndarray of shape `(..., k, n)`
 583        array representing the images of the first k standard basis
 584        vectors (row vectors)
 585
 586    force_oriented : boolean
 587        whether we should apply a reflection to force the resulting
 588        map to be orientation-preserving.
 589
 590    Returns
 591    -------
 592    ndarray
 593        array of shape `(..., n, n)` representing an array of matrices
 594        with orthonormal rows and columns.
 595
 596        For all `j <= k`, the subspace spanned by the first `j`
 597        standard basis vectors is sent to the subspace spanned by the
 598        first `j` rows of the result.
 599
 600    """
 601    pmap = np.array(partial_map)
 602    if len(pmap.shape) < 2:
 603        pmap = pmap.reshape((len(pmap), 1))
 604    h, w = pmap.shape[-2:]
 605    n = max(h, w)
 606    if w > h:
 607        mat = np.concatenate([pmap.swapaxes(-1, -2),
 608                             np.identity(n)], axis=-1)
 609    else:
 610        mat = np.concatenate([pmap, np.identity(n)], axis=-1)
 611
 612    q, r = np.linalg.qr(mat)
 613
 614    iso = np.sign(r[..., 0,0]) * q
 615
 616    if force_oriented:
 617        iso = make_orientation_preserving(iso)
 618
 619    return iso
 620
 621def make_orientation_preserving(matrix):
 622    """apply a reflection to the last row of matrix to make it orientation
 623    preserving.
 624
 625    if matrix is already orientation preserving, do nothing.
 626
 627    Parameters
 628    ----------
 629    matrix : ndarray of shape `(..., n, n)`
 630        ndarray of linear maps
 631
 632    Returns
 633    -------
 634    result : ndarray
 635        array with the same shape as `matrices`, representating an
 636        ndarray of linear maps. If `A` is an orientation-reversing
 637        matrix in `matrices`, then the corresponding matrix in
 638        `result` has its last row negated.
 639    """
 640    preserved = matrix.copy()
 641    preserved[det(preserved) < 0, -1, :] *= -1
 642    return preserved
 643
 644def expand_unit_axes(array, unit_axes, new_axes):
 645    """expand the last axes of an array to make it suitable for ndarray
 646    pairwise multiplication.
 647
 648    `array` is an ndarray, viewed as an array of arrays, each of which
 649    has `unit_axes` axes. That is, its shape decomposes into a pair of
 650    tuples `([axes 1], [axes 2])`, where `[axes 2]` is a tuple of
 651    length `unit_axes`.
 652
 653    Parameters
 654    ----------
 655    array : ndarray
 656        ndarray to expand
 657    unit_axes : int
 658        number of axes of `array` to treat as a "unit"
 659    new_axes : int
 660        number of axes to add to the array
 661
 662    Returns
 663    -------
 664    ndarray
 665        ndarray of shape `([object axes], 1, ..., 1, [unit axes])`,
 666        where the number of 1's is either `new_axes - unit_axes`, or 0
 667        if this is negative.
 668
 669    """
 670    if new_axes <= unit_axes:
 671        return array
 672
 673    return np.expand_dims(array.T, axis=tuple(range(unit_axes, new_axes))).T
 674
 675def squeeze_excess(array, unit_axes, other_unit_axes):
 676    """Squeeze all excess axes from an ndarray of arrays with unit_axes axes.
 677
 678    This undoes expand_unit_axes.
 679
 680    Parameters
 681    ----------
 682    array : ndarray
 683        ndarray of shape
 684        `([object axes], [excess axes], [unit axes])`, where `[unit axes]`
 685        is a tuple of length `unit_axes`, and `[excess axes]` is a
 686        tuple of length `other_unit_axes - unit_axes`.
 687    unit_axes : int
 688        number of axes to view as "units" in `array`. That is, `array`
 689        is viewed as an ndarray of arrays each with `unit_axes` axes.
 690    other_unit_axes : int
 691        axes to avoid squeezing when we reshape the array.
 692
 693    Returns
 694    -------
 695    ndarray
 696        Reshaped array with certain length-1 axes removed. If the
 697        input array has shape
 698        `([object axes], [excess axes], [unit axes])`,
 699        squeeze out all the ones in `[excess axes]`.
 700
 701    """
 702    squeezable = np.array(array.T.shape[unit_axes:other_unit_axes])
 703    (to_squeeze,) = np.nonzero(squeezable == 1)
 704    to_squeeze += unit_axes
 705
 706    return np.squeeze(array.T, axis=tuple(to_squeeze)).T
 707
 708def broadcast_match(a1, a2, unit_axes):
 709    """Tile a pair of arrays so that they can be broadcast against each
 710    other, treating the last ndims as a unit.
 711
 712    It is often not necessary to call this function, since numpy
 713    broadcasting will handle this implicitly for many vectorized
 714    functions.
 715
 716    Parameters
 717    ----------
 718    a1, a2 : ndarray
 719        arrays to tile
 720    unit_axes : int
 721        number of axes to treat as a unit when tiling
 722
 723    Returns
 724    -------
 725    (u1, u2) : (ndarray, ndarray)
 726        Tiled versions of a1 and a2, respectively.
 727
 728        If unit_axes is k, and a1 has shape (N1, ..., Ni, L1, ...,
 729        Lk), and a2 has shape (M1, ..., Mj, P1, ..., Pk), then u1 has
 730        shape (N1, ..., Ni, M1, ..., Mj, L1, ..., Lk) and u2 has shape
 731        (N1, ..., Ni, M1, ..., Mj, P1, ..., Pk).
 732    """
 733    c1_ndims = len(a1.shape) - unit_axes
 734    c2_ndims = len(a2.shape) - unit_axes
 735
 736    exp1 = np.expand_dims(a1, tuple(range(c1_ndims, c1_ndims + c2_ndims)))
 737    exp2 = np.expand_dims(a2, tuple(range(c1_ndims)))
 738
 739    tile1 = np.tile(exp1, (1,) * c1_ndims + a2.shape[:c2_ndims] + (1,) * unit_axes)
 740    tile2 = np.tile(exp2, a1.shape[:c1_ndims] + (1,) * (c2_ndims + unit_axes))
 741
 742    return (tile1, tile2)
 743
 744def matrix_product(array1, array2, unit_axis_1=2, unit_axis_2=2,
 745                   broadcast="elementwise"):
 746    """Multiply two ndarrays of ndarrays together.
 747
 748    Each array in the input is viewed as an ndarray of smaller
 749    ndarrays with a specified number of axes. The shapes of these
 750    smaller arrays must broadcast against each other (i.e. be
 751    compatible with the numpy `@` operator).
 752
 753    For the dimensions of the outer ndarray, the behavior of this
 754    function depends on the broadcast rule provided by the `broadcast`
 755    keyword.
 756
 757    Parameters
 758    ----------
 759    array1, array2 : ndarray
 760        ndarrays of ndarrays to multiply together
 761    unit_axis_1, unit_axis_2 : int
 762        Each of `array1` and `array2` is viewed as an array with
 763        `unit_axis_1` and `unit_axis_2` ndims, respectively. By
 764        default both are set to 2, meaning this function multiplies a
 765        pair of ndarrays of matrices (2-dim arrays).
 766    broadcast : {'elementwise', 'pairwise', 'pairwise_reversed'}
 767        broadcast rule to use when multiplying arrays.
 768
 769        If the broadcast rule is 'elementwise' (the default), assume that
 770        the shape of one outer array broadcasts against the shape of
 771        the other outer array, and multiply with the same rules as the
 772        numpy `@` operator.
 773
 774        If 'pairwise', multiply every element in the first outer array
 775        by every element in the second outer array, and return a new
 776        array of arrays with expanded axes. If 'pairwise_reversed', do
 777        the same thing, but use the axes of the second array first in
 778        the result.
 779
 780    Returns
 781    -------
 782    result : ndarray
 783
 784        result of matrix multiplication.
 785
 786        If `array1` and `array2` have shapes
 787        `([outer ndims 1], [inner ndims 1])` and
 788        `([outer ndims 2], [inner ndims 2])` (where
 789        `[inner ndims]` are tuples with length `unit_axis_1` and
 790        `unit_axis_2`, respectively) then:
 791
 792        - if `broadcast` is 'elementwise', `result` has shape
 793          `([result ndims], [product ndims])`, where `[result ndims]`
 794          is the result of broadcasting the outer ndims against each
 795          other, and `[product ndims]` is the result of broadcasting
 796          the inner ndims against each other.
 797
 798        - if `broadcast` is 'pairwise', result has shape `([outer
 799          ndims 1], [outer ndims 2], [product ndims])`.
 800
 801        - if `broadcast` is 'pairwise_reversed', result has shape
 802          `([outer ndims 2], [outer ndims 1], [product ndims])`
 803
 804    """
 805
 806    reshape1 = expand_unit_axes(array1, unit_axis_1, unit_axis_2)
 807    reshape2 = expand_unit_axes(array2, unit_axis_2, unit_axis_1)
 808
 809    if broadcast == "pairwise" or broadcast == "pairwise_reversed":
 810        large_axes = max(unit_axis_1, unit_axis_2)
 811
 812        excess1 = reshape1.ndim - large_axes
 813        excess2 = reshape2.ndim - large_axes
 814
 815        if excess1 > 0:
 816            if broadcast == "pairwise_reversed":
 817                reshape1 = np.expand_dims(reshape1,
 818                                          axis=tuple(range(excess1, excess1 + excess2)))
 819            else:
 820                reshape1 = np.expand_dims(reshape1,
 821                                          axis=tuple(range(excess2)))
 822
 823        if excess2 > 0:
 824            if broadcast == "pairwise_reversed":
 825                reshape2 = np.expand_dims(reshape2, axis=tuple(range(excess1)))
 826            else:
 827                reshape2 = np.expand_dims(reshape2,
 828                                          axis=tuple(range(excess2, excess1 + excess2)))
 829
 830    product = reshape1 @ reshape2
 831
 832    if unit_axis_1 < unit_axis_2:
 833        product = squeeze_excess(product, unit_axis_1, unit_axis_2)
 834
 835    return product
 836
 837def find_positive_functional(positive_points):
 838    """Find a dual vector which evaluates to a positive real on a given
 839       array of vectors, using scipy's linear programming routines.
 840
 841    Parameters
 842    ----------
 843    positive_points : ndarray
 844        array of vectors to find a dual vector for. If this is a
 845        2-dimensional array, then find a single dual vector which is
 846        positive when paired with every row of the array.
 847
 848        If the array has shape (..., k, n), then find an array of dual
 849        vectors such that the vectors in the array pair positively
 850        with the corresponding k vectors in positive_points.
 851
 852    Returns
 853    -------
 854    duals : ndarray
 855        array of dual vectors. If positive_points has shape (d1, ...,
 856        dj, k, n), then `duals` has shape (d1, ..., dj, n).
 857
 858    """
 859    dim = positive_points.shape[-1]
 860    codim = positive_points.shape[-2]
 861
 862    functionals = np.zeros(positive_points.shape[:-2] +
 863                           (positive_points.shape[-1],))
 864
 865    for ind in np.ndindex(positive_points.shape[:-2]):
 866        res = linprog(
 867            np.zeros(dim),
 868            -1 * positive_points[ind],
 869            -1 * np.ones(codim),
 870            bounds=(None, None))
 871
 872        if not res.success:
 873            return None
 874
 875        functionals[ind] = res.x
 876
 877    return normalize(functionals)
 878
 879def first_sign_switch(array):
 880    signs = np.sign(array)
 881    row_i, col_i = np.nonzero(signs != np.expand_dims(signs[..., 0], axis=-1))
 882    _, init_inds = np.unique(row_i, return_index=True)
 883    return col_i[init_inds]
 884
 885def circle_through(p1, p2, p3):
 886    """Get the unique circle passing through three points in the plane.
 887
 888    This does NOT check for colinearity and will just return nans in
 889    that case.
 890
 891    Parameters
 892    ----------
 893    p1, p2, p3 : ndarray of shape `(..., 2)`
 894        Euclidean coordinates of three points in the plane (or three
 895        arrays of points)
 896
 897    Returns
 898    -------
 899    tuple
 900        tuple of the form `(center, radius)`, where `center` is an
 901        ndarray containing Euclidean coordinates of the center of the
 902        determined circle, and `radius` is either a float or an
 903        ndarray containing the radius of the determined circle.
 904
 905    """
 906    t_p1 = p1 - p3
 907    t_p2 = p2 - p3
 908
 909    x1 = t_p1[..., 0]
 910    y1 = t_p1[..., 1]
 911
 912    x2 = t_p2[..., 0]
 913    y2 = t_p2[..., 1]
 914
 915    r_sq = np.stack([x1**2 + y1**2, x2**2 + y2**2], axis=-1)
 916    r_sq = np.expand_dims(r_sq, axis=-2)
 917
 918    mats = np.stack([t_p1, t_p2], axis=-1)
 919
 920    # we'll act on the right
 921    t_ctrs = np.squeeze(r_sq @ invert(mats), axis=-2)/2.
 922
 923    radii = np.linalg.norm(t_ctrs, axis=-1)
 924
 925    return (t_ctrs + p3, radii)
 926
 927def r_to_c(real_coords):
 928    return real_coords[..., 0] + real_coords[..., 1]*1.0j
 929
 930def c_to_r(cx_array):
 931    return cx_array.astype('complex').view('(2,)float')
 932
 933def order_eigs(eigenvalues, eigenvectors):
 934    """Sort eigenvalue/eigenvector tuples by increasing modulus.
 935
 936    This function accepts the eigenvalue/eigenvector data returned by
 937    the `np.linalg.eig` function.
 938
 939    Parameters
 940    ----------
 941    eigenvalues : ndarray
 942        Array of shape (..., n) representing eigenvalues of some
 943        matrices
 944    eigenvectors : ndarray
 945        Array of shape (..., n, n) representing eigenvectors of some
 946        matrices (as column vectors).
 947
 948    Returns
 949    -------
 950    tuple
 951        Tuple of the form `(eigvals, eigvecs)`, where both eigvals and
 952        eigvecs have the same data as the input arrays, but arranged
 953        so that eigenvalues and eigenvectors are in increasing order
 954        of modulus.
 955
 956    """
 957    eigorder = np.argsort(np.abs(eigenvalues), axis=-1)
 958    eigvecs = np.take_along_axis(eigenvectors, np.expand_dims(eigorder, axis=-2), axis=-1)
 959    eigvals = np.take_along_axis(eigenvalues, eigorder, axis=-1)
 960
 961    return eigvals, eigvecs
 962
 963def affine_disks_contain(cout, rout, cin, rin,
 964                         broadcast="elementwise"):
 965    if broadcast == "pairwise":
 966        pairwise_dists = np.linalg.norm(
 967            np.expand_dims(cout, 0) - np.expand_dims(cin, 1),
 968            axis=-1
 969        )
 970        radial_diffs = np.expand_dims(rout, 0) - np.expand_dims(rin, 1)
 971        return pairwise_dists < radial_diffs
 972
 973    return np.linalg.norm(cout - cin, axis=-1) < (rout - rin)
 974
 975
 976def disk_containments(c1, r1, c2, r2,
 977                      broadcast="elementwise"):
 978    if broadcast == "pairwise":
 979        pairwise_dists = np.linalg.norm(
 980            np.expand_dims(cout, 0) - np.expand_dims(cin, 1),
 981            axis=-1
 982        )
 983        radial_diff1 = np.expand_dims(rout, 0) - np.expand_dims(rin, 1)
 984        return (pairwise_dists < radial_diffs,
 985                pairwise_dists < -radial_diffs)
 986
 987    dists = np.linalg.norm(cout - cin, axis=-1)
 988    return (dists < (rout - rin),
 989            dists < (rin - rout))
 990
 991def disk_interactions(c1, r1, c2, r2,
 992                      broadcast="elementwise"):
 993    # WARNING: this will only work for Nx2 arrays
 994    if broadcast == "pairwise":
 995        pairwise_dists = np.linalg.norm(
 996            np.expand_dims(c1, 1) - np.expand_dims(c2, 0),
 997            axis=-1
 998        )
 999        radial_diff = np.expand_dims(r1, 1) - np.expand_dims(r2, 0)
1000        radial_sum = np.expand_dims(r1, 1) + np.expand_dims(r2, 0)
1001        return (pairwise_dists < radial_diff,
1002                pairwise_dists < -radial_diff,
1003                pairwise_dists < radial_sum)
1004
1005    dists = np.linalg.norm(c1 - c2, axis=-1)
1006
1007    return (dists < (r1 - r2),
1008            dists < (r2 - r1),
1009            dists < (r2 + r1))
1010
1011def indefinite_form(p, q, neg_first=True, **kwargs):
1012    n = p + q
1013
1014    mult = 1
1015    if neg_first:
1016        mult = -1
1017
1018    form_mat = zeros((n,n), **kwargs)
1019    form_mat[:p, :p] = mult * identity(p, **kwargs)
1020    form_mat[p:, p:] = -mult * identity(q, **kwargs)
1021
1022    return form_mat
1023
1024def symplectic_form(n, **kwargs):
1025    if n != n % 2:
1026        raise ValueError("Cannot construct a symplectic form in odd dimensions.")
1027
1028    m = n / 2
1029
1030    form_mat = zeros((n, n), **kwargs)
1031    form_mat[:m, m:] = -identity(m, **kwargs)
1032    form_mat[m:, :m] = identity(m, **kwargs)
1033
1034    return form_mat
1035
1036def symmetric_part(bilinear_form):
1037    return (bilinear_form + bilinear_form.swapaxes(-1, -2)) / 2
1038
1039def antisymmetric_part(bilinear_form):
1040    return (bilinear_form - bilinear_form.swapaxes(-1, -2)) / 2
1041
1042def matrix_func(func):
1043
1044    # get sage_func now, so if there's a name issue we'll throw an
1045    # error when the wrapped function is defined (rather than when
1046    # it's called)
1047    if SAGE_AVAILABLE:
1048        sage_func = getattr(sagewrap, func.__name__)
1049
1050    @functools.wraps(func)
1051    def wrapped(*args, **kwargs):
1052        if not SAGE_AVAILABLE:
1053            return func(*args, **kwargs)
1054
1055        return sage_func(*args, **kwargs)
1056
1057    return wrapped
1058
1059@matrix_func
1060def invert(mat):
1061    return np.linalg.inv(mat)
1062
1063@matrix_func
1064def kernel(mat):
1065    return numerical.svd_kernel(mat)
1066
1067@matrix_func
1068def eig(mat):
1069    return np.linalg.eig(mat)
1070
1071@matrix_func
1072def eigh(mat):
1073    return np.linalg.eigh(mat)
1074
1075@matrix_func
1076def det(mat):
1077    return np.linalg.det(mat)
1078
1079def check_type(base_ring=None, dtype=None, like=None,
1080               default_dtype='float64', integer_type=True):
1081
1082    default_ring = None
1083    if SAGE_AVAILABLE:
1084        default_ring = sagewrap.Integer
1085        if not integer_type:
1086            default_ring = sagewrap.SR
1087
1088    if like is not None:
1089        try:
1090            if dtype is None:
1091                dtype = like.dtype
1092        except AttributeError:
1093            if not types.is_linalg_type(like):
1094                dtype = np.dtype('O')
1095
1096    if base_ring is None and dtype == np.dtype('O') and SAGE_AVAILABLE:
1097        base_ring = default_ring
1098
1099    if base_ring is not None:
1100        if not SAGE_AVAILABLE:
1101            raise EnvironmentError(
1102                "Cannot specify base_ring unless running within sage"
1103            )
1104        if dtype is not None and dtype != np.dtype('object'):
1105            raise TypeError(
1106                "Cannot specify base_ring and dtype unless dtype is dtype('object')"
1107            )
1108        dtype = np.dtype('object')
1109    elif dtype is None:
1110        dtype = default_dtype
1111
1112    if not integer_type and np.can_cast(dtype, int):
1113        dtype = np.dtype('float64')
1114
1115    return (base_ring, dtype)
1116
1117def complex_type(**kwargs):
1118    base_ring, dtype = check_type(integer_type=False, **kwargs)
1119
1120    if np.can_cast(dtype, 'float'):
1121        return base_ring, np.dtype('complex')
1122
1123    return base_ring, dtype
1124
1125def guess_literal_ring(data):
1126    if not SAGE_AVAILABLE:
1127        return None
1128
1129    try:
1130        if data.dtype == np.dtype('O'):
1131            return sage.all.QQ
1132    except AttributeError:
1133        if not types.is_linalg_type(data):
1134            return sage.all.QQ
1135
1136    # this is maybe not Pythonic, but it's also not a mistake.
1137    return None
1138
1139def astype(val, dtype=None):
1140    if dtype is None:
1141        return val
1142
1143    return np.array(val).astype(dtype)
1144
1145def number(val, like=None, dtype=None, base_ring=None):
1146    if like is not None and dtype is not None:
1147        raise UserWarning(
1148            "Passing both 'like' and 'dtype' when specifying a number;"
1149            " 'like' parameter is ignored."
1150        )
1151
1152    if base_ring is not None:
1153        if dtype is not None:
1154            raise UserWarning(
1155                "Passing both 'base_ring' and 'dtype' when specifying a number;"
1156                " 'dtype' is ignored (assumed to be dtype('O'))"
1157            )
1158        dtype = np.dtype('O')
1159
1160        if not SAGE_AVAILABLE:
1161            raise UserWarning( "Specifying base_ring when sage is not"
1162            "available has no effect."  )
1163
1164    if not SAGE_AVAILABLE:
1165        return astype(val, dtype)
1166
1167    if dtype is None and like is not None:
1168        try:
1169            dtype = like.dtype
1170        except AttributeError:
1171            if not types.is_linalg_type(like):
1172                dtype = np.dtype('O')
1173
1174    if dtype == np.dtype('O'):
1175        if isinstance(val, int):
1176            # we use SR instead of Integer here, because numpy
1177            # sometimes silently converts sage Integers
1178            return sage.all.SR(val)
1179        if isinstance(val, float):
1180            return sage.all.SR(sage.all.QQ(val))
1181
1182    return astype(val, dtype)
1183
1184def wrap_elementary(func):
1185    @functools.wraps(func)
1186    def wrapped(*args, like=None, base_ring=None, dtype=None, **kwargs):
1187        if not SAGE_AVAILABLE:
1188            return func(*args, **kwargs)
1189
1190        else:
1191            packaged_args = []
1192            arg_ring = None
1193            if base_ring is not None:
1194                arg_ring = sagewrap.SR
1195
1196            for arg in args:
1197                arg_like = like
1198                if arg_like is None:
1199                    arg_like = arg
1200
1201                packaged_args.append(
1202                    array_like(arg, like=arg_like,
1203                               base_ring=arg_ring,
1204                               dtype=dtype,
1205                               integer_type=False)
1206                )
1207
1208            return sagewrap.change_base_ring(
1209                func(*packaged_args, **kwargs),
1210                base_ring
1211            )
1212
1213    return wrapped
1214
1215@wrap_elementary
1216def power(base, exp):
1217    return np.power(base, exp)
1218
1219@wrap_elementary
1220def cos(arg):
1221    return np.cos(arg)
1222
1223@wrap_elementary
1224def sin(arg):
1225    return np.sin(arg)
1226
1227@wrap_elementary
1228def conjugate(arg):
1229    return np.conjugate(arg)
1230
1231def real(arg):
1232    if not SAGE_AVAILABLE or types.is_linalg_type(arg):
1233        return np.real(arg)
1234
1235    # np.real doesn't work well with sage since it expects the real
1236    # attribute to be a property, not a callable. So we need to wrap
1237    # sage's real function explicitly.
1238    return sagewrap.real(arg)
1239
1240def imag(arg):
1241    if not SAGE_AVAILABLE or types.is_linalg_type(arg):
1242        return np.imag(arg)
1243
1244    return sagewrap.imag(arg)
1245
1246def array_like(array, like=None, base_ring=None, dtype=None,
1247               integer_type=False, **kwargs):
1248
1249    if like is None:
1250        like = array
1251
1252    base_ring, dtype = check_type(base_ring, dtype, like,
1253                                  integer_type=integer_type)
1254
1255    arr = np.array(array, dtype=dtype, **kwargs)
1256
1257    if base_ring is not None:
1258        return sagewrap.change_base_ring(arr, base_ring)
1259    return arr
1260
1261def pi(**kwargs):
1262    if not SAGE_AVAILABLE:
1263        return np.pi
1264
1265    base_ring, dtype = check_type(**kwargs)
1266
1267    if base_ring is not None:
1268        return sagewrap.pi
1269
1270    return np.pi
1271
1272def unit_imag(**kwargs):
1273    if not SAGE_AVAILABLE:
1274        return 1j
1275
1276    base_ring, dtype = check_type(**kwargs)
1277
1278    if base_ring is not None:
1279        return sagewrap.I
1280    return 1j
1281
1282
1283def zeros(shape, base_ring=None, dtype=None, like=None,
1284          integer_type=True, **kwargs):
1285
1286    base_ring, dtype = check_type(base_ring, dtype, like,
1287                                  integer_type=integer_type)
1288
1289    zero_arr = np.zeros(shape, dtype=dtype, **kwargs)
1290    if base_ring is not None:
1291        return sagewrap.change_base_ring(zero_arr, base_ring)
1292    return zero_arr
1293
1294def ones(shape, base_ring=None, dtype=None, like=None,
1295         integer_type=True, **kwargs):
1296    base_ring, dtype = check_type(base_ring, dtype, like,
1297                                  integer_type=integer_type)
1298
1299    ones_arr = np.ones(shape, dtype=dtype, **kwargs)
1300    if base_ring is not None:
1301        return sagewrap.change_base_ring(ones_arr, base_ring)
1302    return ones_arr
1303
1304def identity(n, base_ring=None, dtype=None, like=None,
1305             integer_type=True, **kwargs):
1306    base_ring, dtype = check_type(base_ring, dtype, like,
1307                                  integer_type=integer_type)
1308
1309    identity_arr = np.identity(n, dtype=dtype, **kwargs)
1310
1311    if base_ring is not None:
1312        return sagewrap.change_base_ring(identity_arr, base_ring)
1313    return identity_arr
1314
1315def change_base_ring(array, base_ring=None,
1316                     rational_approx=False):
1317    if base_ring is None:
1318        return array
1319
1320    if not SAGE_AVAILABLE:
1321        raise EnvironmentError(
1322            "Cannot change base ring unless sage is available."
1323        )
1324
1325    return sagewrap.change_base_ring(array, base_ring,
1326                                     rational_approx=rational_approx)