File size: 8,196 Bytes
6ee917b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
import Foundation


struct Array3D {
    var data: [Float]
    
    /// slowest
    let dim0: Int
    let dim1: Int
    /// Fastest dim
    let dim2: Int
    
    var count: Int {
        return dim0 * dim1 * dim2
    }
    
    public init(data: [Float], dim0: Int, dim1: Int, dim2: Int) {
        if (data.count != dim0 * dim1 * dim2) {
            fatalError("Wrong Array3D dimensions. dim0=\(dim0) dim1=\(dim1) dim2=\(dim2) count=\(data.count)")
        }
        self.data = data
        self.dim0 = dim0
        self.dim1 = dim1
        self.dim2 = dim2
    }
    
    public init(repeating: Float, dim0: Int, dim1: Int, dim2: Int) {
        self.data = [Float](repeating: repeating, count: dim0 * dim1 * dim2)
        self.dim0 = dim0
        self.dim1 = dim1
        self.dim2 = dim2
    }
    
    @inlinable subscript(d0: Int, d1: Int, d2: Int) -> Float {
        get {
            assert(d0 < dim0, "dim0 subscript invalid: \(d0) with dim0=\(dim0)")
            assert(d1 < dim1, "dim1 subscript invalid: \(d1) with dim1=\(dim1)")
            assert(d2 < dim2, "dim2 subscript invalid: \(d2) with dim2=\(dim2)")
            return data[d0 * dim1 * dim2 + d1 * dim2 + d2]
        }
        set {
            assert(d0 < dim0, "dim0 subscript invalid: \(d0) with dim0=\(dim0)")
            assert(d1 < dim1, "dim1 subscript invalid: \(d1) with dim1=\(dim1)")
            assert(d2 < dim2, "dim2 subscript invalid: \(d2) with dim2=\(dim2)")
            data[d0 * dim1 * dim2 + d1 * dim2 + d2] = newValue
        }
    }
}

public struct Array3DFastTime {
    public var data: [Float]
    public let nLocations: Int
    public let nLevel: Int
    public let nTime: Int
    
    public init(data: [Float], nLocations: Int, nLevel: Int, nTime: Int) {
        if (data.count != nLocations * nTime * nLevel) {
            fatalError("Wrong Array2DFastTime dimensions. nLocations=\(nLocations) nLevel=\(nLevel) nTime=\(nTime) count=\(data.count)")
        }
        self.data = data
        self.nLocations = nLocations
        self.nLevel = nLevel
        self.nTime = nTime
    }
    
    public init(nLocations: Int, nLevel: Int, nTime: Int) {
        self.data = .init(repeating: .nan, count: nLocations * nTime * nLevel)
        self.nLocations = nLocations
        self.nTime = nTime
        self.nLevel = nLevel
    }
    
    @inlinable subscript(location: Int, level: Int, time: Int) -> Float {
        get {
            precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time < nTime, "time subscript invalid: \(time) with nTime=\(nTime)")
            return data[location * nTime * nLevel + level * nTime + time]
        }
        set {
            precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time < nTime, "time subscript invalid: \(time) with nTime=\(nTime)")
            data[location * nTime * nLevel + level * nTime + time] = newValue
        }
    }
    
    @inlinable subscript(location: Int, level: Int, time: Range<Int>) -> ArraySlice<Float> {
        get {
            precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time.upperBound <= nTime, "time subscript invalid: \(time) with nTime=\(nTime)")
            return data[time.add(location * nTime * nLevel + level * nTime)]
        }
        set {
            precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time.upperBound <= nTime, "time subscript invalid: \(time) with nTime=\(nTime)")
            data[time.add(location * nTime * nLevel + level * nTime)] = newValue
        }
    }
    
    /// One spatial field into time-series array
    @inlinable subscript(location: Range<Int>, level: Int, time: Int) -> [Float] {
        get {
            precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)")
            var out = [Float]()
            out.reserveCapacity(location.count)
            for loc in location {
                out.append(self[loc, level, time])
            }
            return out
        }
        set {
            precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)")
            precondition(newValue.count == location.count, "Array and location count do not match")
            for (loc, value) in zip(location, newValue) {
                data[loc * nTime * nLevel + level * nTime + time] = value
            }
        }
    }
    
    @inlinable subscript(location: Range<Int>, level: Int, time: Range<Int>) -> ArraySlice<Float> {
        get {
            precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time.upperBound <= nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)")
            var out = [Float]()
            out.reserveCapacity(location.count * time.count)
            for loc in location {
                out.append(contentsOf: self[loc, level, time])
            }
            return ArraySlice(out)
        }
        set {
            precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time.upperBound <= nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)")
            precondition(newValue.count == location.count, "Array and location count do not match")
            for loc in location {
                data[time.add(loc * nTime * nLevel + level * nTime)] = newValue[(time.count * loc ..< time.count * (loc+1)).add(newValue.startIndex)]
            }
        }
    }
    
    /// One spatial field into time-series array
    @inlinable subscript(location: Range<Int>, level: Int, time: Int) -> ArraySlice<Float> {
        get {
            precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)")
            var out = [Float]()
            out.reserveCapacity(location.count)
            for loc in location {
                out.append(self[loc, level, time])
            }
            return ArraySlice(out)
        }
        set {
            precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)")
            precondition(level < nLevel, "level subscript invalid: \(level) with nLevel=\(nLevel)")
            precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)")
            precondition(newValue.count == location.count, "Array and location count do not match")
            for (loc, value) in zip(location, newValue) {
                data[loc * nTime * nLevel + level * nTime + time] = value
            }
        }
    }
}