Spaces:
Sleeping
Sleeping
| import Foundation | |
| import SwiftNetCDF | |
| struct Array2D { | |
| /// The underlying data storage for the 2D array. | |
| var data: [Float] | |
| /// The number of elements in x dimension | |
| let nx: Int | |
| /// The number of elements in y dimension | |
| let ny: Int | |
| /// The total number of elements in this array | |
| var count: Int { | |
| return nx * ny | |
| } | |
| init(data: [Float], nx: Int, ny: Int) { | |
| precondition(data.count == nx * ny) | |
| self.data = data | |
| self.nx = nx | |
| self.ny = ny | |
| } | |
| func writeNetcdf(filename: String) throws { | |
| let file = try NetCDF.create(path: filename, overwriteExisting: true) | |
| try file.setAttribute("TITLE", "My data set") | |
| let dimensions = [ | |
| try file.createDimension(name: "LAT", length: ny), | |
| try file.createDimension(name: "LON", length: nx) | |
| ] | |
| var variable = try file.createVariable(name: "data", type: Float.self, dimensions: dimensions) | |
| try variable.write(data) | |
| } | |
| mutating func shift180LongitudeAndFlipLatitude() { | |
| data.shift180LongitudeAndFlipLatitude(nt: 1, ny: ny, nx: nx) | |
| } | |
| mutating func flipEverySecondScanLine() { | |
| // flip every second line | |
| for y in stride(from: 1, to: ny, by: 2) { | |
| for x in 0 ..< nx / 2 { | |
| let temp = data[y * nx + x] | |
| data[y * nx + x] = data[y * nx + nx - x - 1] | |
| data[y * nx + nx - x - 1] = temp | |
| } | |
| } | |
| } | |
| mutating func shift180Longitudee() { | |
| data.shift180Longitude(nt: 1, ny: ny, nx: nx) | |
| } | |
| mutating func flipLatitude() { | |
| data.flipLatitude(nt: 1, ny: ny, nx: nx) | |
| } | |
| } | |
| /** | |
| `Array2DFastSpace` is a struct that represents a 2D array of Float values with fast space indexing. It allows accessing and modifying individual elements using subscript notation. | |
| Data is stored to be accessed quickly for multiple locatios in a row, while accessing timesteps is slower | |
| */ | |
| struct Array2DFastSpace { | |
| /// The underlying data storage for the 2D array. | |
| var data: [Float] | |
| /// The number of spatial locations in the array. | |
| let nLocations: Int | |
| /// The number of time steps in the array. | |
| let nTime: Int | |
| /** | |
| Initializes a new instance of `Array2DFastSpace`. | |
| - Parameters: | |
| - data: The data to be used as the underlying storage of the 2D array. Its count should be equal to `nLocations * nTime`. | |
| - nLocations: The number of spatial locations in the array. | |
| - nTime: The number of time steps in the array. | |
| - Precondition: `data.count` should be equal to `nLocations * nTime`, otherwise the initializer will fatalError. | |
| */ | |
| public init(data: [Float], nLocations: Int, nTime: Int) { | |
| if (data.count != nLocations * nTime) { | |
| fatalError("Wrong Array2DFastTime dimensions. nLocations=\(nLocations) nTime=\(nTime) count=\(data.count)") | |
| } | |
| self.data = data | |
| self.nLocations = nLocations | |
| self.nTime = nTime | |
| } | |
| /** | |
| Initializes a new instance of `Array2DFastSpace` with all elements set to `NaN`. | |
| - Parameters: | |
| - nLocations: The number of spatial locations in the array. | |
| - nTime: The number of time steps in the array. | |
| */ | |
| public init(nLocations: Int, nTime: Int) { | |
| self.data = .init(repeating: .nan, count: nLocations * nTime) | |
| self.nLocations = nLocations | |
| self.nTime = nTime | |
| } | |
| /** | |
| Writes the 2D array data to a NetCDF file. | |
| - Parameters: | |
| - filename: The name of the file to write to. | |
| - nx: The number of grid points in the x (longitude) direction. | |
| - ny: The number of grid points in the y (latitude) direction. | |
| - Throws: An error of type `NetCDFError` if the write operation fails. | |
| */ | |
| func writeNetcdf(filename: String, nx: Int, ny: Int) throws { | |
| let file = try NetCDF.create(path: filename, overwriteExisting: true) | |
| try file.setAttribute("TITLE", "My data set") | |
| let dimensions = [ | |
| try file.createDimension(name: "TIME", length: nTime), | |
| try file.createDimension(name: "LAT", length: ny), | |
| try file.createDimension(name: "LON", length: nx) | |
| ] | |
| var variable = try file.createVariable(name: "MyData", type: Float.self, dimensions: dimensions) | |
| try variable.write(data) | |
| } | |
| /// Accesses the element at the specified time and location in the 2D array. | |
| /// | |
| /// - Parameters: | |
| /// - time: The time index of the element to access. | |
| /// - location: The location index of the element to access. | |
| /// | |
| /// - Precondition: `location` must be less than `nLocations`. | |
| /// - Precondition: `time` must be less than `nTime`. | |
| /// | |
| /// - Returns: The element at the specified time and location in the 2D array. | |
| @inlinable subscript(time: Int, location: Int) -> Float { | |
| get { | |
| precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| return data[time * nLocations + location] | |
| } | |
| set { | |
| precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| data[time * nLocations + location] = newValue | |
| } | |
| } | |
| /// Accesses a range of values in the array for a specific time. | |
| /// | |
| /// Use this subscript to access a range of values from the `Array2DFastSpace` array | |
| /// for a specific time. The range of locations is specified as a `Range` object. | |
| /// | |
| /// - Parameters: | |
| /// - time: The time to access the range of values for. | |
| /// - location: The range of locations to access. | |
| /// | |
| /// - Returns: An array slice that contains the values of the specified range | |
| /// of locations for the specified time. | |
| /// | |
| /// - Precondition: `time` must be less than `nTime` and `location.upperBound` must | |
| /// be less than or equal to `nLocations`. | |
| /// | |
| /// - SeeAlso: `subscript(time: Int, location: Int) -> Float` | |
| @inlinable subscript(time: Int, location: Range<Int>) -> ArraySlice<Float> { | |
| get { | |
| precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| return data[location.add(time * nLocations)] | |
| } | |
| set { | |
| precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| data[location.add(time * nLocations)] = newValue | |
| } | |
| } | |
| /// Transpose to fast time | |
| func transpose() -> Array2DFastTime { | |
| precondition(data.count == nLocations * nTime) | |
| return data.withUnsafeBufferPointer { data in | |
| let out = [Float](unsafeUninitializedCapacity: data.count) { buffer, initializedCount in | |
| for l in 0..<nLocations { | |
| for t in 0..<nTime { | |
| buffer[l * nTime + t] = data[t * nLocations + l] | |
| } | |
| } | |
| initializedCount += data.count | |
| } | |
| return Array2DFastTime(data: out, nLocations: nLocations, nTime: nTime) | |
| } | |
| } | |
| } | |
| /** | |
| `Array2DFastTime` is a struct that represents a 2D array of Float values with fast time indexing. It allows accessing and modifying individual elements using subscript notation. | |
| Data is stored to be accessed quickly for multiple time temps in a row, while accessing locations is slower | |
| */ | |
| public struct Array2DFastTime { | |
| /// The underlying data storage for the 2D array. | |
| public var data: [Float] | |
| /// The number of spatial locations in the array. | |
| public let nLocations: Int | |
| /// The number of time steps in the array. | |
| public let nTime: Int | |
| /** | |
| Initializes a new instance of `Array2DFastTime`. | |
| - Parameters: | |
| - data: The data to be used as the underlying storage of the 2D array. Its count should be equal to `nLocations * nTime`. | |
| - nLocations: The number of spatial locations in the array. | |
| - nTime: The number of time steps in the array. | |
| - Precondition: `data.count` should be equal to `nLocations * nTime`, otherwise the initializer will fatalError. | |
| */ | |
| public init(data: [Float], nLocations: Int, nTime: Int) { | |
| if (data.count != nLocations * nTime) { | |
| fatalError("Wrong Array2DFastTime dimensions. nLocations=\(nLocations) nTime=\(nTime) count=\(data.count)") | |
| } | |
| self.data = data | |
| self.nLocations = nLocations | |
| self.nTime = nTime | |
| } | |
| /** | |
| Initializes a new instance of `Array2DFastTime` with all elements set to `NaN`. | |
| - Parameters: | |
| - nLocations: The number of spatial locations in the array. | |
| - nTime: The number of time steps in the array. | |
| */ | |
| public init(nLocations: Int, nTime: Int) { | |
| self.data = .init(repeating: .nan, count: nLocations * nTime) | |
| self.nLocations = nLocations | |
| self.nTime = nTime | |
| } | |
| /// Accesses the element at the specified time and location in the 2D array. | |
| /// | |
| /// - Parameters: | |
| /// - location: The location index of the element to access. | |
| /// - time: The time index of the element to access. | |
| /// | |
| /// - Precondition: `location` must be less than `nLocations`. | |
| /// - Precondition: `time` must be less than `nTime`. | |
| /// | |
| /// - Returns: The element at the specified time and location in the 2D array. | |
| @inlinable subscript(location: Int, time: Int) -> Float { | |
| get { | |
| precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| return data[location * nTime + time] | |
| } | |
| set { | |
| precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time < nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| data[location * nTime + time] = newValue | |
| } | |
| } | |
| /// Accesses a range of values in the array for a specific location. | |
| /// | |
| /// Use this subscript to access a range of values from the `Array2DFastTime` array | |
| /// for a specific location. The range of time steps is specified as a `Range` object. | |
| /// | |
| /// - Parameters: | |
| /// - location: The location to access the range of values for. | |
| /// - time: The range of time steps to access. | |
| /// | |
| /// - Returns: An array slice that contains the values of the specified range | |
| /// of timesteps for the specified location. | |
| /// | |
| /// - Precondition: `location` must be less than `nLocations` and `time.upperBound` must | |
| /// be less than or equal to `nTime`. | |
| /// | |
| /// - SeeAlso: `subscript(time: Int, location: Int) -> Float` | |
| @inlinable subscript(location: Int, time: Range<Int>) -> ArraySlice<Float> { | |
| get { | |
| precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time.upperBound <= nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| return data[time.add(location * nTime)] | |
| } | |
| set { | |
| precondition(location < nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| precondition(time.upperBound <= nTime, "time subscript invalid: \(nTime) with nTime=\(nTime)") | |
| data[time.add(location * nTime)] = newValue | |
| } | |
| } | |
| /// Accesses a range of values in the array for a specific time. This function is relatively slow, because data needs to be transposed. | |
| /// | |
| /// Use this subscript to access a range of values from the `Array2DFastTime` array | |
| /// for a specific time. The range of locations is specified as a `Range` object. | |
| /// | |
| /// - Parameters: | |
| /// - time: The time to access the range of values for. | |
| /// - location: The range of locations to access. | |
| /// | |
| /// - Returns: An array slice that contains the values of the specified range | |
| /// of locations for the specified time. | |
| /// | |
| /// - Precondition: `time` must be less than `nTime` and `location.upperBound` must | |
| /// be less than or equal to `nLocations`. | |
| /// | |
| /// - SeeAlso: `subscript(time: Int, location: Int) -> Float` | |
| @inlinable subscript(location: Range<Int>, time: Int) -> [Float] { | |
| get { | |
| precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| 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, time]) | |
| } | |
| return out | |
| } | |
| set { | |
| precondition(location.upperBound <= nLocations, "location subscript invalid: \(location) with nLocations=\(nLocations)") | |
| 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 + time] = value | |
| } | |
| } | |
| } | |
| /// Transpose to fast space | |
| func transpose() -> Array2DFastSpace { | |
| precondition(data.count == nLocations * nTime) | |
| return data.withUnsafeBufferPointer { data in | |
| let out = [Float](unsafeUninitializedCapacity: data.count) { buffer, initializedCount in | |
| for t in 0..<nTime { | |
| for l in 0..<nLocations { | |
| buffer[t * nLocations + l] = data[l * nTime + t] | |
| } | |
| } | |
| initializedCount += data.count | |
| } | |
| return Array2DFastSpace(data: out, nLocations: nLocations, nTime: nTime) | |
| } | |
| } | |
| } | |
| /*extension Array { | |
| /// Calculate start positions for cycles | |
| static func transposeCalculateCycles(rows: Int, cols: Int) -> [Int] { | |
| var cycles = [Int]() | |
| let size = rows * cols - 1 | |
| var b = [Bool](repeating: false, count: rows*cols) | |
| b[0] = true | |
| b[size] = true | |
| var i = 1 | |
| while i < size { | |
| let cycleBegin = i | |
| cycles.append(i) | |
| repeat { | |
| b[i] = true | |
| i = (i*rows)%size | |
| } while i != cycleBegin | |
| i = 1 | |
| while i < size && b[i] == true { | |
| i += 1 | |
| } | |
| } | |
| return cycles | |
| } | |
| /// Perform inplace transposition using `following cycles algorithm`. This is 10 slower than double buffer transpose, exclusing the cycle calculation | |
| /// See: https://en.wikipedia.org/wiki/In-place_matrix_transposition | |
| /// See: https://www.geeksforgeeks.org/inplace-m-x-n-size-matrix-transpose/ | |
| mutating func transpose(rows: Int, cols: Int, cycles: [Int]? = nil) { | |
| precondition(count == rows * cols) | |
| let cycles = cycles ?? Self.transposeCalculateCycles(rows: rows, cols: cols) | |
| self.withUnsafeMutableBufferPointer { ptr in | |
| let size = rows * cols - 1 | |
| for cycleBegin in cycles { | |
| var i = cycleBegin | |
| var t = ptr[i] | |
| repeat { | |
| i = (i*rows)%size | |
| swap(&ptr[i], &t) | |
| } while i != cycleBegin | |
| } | |
| } | |
| } | |
| }*/ | |