astroalign.py 文件源码

def _generate_invariants(sources):
    """Return an array of (unique) invariants derived from the array `sources`.
Return an array of the indices of `sources` that correspond to each invariant,
arranged as described in _arrangetriplet.
"""
    from scipy.spatial import KDTree
    from itertools import combinations
    from functools import partial
    arrange = partial(_arrangetriplet, sources=sources)

    inv = []
    triang_vrtx = []
    coordtree = KDTree(sources)
    for asrc in sources:
        __, indx = coordtree.query(asrc, NUM_NEAREST_NEIGHBORS)

        # Generate all possible triangles with the 5 indx provided, and store
        # them with the order (a, b, c) defined in _arrangetriplet
        all_asterism_triang = [arrange(vertex_indices=list(cmb))
                               for cmb in combinations(indx, 3)]
        triang_vrtx.extend(all_asterism_triang)

        inv.extend([_invariantfeatures(*sources[triplet])
                    for triplet in all_asterism_triang])

    # Remove here all possible duplicate triangles
    uniq_ind = [pos for (pos, elem) in enumerate(inv)
                if elem not in inv[pos + 1:]]
    inv_uniq = _np.array(inv)[uniq_ind]
    triang_vrtx_uniq = _np.array(triang_vrtx)[uniq_ind]

    return inv_uniq, triang_vrtx_uniq