astrometry.net-0.85/libkd/kdtree.c

/*
 # This file is part of libkd.
 # Licensed under a 3-clause BSD style license - see LICENSE
 */

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>
#include <math.h>

#include "os-features.h"
#include "kdtree.h"
#include "kdtree_internal.h"
#include "kdtree_internal_common.h"
#include "kdtree_mem.h"
#include "keywords.h"
#include "an-fls.h"
#include "errors.h"
#include "log.h"

kdtree_t* kdtree_build(kdtree_t* kd, void *data, int N, int D, int Nleaf,
                       int treetype, unsigned int options) {
    return kdtree_build_2(kd, data, N, D, Nleaf, treetype, options,
                          NULL, NULL);
}

void kdtree_print(kdtree_t* kd) {
    printf("kdtree:\n");
    printf("  type 0x%x\n", kd->treetype);
    printf("  lr %p\n", kd->lr);
    printf("  perm %p\n", kd->perm);
    printf("  bb %p\n", kd->bb.any);
    printf("  nbb %i\n", kd->n_bb);
    printf("  split %p\n", kd->split.any);
    printf("  splitdim %p\n", kd->splitdim);
    printf("  dimbits %i\n", kd->dimbits);
    printf("  dimmask 0x%x\n", kd->dimmask);
    printf("  splitmask 0x%x\n", kd->splitmask);
    printf("  data %p\n", kd->data.any);
    printf("  free data %i\n", (int)kd->free_data);
    printf("  range");
    if (kd->minval && kd->maxval) {
        int i;
        for (i=0; i<kd->ndim; i++)
            printf(" [%g, %g]", kd->minval[i], kd->maxval[i]);
    } else
        printf(" (none)\n");
    printf("\n");
    printf("  scale %g\n", kd->scale);
    printf("  invscale %g\n", kd->invscale);
    printf("  Ndata %i\n", kd->ndata);
    printf("  Ndim %i\n", kd->ndim);
    printf("  Nnodes %i\n", kd->nnodes);
    printf("  Nbottom %i\n", kd->nbottom);
    printf("  Ninterior %i\n", kd->ninterior);
    printf("  Nlevels %i\n", kd->nlevels);
    printf("  has_linear_lr %i\n", (int)kd->has_linear_lr);
    printf("  name %s\n", kd->name);
}

int kdtree_n(const kdtree_t* kd) {
    return kd->ndata;
}

int kdtree_nnodes(const kdtree_t* kd) {
    return kd->nnodes;
}

int kdtree_has_old_bb(const kdtree_t* kd) {
    return kd->n_bb != kd->nnodes;
}

KD_DECLARE(kdtree_update_funcs, void, (kdtree_t*));

void kdtree_update_funcs(kdtree_t* kd) {
    KD_DISPATCH(kdtree_update_funcs, kd->treetype,, (kd));
}

static int get_tree_size(int treetype) {
    switch (treetype & KDT_TREE_MASK) {
    case KDT_TREE_DOUBLE:
        return sizeof(double);
    case KDT_TREE_FLOAT:
        return sizeof(float);
    case KDT_TREE_U32:
        return sizeof(u32);
    case KDT_TREE_U16:
        return sizeof(u16);
    }
    return -1;
}
static int get_data_size(int treetype) {
    switch (treetype & KDT_DATA_MASK) {
    case KDT_DATA_DOUBLE:
        return sizeof(double);
    case KDT_DATA_FLOAT:
        return sizeof(float);
    case KDT_DATA_U32:
        return sizeof(u32);
    case KDT_DATA_U16:
        return sizeof(u16);
    }
    return -1;
}

size_t kdtree_sizeof_lr(const kdtree_t* kd) {
    return sizeof(int32_t) * kd->nbottom;
}

size_t kdtree_sizeof_perm(const kdtree_t* kd) {
    return sizeof(u32) * kd->ndata;
}

size_t kdtree_sizeof_bb(const kdtree_t* kd) {
    return get_tree_size(kd->treetype) * kd->ndim * kd->nnodes;
}

size_t kdtree_sizeof_split(const kdtree_t* kd) {
    return get_tree_size(kd->treetype) * kd->ninterior;
}

size_t kdtree_sizeof_splitdim(const kdtree_t* kd) {
    return sizeof(u8) * kd->ninterior;
}

size_t kdtree_sizeof_data(const kdtree_t* kd) {
    return get_data_size(kd->treetype) * kd->ndim * kd->ndata;
}

void kdtree_memory_report(kdtree_t* kd) {
    int mem;
    int n, sz;
    int total = 0;

    int tsz = get_tree_size(kd->treetype);
    int dsz = get_data_size(kd->treetype);

    char* fmt = "%-10s:   %12i %10s * %2i = %12i B  (%10.3f MB)\n";

    printf("Memory usage of kdtree (ndata %i, ndim %i, nnodes %i, nleaves %i)\n",
           kd->ndata, kd->ndim, kd->nnodes, kd->nbottom);

    if (kd->lr) {
        n = kd->nbottom;
        sz = sizeof(int32_t);
        mem = n*sz;
        printf(fmt, "lr", n, "leaves", sz, mem, 1e-6*mem);
        total += mem;
    }
    if (kd->perm) {
        n = kd->ndata;
        sz = sizeof(u32);
        mem = n*sz;
        printf(fmt, "perm", n, "points", sz, mem, 1e-6*mem);
        total += mem;
    }
    if (kd->bb.any) {
        n = kd->nnodes;
        sz = tsz * 2 * kd->ndim;
        mem = n*sz;
        printf(fmt, "bbox", n, "nodes", sz, mem, 1e-6*mem);
        total += mem;
    }
    if (kd->split.any) {
        n = kd->ninterior;
        sz = tsz;
        mem = n*sz;
        printf(fmt, "split", n, "splits", sz, mem, 1e-6*mem);
        total += mem;
    }
    if (kd->splitdim) {
        n = kd->ninterior;
        sz = sizeof(u8);
        mem = n*sz;
        printf(fmt, "splitdim", n, "splits", sz, mem, 1e-6*mem);
        total += mem;
    }

    printf("Total without data:                          %12i B  (%10.3f MB)\n", total, total*1e-6);

    if (kd->data.any) {
        n = kd->ndata;
        sz = dsz * kd->ndim;
        mem = n*sz;
        printf(fmt, "data", n, "points", sz, mem, 1e-6*mem);
        total += mem;
    }
    printf("Total including data:                        %12i B  (%10.3f MB)\n", total, total*1e-6);

}

int kdtree_level_start(const kdtree_t* kd, int level) {
    return (1 << level) - 1;
}

int kdtree_level_end(const kdtree_t* kd, int level) {
    return (1 << (level+1)) - 2;
}

static inline u8 node_level(int nodeid) {
    return an_flsB(nodeid + 1);
}

int kdtree_get_level(const kdtree_t* kd, int nodeid) {
    return node_level(nodeid);
}

int kdtree_get_splitdim(const kdtree_t* kd, int nodeid) {
    u32 tmpsplit;
    if (kd->splitdim)
        return kd->splitdim[nodeid];

    switch (kdtree_treetype(kd)) {
    case KDT_TREE_U32:
        tmpsplit = kd->split.u[nodeid];
        break;
    case KDT_TREE_U16:
        tmpsplit = kd->split.s[nodeid];
        break;
    default:
        return -1;
    }
    return tmpsplit & kd->dimmask;
}

double kdtree_get_splitval(const kdtree_t* kd, int nodeid) {
    assert(kd->fun.get_splitval);
    return kd->fun.get_splitval(kd, nodeid);
}

void* kdtree_get_data(const kdtree_t* kd, int i) {
    switch (kdtree_datatype(kd)) {
    case KDT_DATA_DOUBLE:
        return kd->data.d + kd->ndim * i;
    case KDT_DATA_FLOAT:
        return kd->data.f + kd->ndim * i;
    case KDT_DATA_U32:
        return kd->data.u + kd->ndim * i;
    case KDT_DATA_U16:
        return kd->data.s + kd->ndim * i;
    default:
        ERROR("kdtree_get_data: invalid data type %i", kdtree_datatype(kd));
        return NULL;
    }
}

void kdtree_copy_data_double(const kdtree_t* kd, int start, int N, double* dest) {
    int i;
    int d, D;
    D = kd->ndim;
    switch (kdtree_datatype(kd)) {
    case KDT_DATA_DOUBLE:
        memcpy(dest, kd->data.d + start*D,
               N * D * sizeof(double));
        break;
    case KDT_DATA_FLOAT:
        for (i=0; i<(N * D); i++)
            dest[i] = kd->data.f[start*D + i];
        break;
    case KDT_DATA_U32:
        for (i=0; i<N; i++)
            for (d=0; d<D; d++)
                dest[i*D + d] = POINT_INVSCALE(kd, d, kd->data.u[(start + i)*D + d]);
        break;
    case KDT_DATA_U16:
        for (i=0; i<N; i++)
            for (d=0; d<D; d++)
                dest[i*D + d] = POINT_INVSCALE(kd, d, kd->data.s[(start + i)*D + d]);
        break;
    default:
        ERROR("kdtree_copy_data_double: invalid data type %i", kdtree_datatype(kd));
        return;
    }
}

int kdtree_permute(const kdtree_t* tree, int ind) {
    if (!tree->perm)
        return ind;
    return tree->perm[ind];
}

void kdtree_inverse_permutation(const kdtree_t* tree, int* invperm) {
    int i;
    if (!tree->perm) {
        for (i=0; i<tree->ndata; i++)
            invperm[i] = i;
    } else {
        for (i=0; i<tree->ndata; i++) {
            assert((int)(tree->perm[i]) >= 0);
            assert((int)(tree->perm[i]) < tree->ndata);
            invperm[tree->perm[i]] = i;
        }
    }
}

const char* kdtree_build_options_to_string(int opts) {
    static char buf[256];
    sprintf(buf, "%s%s%s%s%s",
            (opts & KD_BUILD_BBOX) ? "BBOX ":"",
            (opts & KD_BUILD_SPLIT) ? "SPLIT ":"",
            (opts & KD_BUILD_SPLITDIM) ? "SPLITDIM ":"",
            (opts & KD_BUILD_NO_LR) ? "NOLR ":"",
            (opts & KD_BUILD_LINEAR_LR) ? "LINEARLR ":"");
    return buf;
}

const char* kdtree_kdtype_to_string(int kdtype) {
    switch (kdtype) {
    case KDT_DATA_DOUBLE:
    case KDT_TREE_DOUBLE:
    case KDT_EXT_DOUBLE:
        return "double";
    case KDT_DATA_FLOAT:
    case KDT_TREE_FLOAT:
    case KDT_EXT_FLOAT:
        return "float";
    case KDT_DATA_U32:
    case KDT_TREE_U32:
        return "u32";
    case KDT_DATA_U16:
    case KDT_TREE_U16:
        return "u16";
    default:
        return NULL;
    }
}

int kdtree_kdtype_parse_data_string(const char* str) {
    if (!str) return KDT_DATA_NULL;
    if (!strcmp(str, "double")) {
        return KDT_DATA_DOUBLE;
    } else if (!strcmp(str, "float")) {
        return KDT_DATA_FLOAT;
    } else if (!strcmp(str, "u32")) {
        return KDT_DATA_U32;
    } else if (!strcmp(str, "u16")) {
        return KDT_DATA_U16;
    } else
        return KDT_DATA_NULL;
}

int kdtree_kdtype_parse_tree_string(const char* str) {
    if (!str) return KDT_TREE_NULL;
    if (!strcmp(str, "double")) {
        return KDT_TREE_DOUBLE;
    } else if (!strcmp(str, "float")) {
        return KDT_TREE_FLOAT;
    } else if (!strcmp(str, "u32")) {
        return KDT_TREE_U32;
    } else if (!strcmp(str, "u16")) {
        return KDT_TREE_U16;
    } else
        return KDT_TREE_NULL;
}

int kdtree_kdtype_parse_ext_string(const char* str) {
    if (!str) return KDT_EXT_NULL;
    if (!strcmp(str, "double")) {
        return KDT_EXT_DOUBLE;
    } else if (!strcmp(str, "float")) {
        return KDT_EXT_FLOAT;
    } else
        return KDT_EXT_NULL;
}

int kdtree_kdtypes_to_treetype(int exttype, int treetype, int datatype) {
    // HACK - asserts here...
    return exttype | treetype | datatype;
}

kdtree_t* kdtree_new(int N, int D, int Nleaf) {
    kdtree_t* kd;
    int maxlevel, nnodes;
    maxlevel = kdtree_compute_levels(N, Nleaf);
    kd = CALLOC(1, sizeof(kdtree_t));
    nnodes = (1 << maxlevel) - 1;
    kd->nlevels = maxlevel;
    kd->ndata = N;
    kd->ndim = D;
    kd->nnodes = nnodes;
    kd->nbottom = 1 << (maxlevel - 1);
    kd->ninterior = kd->nbottom - 1;
    assert(kd->nbottom + kd->ninterior == kd->nnodes);
    return kd;
}

void kdtree_set_limits(kdtree_t* kd, double* low, double* high) {
    int D = kd->ndim;
    if (!kd->minval) {
        kd->minval = MALLOC(D * sizeof(double));
    }
    if (!kd->maxval) {
        kd->maxval = MALLOC(D * sizeof(double));
    }
    memcpy(kd->minval, low,  D * sizeof(double));
    memcpy(kd->maxval, high, D * sizeof(double));
}

double kdtree_get_conservative_query_radius(const kdtree_t* kd, double radius) {
    if (!kd->minval) {
        return radius;
    }
    return sqrt(radius* radius + kd->ndim * 0.25 * kd->invscale*kd->invscale);
}

int kdtree_compute_levels(int N, int Nleaf) {
    int nnodes = N / Nleaf;
    int maxlevel = 1;
    while (nnodes) {
        nnodes = nnodes >> 1;
        maxlevel++;
    }
    return maxlevel;
}

int kdtree_nnodes_to_nlevels(int Nnodes) {
    return an_flsB(Nnodes + 1);
}

/* This returns the NODE id (not leaf index) */
int kdtree_first_leaf(const kdtree_t* kd, int nodeid) {
    int dlevel;
    dlevel = (kd->nlevels - 1) - node_level(nodeid);
    return ((nodeid+1) << dlevel) - 1;
}

/* This returns the NODE id (not leaf index) */
int kdtree_last_leaf(const kdtree_t* kd, int nodeid) {
    int dlevel, twodl, nodeid_twodl;
    dlevel = (kd->nlevels - 1) - node_level(nodeid);
    twodl = (1 << dlevel);
    nodeid_twodl = (nodeid << dlevel);
    return nodeid_twodl + (twodl - 1)*2;
}

static int calculate_R(const kdtree_t* kd, int leafid) {
    int l;
    unsigned int mask, N, L;

    mask = (1 << (kd->nlevels-1));
    N = kd->ndata;
    L = 0;
    // Compute the L index of the node one to the right of this node.
    int nextguy = leafid + 1;
    if (nextguy == kd->nbottom)
        return kd->ndata-1;
    for (l=0; l<(kd->nlevels-1); l++) {
        mask /= 2;
        if (nextguy & mask) {
            L += N/2;
            N = (N+1)/2;
        } else {
            N = N/2;
        }
    }
    L--;
    return L;
}

static int linear_lr(const kdtree_t* kd, int leafid) {
    return ((u64)leafid * kd->ndata) / kd->nbottom;
}

int kdtree_leaf_left(const kdtree_t* kd, int nodeid) {
    int leafid = nodeid - kd->ninterior;
    if (!leafid)
        return 0;
    if (kd->has_linear_lr)
        return linear_lr(kd, leafid);
    if (kd->lr)
        return kd->lr[leafid-1] + 1;
    return calculate_R(kd, leafid-1) + 1;
}

int kdtree_leaf_right(const kdtree_t* kd, int nodeid) {
    int leafid = nodeid - kd->ninterior;
    if (kd->has_linear_lr)
        return linear_lr(kd, leafid+1) - 1;
    if (kd->lr)
        return kd->lr[leafid];
    return calculate_R(kd, leafid);
}


int kdtree_left(const kdtree_t* kd, int nodeid) {
    if (KD_IS_LEAF(kd, nodeid)) {
        return kdtree_leaf_left(kd, nodeid);
    } else {
        // leftmost child's L.
        int leftmost = kdtree_first_leaf(kd, nodeid);
        return kdtree_leaf_left(kd, leftmost);
    }
}

int kdtree_right(const kdtree_t* kd, int nodeid) {
    if (KD_IS_LEAF(kd, nodeid)) {
        return kdtree_leaf_right(kd, nodeid);
    } else {
        // rightmost child's R.
        int rightmost = kdtree_last_leaf(kd, nodeid);
        return kdtree_leaf_right(kd, rightmost);
    }
}

int kdtree_is_leaf_node_empty(const kdtree_t* kd, int nodeid) {
    int L, R;
    L = kdtree_leaf_left (kd, nodeid);
    R = kdtree_leaf_right(kd, nodeid);
    return (L == R+1);
}

int kdtree_is_node_empty(const kdtree_t* kd, int nodeid) {
    int L, R;
    L = kdtree_left (kd, nodeid);
    R = kdtree_right(kd, nodeid);
    return (L == R+1);
}

int kdtree_npoints(const kdtree_t* kd, int nodeid) {
    return 1 + kdtree_right(kd, nodeid) - kdtree_left(kd, nodeid);
}

void kdtree_free_query(kdtree_qres_t *kq) {
    if (!kq) return;
    FREE(kq->results.any);
    FREE(kq->sdists);
    FREE(kq->inds);
    FREE(kq);
}

void kdtree_free(kdtree_t *kd) {
    if (!kd) return;
    FREE(kd->name);
    FREE(kd->lr);
    FREE(kd->perm);
    FREE(kd->bb.any);
    FREE(kd->split.any);
    FREE(kd->splitdim);
    if (kd->free_data)
        FREE(kd->data.any);
    FREE(kd->minval);
    FREE(kd->maxval);
    //FREE(kd->fun);
    FREE(kd);
}

int kdtree_nearest_neighbour(const kdtree_t* kd, const void* pt, double* p_mindist2) {
    return kdtree_nearest_neighbour_within(kd, pt, HUGE_VAL, p_mindist2);
}

int kdtree_check(const kdtree_t* kd) {
    assert(kd->fun.check);
    return kd->fun.check(kd);
}

int kdtree_nearest_neighbour_within(const kdtree_t* kd, const void *pt, double maxd2,
                                    double* p_mindist2) {
    double bestd2 = maxd2;
    int ibest = -1;

    assert(kd->fun.nearest_neighbour_internal);
    kd->fun.nearest_neighbour_internal(kd, pt, &bestd2, &ibest);

    if (p_mindist2 && (ibest != -1))
        *p_mindist2 = bestd2;
    return ibest;
}

KD_DECLARE(kdtree_node_node_mindist2, double, (const kdtree_t* kd1, int node1, const kdtree_t* kd2, int node2));

double kdtree_node_node_mindist2(const kdtree_t* kd1, int node1,
                                 const kdtree_t* kd2, int node2) {
    double res = HUGE_VAL;
    KD_DISPATCH(kdtree_node_node_mindist2, kd1->treetype, res=, (kd1, node1, kd2, node2));
    return res;
}

KD_DECLARE(kdtree_node_node_maxdist2, double, (const kdtree_t* kd1, int node1, const kdtree_t* kd2, int node2));

double kdtree_node_node_maxdist2(const kdtree_t* kd1, int node1,
                                 const kdtree_t* kd2, int node2) {
    double res = HUGE_VAL;
    KD_DISPATCH(kdtree_node_node_maxdist2, kd1->treetype, res=, (kd1, node1, kd2, node2));
    return res;
}

KD_DECLARE(kdtree_node_node_mindist2_exceeds, int, (const kdtree_t* kd1, int node1, const kdtree_t* kd2, int node2, double maxd2));

int kdtree_node_node_mindist2_exceeds(const kdtree_t* kd1, int node1,
                                      const kdtree_t* kd2, int node2,
                                      double dist2) {
    int res = FALSE;
    KD_DISPATCH(kdtree_node_node_mindist2_exceeds, kd1->treetype, res=, (kd1, node1, kd2, node2, dist2));
    return res;
}

KD_DECLARE(kdtree_node_node_maxdist2_exceeds, int, (const kdtree_t* kd1, int node1, const kdtree_t* kd2, int node2, double maxd2));

int kdtree_node_node_maxdist2_exceeds(const kdtree_t* kd1, int node1,
                                      const kdtree_t* kd2, int node2,
                                      double dist2) {
    int res = FALSE;
    KD_DISPATCH(kdtree_node_node_maxdist2_exceeds, kd1->treetype, res=, (kd1, node1, kd2, node2, dist2));
    return res;
}

KD_DECLARE(kdtree_node_point_mindist2, double, (const kdtree_t* kd, int node, const void* query));

double kdtree_node_point_mindist2(const kdtree_t* kd, int node, const void* pt) {
    double res = HUGE_VAL;
    KD_DISPATCH(kdtree_node_point_mindist2, kd->treetype, res=, (kd, node, pt));
    return res;
}

KD_DECLARE(kdtree_node_point_maxdist2, double, (const kdtree_t* kd, int node, const void* query));

double kdtree_node_point_maxdist2(const kdtree_t* kd, int node, const void* pt) {
    double res = HUGE_VAL;
    KD_DISPATCH(kdtree_node_point_maxdist2, kd->treetype, res=, (kd, node, pt));
    return res;
}

KD_DECLARE(kdtree_node_point_mindist2_exceeds, int, (const kdtree_t* kd, int node, const void* query, double maxd2));

int kdtree_node_point_mindist2_exceeds(const kdtree_t* kd, int node, const void* pt,
                                       double dist2) {
    int res = FALSE;
    KD_DISPATCH(kdtree_node_point_mindist2_exceeds, kd->treetype, res=, (kd, node, pt, dist2));
    return res;
}

KD_DECLARE(kdtree_node_point_maxdist2_exceeds, int, (const kdtree_t* kd, int node, const void* query, double maxd2));

int kdtree_node_point_maxdist2_exceeds(const kdtree_t* kd, int node, const void* pt,
                                       double dist2) {
    int res = FALSE;
    KD_DISPATCH(kdtree_node_point_maxdist2_exceeds, kd->treetype, res=, (kd, node, pt, dist2));
    return res;
}

void kdtree_nodes_contained(const kdtree_t* kd,
                            const void* querylow, const void* queryhi,
                            void (*callback_contained)(const kdtree_t* kd, int node, void* extra),
                            void (*callback_overlap)(const kdtree_t* kd, int node, void* extra),
                            void* cb_extra) {
    assert(kd->fun.nodes_contained);
    kd->fun.nodes_contained(kd, querylow, queryhi, callback_contained, callback_overlap, cb_extra);
}

int kdtree_get_bboxes(const kdtree_t* kd, int node, void* bblo, void* bbhi) {
    assert(kd->fun.get_bboxes);
    return kd->fun.get_bboxes(kd, node, bblo, bbhi);
}

void kdtree_fix_bounding_boxes(kdtree_t* kd) {
    assert(kd->fun.fix_bounding_boxes);
    kd->fun.fix_bounding_boxes(kd);
}