Multi_NNGP/util2.j at master · LuZhangstat/Multi_NNGP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
# build nearest neighbor for observed location set S #
function BuildNN(coords, m, nThreads = 1.0, brute = 0.0)
    # coords: n by 2 array
    # m:      number of neighbors
    n = size(coords)[1]
    # coords_fit = [coords[1, :] coords[2, :]]
    nIndx = trunc(Cint,(1 + m) / 2 * m + (n - m - 1) * m);
    nnIndx = zeros(Cint, nIndx);
    nnDist = zeros(Float64, nIndx);
    ##first of nnIndxLU column holds the nnIndx index for the i-th location and the second columns holds the number of neighbors the i-th location has (the second column is a bit of a waste but simplifies my life in the spNNGP)
    nnIndxLU = zeros(Cint, 2 * n);
    ##Note I need to convert n and m to pointers only because the R example of .C wants only pointers so that's how I worte the mkNNIndx and mkNNIndxTree0
    nPtr = Cint[n];
    mPtr = Cint[m];
    nThreadsPtr = Cint[nThreads];

    if brute == 1.0
        ccall((:mkNNIndx, "../../../julia-R-nn-ccall2/nn.so"), Cvoid, (Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble},
            Ptr{Cint}, Ptr{Cdouble}, Ptr{Cint}, Ptr{Cint}), nPtr, mPtr, coords, nnIndx, nnDist, nnIndxLU, nThreadsPtr)
    else
        ccall((:mkNNIndxCB, "../../../julia-R-nn-ccall2/nn.so"), Cvoid, (Ptr{Cint}, Ptr{Cint}, Ptr{Cdouble},
            Ptr{Cint}, Ptr{Cdouble}, Ptr{Cint}, Ptr{Cint}), nPtr, mPtr, coords, nnIndx, nnDist, nnIndxLU, nThreadsPtr)
    end

    nnIndx = nnIndx .+ 1
    nnIndxLU = nnIndxLU[1:(n + 1)]  .+ 1
    nnIndxLU[n + 1] = nIndx + 1
    return (nnIndx = nnIndx, nnDist = nnDist, nnIndxLU = nnIndxLU[1:(n + 1)])
end

function getAD(coords, nnIndx, nnDist, nnIndxLU, ϕ, ν, A, D)
    # coords: n by 2 array,
    # nnIndx, nnDist, nnIndxLU output of the mkNNIndx function

    # return: A D
    # A = []
    # D = ones(size(coords)[1])

    # initialize memory
    NNdistM = []
    NNdistv = []
    n = length(nnIndxLU) - 1;
    if nnIndxLU[1] == nnIndxLU[2]
        D[1] = 1.0
        l = 2
    else
        l = 1
    end
    for i in l:n
        if ν == 0.5
            NNdistM = cholesky(exp.(-ϕ .* pairwise(Euclidean(),
                        coords[nnIndx[nnIndxLU[i]:(nnIndxLU[i + 1] - 1)], :],
                        dims = 1)))
            NNdistv = NNdistM.L \ (exp.(-ϕ .* nnDist[nnIndxLU[i]:(nnIndxLU[i + 1] - 1)]))
            D[i] = 1.0 - NNdistv⋅NNdistv
            A[nnIndxLU[i]:(nnIndxLU[i + 1] - 1)] = NNdistM.U \ NNdistv
        end
    end
    # return (A = A, D = D)
end

function getAD_collapse(coords, nnIndx, nnDist, nnIndxLU, ϕ, ν, α, A, D)
    # coords: n by 2 array,
    # nnIndx, nnDist, nnIndxLU output of the mkNNIndx function
    # ϕ, ν, α: covariance parameter set
    # hold: whether AD is for holded locations or not

    # return: A D
    # A = []
    # D = ones(size(coords)[1])

    # initialize memory
    NNdistM = []
    NNdistv = []
    n = length(nnIndxLU) - 1;
    if nnIndxLU[1] == nnIndxLU[2]
        D[1] = 1.0 / α
        l = 2
    else
        l = 1
    end
    for i in l:n
        if ν == 0.5
            NNdistM = cholesky(exp.(-ϕ .* pairwise(Euclidean(),
                        coords[nnIndx[nnIndxLU[i]:(nnIndxLU[i + 1] - 1)], :],
                        dims = 1)) + (1.0 / α - 1.0) * I)
            NNdistv = NNdistM.L \ (exp.(-ϕ .* nnDist[nnIndxLU[i]:(nnIndxLU[i + 1] - 1)]))
            D[i] = 1.0 / α - NNdistv⋅NNdistv
            A[nnIndxLU[i]:(nnIndxLU[i + 1] - 1)] = NNdistM.U \ NNdistv
        end
    end

    # return (A = A, D = D)
end


function kfoldperm(N,k)

    # k folder split of 1:N

    n,r = divrem(N,k)
    b = collect(1:n:N+1)
    for i in 1:length(b)
        b[i] += i > r ? r : i-1
    end
    p = randperm(N)
    return [p[r] for r in [b[i]:b[i+1]-1 for i=1:k]]
end

colnorm(A) = [norm(A[:,i], 2) for i=1:size(A,2)]