rusty_tree/
morton.rs

1//! Routines for Morton encoding and decoding.
2
3use ndarray::{Array1, Array2, ArrayView1, ArrayView2, Axis, Zip};
4use rusty_kernel_tools::RealType;
5use std::collections::HashSet;
6
7const X_LOOKUP_ENCODE: [usize; 256] = [
8    0x00000000, 0x00000001, 0x00000008, 0x00000009, 0x00000040, 0x00000041, 0x00000048, 0x00000049,
9    0x00000200, 0x00000201, 0x00000208, 0x00000209, 0x00000240, 0x00000241, 0x00000248, 0x00000249,
10    0x00001000, 0x00001001, 0x00001008, 0x00001009, 0x00001040, 0x00001041, 0x00001048, 0x00001049,
11    0x00001200, 0x00001201, 0x00001208, 0x00001209, 0x00001240, 0x00001241, 0x00001248, 0x00001249,
12    0x00008000, 0x00008001, 0x00008008, 0x00008009, 0x00008040, 0x00008041, 0x00008048, 0x00008049,
13    0x00008200, 0x00008201, 0x00008208, 0x00008209, 0x00008240, 0x00008241, 0x00008248, 0x00008249,
14    0x00009000, 0x00009001, 0x00009008, 0x00009009, 0x00009040, 0x00009041, 0x00009048, 0x00009049,
15    0x00009200, 0x00009201, 0x00009208, 0x00009209, 0x00009240, 0x00009241, 0x00009248, 0x00009249,
16    0x00040000, 0x00040001, 0x00040008, 0x00040009, 0x00040040, 0x00040041, 0x00040048, 0x00040049,
17    0x00040200, 0x00040201, 0x00040208, 0x00040209, 0x00040240, 0x00040241, 0x00040248, 0x00040249,
18    0x00041000, 0x00041001, 0x00041008, 0x00041009, 0x00041040, 0x00041041, 0x00041048, 0x00041049,
19    0x00041200, 0x00041201, 0x00041208, 0x00041209, 0x00041240, 0x00041241, 0x00041248, 0x00041249,
20    0x00048000, 0x00048001, 0x00048008, 0x00048009, 0x00048040, 0x00048041, 0x00048048, 0x00048049,
21    0x00048200, 0x00048201, 0x00048208, 0x00048209, 0x00048240, 0x00048241, 0x00048248, 0x00048249,
22    0x00049000, 0x00049001, 0x00049008, 0x00049009, 0x00049040, 0x00049041, 0x00049048, 0x00049049,
23    0x00049200, 0x00049201, 0x00049208, 0x00049209, 0x00049240, 0x00049241, 0x00049248, 0x00049249,
24    0x00200000, 0x00200001, 0x00200008, 0x00200009, 0x00200040, 0x00200041, 0x00200048, 0x00200049,
25    0x00200200, 0x00200201, 0x00200208, 0x00200209, 0x00200240, 0x00200241, 0x00200248, 0x00200249,
26    0x00201000, 0x00201001, 0x00201008, 0x00201009, 0x00201040, 0x00201041, 0x00201048, 0x00201049,
27    0x00201200, 0x00201201, 0x00201208, 0x00201209, 0x00201240, 0x00201241, 0x00201248, 0x00201249,
28    0x00208000, 0x00208001, 0x00208008, 0x00208009, 0x00208040, 0x00208041, 0x00208048, 0x00208049,
29    0x00208200, 0x00208201, 0x00208208, 0x00208209, 0x00208240, 0x00208241, 0x00208248, 0x00208249,
30    0x00209000, 0x00209001, 0x00209008, 0x00209009, 0x00209040, 0x00209041, 0x00209048, 0x00209049,
31    0x00209200, 0x00209201, 0x00209208, 0x00209209, 0x00209240, 0x00209241, 0x00209248, 0x00209249,
32    0x00240000, 0x00240001, 0x00240008, 0x00240009, 0x00240040, 0x00240041, 0x00240048, 0x00240049,
33    0x00240200, 0x00240201, 0x00240208, 0x00240209, 0x00240240, 0x00240241, 0x00240248, 0x00240249,
34    0x00241000, 0x00241001, 0x00241008, 0x00241009, 0x00241040, 0x00241041, 0x00241048, 0x00241049,
35    0x00241200, 0x00241201, 0x00241208, 0x00241209, 0x00241240, 0x00241241, 0x00241248, 0x00241249,
36    0x00248000, 0x00248001, 0x00248008, 0x00248009, 0x00248040, 0x00248041, 0x00248048, 0x00248049,
37    0x00248200, 0x00248201, 0x00248208, 0x00248209, 0x00248240, 0x00248241, 0x00248248, 0x00248249,
38    0x00249000, 0x00249001, 0x00249008, 0x00249009, 0x00249040, 0x00249041, 0x00249048, 0x00249049,
39    0x00249200, 0x00249201, 0x00249208, 0x00249209, 0x00249240, 0x00249241, 0x00249248, 0x00249249,
40];
41
42const Y_LOOKUP_ENCODE: [usize; 256] = [
43    0x00000000, 0x00000002, 0x00000010, 0x00000012, 0x00000080, 0x00000082, 0x00000090, 0x00000092,
44    0x00000400, 0x00000402, 0x00000410, 0x00000412, 0x00000480, 0x00000482, 0x00000490, 0x00000492,
45    0x00002000, 0x00002002, 0x00002010, 0x00002012, 0x00002080, 0x00002082, 0x00002090, 0x00002092,
46    0x00002400, 0x00002402, 0x00002410, 0x00002412, 0x00002480, 0x00002482, 0x00002490, 0x00002492,
47    0x00010000, 0x00010002, 0x00010010, 0x00010012, 0x00010080, 0x00010082, 0x00010090, 0x00010092,
48    0x00010400, 0x00010402, 0x00010410, 0x00010412, 0x00010480, 0x00010482, 0x00010490, 0x00010492,
49    0x00012000, 0x00012002, 0x00012010, 0x00012012, 0x00012080, 0x00012082, 0x00012090, 0x00012092,
50    0x00012400, 0x00012402, 0x00012410, 0x00012412, 0x00012480, 0x00012482, 0x00012490, 0x00012492,
51    0x00080000, 0x00080002, 0x00080010, 0x00080012, 0x00080080, 0x00080082, 0x00080090, 0x00080092,
52    0x00080400, 0x00080402, 0x00080410, 0x00080412, 0x00080480, 0x00080482, 0x00080490, 0x00080492,
53    0x00082000, 0x00082002, 0x00082010, 0x00082012, 0x00082080, 0x00082082, 0x00082090, 0x00082092,
54    0x00082400, 0x00082402, 0x00082410, 0x00082412, 0x00082480, 0x00082482, 0x00082490, 0x00082492,
55    0x00090000, 0x00090002, 0x00090010, 0x00090012, 0x00090080, 0x00090082, 0x00090090, 0x00090092,
56    0x00090400, 0x00090402, 0x00090410, 0x00090412, 0x00090480, 0x00090482, 0x00090490, 0x00090492,
57    0x00092000, 0x00092002, 0x00092010, 0x00092012, 0x00092080, 0x00092082, 0x00092090, 0x00092092,
58    0x00092400, 0x00092402, 0x00092410, 0x00092412, 0x00092480, 0x00092482, 0x00092490, 0x00092492,
59    0x00400000, 0x00400002, 0x00400010, 0x00400012, 0x00400080, 0x00400082, 0x00400090, 0x00400092,
60    0x00400400, 0x00400402, 0x00400410, 0x00400412, 0x00400480, 0x00400482, 0x00400490, 0x00400492,
61    0x00402000, 0x00402002, 0x00402010, 0x00402012, 0x00402080, 0x00402082, 0x00402090, 0x00402092,
62    0x00402400, 0x00402402, 0x00402410, 0x00402412, 0x00402480, 0x00402482, 0x00402490, 0x00402492,
63    0x00410000, 0x00410002, 0x00410010, 0x00410012, 0x00410080, 0x00410082, 0x00410090, 0x00410092,
64    0x00410400, 0x00410402, 0x00410410, 0x00410412, 0x00410480, 0x00410482, 0x00410490, 0x00410492,
65    0x00412000, 0x00412002, 0x00412010, 0x00412012, 0x00412080, 0x00412082, 0x00412090, 0x00412092,
66    0x00412400, 0x00412402, 0x00412410, 0x00412412, 0x00412480, 0x00412482, 0x00412490, 0x00412492,
67    0x00480000, 0x00480002, 0x00480010, 0x00480012, 0x00480080, 0x00480082, 0x00480090, 0x00480092,
68    0x00480400, 0x00480402, 0x00480410, 0x00480412, 0x00480480, 0x00480482, 0x00480490, 0x00480492,
69    0x00482000, 0x00482002, 0x00482010, 0x00482012, 0x00482080, 0x00482082, 0x00482090, 0x00482092,
70    0x00482400, 0x00482402, 0x00482410, 0x00482412, 0x00482480, 0x00482482, 0x00482490, 0x00482492,
71    0x00490000, 0x00490002, 0x00490010, 0x00490012, 0x00490080, 0x00490082, 0x00490090, 0x00490092,
72    0x00490400, 0x00490402, 0x00490410, 0x00490412, 0x00490480, 0x00490482, 0x00490490, 0x00490492,
73    0x00492000, 0x00492002, 0x00492010, 0x00492012, 0x00492080, 0x00492082, 0x00492090, 0x00492092,
74    0x00492400, 0x00492402, 0x00492410, 0x00492412, 0x00492480, 0x00492482, 0x00492490, 0x00492492,
75];
76
77const Z_LOOKUP_ENCODE: [usize; 256] = [
78    0x00000000, 0x00000004, 0x00000020, 0x00000024, 0x00000100, 0x00000104, 0x00000120, 0x00000124,
79    0x00000800, 0x00000804, 0x00000820, 0x00000824, 0x00000900, 0x00000904, 0x00000920, 0x00000924,
80    0x00004000, 0x00004004, 0x00004020, 0x00004024, 0x00004100, 0x00004104, 0x00004120, 0x00004124,
81    0x00004800, 0x00004804, 0x00004820, 0x00004824, 0x00004900, 0x00004904, 0x00004920, 0x00004924,
82    0x00020000, 0x00020004, 0x00020020, 0x00020024, 0x00020100, 0x00020104, 0x00020120, 0x00020124,
83    0x00020800, 0x00020804, 0x00020820, 0x00020824, 0x00020900, 0x00020904, 0x00020920, 0x00020924,
84    0x00024000, 0x00024004, 0x00024020, 0x00024024, 0x00024100, 0x00024104, 0x00024120, 0x00024124,
85    0x00024800, 0x00024804, 0x00024820, 0x00024824, 0x00024900, 0x00024904, 0x00024920, 0x00024924,
86    0x00100000, 0x00100004, 0x00100020, 0x00100024, 0x00100100, 0x00100104, 0x00100120, 0x00100124,
87    0x00100800, 0x00100804, 0x00100820, 0x00100824, 0x00100900, 0x00100904, 0x00100920, 0x00100924,
88    0x00104000, 0x00104004, 0x00104020, 0x00104024, 0x00104100, 0x00104104, 0x00104120, 0x00104124,
89    0x00104800, 0x00104804, 0x00104820, 0x00104824, 0x00104900, 0x00104904, 0x00104920, 0x00104924,
90    0x00120000, 0x00120004, 0x00120020, 0x00120024, 0x00120100, 0x00120104, 0x00120120, 0x00120124,
91    0x00120800, 0x00120804, 0x00120820, 0x00120824, 0x00120900, 0x00120904, 0x00120920, 0x00120924,
92    0x00124000, 0x00124004, 0x00124020, 0x00124024, 0x00124100, 0x00124104, 0x00124120, 0x00124124,
93    0x00124800, 0x00124804, 0x00124820, 0x00124824, 0x00124900, 0x00124904, 0x00124920, 0x00124924,
94    0x00800000, 0x00800004, 0x00800020, 0x00800024, 0x00800100, 0x00800104, 0x00800120, 0x00800124,
95    0x00800800, 0x00800804, 0x00800820, 0x00800824, 0x00800900, 0x00800904, 0x00800920, 0x00800924,
96    0x00804000, 0x00804004, 0x00804020, 0x00804024, 0x00804100, 0x00804104, 0x00804120, 0x00804124,
97    0x00804800, 0x00804804, 0x00804820, 0x00804824, 0x00804900, 0x00804904, 0x00804920, 0x00804924,
98    0x00820000, 0x00820004, 0x00820020, 0x00820024, 0x00820100, 0x00820104, 0x00820120, 0x00820124,
99    0x00820800, 0x00820804, 0x00820820, 0x00820824, 0x00820900, 0x00820904, 0x00820920, 0x00820924,
100    0x00824000, 0x00824004, 0x00824020, 0x00824024, 0x00824100, 0x00824104, 0x00824120, 0x00824124,
101    0x00824800, 0x00824804, 0x00824820, 0x00824824, 0x00824900, 0x00824904, 0x00824920, 0x00824924,
102    0x00900000, 0x00900004, 0x00900020, 0x00900024, 0x00900100, 0x00900104, 0x00900120, 0x00900124,
103    0x00900800, 0x00900804, 0x00900820, 0x00900824, 0x00900900, 0x00900904, 0x00900920, 0x00900924,
104    0x00904000, 0x00904004, 0x00904020, 0x00904024, 0x00904100, 0x00904104, 0x00904120, 0x00904124,
105    0x00904800, 0x00904804, 0x00904820, 0x00904824, 0x00904900, 0x00904904, 0x00904920, 0x00904924,
106    0x00920000, 0x00920004, 0x00920020, 0x00920024, 0x00920100, 0x00920104, 0x00920120, 0x00920124,
107    0x00920800, 0x00920804, 0x00920820, 0x00920824, 0x00920900, 0x00920904, 0x00920920, 0x00920924,
108    0x00924000, 0x00924004, 0x00924020, 0x00924024, 0x00924100, 0x00924104, 0x00924120, 0x00924124,
109    0x00924800, 0x00924804, 0x00924820, 0x00924824, 0x00924900, 0x00924904, 0x00924920, 0x00924924,
110];
111
112const X_LOOKUP_DECODE: [usize; 512] = [
113    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
114    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
115    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
116    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
117    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
118    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
119    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
120    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
121    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
122    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
123    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
124    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
125    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
126    0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3,
127    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
128    4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7,
129];
130
131const Y_LOOKUP_DECODE: [usize; 512] = [
132    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
133    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
134    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
135    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
136    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
137    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
138    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
139    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
140    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
141    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
142    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
143    0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3, 2, 2, 3, 3,
144    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
145    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
146    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
147    4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 4, 4, 5, 5, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7, 6, 6, 7, 7,
148];
149
150const Z_LOOKUP_DECODE: [usize; 512] = [
151    0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1,
152    2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3,
153    0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1,
154    2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3,
155    0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1,
156    2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3,
157    0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1,
158    2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3, 2, 2, 2, 2, 3, 3, 3, 3,
159    4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5,
160    6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7,
161    4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5,
162    6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7,
163    4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5,
164    6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7,
165    4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5, 4, 4, 4, 4, 5, 5, 5, 5,
166    6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 6, 6, 7, 7, 7, 7,
167];
168
169// Number of bits used for Level information.
170const LEVEL_DISPLACEMENT: usize = 15;
171
172// Mask for the last 15 bits.
173const LEVEL_MASK: usize = 0x7FFF;
174
175// Mask for lowest order byte.
176const BYTE_MASK: usize = 0xFF;
177const BYTE_DISPLACEMENT: usize = 8;
178
179// Mask encapsulating a bit.
180const NINE_BIT_MASK: usize = 0x1FF;
181
182/// Return the level associated with a key.
183pub fn find_level(key: usize) -> usize {
184    return key & LEVEL_MASK;
185}
186
187/// Map a point to the integer coordinates of its enclosing box.
188///
189/// Returns the 3 integeger coordinates of the enclosing box.
190///
191/// # Arguments
192/// `point` - The (x, y, z) coordinates of the point to map.
193/// `level` - The level of the tree at which the point will be mapped.
194/// `origin` - The origin of the bounding box.
195/// `diameter` - The diameter of the bounding box in each dimension.
196pub fn point_to_anchor(
197    point: &[f64; 3],
198    level: usize,
199    origin: &[f64; 3],
200    diameter: &[f64; 3],
201) -> [usize; 4] {
202    use itertools::izip;
203    let mut anchor: [usize; 4] = [0, 0, 0, 0];
204    anchor[3] = level;
205
206    let level_size = (1 << level) as f64;
207
208    for (anchor_value, point_value, &origin_value, &diameter_value) in
209        izip!(&mut anchor, point, origin, diameter)
210    {
211        *anchor_value =
212            ((point_value - origin_value) * level_size / diameter_value).floor() as usize
213    }
214
215    anchor
216}
217
218/// Encode an anchor.
219///
220/// Returns the Morton key associated with the given anchor.
221///
222/// # Arguments
223/// `anchor` - A vector with 4 elements defining the integer coordinates and level.
224pub fn encode_anchor(anchor: &[usize; 4]) -> usize {
225    let x = anchor[0];
226    let y = anchor[1];
227    let z = anchor[2];
228    let level = anchor[3];
229
230    let key: usize = Z_LOOKUP_ENCODE[(z >> BYTE_DISPLACEMENT) & BYTE_MASK]
231        | Y_LOOKUP_ENCODE[(y >> BYTE_DISPLACEMENT) & BYTE_MASK]
232        | X_LOOKUP_ENCODE[(x >> BYTE_DISPLACEMENT) & BYTE_MASK];
233
234    let key = (key << 24)
235        | Z_LOOKUP_ENCODE[z & BYTE_MASK]
236        | Y_LOOKUP_ENCODE[y & BYTE_MASK]
237        | X_LOOKUP_ENCODE[x & BYTE_MASK];
238
239    let key = key << LEVEL_DISPLACEMENT;
240    key | level
241}
242
243/// Encode a point.
244///
245/// Return the Morton key of a point for a given level.
246///
247/// # Arguments
248/// `point` - The (x, y, z) coordinates of the point to map.
249/// `level` - The level of the tree at which the point will be mapped.
250/// `origin` - The origin of the bounding box.
251/// `diameter` - The diameter of the bounding box in each dimension.
252pub fn encode_point(
253    point: &[f64; 3],
254    level: usize,
255    origin: &[f64; 3],
256    diameter: &[f64; 3],
257) -> usize {
258    let anchor = point_to_anchor(point, level, origin, diameter);
259    encode_anchor(&anchor)
260}
261
262/// Given an anchor, return the corresponding x,y,z coordinates.
263///
264/// The result is the coordinates of the lower left corner of the box described by the given anchor.
265///
266/// # Arguments
267/// `anchor` - The three indices describing the box and the level information.
268/// `origin` - The origin of the bounding box.
269/// `diameter` - The diameter of the bounding box in each dimension.
270pub fn anchor_to_coordinates(
271    anchor: &[usize; 4],
272    origin: &[f64; 3],
273    diameter: &[f64; 3],
274) -> [f64; 3] {
275    use itertools::izip;
276    let mut coord: [f64; 3] = [0.0; 3];
277
278    let level = anchor[3];
279    let level_size = (1 << level) as f64;
280
281    for (&anchor_value, coord_ref, &origin_value, &diameter_value) in
282        izip!(anchor, &mut coord, origin, diameter)
283    {
284        *coord_ref = origin_value + diameter_value * (anchor_value as f64) / level_size;
285    }
286
287    coord
288}
289
290/// Serialized representation of a box associated with a key.
291///
292/// Returns a vector with 24 f64 entries, associated with the 8 x,y,z coordinates
293/// of the box associated with the key.
294/// If the lower left corner of the box is (0, 0, 0). Then the points are numbered in the
295/// following order.
296/// 1. (0, 0, 0)
297/// 2. (1, 0, 0)
298/// 3. (0, 1, 0)
299/// 4. (0, 1, 1)
300/// 5. (0, 0, 1)
301/// 6. (1, 0, 1)
302/// 7. (0, 1, 1)
303/// 8. (1, 1, 1)
304/// 
305/// # Arguments
306/// * `key` - The key for which the box is taken.
307/// `origin` - The origin of the bounding box.
308/// `diameter` - The diameter of the bounding box in each dimension.
309pub fn serialize_box_from_key(key: usize, origin: &[f64; 3], diameter: &[f64; 3]) -> Vec<f64> {
310    let anchor = decode_key(key);
311
312    let mut serialized = Vec::<f64>::with_capacity(24);
313
314    let anchors = [
315        [anchor[0], anchor[1], anchor[2], anchor[3]],
316        [1 + anchor[0], anchor[1], anchor[2], anchor[3]],
317        [anchor[0], 1 + anchor[1], anchor[2], anchor[3]],
318        [1 + anchor[0], 1 + anchor[1], anchor[2], anchor[3]],
319        [anchor[0], anchor[1], 1 + anchor[2], anchor[3]],
320        [1 + anchor[0], anchor[1], 1 + anchor[2], anchor[3]],
321        [anchor[0], 1 + anchor[1], 1 + anchor[2], anchor[3]],
322        [1 + anchor[0], 1 + anchor[1], 1 + anchor[2], anchor[3]],
323    ];
324
325    for anchor in anchors.iter() {
326        let coords = anchor_to_coordinates(anchor, origin, diameter);
327        for index in 0..3 {
328            serialized.push(coords[index]);
329        }
330    }
331
332    serialized
333}
334
335/// Encode many points.
336///
337/// Return an array containing all Morton keys of a given array of points.
338///
339/// # Arguments
340/// `point` - A (3 ,N) array of N points of type f32 or f64.
341/// `level` - The level of the tree at which the point will be mapped.
342/// `origin` - The origin of the bounding box.
343/// `diameter` - The diameter of the bounding box in each dimension.
344pub fn encode_points<T: RealType>(
345    points: ArrayView2<T>,
346    level: usize,
347    origin: &[f64; 3],
348    diameter: &[f64; 3],
349) -> Array1<usize> {
350    let npoints = points.len_of(Axis(1));
351    let mut box_coordinates = Array2::<usize>::zeros((3, npoints));
352    let mut keys = Array1::<usize>::zeros(npoints);
353
354    let level_size = (1 << level) as f64;
355
356    let origin_view = ArrayView1::from(origin);
357    let diameter_view = ArrayView1::from(diameter);
358
359    Zip::from(points.axis_iter(Axis(0)))
360        .and(box_coordinates.axis_iter_mut(Axis(0)))
361        .and(origin_view)
362        .and(diameter_view)
363        .apply(
364            |points_row, box_coordinates_row, &origin_value, &diameter_value| {
365                Zip::from(points_row).and(box_coordinates_row).par_apply(
366                    |&point_value, box_coordinate_value| {
367                        let tmp = (point_value.to_f64().unwrap() - origin_value) * level_size
368                            / diameter_value;
369                        *box_coordinate_value = tmp.floor() as usize;
370                    },
371                )
372            },
373        );
374
375    Zip::from(keys.view_mut())
376        .and(box_coordinates.axis_iter(Axis(1)))
377        .par_apply(|key, box_coordinate| {
378            let anchor: [usize; 4] = [
379                box_coordinate[0],
380                box_coordinate[1],
381                box_coordinate[2],
382                level,
383            ];
384            *key = encode_anchor(&anchor);
385        });
386
387    keys
388}
389
390/// Helper function for decoding keys.
391fn decode_key_helper(key: usize, lookup_table: &[usize; 512]) -> usize {
392    const N_LOOPS: usize = 7; // 8 bytes in 64 bit key
393    let mut coord: usize = 0;
394
395    for index in 0..N_LOOPS {
396        coord |= lookup_table[(key >> (index * 9)) & NINE_BIT_MASK] << (3 * index);
397    }
398
399    coord
400}
401
402/// Decode a given key.
403///
404/// Returns an array containing the three coordinates and level of the key.
405pub fn decode_key(key: usize) -> [usize; 4] {
406    let level = find_level(key);
407    let key = key >> LEVEL_DISPLACEMENT;
408
409    let x = decode_key_helper(key, &X_LOOKUP_DECODE);
410    let y = decode_key_helper(key, &Y_LOOKUP_DECODE);
411    let z = decode_key_helper(key, &Z_LOOKUP_DECODE);
412
413    [x, y, z, level]
414}
415
416/// Return the key of the parent node.
417pub fn find_parent(key: usize) -> usize {
418    let level = find_level(key);
419    let key = key >> LEVEL_DISPLACEMENT;
420
421    let parent_level = level - 1;
422
423    ((key >> 3) << LEVEL_DISPLACEMENT) | parent_level
424}
425
426/// Return all children of a given key.
427pub fn find_children(key: usize) -> [usize; 8] {
428    let level = find_level(key);
429    let key = key >> LEVEL_DISPLACEMENT;
430
431    let mut children: [usize; 8] = [0; 8];
432
433    let root = (key >> 3) << 3;
434
435    for (index, item) in children.iter_mut().enumerate() {
436        *item = ((root | index) << LEVEL_DISPLACEMENT) | (level + 1);
437    }
438
439    children
440}
441
442/// Return all siblings of a key.
443///
444/// For a given key this function returns all 8 children
445/// of the parent of the key. Hence, the key itself is
446/// returned as well.
447pub fn find_siblings(key: usize) -> [usize; 8] {
448    let parent = find_parent(key);
449    find_children(parent)
450}
451
452/// Find key in a given direction.
453///
454/// Returns the key obtained by moving direction\[j\] boxes into direction j
455/// starting from the anchor associated with the given key.
456/// Negative steps are possible. If the result is out of bounds,
457/// i.e. anchor\[j\] + direction\[j\] is negative or larger than the number of boxes
458/// across each dimension, `None` is returned. Otherwise, `Some(new_key)` is returned,
459/// where `new_key` is the Morton key after moving into the given direction.
460///
461/// # Arguments
462/// `key` - The starting key.
463/// `direction` - A vector describing how many boxes we move along each coordinate direction.
464///               Negative values are possible (meaning that we move backwards).
465pub fn find_key_in_direction(key: usize, direction: &[i64; 3]) -> Option<usize> {
466    let anchor = decode_key(key);
467
468    let level = anchor[3];
469
470    let max_number_of_boxes: i64 = 1 << level;
471
472    let x: i64 = anchor[0] as i64;
473    let y: i64 = anchor[1] as i64;
474    let z: i64 = anchor[2] as i64;
475
476    let x = x + direction[0];
477    let y = y + direction[1];
478    let z = z + direction[2];
479
480    if (x >= 0)
481        & (y >= 0)
482        & (z >= 0)
483        & (x < max_number_of_boxes)
484        & (y < max_number_of_boxes)
485        & (z < max_number_of_boxes)
486    {
487        let new_anchor: [usize; 4] = [x as usize, y as usize, z as usize, level];
488        Some(encode_anchor(&new_anchor))
489    } else {
490        None
491    }
492}
493
494/// Compute near field.
495///
496/// The near field is the set of all boxes that are bordering the current box, including the box itself.
497///
498/// # Arguments
499/// `key` - The key for which we want to compute the neighbours.
500pub fn compute_near_field(key: usize) -> HashSet<usize> {
501    let mut near_field = HashSet::<usize>::new();
502
503    use itertools::iproduct;
504
505    for (i, j, k) in iproduct!(0..3, 0..3, 0..3) {
506        let direction: [i64; 3] = [i - 1, j - 1, k - 1];
507        if let Some(key) = find_key_in_direction(key, &direction) {
508            near_field.insert(key);
509        }
510    }
511    near_field
512}
513
514/// Compute interaction list.
515///
516/// The interaction list of a key consists of all the children of the near field of the
517/// parent that are not themselves in the near field of the key.
518/// The function returns a set of all keys that form the interaction list of the
519/// current key.
520pub fn compute_interaction_list(key: usize) -> HashSet<usize> {
521    let mut interaction_list = HashSet::<usize>::new();
522    let level = find_level(key);
523
524    if level < 2 {
525        // Levels zero and one always have empty interaction lists.
526        return interaction_list;
527    }
528
529    let near_field = compute_near_field(key);
530
531    let parent = find_parent(key);
532    let parent_near_field = compute_near_field(parent);
533
534    for &parent_neighbour in parent_near_field.iter() {
535        let children = find_children(parent_neighbour);
536        for &child in children.iter() {
537            if !near_field.contains(&child) {
538                interaction_list.insert(child);
539            }
540        }
541    }
542
543    interaction_list
544}
545
546#[cfg(test)]
547mod tests {
548    use super::*;
549
550    /// Test the encoding table for the x-coordinate.
551    #[test]
552    fn test_x_encode_table() {
553        for (mut index, actual) in X_LOOKUP_ENCODE.iter().enumerate() {
554            let mut sum: usize = 0;
555
556            for shift in 0..8 {
557                sum |= (index & 1) << (3 * shift);
558                index = index >> 1;
559            }
560
561            assert_eq!(sum, *actual);
562        }
563    }
564
565    /// Test the encoding table for the y-coordinate.
566    #[test]
567    fn test_y_encode_table() {
568        for (mut index, actual) in Y_LOOKUP_ENCODE.iter().enumerate() {
569            let mut sum: usize = 0;
570
571            for shift in 0..8 {
572                sum |= (index & 1) << (3 * shift + 1);
573                index = index >> 1;
574            }
575
576            assert_eq!(sum, *actual);
577        }
578    }
579
580    /// Test the encoding table for the z-coordinate.
581    #[test]
582    fn test_z_encode_table() {
583        for (mut index, actual) in Z_LOOKUP_ENCODE.iter().enumerate() {
584            let mut sum: usize = 0;
585
586            for shift in 0..8 {
587                sum |= (index & 1) << (3 * shift + 2);
588                index = index >> 1;
589            }
590
591            assert_eq!(sum, *actual);
592        }
593    }
594
595    /// Test the decoding table for the x-coordinate.
596    #[test]
597    fn test_x_decode_table() {
598        for (index, &actual) in X_LOOKUP_DECODE.iter().enumerate() {
599            let mut expected: usize = index & 1;
600            expected |= ((index >> 3) & 1) << 1;
601            expected |= ((index >> 6) & 1) << 2;
602
603            assert_eq!(actual, expected);
604        }
605    }
606
607    /// Test the decoding table for the y-coordinate.
608    #[test]
609    fn test_y_decode_table() {
610        for (index, &actual) in Y_LOOKUP_DECODE.iter().enumerate() {
611            let mut expected: usize = (index >> 1) & 1;
612            expected |= ((index >> 4) & 1) << 1;
613            expected |= ((index >> 7) & 1) << 2;
614
615            assert_eq!(actual, expected);
616        }
617    }
618
619    /// Test the decoding table for the z-coordinate.
620    #[test]
621    fn test_z_decode_table() {
622        for (index, &actual) in Z_LOOKUP_DECODE.iter().enumerate() {
623            let mut expected: usize = (index >> 2) & 1;
624            expected |= ((index >> 5) & 1) << 1;
625            expected |= ((index >> 8) & 1) << 2;
626
627            assert_eq!(actual, expected);
628        }
629    }
630
631    /// Test encoding and decoding an anchor
632    #[test]
633    fn test_encoding_decodiing() {
634        let anchor: [usize; 4] = [65535, 65535, 65535, 16];
635
636        let actual = decode_key(encode_anchor(&anchor));
637
638        assert_eq!(anchor, actual);
639    }
640
641    /// Test encoding many points
642    #[test]
643    fn test_encode_many_points() {
644        use rand::prelude::*;
645
646        const NPOINTS: usize = 100;
647        const LEVEL: usize = 4;
648
649        let mut rng = rand::thread_rng();
650
651        let mut points = Array2::<f64>::zeros((3, NPOINTS));
652
653        points.iter_mut().for_each(|item| *item = rng.gen::<f64>());
654
655        let origin = [-0.1, -0.2, 0.0];
656        let diameter = [1.2, 1.3, 1.01];
657        let keys = encode_points(points.view(), LEVEL, &origin, &diameter);
658
659        let level_size = (1 << LEVEL) as f64;
660
661        for (point, &key) in points.axis_iter(Axis(1)).zip(keys.iter()) {
662            let point_arr = [point[0], point[1], point[2]];
663            let single_key = encode_point(&point_arr, LEVEL, &origin, &diameter);
664
665            // Check that the key via encode_point is the same as the key via encode_points
666
667            assert_eq!(single_key, key);
668
669            // Check if box is close to the point.
670
671            let box_diameter = [
672                diameter[0] / level_size,
673                diameter[1] / level_size,
674                diameter[2] / level_size,
675            ];
676
677            let anchor = decode_key(single_key);
678            let coords = anchor_to_coordinates(&anchor, &origin, &diameter);
679            for dim in 0..3 {
680                assert!(coords[dim] <= point_arr[dim]);
681                assert!(point_arr[dim] < coords[dim] + box_diameter[dim]);
682            }
683        }
684    }
685}