1 /*
2 # This file is part of the Astrometry.net suite.
3 # Licensed under a 3-clause BSD style license - see LICENSE
4 */
5
6 #include <stdio.h>
7 #include <stdlib.h>
8 #include <assert.h>
9
10 #include "starkd.h"
11 #include "kdtree.h"
12 #include "kdtree_fits_io.h"
13 #include "starutil.h"
14 #include "fitsbin.h"
15 #include "fitstable.h"
16 #include "errors.h"
17 #include "tic.h"
18 #include "log.h"
19 #include "ioutils.h"
20 #include "fitsioutils.h"
21
startree_alloc()22 static startree_t* startree_alloc() {
23 startree_t* s = calloc(1, sizeof(startree_t));
24 if (!s) {
25 fprintf(stderr, "Failed to allocate a star kdtree struct.\n");
26 return NULL;
27 }
28 return s;
29 }
30
startree_get_tagalong_column_names(startree_t * s,sl * lst)31 sl* startree_get_tagalong_column_names(startree_t* s, sl* lst) {
32 if (!startree_has_tagalong(s))
33 return NULL;
34 return fitstable_get_fits_column_names(startree_get_tagalong(s), lst);
35 }
36
startree_get_tagalong_N_columns(startree_t * s)37 int startree_get_tagalong_N_columns(startree_t* s) {
38 if (!startree_has_tagalong(s))
39 return 0;
40 return fitstable_get_N_fits_columns(startree_get_tagalong(s));
41 }
42
43 /**
44 Returns the name of the 'i'th column in the tagalong table.
45 */
startree_get_tagalong_column_name(startree_t * s,int i)46 const char* startree_get_tagalong_column_name(startree_t* s, int i) {
47 if (!startree_has_tagalong(s))
48 return NULL;
49 return fitstable_get_fits_column_name(startree_get_tagalong(s), i);
50 }
51
startree_get_tagalong_column_fits_type(startree_t * s,int i)52 tfits_type startree_get_tagalong_column_fits_type(startree_t* s, int i) {
53 if (!startree_has_tagalong(s))
54 return TFITS_BIN_TYPE_UNKNOWN;
55 return fitstable_get_fits_column_type(startree_get_tagalong(s), i);
56 }
57
startree_get_tagalong_column_array_size(startree_t * s,int i)58 int startree_get_tagalong_column_array_size(startree_t* s, int i) {
59 if (!startree_has_tagalong(s))
60 return -1;
61 return fitstable_get_fits_column_array_size(startree_get_tagalong(s), i);
62 }
63
64
get_data_column(startree_t * s,const char * colname,const int * inds,int N,tfits_type tt)65 static void* get_data_column(startree_t* s, const char* colname, const int* inds, int N, tfits_type tt) {
66 fitstable_t* table;
67 void* arr;
68 if (N == 0) {
69 logmsg("Warning: zero stars (elements) in your request for data column \"%s\"\n", colname);
70 return NULL;
71 }
72 table = startree_get_tagalong(s);
73 if (!table) {
74 ERROR("No tag-along data found");
75 return NULL;
76 }
77 arr = fitstable_read_column_inds(table, colname, tt, inds, N);
78 if (!arr) {
79 ERROR("Failed to read tag-along data column \"%s\"", colname);
80 return NULL;
81 }
82 return arr;
83 }
84
85
startree_get_data_column(startree_t * s,const char * colname,const int * inds,int N)86 double* startree_get_data_column(startree_t* s, const char* colname, const int* inds, int N) {
87 return get_data_column(s, colname, inds, N, fitscolumn_double_type());
88 }
89
startree_get_data_column_int64(startree_t * s,const char * colname,const int * inds,int N)90 int64_t* startree_get_data_column_int64(startree_t* s, const char* colname, const int* inds, int N) {
91 return get_data_column(s, colname, inds, N, fitscolumn_i64_type());
92 }
93
startree_get_data_column_array(startree_t * s,const char * colname,const int * indices,int N,int * arraysize)94 double* startree_get_data_column_array(startree_t* s, const char* colname, const int* indices, int N, int* arraysize) {
95 fitstable_t* table;
96 tfits_type dubl = fitscolumn_double_type();
97 double* arr;
98 table = startree_get_tagalong(s);
99 if (!table) {
100 ERROR("No tag-along data found");
101 return NULL;
102 }
103 arr = fitstable_read_column_array_inds(table, colname, dubl, indices, N, arraysize);
104 if (!arr) {
105 ERROR("Failed to read tag-along data");
106 return NULL;
107 }
108 return arr;
109 }
110
startree_free_data_column(startree_t * s,double * d)111 void startree_free_data_column(startree_t* s, double* d) {
112 free(d);
113 }
114
startree_search_for_radec(const startree_t * s,double ra,double dec,double radius,double ** xyzresults,double ** radecresults,int ** starinds,int * nresults)115 void startree_search_for_radec(const startree_t* s, double ra, double dec, double radius,
116 double** xyzresults, double** radecresults,
117 int** starinds, int* nresults) {
118 double xyz[3];
119 double r2;
120 radecdeg2xyzarr(ra, dec, xyz);
121 r2 = deg2distsq(radius);
122 startree_search_for(s, xyz, r2, xyzresults, radecresults, starinds, nresults);
123 }
124
startree_search_for(const startree_t * s,const double * xyzcenter,double radius2,double ** xyzresults,double ** radecresults,int ** starinds,int * nresults)125 void startree_search_for(const startree_t* s, const double* xyzcenter, double radius2,
126 double** xyzresults, double** radecresults,
127 int** starinds, int* nresults) {
128 kdtree_qres_t* res = NULL;
129 int opts;
130 double* xyz;
131 int i, N;
132
133 opts = KD_OPTIONS_SMALL_RADIUS;
134 if (xyzresults || radecresults)
135 opts |= KD_OPTIONS_RETURN_POINTS;
136
137 res = kdtree_rangesearch_options(s->tree, xyzcenter, radius2, opts);
138
139 if (!res || !res->nres) {
140 if (xyzresults)
141 *xyzresults = NULL;
142 if (radecresults)
143 *radecresults = NULL;
144 if (starinds)
145 *starinds = NULL;
146 *nresults = 0;
147 if (res)
148 kdtree_free_query(res);
149 return;
150 }
151
152 xyz = res->results.d;
153 N = res->nres;
154 *nresults = N;
155
156 if (radecresults) {
157 *radecresults = malloc(N * 2 * sizeof(double));
158 for (i=0; i<N; i++)
159 xyzarr2radecdegarr(xyz + i*3, (*radecresults) + i*2);
160 }
161 if (xyzresults) {
162 // Steal the results array.
163 *xyzresults = xyz;
164 res->results.d = NULL;
165 }
166 if (starinds) {
167 *starinds = malloc(res->nres * sizeof(int));
168 for (i=0; i<N; i++)
169 (*starinds)[i] = res->inds[i];
170 }
171 kdtree_free_query(res);
172 }
173
174
startree_search(const startree_t * s,const double * xyzcenter,double radius2,double ** xyzresults,double ** radecresults,int * nresults)175 void startree_search(const startree_t* s, const double* xyzcenter, double radius2,
176 double** xyzresults, double** radecresults, int* nresults) {
177 startree_search_for(s, xyzcenter, radius2, xyzresults, radecresults, NULL, nresults);
178 }
179
startree_N(const startree_t * s)180 int startree_N(const startree_t* s) {
181 return s->tree->ndata;
182 }
183
startree_nodes(const startree_t * s)184 int startree_nodes(const startree_t* s) {
185 return s->tree->nnodes;
186 }
187
startree_D(const startree_t * s)188 int startree_D(const startree_t* s) {
189 return s->tree->ndim;
190 }
191
startree_header(const startree_t * s)192 qfits_header* startree_header(const startree_t* s) {
193 return s->header;
194 }
195
get_chunks(startree_t * s,il * wordsizes)196 static bl* get_chunks(startree_t* s, il* wordsizes) {
197 bl* chunks = bl_new(4, sizeof(fitsbin_chunk_t));
198 fitsbin_chunk_t chunk;
199 kdtree_t* kd = s->tree;
200
201 fitsbin_chunk_init(&chunk);
202 chunk.tablename = "sweep";
203 chunk.forced_type = fitscolumn_u8_type();
204 chunk.itemsize = sizeof(uint8_t);
205 chunk.nrows = kd->ndata;
206 chunk.data = s->sweep;
207 chunk.userdata = &(s->sweep);
208 chunk.required = FALSE;
209 bl_append(chunks, &chunk);
210 if (wordsizes)
211 il_append(wordsizes, sizeof(uint8_t));
212
213 fitsbin_chunk_clean(&chunk);
214 return chunks;
215 }
216
my_open(const char * fn,anqfits_t * fits)217 static startree_t* my_open(const char* fn, anqfits_t* fits) {
218 struct timeval tv1, tv2;
219 startree_t* s;
220 bl* chunks;
221 int i;
222 kdtree_fits_t* io;
223 char* treename = STARTREE_NAME;
224 const char* thefn = fn;
225
226 assert(fn || fits);
227
228 if (!thefn)
229 thefn = fits->filename;
230
231 s = startree_alloc();
232 if (!s)
233 return NULL;
234
235 gettimeofday(&tv1, NULL);
236 if (fn)
237 io = kdtree_fits_open(fn);
238 else
239 io = kdtree_fits_open_fits(fits);
240
241 gettimeofday(&tv2, NULL);
242 debug("kdtree_fits_open() took %g ms\n", millis_between(&tv1, &tv2));
243 if (!io) {
244 ERROR("Failed to open FITS file \"%s\"", thefn);
245 goto bailout;
246 }
247
248 gettimeofday(&tv1, NULL);
249 if (!kdtree_fits_contains_tree(io, treename))
250 treename = NULL;
251 gettimeofday(&tv2, NULL);
252 debug("kdtree_fits_contains_tree() took %g ms\n", millis_between(&tv1, &tv2));
253
254 gettimeofday(&tv1, NULL);
255 s->tree = kdtree_fits_read_tree(io, treename, &s->header);
256 gettimeofday(&tv2, NULL);
257 debug("kdtree_fits_read_tree() took %g ms\n", millis_between(&tv1, &tv2));
258 if (!s->tree) {
259 ERROR("Failed to read kdtree from file \"%s\"", thefn);
260 goto bailout;
261 }
262
263 // Check the tree dimensionality.
264 // (because code trees can be confused...)
265 if (s->tree->ndim != 3) {
266 logverb("File %s contains a kd-tree with dim %i (not 3), named %s\n",
267 thefn, s->tree->ndim, treename);
268 s->tree->io = NULL;
269 goto bailout;
270 }
271
272 gettimeofday(&tv1, NULL);
273 chunks = get_chunks(s, NULL);
274 for (i=0; i<bl_size(chunks); i++) {
275 fitsbin_chunk_t* chunk = bl_access(chunks, i);
276 void** dest = chunk->userdata;
277 kdtree_fits_read_chunk(io, chunk);
278 *dest = chunk->data;
279 }
280 bl_free(chunks);
281 gettimeofday(&tv2, NULL);
282 debug("reading chunks took %g ms\n", millis_between(&tv1, &tv2));
283
284 // kdtree_fits_t is a typedef of fitsbin_t
285 fitsbin_close_fd(io);
286
287 return s;
288
289 bailout:
290 kdtree_fits_io_close(io);
291 startree_close(s);
292 return NULL;
293 }
294
startree_open_fits(anqfits_t * fits)295 startree_t* startree_open_fits(anqfits_t* fits) {
296 return my_open(NULL, fits);
297 }
298
startree_open(const char * fn)299 startree_t* startree_open(const char* fn) {
300 return my_open(fn, NULL);
301 }
302
303 /*
304 uint64_t startree_get_starid(const startree_t* s, int ind) {
305 if (!s->starids)
306 return 0;
307 return s->starids[ind];
308 }
309 */
startree_close(startree_t * s)310 int startree_close(startree_t* s) {
311 if (!s) return 0;
312 if (s->inverse_perm)
313 free(s->inverse_perm);
314 if (s->header)
315 qfits_header_destroy(s->header);
316 if (s->tree) {
317 if (s->writing) {
318 free(s->tree->data.any);
319 s->tree->data.any = NULL;
320 kdtree_free(s->tree);
321 free(s->sweep);
322 }
323 else
324 kdtree_fits_close(s->tree);
325 }
326 if (s->tagalong)
327 fitstable_close(s->tagalong);
328 free(s);
329 return 0;
330 }
331
get_tagalong(startree_t * s,anbool report_errs)332 static fitstable_t* get_tagalong(startree_t* s, anbool report_errs) {
333 char* fn;
334 int next;
335 int i;
336 int ext = -1;
337 fitstable_t* tag;
338
339 if (!s->tree->io)
340 return NULL;
341 fn = fitsbin_get_filename(s->tree->io);
342 if (!fn) {
343 if (report_errs)
344 ERROR("No filename");
345 return NULL;
346 }
347 tag = fitstable_open(fn);
348 if (!tag) {
349 if (report_errs)
350 ERROR("Failed to open FITS table from %s", fn);
351 return NULL;
352 }
353 next = fitstable_n_extensions(tag);
354 for (i=1; i<next; i++) {
355 char* type;
356 anbool eq;
357 const qfits_header* hdr;
358 hdr = anqfits_get_header_const(tag->anq, i);
359 if (!hdr) {
360 if (report_errs)
361 ERROR("Failed to read FITS header for ext %i in %s", i, fn);
362 return NULL;
363 }
364 type = fits_get_dupstring(hdr, "AN_FILE");
365 eq = streq(type, AN_FILETYPE_TAGALONG);
366 free(type);
367 if (!eq)
368 continue;
369 ext = i;
370 break;
371 }
372 if (ext == -1) {
373 if (report_errs)
374 ERROR("Failed to find a FITS header with the card AN_FILE = TAGALONG");
375 return NULL;
376 }
377 fitstable_open_extension(tag, ext);
378 return tag;
379 }
380
startree_get_tagalong(startree_t * s)381 fitstable_t* startree_get_tagalong(startree_t* s) {
382 if (s->tagalong)
383 return s->tagalong;
384 s->tagalong = get_tagalong(s, TRUE);
385 return s->tagalong;
386 }
387
startree_has_tagalong(startree_t * s)388 anbool startree_has_tagalong(startree_t* s) {
389 return (startree_get_tagalong(s) != NULL);
390 }
391
Ndata(const startree_t * s)392 static int Ndata(const startree_t* s) {
393 return s->tree->ndata;
394 }
395
startree_check_inverse_perm(startree_t * s)396 int startree_check_inverse_perm(startree_t* s) {
397 // ensure that each value appears exactly once.
398 int i, N;
399 uint8_t* counts;
400 N = Ndata(s);
401 counts = calloc(Ndata(s), sizeof(uint8_t));
402 for (i=0; i<N; i++) {
403 assert(s->inverse_perm[i] >= 0);
404 assert(s->inverse_perm[i] < N);
405 counts[s->inverse_perm[i]]++;
406 }
407 for (i=0; i<N; i++) {
408 assert(counts[i] == 1);
409 }
410 return 0;
411 }
412
startree_compute_inverse_perm(startree_t * s)413 void startree_compute_inverse_perm(startree_t* s) {
414 if (s->inverse_perm)
415 return;
416 // compute inverse permutation vector.
417 s->inverse_perm = malloc(Ndata(s) * sizeof(int));
418 if (!s->inverse_perm) {
419 fprintf(stderr, "Failed to allocate star kdtree inverse permutation vector.\n");
420 return;
421 }
422 #ifndef NDEBUG
423 {
424 int i;
425 for (i=0; i<Ndata(s); i++)
426 s->inverse_perm[i] = -1;
427 }
428 #endif
429 kdtree_inverse_permutation(s->tree, s->inverse_perm);
430 #ifndef NDEBUG
431 {
432 int i;
433 for (i=0; i<Ndata(s); i++)
434 assert(s->inverse_perm[i] != -1);
435 }
436 #endif
437 }
438
startree_get_cut_nside(const startree_t * s)439 int startree_get_cut_nside(const startree_t* s) {
440 return qfits_header_getint(s->header, "CUTNSIDE", -1);
441 }
442
startree_get_cut_nsweeps(const startree_t * s)443 int startree_get_cut_nsweeps(const startree_t* s) {
444 return qfits_header_getint(s->header, "CUTNSWEP", -1);
445 }
446
startree_get_cut_dedup(const startree_t * s)447 double startree_get_cut_dedup(const startree_t* s) {
448 return qfits_header_getdouble(s->header, "CUTDEDUP", 0.0);
449 }
450
startree_get_cut_band(const startree_t * s)451 char* startree_get_cut_band(const startree_t* s) {
452 static char* bands[] = { "R", "B", "J" };
453 int i;
454 char* str = fits_get_dupstring(s->header, "CUTBAND");
455 char* rtn = NULL;
456 if (!str)
457 return NULL;
458 for (i=0; i<sizeof(bands) / sizeof(char*); i++) {
459 if (streq(str, bands[i])) {
460 rtn = bands[i];
461 break;
462 }
463 }
464 free(str);
465 return rtn;
466 }
467
startree_get_cut_margin(const startree_t * s)468 int startree_get_cut_margin(const startree_t* s) {
469 return qfits_header_getint(s->header, "CUTMARG", -1);
470 }
471
startree_get_jitter(const startree_t * s)472 double startree_get_jitter(const startree_t* s) {
473 return qfits_header_getdouble(s->header, "JITTER", 0.0);
474 }
475
startree_set_jitter(startree_t * s,double jitter_arcsec)476 void startree_set_jitter(startree_t* s, double jitter_arcsec) {
477 fits_header_set_double(s->header, "JITTER", jitter_arcsec, "Positional error of stars [arcsec]");
478 }
479
startree_get_sweep(const startree_t * s,int ind)480 int startree_get_sweep(const startree_t* s, int ind) {
481 if (ind < 0 || ind >= Ndata(s) || !s->sweep)
482 return -1;
483 return s->sweep[ind];
484 }
485
startree_get(startree_t * s,int starid,double * posn)486 int startree_get(startree_t* s, int starid, double* posn) {
487 if (s->tree->perm && !s->inverse_perm) {
488 startree_compute_inverse_perm(s);
489 if (!s->inverse_perm)
490 return -1;
491 }
492 if (starid >= Ndata(s)) {
493 fprintf(stderr, "Invalid star ID: %u >= %u.\n", starid, Ndata(s));
494 assert(0);
495 return -1;
496 }
497 if (s->inverse_perm) {
498 kdtree_copy_data_double(s->tree, s->inverse_perm[starid], 1, posn);
499 } else {
500 kdtree_copy_data_double(s->tree, starid, 1, posn);
501 }
502 return 0;
503 }
504
startree_get_radec(startree_t * s,int starid,double * ra,double * dec)505 int startree_get_radec(startree_t* s, int starid, double* ra, double* dec) {
506 double xyz[3];
507 int rtn;
508 rtn = startree_get(s, starid, xyz);
509 if (rtn)
510 return rtn;
511 xyzarr2radecdeg(xyz, ra, dec);
512 return rtn;
513 }
514
startree_new()515 startree_t* startree_new() {
516 startree_t* s = startree_alloc();
517 s->header = qfits_header_default();
518 if (!s->header) {
519 fprintf(stderr, "Failed to create a qfits header for star kdtree.\n");
520 free(s);
521 return NULL;
522 }
523 qfits_header_add(s->header, "AN_FILE", AN_FILETYPE_STARTREE, "This file is a star kdtree.", NULL);
524 s->writing = TRUE;
525 return s;
526 }
527
write_to_file(startree_t * s,const char * fn,anbool flipped,FILE * fid)528 static int write_to_file(startree_t* s, const char* fn, anbool flipped,
529 FILE* fid) {
530 bl* chunks;
531 il* wordsizes = NULL;
532 int i;
533 kdtree_fits_t* io = NULL;
534
535 // just haven't bothered...
536 assert(!(flipped && fid));
537
538 if (fn) {
539 io = kdtree_fits_open_for_writing(fn);
540 if (!io) {
541 ERROR("Failed to open file \"%s\" for writing kdtree", fn);
542 return -1;
543 }
544 }
545 if (flipped) {
546 if (kdtree_fits_write_tree_flipped(io, s->tree, s->header)) {
547 ERROR("Failed to write (flipped) kdtree to file \"%s\"", fn);
548 return -1;
549 }
550 } else {
551 if (fid) {
552 if (kdtree_fits_append_tree_to(s->tree, s->header, fid)) {
553 ERROR("Failed to write star kdtree");
554 return -1;
555 }
556 } else {
557 if (kdtree_fits_write_tree(io, s->tree, s->header)) {
558 ERROR("Failed to write kdtree to file \"%s\"", fn);
559 return -1;
560 }
561 }
562 }
563
564 if (flipped)
565 wordsizes = il_new(4);
566
567 chunks = get_chunks(s, wordsizes);
568 for (i=0; i<bl_size(chunks); i++) {
569 fitsbin_chunk_t* chunk = bl_access(chunks, i);
570 if (!chunk->data)
571 continue;
572 if (flipped)
573 kdtree_fits_write_chunk_flipped(io, chunk, il_get(wordsizes, i));
574 else {
575 if (fid) {
576 kdtree_fits_write_chunk_to(chunk, fid);
577 } else {
578 kdtree_fits_write_chunk(io, chunk);
579 }
580 }
581 fitsbin_chunk_clean(chunk);
582 }
583 bl_free(chunks);
584
585 if (flipped)
586 il_free(wordsizes);
587
588 if (io)
589 kdtree_fits_io_close(io);
590 return 0;
591 }
592
startree_write_to_file(startree_t * s,const char * fn)593 int startree_write_to_file(startree_t* s, const char* fn) {
594 return write_to_file(s, fn, FALSE, NULL);
595 }
596
startree_write_to_file_flipped(startree_t * s,const char * fn)597 int startree_write_to_file_flipped(startree_t* s, const char* fn) {
598 return write_to_file(s, fn, TRUE, NULL);
599 }
600
startree_append_to(startree_t * s,FILE * fid)601 int startree_append_to(startree_t* s, FILE* fid) {
602 return write_to_file(s, NULL, FALSE, fid);
603 }
604
605