1!
2!(c) Matthew Kennel, Institute for Nonlinear Science (2004)
3!
4! Licensed under the Academic Free License version 1.1 found in file LICENSE_kdTree2.txt
5! with additional provisions found in that same file.
6!
7module kdtree2_precision_mod
8 ! MODULE kdtree2_precision_mod (prefix='none' category='8. Low-level utilities and constants')
9 !
10 !: Purpose: Module used by the kdTree2_mod module.
11 !
12 integer, parameter :: sp = kind(0.0)
13 integer, parameter :: dp = kind(0.0d0)
14
15 private :: sp, dp
16
17 !
18 ! You must comment out exactly one
19 ! of the two lines. If you comment
20 ! out kdkind = sp then you get single precision
21 ! and if you comment out kdkind = dp
22 ! you get double precision.
23 !
24
25 !integer, parameter :: kdkind = sp
26 integer, parameter :: kdkind = dp
27 public :: kdkind
28
29end module kdtree2_precision_mod
30
31module kdtree2_priority_queue_mod
32 ! MODULE kdtree2_priority_queue_mod (prefix='none' category='8. Low-level utilities and constants')
33 !
34 !: Purpose: Module used by the kdTree2_mod module.
35 !
36 use utilities_mod
37 use kdtree2_precision_mod
38 !
39 ! maintain a priority queue (PQ) of data, pairs of 'priority/payload',
40 ! implemented with a binary heap. This is the type, and the 'dis' field
41 ! is the priority.
42 !
43 type kdtree2_result
44 ! a pair of distances, indexes
45 real(kdkind) :: dis!=0.0
46 integer :: idx!=-1 Initializers cause some bugs in compilers.
47 end type kdtree2_result
48 !
49 ! A heap-based priority queue lets one efficiently implement the following
50 ! operations, each in log(N) time, as opposed to linear time.
51 !
52 ! 1) add a datum (push a datum onto the queue, increasing its length)
53 ! 2) return the priority value of the maximum priority element
54 ! 3) pop-off (and delete) the element with the maximum priority, decreasing
55 ! the size of the queue.
56 ! 4) replace the datum with the maximum priority with a supplied datum
57 ! (of either higher or lower priority), maintaining the size of the
58 ! queue.
59 !
60 !
61 ! In the k-d tree case, the 'priority' is the square distance of a point in
62 ! the data set to a reference point. The goal is to keep the smallest M
63 ! distances to a reference point. The tree algorithm searches terminal
64 ! nodes to decide whether to add points under consideration.
65 !
66 ! A priority queue is useful here because it lets one quickly return the
67 ! largest distance currently existing in the list. If a new candidate
68 ! distance is smaller than this, then the new candidate ought to replace
69 ! the old candidate. In priority queue terms, this means removing the
70 ! highest priority element, and inserting the new one.
71 !
72 ! Algorithms based on Cormen, Leiserson, Rivest, _Introduction
73 ! to Algorithms_, 1990, with further optimization by the author.
74 !
75 ! Originally informed by a C implementation by Sriranga Veeraraghavan.
76 !
77 ! This module is not written in the most clear way, but is implemented such
78 ! for speed, as it its operations will be called many times during searches
79 ! of large numbers of neighbors.
80 !
81 type pq
82 !
83 ! The priority queue consists of elements
84 ! priority(1:heap_size), with associated payload(:).
85 !
86 ! There are heap_size active elements.
87 ! Assumes the allocation is always sufficient. Will NOT increase it
88 ! to match.
89 integer :: heap_size = 0
90 type(kdtree2_result), pointer :: elems(:)
91 end type pq
92
93 public :: kdtree2_result
94
95 public :: pq
96 public :: pq_create
97 public :: pq_delete, pq_insert
98 public :: pq_extract_max, pq_max, pq_replace_max, pq_maxpri
99 private
100
101contains
102
103
104 function pq_create(results_in) result(res)
105 !:Purpose: Create a priority queue from ALREADY allocated
106 ! array pointers for storage. NOTE! It will NOT
107 ! add any alements to the heap, i.e. any existing
108 ! data in the input arrays will NOT be used and may
109 ! be overwritten.
110 !
111
112 ! usage:
113 ! real(kdkind), pointer :: x(:)
114 ! integer, pointer :: k(:)
115 ! allocate(x(1000),k(1000))
116 ! pq => pq_create(x,k)
117 !
118 type(kdtree2_result), target:: results_in(:)
119 type(pq) :: res
120 !
121 !
122 integer :: nalloc
123
124 nalloc = size(results_in,1)
125 if (nalloc .lt. 1) then
126 write (*,*) 'PQ_CREATE: error, input arrays must be allocated.'
127 end if
128 res%elems => results_in
129 res%heap_size = 0
130 return
131 end function pq_create
132
133
134 subroutine heapify(a,i_in)
135 !:Purpose: take a heap rooted at 'i' and force it to be in the
136 ! heap canonical form. This is performance critical
137 ! and has been tweaked a little to reflect this.
138 !
139 type(pq),pointer :: a
140 integer, intent(in) :: i_in
141 !
142 integer :: i, l, r, largest
143
144 real(kdkind) :: pri_i, pri_l, pri_r, pri_largest
145
146
147 type(kdtree2_result) :: temp
148
149 i = i_in
150
151bigloop: do
152 l = 2*i ! left(i)
153 r = l+1 ! right(i)
154 !
155 ! set 'largest' to the index of either i, l, r
156 ! depending on whose priority is largest.
157 !
158 ! note that l or r can be larger than the heap size
159 ! in which case they do not count.
160
161
162 ! does left child have higher priority?
163 if (l .gt. a%heap_size) then
164 ! we know that i is the largest as both l and r are invalid.
165 exit
166 else
167 pri_i = a%elems(i)%dis
168 pri_l = a%elems(l)%dis
169 if (pri_l .gt. pri_i) then
170 largest = l
171 pri_largest = pri_l
172 else
173 largest = i
174 pri_largest = pri_i
175 endif
176
177 !
178 ! between i and l we have a winner
179 ! now choose between that and r.
180 !
181 if (r .le. a%heap_size) then
182 pri_r = a%elems(r)%dis
183 if (pri_r .gt. pri_largest) then
184 largest = r
185 endif
186 endif
187 endif
188
189 if (largest .ne. i) then
190 ! swap data in nodes largest and i, then heapify
191
192 temp = a%elems(i)
193 a%elems(i) = a%elems(largest)
194 a%elems(largest) = temp
195 !
196 ! Canonical heapify() algorithm has tail-ecursive call:
197 !
198 ! call heapify(a,largest)
199 ! we will simulate with cycle
200 !
201 i = largest
202 cycle bigloop ! continue the loop
203 else
204 return ! break from the loop
205 end if
206 enddo bigloop
207 return
208 end subroutine heapify
209
210 subroutine pq_max(a,e)
211 !:Purpose: return the priority and its payload of the maximum priority element
212 ! on the queue, which should be the first one, if it is
213 ! in heapified form.
214 !
215 type(pq),pointer :: a
216 type(kdtree2_result),intent(out) :: e
217
218 if (a%heap_size .gt. 0) then
219 e = a%elems(1)
220 else
221 call utl_abort('kdtree2_mod-PQ_MAX: ERROR, heap_size < 1')
222 endif
223 return
224 end subroutine pq_max
225
226 real(kdkind) function pq_maxpri(a)
227 type(pq), pointer :: a
228
229 if (a%heap_size .gt. 0) then
230 pq_maxpri = a%elems(1)%dis
231 else
232 call utl_abort('kdtrees_mod-PQ_MAX_PRI: ERROR, heapsize < 1')
233 endif
234 return
235 end function pq_maxpri
236
237 subroutine pq_extract_max(a,e)
238 !:Purpose: return the priority and payload of maximum priority
239 ! element, and remove it from the queue.
240 ! (equivalent to 'pop()' on a stack)
241 !
242 type(pq),pointer :: a
243 type(kdtree2_result), intent(out) :: e
244
245 if (a%heap_size .ge. 1) then
246 !
247 ! return max as first element
248 !
249 e = a%elems(1)
250
251 !
252 ! move last element to first
253 !
254 a%elems(1) = a%elems(a%heap_size)
255 a%heap_size = a%heap_size-1
256 call heapify(a,1)
257 return
258 else
259 call utl_abort('kdtree2_mod-PQ_EXTRACT_MAX: error, attempted to pop non-positive PQ')
260 end if
261
262 end subroutine pq_extract_max
263
264
265 real(kdkind) function pq_insert(a,dis,idx)
266 !:Purpose: Insert a new element and return the new maximum priority,
267 ! which may or may not be the same as the old maximum priority.
268 !
269 type(pq),pointer :: a
270 real(kdkind), intent(in) :: dis
271 integer, intent(in) :: idx
272 ! type(kdtree2_result), intent(in) :: e
273 !
274 integer :: i, isparent
275 real(kdkind) :: parentdis
276 !
277
278 ! if (a%heap_size .ge. a%max_elems) then
279 ! write (*,*) 'PQ_INSERT: error, attempt made to insert element on full PQ'
280 ! stop
281 ! else
282 a%heap_size = a%heap_size + 1
283 i = a%heap_size
284
285 do while (i .gt. 1)
286 isparent = int(i/2)
287 parentdis = a%elems(isparent)%dis
288 if (dis .gt. parentdis) then
289 ! move what was in i's parent into i.
290 a%elems(i)%dis = parentdis
291 a%elems(i)%idx = a%elems(isparent)%idx
292 i = isparent
293 else
294 exit
295 endif
296 end do
297
298 ! insert the element at the determined position
299 a%elems(i)%dis = dis
300 a%elems(i)%idx = idx
301
302 pq_insert = a%elems(1)%dis
303 return
304 ! end if
305
306 end function pq_insert
307
308 subroutine pq_adjust_heap(a,i)
309 type(pq),pointer :: a
310 integer, intent(in) :: i
311 !
312 ! nominally arguments (a,i), but specialize for a=1
313 !
314 ! This routine assumes that the trees with roots 2 and 3 are already heaps, i.e.
315 ! the children of '1' are heaps. When the procedure is completed, the
316 ! tree rooted at 1 is a heap.
317 real(kdkind) :: prichild
318 integer :: parent, child, N
319
320 type(kdtree2_result) :: e
321
322 e = a%elems(i)
323
324 parent = i
325 child = 2*i
326 N = a%heap_size
327
328 do while (child .le. N)
329 if (child .lt. N) then
330 if (a%elems(child)%dis .lt. a%elems(child+1)%dis) then
331 child = child+1
332 endif
333 endif
334 prichild = a%elems(child)%dis
335 if (e%dis .ge. prichild) then
336 exit
337 else
338 ! move child into parent.
339 a%elems(parent) = a%elems(child)
340 parent = child
341 child = 2*parent
342 end if
343 end do
344 a%elems(parent) = e
345 return
346 end subroutine pq_adjust_heap
347
348
349 real(kdkind) function pq_replace_max(a,dis,idx)
350 !:Purpose: Replace the extant maximum priority element
351 ! in the PQ with (dis,idx). Return
352 ! the new maximum priority, which may be larger
353 ! or smaller than the old one.
354 !
355 type(pq),pointer :: a
356 real(kdkind), intent(in) :: dis
357 integer, intent(in) :: idx
358! type(kdtree2_result), intent(in) :: e
359 ! not tested as well!
360
361 integer :: parent, child, N
362 real(kdkind) :: prichild, prichildp1
363
364 type(kdtree2_result) :: etmp
365
366 if (.true.) then
367 N=a%heap_size
368 if (N .ge. 1) then
369 parent =1
370 child=2
371
372 loop: do while (child .le. N)
373 prichild = a%elems(child)%dis
374
375 !
376 ! posibly child+1 has higher priority, and if
377 ! so, get it, and increment child.
378 !
379
380 if (child .lt. N) then
381 prichildp1 = a%elems(child+1)%dis
382 if (prichild .lt. prichildp1) then
383 child = child+1
384 prichild = prichildp1
385 endif
386 endif
387
388 if (dis .ge. prichild) then
389 exit loop
390 ! we have a proper place for our new element,
391 ! bigger than either children's priority.
392 else
393 ! move child into parent.
394 a%elems(parent) = a%elems(child)
395 parent = child
396 child = 2*parent
397 end if
398 end do loop
399 a%elems(parent)%dis = dis
400 a%elems(parent)%idx = idx
401 pq_replace_max = a%elems(1)%dis
402 else
403 a%elems(1)%dis = dis
404 a%elems(1)%idx = idx
405 pq_replace_max = dis
406 endif
407 else
408 !
409 ! slower version using elementary pop and push operations.
410 !
411 call pq_extract_max(a,etmp)
412 etmp%dis = dis
413 etmp%idx = idx
414 pq_replace_max = pq_insert(a,dis,idx)
415 endif
416 return
417 end function pq_replace_max
418
419 subroutine pq_delete(a,i)
420 !:Purpose: delete item with index 'i'
421 !
422 type(pq),pointer :: a
423 integer :: i
424
425 if ((i .lt. 1) .or. (i .gt. a%heap_size)) then
426 call utl_abort('kdtree2_mod-PQ_DELETE: error, attempt to remove out of bounds element.')
427 endif
428
429 ! swap the item to be deleted with the last element
430 ! and shorten heap by one.
431 a%elems(i) = a%elems(a%heap_size)
432 a%heap_size = a%heap_size - 1
433
434 call heapify(a,i)
435
436 end subroutine pq_delete
437
438end module kdtree2_priority_queue_mod
439
440
441module kdTree2_mod
442 ! MODULE kdTree2_mod (prefix='kdtree2' category='8. Low-level utilities and constants')
443 !
444 !:Purpose: An implementation of the kdtree algorithm for efficiently finding nearby
445 ! locations (first used in LETKF to find all obs near an analysis grid point).
446 ! Written by Matt Kennel.
447 !
448 use utilities_mod
449 use kdtree2_precision_mod
450 use kdtree2_priority_queue_mod
451 use earthConstants_mod
452 ! K-D tree routines in Fortran 90 by Matt Kennel.
453 ! Original program was written in Sather by Steve Omohundro and
454 ! Matt Kennel. Only the Euclidean metric is supported.
455 !
456 !
457 ! This module is identical to 'kd_tree', except that the order
458 ! of subscripts is reversed in the data file.
459 ! In otherwords for an embedding of N D-dimensional vectors, the
460 ! data file is here, in natural Fortran order data(1:D, 1:N)
461 ! because Fortran lays out columns first,
462 !
463 ! whereas conventionally (C-style) it is data(1:N,1:D)
464 ! as in the original kd_tree module.
465 !
466 !-------------DATA TYPE, CREATION, DELETION---------------------
467 public :: kdkind
468 public :: kdtree2, kdtree2_result, tree_node, kdtree2_create, kdtree2_destroy
469 !---------------------------------------------------------------
470 !-------------------SEARCH ROUTINES-----------------------------
471 public :: kdtree2_n_nearest,kdtree2_n_nearest_around_point
472 ! Return fixed number of nearest neighbors around arbitrary vector,
473 ! or extant point in dataset, with decorrelation window.
474 !
475 public :: kdtree2_r_nearest, kdtree2_r_nearest_around_point
476 ! Return points within a fixed ball of arb vector/extant point
477 !
478 public :: kdtree2_sort_results
479 ! Sort, in order of increasing distance, rseults from above.
480 !
481 public :: kdtree2_r_count, kdtree2_r_count_around_point
482 ! Count points within a fixed ball of arb vector/extant point
483 !
484 public :: kdtree2_n_nearest_brute_force, kdtree2_r_nearest_brute_force
485 ! brute force of kdtree2_[n|r]_nearest
486 !----------------------------------------------------------------
487
488 public :: kdtree2_3dPosition
489
490 integer, parameter :: bucket_size = 12
491 ! The maximum number of points to keep in a terminal node.
492
493 type interval
494 real(kdkind) :: lower,upper
495 end type interval
496
497 type :: tree_node
498 ! an internal tree node
499 private
500 integer :: cut_dim
501 ! the dimension to cut
502 real(kdkind) :: cut_val
503 ! where to cut the dimension
504 real(kdkind) :: cut_val_left, cut_val_right
505 ! improved cutoffs knowing the spread in child boxes.
506 integer :: l, u
507 type (tree_node), pointer :: left, right
508 type(interval), pointer :: box(:) => null()
509 ! child pointers
510 ! Points included in this node are indexes[k] with k \in [l,u]
511
512
513 end type tree_node
514
515 type :: kdtree2
516 ! Global information about the tree, one per tree
517 integer :: dimen=0, n=0
518 ! dimensionality and total # of points
519 real(kdkind), pointer :: the_data(:,:) => null()
520 ! pointer to the actual data array
521 !
522 ! IMPORTANT NOTE: IT IS DIMENSIONED the_data(1:d,1:N)
523 ! which may be opposite of what may be conventional.
524 ! This is, because in Fortran, the memory layout is such that
525 ! the first dimension is in sequential order. Hence, with
526 ! (1:d,1:N), all components of the vector will be in consecutive
527 ! memory locations. The search time is dominated by the
528 ! evaluation of distances in the terminal nodes. Putting all
529 ! vector components in consecutive memory location improves
530 ! memory cache locality, and hence search speed, and may enable
531 ! vectorization on some processors and compilers.
532
533 integer, pointer :: ind(:) => null()
534 ! permuted index into the data, so that indexes[l..u] of some
535 ! bucket represent the indexes of the actual points in that
536 ! bucket.
537 logical :: sort = .false.
538 ! do we always sort output results?
539 logical :: rearrange = .false.
540 real(kdkind), pointer :: rearranged_data(:,:) => null()
541 ! if (rearrange .eqv. .true.) then rearranged_data has been
542 ! created so that rearranged_data(:,i) = the_data(:,ind(i)),
543 ! permitting search to use more cache-friendly rearranged_data, at
544 ! some initial computation and storage cost.
545 type (tree_node), pointer :: root => null()
546 ! root pointer of the tree
547 end type kdtree2
548
549
550 type :: tree_search_record
551 !
552 ! One of these is created for each search.
553 !
554 private
555 !
556 ! Many fields are copied from the tree structure, in order to
557 ! speed up the search.
558 !
559 integer :: dimen
560 integer :: nn, nfound
561 real(kdkind) :: ballsize
562 integer :: centeridx=999, correltime=9999
563 ! exclude points within 'correltime' of 'centeridx', iff centeridx >= 0
564 integer :: nalloc ! how much allocated for results(:)?
565 logical :: rearrange ! are the data rearranged or original?
566 ! did the # of points found overflow the storage provided?
567 logical :: overflow
568 real(kdkind), pointer :: qv(:) ! query vector
569 type(kdtree2_result), pointer :: results(:) ! results
570 type(pq) :: pq
571 real(kdkind), pointer :: data(:,:) ! temp pointer to data
572 integer, pointer :: ind(:) ! temp pointer to indexes
573 end type tree_search_record
574
575 private
576 ! everything else is private.
577
578 type(tree_search_record), save, allocatable, target :: sr(:) ! A GLOBAL VARIABLE for search
579
580contains
581
582 function kdtree2_create(input_data,dim,sort,rearrange) result (mr)
583 !:Purpose: Create the actual tree structure, given an input array of data.
584
585 ! Note, input data is input_data(1:d,1:N), NOT the other way around.
586 ! THIS IS THE REVERSE OF THE PREVIOUS VERSION OF THIS MODULE.
587 ! The reason for it is cache friendliness, improving performance.
588 !
589 ! Optional arguments: If 'dim' is specified, then the tree
590 ! will only search the first 'dim' components
591 ! of input_data, otherwise, dim is inferred
592 ! from SIZE(input_data,1).
593 !
594 ! if sort .eqv. .true. then output results
595 ! will be sorted by increasing distance.
596 ! default=.false., as it is faster to not sort.
597 !
598 ! if rearrange .eqv. .true. then an internal
599 ! copy of the data, rearranged by terminal node,
600 ! will be made for cache friendliness.
601 ! default=.true., as it speeds searches, but
602 ! building takes longer, and extra memory is used.
603 !
604 ! .. Function Return Cut_value ..
605 type (kdtree2), pointer :: mr
606 integer, intent(in), optional :: dim
607 logical, intent(in), optional :: sort
608 logical, intent(in), optional :: rearrange
609 ! ..
610 ! .. Array Arguments ..
611 real(kdkind), target :: input_data(:,:)
612 !
613 integer :: i, numthread, mythread
614 integer, external :: omp_get_thread_num, omp_get_num_threads
615 ! ..
616 allocate (mr)
617 mr%the_data => input_data
618 ! pointer assignment
619
620 if (present(dim)) then
621 mr%dimen = dim
622 else
623 mr%dimen = size(input_data,1)
624 end if
625 mr%n = size(input_data,2)
626
627 if (mr%n < 1) then
628 ! unlikely to be correct
629 write (*,*) 'KD_TREE_TRANS: likely user error.'
630 write (*,*) 'KD_TREE_TRANS: You passed in matrix with D=',mr%dimen
631 write (*,*) 'KD_TREE_TRANS: and N=',mr%n
632 write (*,*) 'KD_TREE_TRANS: note, that new format is data(1:D,1:N)'
633 write (*,*) 'KD_TREE_TRANS: with usually N >> D. If N =approx= D, then a k-d tree'
634 write (*,*) 'KD_TREE_TRANS: is not an appropriate data structure.'
635 call utl_abort('kdtree2_create: supplied array has no data')
636 end if
637
638 !$OMP PARALLEL PRIVATE(mythread)
639 mythread = omp_get_thread_num()
640 if ( mythread == 0 ) then
641 numthread = omp_get_num_threads()
642 end if
643 !$OMP END PARALLEL
644
645 if ( .not. allocated(sr) ) allocate(sr(0:numthread-1))
646
647 call build_tree(mr)
648
649 if (present(sort)) then
650 mr%sort = sort
651 else
652 mr%sort = .false.
653 endif
654
655 if (present(rearrange)) then
656 mr%rearrange = rearrange
657 else
658 mr%rearrange = .true.
659 endif
660
661 if (mr%rearrange) then
662 allocate(mr%rearranged_data(mr%dimen,mr%n))
663 do i=1,mr%n
664 mr%rearranged_data(:,i) = mr%the_data(:, &
665 mr%ind(i))
666 enddo
667 else
668 nullify(mr%rearranged_data)
669 endif
670
671 end function kdtree2_create
672
673 subroutine build_tree(tp)
674 type (kdtree2), pointer :: tp
675 ! ..
676 integer :: j
677 type(tree_node), pointer :: dummy => null()
678 ! ..
679 allocate (tp%ind(tp%n))
680 forall (j=1:tp%n)
681 tp%ind(j) = j
682 end forall
683 tp%root => build_tree_for_range(tp,1,tp%n, dummy)
684 end subroutine build_tree
685
686 recursive function build_tree_for_range(tp,l,u,parent) result (res)
687 ! .. Function Return Cut_value ..
688 type (tree_node), pointer :: res
689 ! ..
690 ! .. Structure Arguments ..
691 type (kdtree2), pointer :: tp
692 type (tree_node),pointer :: parent
693 ! ..
694 ! .. Scalar Arguments ..
695 integer, intent (In) :: l, u
696 ! ..
697 ! .. Local Scalars ..
698 integer :: i, c, m, dimen
699 logical :: recompute
700 real(kdkind) :: average
701
702 ! first compute min and max
703 dimen = tp%dimen
704 allocate (res)
705 allocate(res%box(dimen))
706
707 ! First, compute an APPROXIMATE bounding box of all points associated with this node.
708 if ( u < l ) then
709 ! no points in this box
710 nullify(res)
711 return
712 end if
713
714 if ((u-l)<=bucket_size) then
715 !
716 ! always compute true bounding box for terminal nodes.
717 !
718 do i=1,dimen
719 call spread_in_coordinate(tp,i,l,u,res%box(i))
720 end do
721 res%cut_dim = 0
722 res%cut_val = 0.0
723 res%l = l
724 res%u = u
725 res%left =>null()
726 res%right => null()
727 else
728 !
729 ! modify approximate bounding box. This will be an
730 ! overestimate of the true bounding box, as we are only recomputing
731 ! the bounding box for the dimension that the parent split on.
732 !
733 ! Going to a true bounding box computation would significantly
734 ! increase the time necessary to build the tree, and usually
735 ! has only a very small difference. This box is not used
736 ! for searching but only for deciding which coordinate to split on.
737 !
738 do i=1,dimen
739 recompute=.true.
740 if (associated(parent)) then
741 if (i .ne. parent%cut_dim) then
742 recompute=.false.
743 end if
744 endif
745 if (recompute) then
746 call spread_in_coordinate(tp,i,l,u,res%box(i))
747 else
748 res%box(i) = parent%box(i)
749 endif
750 end do
751
752
753 c = maxloc(res%box(1:dimen)%upper-res%box(1:dimen)%lower,1)
754 !
755 ! c is the identity of which coordinate has the greatest spread.
756 !
757
758 if (.false.) then
759 ! select exact median to have fully balanced tree.
760 m = (l+u)/2
761 call select_on_coordinate(tp%the_data,tp%ind,c,m,l,u)
762 else
763 !
764 ! select point halfway between min and max, as per A. Moore,
765 ! who says this helps in some degenerate cases, or
766 ! actual arithmetic average.
767 !
768 if (.true.) then
769 ! actually compute average
770 average = sum(tp%the_data(c,tp%ind(l:u))) / real(u-l+1,kdkind)
771 else
772 average = (res%box(c)%upper + res%box(c)%lower)/2.0
773 endif
774
775 res%cut_val = average
776 m = select_on_coordinate_value(tp%the_data,tp%ind,c,average,l,u)
777 endif
778
779 ! moves indexes around
780 res%cut_dim = c
781 res%l = l
782 res%u = u
783! res%cut_val = tp%the_data(c,tp%ind(m))
784
785 res%left => build_tree_for_range(tp,l,m,res)
786 res%right => build_tree_for_range(tp,m+1,u,res)
787
788 if (associated(res%right) .eqv. .false.) then
789 res%box = res%left%box
790 res%cut_val_left = res%left%box(c)%upper
791 res%cut_val = res%cut_val_left
792 elseif (associated(res%left) .eqv. .false.) then
793 res%box = res%right%box
794 res%cut_val_right = res%right%box(c)%lower
795 res%cut_val = res%cut_val_right
796 else
797 res%cut_val_right = res%right%box(c)%lower
798 res%cut_val_left = res%left%box(c)%upper
799 res%cut_val = (res%cut_val_left + res%cut_val_right)/2
800
801
802 ! now remake the true bounding box for self.
803 ! Since we are taking unions (in effect) of a tree structure,
804 ! this is much faster than doing an exhaustive
805 ! search over all points
806 res%box%upper = max(res%left%box%upper,res%right%box%upper)
807 res%box%lower = min(res%left%box%lower,res%right%box%lower)
808 endif
809 end if
810 end function build_tree_for_range
811
812 integer function select_on_coordinate_value(v,ind,c,alpha,li,ui) &
813 result(res)
814 !:Purpose: Move elts of ind around between l and u, so that all points
815 ! <= than alpha (in c cooordinate) are first, and then
816 ! all points > alpha are second.
817
818 !
819 ! Algorithm (matt kennel).
820 !
821 ! Consider the list as having three parts: on the left,
822 ! the points known to be <= alpha. On the right, the points
823 ! known to be > alpha, and in the middle, the currently unknown
824 ! points. The algorithm is to scan the unknown points, starting
825 ! from the left, and swapping them so that they are added to
826 ! the left stack or the right stack, as appropriate.
827 !
828 ! The algorithm finishes when the unknown stack is empty.
829 !
830 ! .. Scalar Arguments ..
831 integer, intent (In) :: c, li, ui
832 real(kdkind), intent(in) :: alpha
833 ! ..
834 real(kdkind) :: v(1:,1:)
835 integer :: ind(1:)
836 integer :: tmp
837 ! ..
838 integer :: lb, rb
839 !
840 ! The points known to be <= alpha are in
841 ! [l,lb-1]
842 !
843 ! The points known to be > alpha are in
844 ! [rb+1,u].
845 !
846 ! Therefore we add new points into lb or
847 ! rb as appropriate. When lb=rb
848 ! we are done. We return the location of the last point <= alpha.
849 !
850 !
851 lb = li; rb = ui
852
853 do while (lb < rb)
854 if ( v(c,ind(lb)) <= alpha ) then
855 ! it is good where it is.
856 lb = lb+1
857 else
858 ! swap it with rb.
859 tmp = ind(lb); ind(lb) = ind(rb); ind(rb) = tmp
860 rb = rb-1
861 endif
862 end do
863
864 ! now lb .eq. ub
865 if (v(c,ind(lb)) <= alpha) then
866 res = lb
867 else
868 res = lb-1
869 endif
870
871 ! add a check to prevent something going wrong (Buehner)
872 if (res < li .or. res >= ui) then
873 !write(*,*) 'patch to handle many equal positions:',li,ui,res,v(c,ind(li:ui))
874 res = (li + ui)/2
875 end if
876
877 end function select_on_coordinate_value
878
879 subroutine select_on_coordinate(v,ind,c,k,li,ui)
880 !:Purpose: Move elts of ind around between l and u, so that the kth
881 ! element is >= those below, <= those above, in the coordinate c.
882
883 ! .. Scalar Arguments ..
884 integer, intent (In) :: c, k, li, ui
885 ! ..
886 integer :: i, l, m, s, t, u
887 ! ..
888 real(kdkind) :: v(:,:)
889 integer :: ind(:)
890 ! ..
891 l = li
892 u = ui
893 do while (l<u)
894 t = ind(l)
895 m = l
896 do i = l + 1, u
897 if (v(c,ind(i))<v(c,t)) then
898 m = m + 1
899 s = ind(m)
900 ind(m) = ind(i)
901 ind(i) = s
902 end if
903 end do
904 s = ind(l)
905 ind(l) = ind(m)
906 ind(m) = s
907 if (m<=k) l = m + 1
908 if (m>=k) u = m - 1
909 end do
910 end subroutine select_on_coordinate
911
912 subroutine spread_in_coordinate(tp,c,l,u,interv)
913 !:Purpose: the spread in coordinate 'c', between l and u.
914 ! Return lower bound in 'smin', and upper in 'smax',
915
916 ! ..
917 ! .. Structure Arguments ..
918 type (kdtree2), pointer :: tp
919 type(interval), intent(out) :: interv
920 ! ..
921 ! .. Scalar Arguments ..
922 integer, intent (In) :: c, l, u
923 ! ..
924 ! .. Local Scalars ..
925 real(kdkind) :: last, lmax, lmin, t, smin,smax
926 integer :: i, ulocal
927 ! ..
928 ! .. Local Arrays ..
929 real(kdkind), pointer :: v(:,:)
930 integer, pointer :: ind(:)
931 ! ..
932 v => tp%the_data(1:,1:)
933 ind => tp%ind(1:)
934 smin = v(c,ind(l))
935 smax = smin
936
937 ulocal = u
938
939 do i = l + 2, ulocal, 2
940 lmin = v(c,ind(i-1))
941 lmax = v(c,ind(i))
942 if (lmin>lmax) then
943 t = lmin
944 lmin = lmax
945 lmax = t
946 end if
947 if (smin>lmin) smin = lmin
948 if (smax<lmax) smax = lmax
949 end do
950 if (i==ulocal+1) then
951 last = v(c,ind(ulocal))
952 if (smin>last) smin = last
953 if (smax<last) smax = last
954 end if
955
956 interv%lower = smin
957 interv%upper = smax
958
959 end subroutine spread_in_coordinate
960
961
962 subroutine kdtree2_destroy(tp)
963 !:Purpose: Deallocates all memory for the tree, except input data matrix
964
965 ! .. Structure Arguments ..
966 type (kdtree2), pointer :: tp
967 ! ..
968 call destroy_node(tp%root)
969
970 deallocate (tp%ind)
971 nullify (tp%ind)
972
973 if (tp%rearrange) then
974 deallocate(tp%rearranged_data)
975 nullify(tp%rearranged_data)
976 endif
977
978 nullify(tp%the_data)
979
980 deallocate(tp)
981 nullify(tp)
982 return
983
984 contains
985 recursive subroutine destroy_node(np)
986 ! .. Structure Arguments ..
987 type (tree_node), pointer :: np
988 ! ..
989 ! .. Intrinsic Functions ..
990 intrinsic ASSOCIATED
991 ! ..
992 if (associated(np%left)) then
993 call destroy_node(np%left)
994 nullify (np%left)
995 end if
996 if (associated(np%right)) then
997 call destroy_node(np%right)
998 nullify (np%right)
999 end if
1000 if (associated(np%box)) deallocate(np%box)
1001 deallocate(np)
1002 return
1003
1004 end subroutine destroy_node
1005
1006 end subroutine kdtree2_destroy
1007
1008 subroutine kdtree2_n_nearest(tp,qv,nn,results)
1009 !:Purpose: Find the 'nn' vectors in the tree nearest to 'qv' in euclidean norm
1010 ! returning their indexes and distances in 'indexes' and 'distances'
1011 ! arrays already allocated passed to this subroutine.
1012 type (kdtree2), pointer :: tp
1013 real(kdkind), target, intent (In) :: qv(:)
1014 integer, intent (In) :: nn
1015 type(kdtree2_result), target :: results(:)
1016 integer :: mythread
1017 integer, external :: omp_get_thread_num
1018
1019 mythread = omp_get_thread_num()
1020
1021 sr(mythread)%ballsize = huge(1.0)
1022 sr(mythread)%qv => qv
1023 sr(mythread)%nn = nn
1024 sr(mythread)%nfound = 0
1025 sr(mythread)%centeridx = -1
1026 sr(mythread)%correltime = 0
1027 sr(mythread)%overflow = .false.
1028
1029 sr(mythread)%results => results
1030
1031 sr(mythread)%nalloc = nn ! will be checked
1032
1033 sr(mythread)%ind => tp%ind
1034 sr(mythread)%rearrange = tp%rearrange
1035 if (tp%rearrange) then
1036 sr(mythread)%Data => tp%rearranged_data
1037 else
1038 sr(mythread)%Data => tp%the_data
1039 endif
1040 sr(mythread)%dimen = tp%dimen
1041
1042 call validate_query_storage(nn)
1043 sr(mythread)%pq = pq_create(results)
1044
1045 call search(tp%root)
1046
1047 if (tp%sort) then
1048 call kdtree2_sort_results(nn, results)
1049 endif
1050! deallocate(sr(mythread)%pqp)
1051 return
1052 end subroutine kdtree2_n_nearest
1053
1054 subroutine kdtree2_n_nearest_around_point(tp,idxin,correltime,nn,results)
1055 !:Purpose: Find the 'nn' vectors in the tree nearest to point 'idxin',
1056 ! with correlation window 'correltime', returing results in
1057 ! results(:), which must be pre-allocated upon entry.
1058 type (kdtree2), pointer :: tp
1059 integer, intent (In) :: idxin, correltime, nn
1060 type(kdtree2_result), target :: results(:)
1061 integer :: mythread
1062 integer, external :: omp_get_thread_num
1063
1064 mythread = omp_get_thread_num()
1065
1066 allocate (sr(mythread)%qv(tp%dimen))
1067 sr(mythread)%qv = tp%the_data(:,idxin) ! copy the vector
1068 sr(mythread)%ballsize = huge(1.0) ! the largest real(kdkind) number
1069 sr(mythread)%centeridx = idxin
1070 sr(mythread)%correltime = correltime
1071
1072 sr(mythread)%nn = nn
1073 sr(mythread)%nfound = 0
1074
1075 sr(mythread)%dimen = tp%dimen
1076 sr(mythread)%nalloc = nn
1077
1078 sr(mythread)%results => results
1079
1080 sr(mythread)%ind => tp%ind
1081 sr(mythread)%rearrange = tp%rearrange
1082
1083 if (sr(mythread)%rearrange) then
1084 sr(mythread)%Data => tp%rearranged_data
1085 else
1086 sr(mythread)%Data => tp%the_data
1087 endif
1088
1089 call validate_query_storage(nn)
1090 sr(mythread)%pq = pq_create(results)
1091
1092 call search(tp%root)
1093
1094 if (tp%sort) then
1095 call kdtree2_sort_results(nn, results)
1096 endif
1097 deallocate (sr(mythread)%qv)
1098 return
1099 end subroutine kdtree2_n_nearest_around_point
1100
1101 subroutine kdtree2_r_nearest(tp,qv,r2,nfound,nalloc,results)
1102 !:Purpose: find the nearest neighbors to point 'idxin', within SQUARED
1103 ! Euclidean distance 'r2'. Upon ENTRY, nalloc must be the
1104 ! size of memory allocated for results(1:nalloc). Upon
1105 ! EXIT, nfound is the number actually found within the ball.
1106 !
1107 ! Note that if nfound .gt. nalloc then more neighbors were found
1108 ! than there were storage to store. The resulting list is NOT
1109 ! the smallest ball inside norm r^2
1110 !
1111 ! Results are NOT sorted unless tree was created with sort option.
1112 type (kdtree2), pointer :: tp
1113 real(kdkind), target, intent (In) :: qv(:)
1114 real(kdkind), intent(in) :: r2
1115 integer, intent(out) :: nfound
1116 integer, intent (In) :: nalloc
1117 type(kdtree2_result), target :: results(:)
1118 integer :: mythread
1119 integer, external :: omp_get_thread_num
1120
1121 mythread = omp_get_thread_num()
1122 !
1123 sr(mythread)%qv => qv
1124 sr(mythread)%ballsize = r2
1125 sr(mythread)%nn = 0 ! flag for fixed ball search
1126 sr(mythread)%nfound = 0
1127 sr(mythread)%centeridx = -1
1128 sr(mythread)%correltime = 0
1129
1130 sr(mythread)%results => results
1131
1132 call validate_query_storage(nalloc)
1133 sr(mythread)%nalloc = nalloc
1134 sr(mythread)%overflow = .false.
1135 sr(mythread)%ind => tp%ind
1136 sr(mythread)%rearrange= tp%rearrange
1137
1138 if (tp%rearrange) then
1139 sr(mythread)%Data => tp%rearranged_data
1140 else
1141 sr(mythread)%Data => tp%the_data
1142 endif
1143 sr(mythread)%dimen = tp%dimen
1144
1145 !
1146 !sr(mythread)%dsl = Huge(sr(mythread)%dsl) ! set to huge positive values
1147 !sr(mythread)%il = -1 ! set to invalid indexes
1148 !
1149
1150 call search(tp%root)
1151 nfound = sr(mythread)%nfound
1152
1153 if (sr(mythread)%overflow) then
1154 write (*,*) 'KD_TREE_TRANS: warning! return from kdtree2_r_nearest found more neighbors'
1155 write (*,*) 'KD_TREE_TRANS: than storage was provided for. Answer is NOT smallest ball'
1156 write (*,*) 'KD_TREE_TRANS: with that number of neighbors! I.e. it is wrong.'
1157 write (*,*) 'KD_TREE_TRANS: nfound =', nfound, ', nalloc =', nalloc
1158 endif
1159
1160 if (tp%sort .and. .not.sr(mythread)%overflow) then
1161 call kdtree2_sort_results(nfound, results)
1162 endif
1163
1164 return
1165 end subroutine kdtree2_r_nearest
1166
1167 subroutine kdtree2_r_nearest_around_point(tp,idxin,correltime,r2,&
1168 nfound,nalloc,results)
1169 !:Purpose: Like kdtree2_r_nearest, but around a point 'idxin' already existing
1170 ! in the data set.
1171 !
1172 ! Results are NOT sorted unless tree was created with sort option.
1173 !
1174 type (kdtree2), pointer :: tp
1175 integer, intent (In) :: idxin, correltime, nalloc
1176 real(kdkind), intent(in) :: r2
1177 integer, intent(out) :: nfound
1178 type(kdtree2_result), target :: results(:)
1179 integer :: mythread
1180 integer, external :: omp_get_thread_num
1181 ! ..
1182 ! .. Intrinsic Functions ..
1183 intrinsic HUGE
1184 ! ..
1185
1186 mythread = omp_get_thread_num()
1187
1188 allocate (sr(mythread)%qv(tp%dimen))
1189 sr(mythread)%qv = tp%the_data(:,idxin) ! copy the vector
1190 sr(mythread)%ballsize = r2
1191 sr(mythread)%nn = 0 ! flag for fixed r search
1192 sr(mythread)%nfound = 0
1193 sr(mythread)%centeridx = idxin
1194 sr(mythread)%correltime = correltime
1195
1196 sr(mythread)%results => results
1197
1198 sr(mythread)%nalloc = nalloc
1199 sr(mythread)%overflow = .false.
1200
1201 call validate_query_storage(nalloc)
1202
1203 ! sr(mythread)%dsl = HUGE(sr(mythread)%dsl) ! set to huge positive values
1204 ! sr(mythread)%il = -1 ! set to invalid indexes
1205
1206 sr(mythread)%ind => tp%ind
1207 sr(mythread)%rearrange = tp%rearrange
1208
1209 if (tp%rearrange) then
1210 sr(mythread)%Data => tp%rearranged_data
1211 else
1212 sr(mythread)%Data => tp%the_data
1213 endif
1214 sr(mythread)%rearrange = tp%rearrange
1215 sr(mythread)%dimen = tp%dimen
1216
1217 !
1218 !sr(mythread)%dsl = Huge(sr(mythread)%dsl) ! set to huge positive values
1219 !sr(mythread)%il = -1 ! set to invalid indexes
1220 !
1221
1222 call search(tp%root)
1223 nfound = sr(mythread)%nfound
1224 if (tp%sort) then
1225 call kdtree2_sort_results(nfound,results)
1226 endif
1227
1228 if (sr(mythread)%overflow) then
1229 write (*,*) 'KD_TREE_TRANS: warning! return from kdtree2_r_nearest found more neighbors'
1230 write (*,*) 'KD_TREE_TRANS: than storage was provided for. Answer is NOT smallest ball'
1231 write (*,*) 'KD_TREE_TRANS: with that number of neighbors! I.e. it is wrong.'
1232 endif
1233
1234 deallocate (sr(mythread)%qv)
1235 return
1236 end subroutine kdtree2_r_nearest_around_point
1237
1238 function kdtree2_r_count(tp,qv,r2) result(nfound)
1239 !:Purpose: Count the number of neighbors within square distance 'r2'.
1240 type (kdtree2), pointer :: tp
1241 real(kdkind), target, intent (In) :: qv(:)
1242 real(kdkind), intent(in) :: r2
1243 integer :: nfound
1244 integer :: mythread
1245 integer, external :: omp_get_thread_num
1246 ! ..
1247 ! .. Intrinsic Functions ..
1248 intrinsic HUGE
1249 ! ..
1250
1251 mythread = omp_get_thread_num()
1252
1253 sr(mythread)%qv => qv
1254 sr(mythread)%ballsize = r2
1255
1256 sr(mythread)%nn = 0 ! flag for fixed r search
1257 sr(mythread)%nfound = 0
1258 sr(mythread)%centeridx = -1
1259 sr(mythread)%correltime = 0
1260
1261 nullify(sr(mythread)%results) ! for some reason, FTN 95 chokes on '=> null()'
1262
1263 sr(mythread)%nalloc = 0 ! we do not allocate any storage but that's OK
1264 ! for counting.
1265 sr(mythread)%ind => tp%ind
1266 sr(mythread)%rearrange = tp%rearrange
1267 if (tp%rearrange) then
1268 sr(mythread)%Data => tp%rearranged_data
1269 else
1270 sr(mythread)%Data => tp%the_data
1271 endif
1272 sr(mythread)%dimen = tp%dimen
1273
1274 !
1275 !sr(mythread)%dsl = Huge(sr(mythread)%dsl) ! set to huge positive values
1276 !sr(mythread)%il = -1 ! set to invalid indexes
1277 !
1278 sr(mythread)%overflow = .false.
1279
1280 call search(tp%root)
1281
1282 nfound = sr(mythread)%nfound
1283
1284 return
1285 end function kdtree2_r_count
1286
1287 function kdtree2_r_count_around_point(tp,idxin,correltime,r2) &
1288 result(nfound)
1289 !:Purpose: Count the number of neighbors within square distance 'r2' around
1290 ! point 'idxin' with decorrelation time 'correltime'.
1291 !
1292 type (kdtree2), pointer :: tp
1293 integer, intent (In) :: correltime, idxin
1294 real(kdkind), intent(in) :: r2
1295 integer :: nfound
1296 integer :: mythread
1297 integer, external :: omp_get_thread_num
1298 ! ..
1299 ! ..
1300 ! .. Intrinsic Functions ..
1301 intrinsic HUGE
1302 ! ..
1303
1304 mythread = omp_get_thread_num()
1305
1306 allocate (sr(mythread)%qv(tp%dimen))
1307 sr(mythread)%qv = tp%the_data(:,idxin)
1308 sr(mythread)%ballsize = r2
1309
1310 sr(mythread)%nn = 0 ! flag for fixed r search
1311 sr(mythread)%nfound = 0
1312 sr(mythread)%centeridx = idxin
1313 sr(mythread)%correltime = correltime
1314 nullify(sr(mythread)%results)
1315
1316 sr(mythread)%nalloc = 0 ! we do not allocate any storage but that's OK
1317 ! for counting.
1318
1319 sr(mythread)%ind => tp%ind
1320 sr(mythread)%rearrange = tp%rearrange
1321
1322 if (sr(mythread)%rearrange) then
1323 sr(mythread)%Data => tp%rearranged_data
1324 else
1325 sr(mythread)%Data => tp%the_data
1326 endif
1327 sr(mythread)%dimen = tp%dimen
1328
1329 !
1330 !sr(mythread)%dsl = Huge(sr(mythread)%dsl) ! set to huge positive values
1331 !sr(mythread)%il = -1 ! set to invalid indexes
1332 !
1333 sr(mythread)%overflow = .false.
1334
1335 call search(tp%root)
1336
1337 nfound = sr(mythread)%nfound
1338
1339 return
1340 end function kdtree2_r_count_around_point
1341
1342
1343 subroutine validate_query_storage(n)
1344 !:Purpose: make sure we have enough storage for n
1345 integer, intent(in) :: n
1346 integer :: mythread
1347 integer, external :: omp_get_thread_num
1348
1349 mythread = omp_get_thread_num()
1350
1351 if (size(sr(mythread)%results,1) .lt. n) then
1352 call utl_abort('kdtrees_mod-validate_query_storage: not enough storage for results')
1353 endif
1354
1355 return
1356 end subroutine validate_query_storage
1357
1358 function square_distance(d, iv,qv) result (res)
1359 !:Purpose: distance between iv[1:n] and qv[1:n]
1360
1361 ! .. Function Return Value ..
1362 ! re-implemented to improve vectorization.
1363 real(kdkind) :: res
1364 ! ..
1365 ! ..
1366 ! .. Scalar Arguments ..
1367 integer :: d
1368 ! ..
1369 ! .. Array Arguments ..
1370 real(kdkind) :: iv(:),qv(:)
1371 ! ..
1372 ! ..
1373 res = sum( (iv(1:d)-qv(1:d))**2 )
1374 end function square_distance
1375
1376 recursive subroutine search(node)
1377 !:Purpose: This is the innermost core routine of the kd-tree search. Along
1378 ! with "process_terminal_node", it is the performance bottleneck.
1379 !
1380 ! This version uses a logically complete secondary search of
1381 ! "box in bounds", whether the sear
1382 !
1383 type (Tree_node), pointer :: node
1384 ! ..
1385 type(tree_node),pointer :: ncloser, nfarther
1386 !
1387 integer :: cut_dim, i
1388 ! ..
1389 real(kdkind) :: qval, dis
1390 real(kdkind) :: ballsize
1391 real(kdkind), pointer :: qv(:)
1392 type(interval), pointer :: box(:)
1393 integer :: mythread
1394 integer, external :: omp_get_thread_num
1395
1396 mythread = omp_get_thread_num()
1397
1398 if ((associated(node%left) .and. associated(node%right)) .eqv. .false.) then
1399 ! we are on a terminal node
1400 if (sr(mythread)%nn .eq. 0) then
1401 call process_terminal_node_fixedball(node)
1402 else
1403 call process_terminal_node(node)
1404 endif
1405 else
1406 ! we are not on a terminal node
1407 qv => sr(mythread)%qv(1:)
1408 cut_dim = node%cut_dim
1409 qval = qv(cut_dim)
1410
1411 if (qval < node%cut_val) then
1412 ncloser => node%left
1413 nfarther => node%right
1414 dis = (node%cut_val_right - qval)**2
1415! extra = node%cut_val - qval
1416 else
1417 ncloser => node%right
1418 nfarther => node%left
1419 dis = (node%cut_val_left - qval)**2
1420! extra = qval- node%cut_val_left
1421 endif
1422
1423 if (associated(ncloser)) call search(ncloser)
1424
1425 ! we may need to search the second node.
1426 if (associated(nfarther)) then
1427 ballsize = sr(mythread)%ballsize
1428! dis=extra**2
1429 if (dis <= ballsize) then
1430 !
1431 ! we do this separately as going on the first cut dimen is often
1432 ! a good idea.
1433 ! note that if extra**2 < sr(mythread)%ballsize, then the next
1434 ! check will also be false.
1435 !
1436 box => node%box(1:)
1437 do i=1,sr(mythread)%dimen
1438 if (i .ne. cut_dim) then
1439 dis = dis + dis2_from_bnd(qv(i),box(i)%lower,box(i)%upper)
1440 if (dis > ballsize) then
1441 return
1442 endif
1443 endif
1444 end do
1445
1446 !
1447 ! if we are still here then we need to search mroe.
1448 !
1449 call search(nfarther)
1450 endif
1451 endif
1452 end if
1453 end subroutine search
1454
1455
1456 real(kdkind) function dis2_from_bnd(x,amin,amax) result (res)
1457 real(kdkind), intent(in) :: x, amin,amax
1458
1459 if (x > amax) then
1460 res = (x-amax)**2;
1461 return
1462 else
1463 if (x < amin) then
1464 res = (amin-x)**2;
1465 return
1466 else
1467 res = 0.0
1468 return
1469 endif
1470 endif
1471 return
1472 end function dis2_from_bnd
1473
1474 logical function box_in_search_range(node, sr) result(res)
1475 !:Purpose: Return the distance from 'qv' to the CLOSEST corner of node's
1476 ! bounding box for all coordinates outside the box. Coordinates
1477 ! inside the box contribute nothing to the distance.
1478 !
1479 type (tree_node), pointer :: node
1480 type (tree_search_record), pointer :: sr(:)
1481
1482 integer :: dimen, i
1483 real(kdkind) :: dis, ballsize
1484 real(kdkind) :: l, u
1485 integer :: mythread
1486 integer, external :: omp_get_thread_num
1487
1488 mythread = omp_get_thread_num()
1489
1490 dimen = sr(mythread)%dimen
1491 ballsize = sr(mythread)%ballsize
1492 dis = 0.0
1493 res = .true.
1494 do i=1,dimen
1495 l = node%box(i)%lower
1496 u = node%box(i)%upper
1497 dis = dis + (dis2_from_bnd(sr(mythread)%qv(i),l,u))
1498 if (dis > ballsize) then
1499 res = .false.
1500 return
1501 endif
1502 end do
1503 res = .true.
1504 return
1505 end function box_in_search_range
1506
1507
1508 subroutine process_terminal_node(node)
1509 !:Purpose: Look for actual near neighbors in 'node', and update
1510 ! the search results on the sr(:) data structure.
1511 !
1512 type (tree_node), pointer :: node
1513 !
1514 real(kdkind), pointer :: qv(:)
1515 integer, pointer :: ind(:)
1516 real(kdkind), pointer :: data(:,:)
1517 !
1518 integer :: dimen, i, indexofi, k, centeridx, correltime
1519 real(kdkind) :: ballsize, sd, newpri
1520 logical :: rearrange
1521 type(pq), pointer :: pqp
1522 integer :: mythread
1523 integer, external :: omp_get_thread_num
1524
1525 mythread = omp_get_thread_num()
1526
1527 !
1528 ! copy values from sr(:) to local variables
1529 !
1530 !
1531 ! Notice, making local pointers with an EXPLICIT lower bound
1532 ! seems to generate faster code.
1533 ! why? I don't know.
1534 qv => sr(mythread)%qv(1:)
1535 pqp => sr(mythread)%pq
1536 dimen = sr(mythread)%dimen
1537 ballsize = sr(mythread)%ballsize
1538 rearrange = sr(mythread)%rearrange
1539 ind => sr(mythread)%ind(1:)
1540 data => sr(mythread)%Data(1:,1:)
1541 centeridx = sr(mythread)%centeridx
1542 correltime = sr(mythread)%correltime
1543
1544 ! doing_correl = (centeridx >= 0) ! Do we have a decorrelation window?
1545 ! include_point = .true. ! by default include all points
1546 ! search through terminal bucket.
1547
1548 mainloop: do i = node%l, node%u
1549 if (rearrange) then
1550 sd = 0.0
1551 do k = 1,dimen
1552 sd = sd + (data(k,i) - qv(k))**2
1553 if (sd>ballsize) cycle mainloop
1554 end do
1555 indexofi = ind(i) ! only read it if we have not broken out
1556 else
1557 indexofi = ind(i)
1558 sd = 0.0
1559 do k = 1,dimen
1560 sd = sd + (data(k,indexofi) - qv(k))**2
1561 if (sd>ballsize) cycle mainloop
1562 end do
1563 endif
1564
1565 if (centeridx > 0) then ! doing correlation interval?
1566 if (abs(indexofi-centeridx) < correltime) cycle mainloop
1567 endif
1568
1569
1570 !
1571 ! two choices for any point. The list so far is either undersized,
1572 ! or it is not.
1573 !
1574 ! If it is undersized, then add the point and its distance
1575 ! unconditionally. If the point added fills up the working
1576 ! list then set the sr(mythread)%ballsize, maximum distance bound (largest distance on
1577 ! list) to be that distance, instead of the initialized +infinity.
1578 !
1579 ! If the running list is full size, then compute the
1580 ! distance but break out immediately if it is larger
1581 ! than sr(mythread)%ballsize, "best squared distance" (of the largest element),
1582 ! as it cannot be a good neighbor.
1583 !
1584 ! Once computed, compare to best_square distance.
1585 ! if it is smaller, then delete the previous largest
1586 ! element and add the new one.
1587
1588 if (sr(mythread)%nfound .lt. sr(mythread)%nn) then
1589 !
1590 ! add this point unconditionally to fill list.
1591 !
1592 sr(mythread)%nfound = sr(mythread)%nfound +1
1593 newpri = pq_insert(pqp,sd,indexofi)
1594 if (sr(mythread)%nfound .eq. sr(mythread)%nn) ballsize = newpri
1595 ! we have just filled the working list.
1596 ! put the best square distance to the maximum value
1597 ! on the list, which is extractable from the PQ.
1598
1599 else
1600 !
1601 ! now, if we get here,
1602 ! we know that the current node has a squared
1603 ! distance smaller than the largest one on the list, and
1604 ! belongs on the list.
1605 ! Hence we replace that with the current one.
1606 !
1607 ballsize = pq_replace_max(pqp,sd,indexofi)
1608 endif
1609 end do mainloop
1610 !
1611 ! Reset sr(mythread) variables which may have changed during loop
1612 !
1613 sr(mythread)%ballsize = ballsize
1614
1615 end subroutine process_terminal_node
1616
1617 subroutine process_terminal_node_fixedball(node)
1618 !:Purpose: Look for actual near neighbors in 'node', and update
1619 ! the search results on the sr(:) data structure, i.e.
1620 ! save all within a fixed ball.
1621 !
1622 type (tree_node), pointer :: node
1623 !
1624 real(kdkind), pointer :: qv(:)
1625 integer, pointer :: ind(:)
1626 real(kdkind), pointer :: data(:,:)
1627 !
1628 integer :: nfound
1629 integer :: dimen, i, indexofi, k
1630 integer :: centeridx, correltime, nn
1631 real(kdkind) :: ballsize, sd
1632 logical :: rearrange
1633 integer :: mythread
1634 integer, external :: omp_get_thread_num
1635
1636 mythread = omp_get_thread_num()
1637
1638 !
1639 ! copy values from sr(:) to local variables
1640 !
1641 qv => sr(mythread)%qv(1:)
1642 dimen = sr(mythread)%dimen
1643 ballsize = sr(mythread)%ballsize
1644 rearrange = sr(mythread)%rearrange
1645 ind => sr(mythread)%ind(1:)
1646 data => sr(mythread)%Data(1:,1:)
1647 centeridx = sr(mythread)%centeridx
1648 correltime = sr(mythread)%correltime
1649 nn = sr(mythread)%nn ! number to search for
1650 nfound = sr(mythread)%nfound
1651
1652 ! search through terminal bucket.
1653 mainloop: do i = node%l, node%u
1654
1655 !
1656 ! two choices for any point. The list so far is either undersized,
1657 ! or it is not.
1658 !
1659 ! If it is undersized, then add the point and its distance
1660 ! unconditionally. If the point added fills up the working
1661 ! list then set the sr(mythread)%ballsize, maximum distance bound (largest distance on
1662 ! list) to be that distance, instead of the initialized +infinity.
1663 !
1664 ! If the running list is full size, then compute the
1665 ! distance but break out immediately if it is larger
1666 ! than sr(mythread)%ballsize, "best squared distance" (of the largest element),
1667 ! as it cannot be a good neighbor.
1668 !
1669 ! Once computed, compare to best_square distance.
1670 ! if it is smaller, then delete the previous largest
1671 ! element and add the new one.
1672
1673 ! which index to the point do we use?
1674
1675 if (rearrange) then
1676 sd = 0.0
1677 do k = 1,dimen
1678 sd = sd + (data(k,i) - qv(k))**2
1679 if (sd>ballsize) cycle mainloop
1680 end do
1681 indexofi = ind(i) ! only read it if we have not broken out
1682 else
1683 indexofi = ind(i)
1684 sd = 0.0
1685 do k = 1,dimen
1686 sd = sd + (data(k,indexofi) - qv(k))**2
1687 if (sd>ballsize) cycle mainloop
1688 end do
1689 endif
1690
1691 if (centeridx > 0) then ! doing correlation interval?
1692 if (abs(indexofi-centeridx)<correltime) cycle mainloop
1693 endif
1694
1695 nfound = nfound+1
1696 if (nfound .gt. sr(mythread)%nalloc) then
1697 ! oh nuts, we have to add another one to the tree but
1698 ! there isn't enough room.
1699 sr(mythread)%overflow = .true.
1700 else
1701 sr(mythread)%results(nfound)%dis = sd
1702 sr(mythread)%results(nfound)%idx = indexofi
1703 endif
1704 end do mainloop
1705 !
1706 ! Reset sr(mythread) variables which may have changed during loop
1707 !
1708 sr(mythread)%nfound = nfound
1709 end subroutine process_terminal_node_fixedball
1710
1711 subroutine kdtree2_n_nearest_brute_force(tp,qv,nn,results)
1712 !:Purpose: find the 'n' nearest neighbors to 'qv' by exhaustive search.
1713 ! only use this subroutine for testing, as it is SLOW! The
1714 ! whole point of a k-d tree is to avoid doing what this subroutine
1715 ! does.
1716 type (kdtree2), pointer :: tp
1717 real(kdkind), intent (In) :: qv(:)
1718 integer, intent (In) :: nn
1719 type(kdtree2_result) :: results(:)
1720
1721 integer :: i, j, k
1722 real(kdkind), allocatable :: all_distances(:)
1723 ! ..
1724 allocate (all_distances(tp%n))
1725 do i = 1, tp%n
1726 all_distances(i) = square_distance(tp%dimen,qv,tp%the_data(:,i))
1727 end do
1728 ! now find 'n' smallest distances
1729 do i = 1, nn
1730 results(i)%dis = huge(1.0)
1731 results(i)%idx = -1
1732 end do
1733 do i = 1, tp%n
1734 if (all_distances(i)<results(nn)%dis) then
1735 ! insert it somewhere on the list
1736 do j = 1, nn
1737 if (all_distances(i)<results(j)%dis) exit
1738 end do
1739 ! now we know 'j'
1740 do k = nn - 1, j, -1
1741 results(k+1) = results(k)
1742 end do
1743 results(j)%dis = all_distances(i)
1744 results(j)%idx = i
1745 end if
1746 end do
1747 deallocate (all_distances)
1748 end subroutine kdtree2_n_nearest_brute_force
1749
1750
1751 subroutine kdtree2_r_nearest_brute_force(tp,qv,r2,nfound,results)
1752 !:Purpose: find the nearest neighbors to 'qv' with distance**2 <= r2 by exhaustive search.
1753 ! only use this subroutine for testing, as it is SLOW! The
1754 ! whole point of a k-d tree is to avoid doing what this subroutine
1755 ! does.
1756 type (kdtree2), pointer :: tp
1757 real(kdkind), intent (In) :: qv(:)
1758 real(kdkind), intent (In) :: r2
1759 integer, intent(out) :: nfound
1760 type(kdtree2_result) :: results(:)
1761
1762 integer :: i, nalloc
1763 real(kdkind), allocatable :: all_distances(:)
1764 ! ..
1765 allocate (all_distances(tp%n))
1766 do i = 1, tp%n
1767 all_distances(i) = square_distance(tp%dimen,qv,tp%the_data(:,i))
1768 end do
1769
1770 nfound = 0
1771 nalloc = size(results,1)
1772
1773 do i = 1, tp%n
1774 if (all_distances(i)< r2) then
1775 ! insert it somewhere on the list
1776 if (nfound .lt. nalloc) then
1777 nfound = nfound+1
1778 results(nfound)%dis = all_distances(i)
1779 results(nfound)%idx = i
1780 endif
1781 end if
1782 enddo
1783 deallocate (all_distances)
1784
1785 call kdtree2_sort_results(nfound,results)
1786
1787
1788 end subroutine kdtree2_r_nearest_brute_force
1789
1790 subroutine kdtree2_sort_results(nfound,results)
1791 !:Purpose: Use after search to sort results(1:nfound) in order of increasing
1792 ! distance.
1793 integer, intent(in) :: nfound
1794 type(kdtree2_result), target :: results(:)
1795 !
1796 !
1797
1798 !THIS IS BUGGY WITH INTEL FORTRAN
1799 ! If (nfound .Gt. 1) Call heapsort(results(1:nfound)%dis,results(1:nfound)%ind,nfound)
1800 !
1801 if (nfound .gt. 1) call heapsort_struct(results,nfound)
1802
1803 return
1804 end subroutine kdtree2_sort_results
1805
1806 subroutine heapsort(a,ind,n)
1807 !:Purpose: Sort a(1:n) in ascending order, permuting ind(1:n) similarly.
1808 !
1809 ! If ind(k) = k upon input, then it will give a sort index upon output.
1810 !
1811 integer,intent(in) :: n
1812 real(kdkind), intent(inout) :: a(:)
1813 integer, intent(inout) :: ind(:)
1814
1815 !
1816 !
1817 real(kdkind) :: value ! temporary for a value from a()
1818 integer :: ivalue ! temporary for a value from ind()
1819
1820 integer :: i,j
1821 integer :: ileft,iright
1822
1823 ileft=n/2+1
1824 iright=n
1825
1826 ! do i=1,n
1827 ! ind(i)=i
1828 ! Generate initial idum array
1829 ! end do
1830
1831 if(n.eq.1) return
1832
1833 do
1834 if(ileft > 1)then
1835 ileft=ileft-1
1836 value=a(ileft); ivalue=ind(ileft)
1837 else
1838 value=a(iright); ivalue=ind(iright)
1839 a(iright)=a(1); ind(iright)=ind(1)
1840 iright=iright-1
1841 if (iright == 1) then
1842 a(1)=value;ind(1)=ivalue
1843 return
1844 endif
1845 endif
1846 i=ileft
1847 j=2*ileft
1848 do while (j <= iright)
1849 if(j < iright) then
1850 if(a(j) < a(j+1)) j=j+1
1851 endif
1852 if(value < a(j)) then
1853 a(i)=a(j); ind(i)=ind(j)
1854 i=j
1855 j=j+j
1856 else
1857 j=iright+1
1858 endif
1859 end do
1860 a(i)=value; ind(i)=ivalue
1861 end do
1862 end subroutine heapsort
1863
1864 subroutine heapsort_struct(a,n)
1865 !:Purpose: Sort a(1:n) in ascending order
1866 !
1867 integer,intent(in) :: n
1868 type(kdtree2_result),intent(inout) :: a(:)
1869
1870 !
1871 !
1872 type(kdtree2_result) :: value ! temporary value
1873
1874 integer :: i,j
1875 integer :: ileft,iright
1876
1877 ileft=n/2+1
1878 iright=n
1879
1880 ! do i=1,n
1881 ! ind(i)=i
1882 ! Generate initial idum array
1883 ! end do
1884
1885 if(n.eq.1) return
1886
1887 do
1888 if(ileft > 1)then
1889 ileft=ileft-1
1890 value=a(ileft)
1891 else
1892 value=a(iright)
1893 a(iright)=a(1)
1894 iright=iright-1
1895 if (iright == 1) then
1896 a(1) = value
1897 return
1898 endif
1899 endif
1900 i=ileft
1901 j=2*ileft
1902 do while (j <= iright)
1903 if(j < iright) then
1904 if(a(j)%dis < a(j+1)%dis) j=j+1
1905 endif
1906 if(value%dis < a(j)%dis) then
1907 a(i)=a(j);
1908 i=j
1909 j=j+j
1910 else
1911 j=iright+1
1912 endif
1913 end do
1914 a(i)=value
1915 end do
1916 end subroutine heapsort_struct
1917
1918 !--------------------------------------------------------------------------
1919 ! kdtree2_3dPosition
1920 !--------------------------------------------------------------------------
1921 function kdtree2_3dPosition(lon_rad, lat_rad) result(positionArray)
1922 !
1923 !:Purpose: Calculate the 3 dimensional coordinates of a point
1924 ! on the sphere of Earth radius given
1925 ! the longitude-latitude position on the surface of the sphere
1926 ! in radian.
1927 !
1928 implicit none
1929
1930 ! arguments
1931 real(kdkind) :: positionArray(3) ! returned values of function
1932 real(8), intent(in) :: lon_rad, lat_rad
1933
1934 positionArray(1) = ec_ra * sin(lon_rad) * cos(lat_rad)
1935 positionArray(2) = ec_ra * cos(lon_rad) * cos(lat_rad)
1936 positionArray(3) = ec_ra * sin(lat_rad)
1937
1938 end function kdtree2_3dPosition
1939
1940end module kdTree2_mod
