From c69bf3903dcf61d292823e5da82b78b6e7086d71 Mon Sep 17 00:00:00 2001 From: Boris Gromov <0xff.root@gmail.com> Date: Sun, 24 May 2020 12:03:10 +0200 Subject: [PATCH] Add range-based Matrix subscripts (#11) --- Example/Tests/MatrixTests.swift | 54 +++++++++++++- Sources/Matrix.swift | 123 ++++++++++++++++++++++++++++++++ 2 files changed, 176 insertions(+), 1 deletion(-) diff --git a/Example/Tests/MatrixTests.swift b/Example/Tests/MatrixTests.swift index d9301ce..87feb55 100644 --- a/Example/Tests/MatrixTests.swift +++ b/Example/Tests/MatrixTests.swift @@ -322,7 +322,59 @@ class MatrixSpec: QuickSpec { expect(m1[1, 0]) == 30.0 } } - + + describe("Matrix range-based subscript") { + it("[i,j] -> Matrix") { + let m1 = Matrix([[1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0]]) + let m2 = Matrix([[10.0, 11.0], + [12.0, 13.0]]) + expect(m1[0, 1].flat) == [2.0, 3.0, 5.0, 6.0, 8.0, 9.0] + m1[1, 0] = m2 + expect(m1[1..<3, 0...1].flat) == [10.0, 11.0, 12.0, 13.0] + } + it("closed") { + let m1 = Matrix([[1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0]]) + let m2 = Matrix([[10.0, 11.0], + [12.0, 13.0]]) + expect(m1[1...2, 0...1].flat) == [4.0, 5.0, 7.0, 8.0] + m1[0...1, 1...2] = m2 + expect(m1.flat) == [1.0, 10.0, 11.0, 4.0, 12.0, 13.0, 7.0, 8.0, 9.0] + } + it("open") { + let m1 = Matrix([[1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0]]) + let m2 = Matrix([[10.0, 11.0], + [12.0, 13.0]]) + expect(m1[1..<3, 0..<2].flat) == [4.0, 5.0, 7.0, 8.0] + m1[0..<2, 1..<3] = m2 + expect(m1.flat) == [1.0, 10.0, 11.0, 4.0, 12.0, 13.0, 7.0, 8.0, 9.0] + } + it("partial") { + let m1 = Matrix([[1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0]]) + let m2 = Matrix([[10.0, 11.0], + [12.0, 13.0]]) + expect(m1[1..., ..<2].flat) == [4.0, 5.0, 7.0, 8.0] + m1[1..., 1...] = m2 + expect(m1.flat) == [1.0, 2.0, 3.0, 4.0, 10.0, 11.0, 7.0, 12.0, 13.0] + } + it("unbounded") { + let m1 = Matrix([[1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0]]) + let m2 = Matrix([[10.0, 11.0, 12.0]]) + expect(m1[..., 2...2].flat) == [3.0, 6.0, 9.0] + m1[1...1, ...] = m2 + expect(m1.flat) == [1.0, 2.0, 3.0, 10.0, 11.0, 12.0, 7.0, 8.0, 9.0] + } + } + describe("Matrix map/reduce") { it("map") { let m1 = Matrix([[1.0, 2.0], [3.0, 4.0]]) diff --git a/Sources/Matrix.swift b/Sources/Matrix.swift index a781331..d57243f 100644 --- a/Sources/Matrix.swift +++ b/Sources/Matrix.swift @@ -273,6 +273,129 @@ extension Matrix { } } } + + /// Get and set M(row, col) submatrix of Matrix. + /// + /// The range-based subscript methods for getting and setting submatricies. + /// + /// var M = Matrix([[1, 2, 3, 4], + /// [5, 6, 7, 8], + /// [9, 10, 11, 12]) + /// + /// var K = Matrix([[1, 0], + /// [0, 1]) + /// + /// - Using bounded ranges, including partial (e.g. `..<3`): + /// + /// M[1..<3, 0..1] = K + /// // M is now: + /// // [[1, 2, 3, 4], + /// // [1, 0, 7, 8], + /// // [0, 1, 11, 12]] + /// + /// K = M[0...1, 2...] + /// // K is now: + /// // [[3, 4] + /// // [7, 8]] + /// + /// - Using unbounded ranges: + /// + /// K = M[..., 1..2] + /// // K is now: + /// // [[ 2, 3], + /// // [ 6, 7], + /// // [10, 11]] + /// + /// - Parameters: + /// - row: Range for rows (0-based) + /// - col: Range for cols (0-based) + /// + /// - Returns: submatrix of size `row.count` by `col.count` + /// + public subscript(_ row: A, _ col: B) -> Matrix where A.Bound == Int, B.Bound == Int { + get { + return self[ClosedRange(row.relative(to: self[row: 0])), ClosedRange(col.relative(to: self[col: 0]))] + } + + set { + self[ClosedRange(row.relative(to: self[row: 0])), ClosedRange(col.relative(to: self[col: 0]))] = newValue + } + } + + public subscript(_ : UnboundedRange, _ col: B) -> Matrix where B.Bound == Int { + get { return self[0..(col.relative(to: self[col: 0]))] } + set { self[0..(col.relative(to: self[col: 0]))] = newValue } + } + + public subscript(_ row: A, _ : UnboundedRange) -> Matrix where A.Bound == Int { + get { return self[ClosedRange(row.relative(to: self[row: 0])), 0..(row.relative(to: self[row: 0])), 0.. Matrix { + get { return self} + set { self[0.. Matrix { + get { return self[row..., col...]} + set { self[row..<(row + newValue.rows), col..<(col + newValue.cols)] = newValue} + } + + public subscript(_ row: ClosedRange, _ col: ClosedRange) -> Matrix { + get { + precondition(indexIsValidForRow(row.lowerBound, col.lowerBound), "Invalid range") + precondition(indexIsValidForRow(row.upperBound, col.upperBound), "Invalid range") + + let dst = Matrix(row.count, col.count) + + flat.withUnsafeBufferPointer { srcBuf in + let srcPtr = srcBuf.baseAddress! + row.lowerBound * rows + col.lowerBound + vDSP_mmovD(srcPtr, &dst.flat, + vDSP_Length(col.count), vDSP_Length(row.count), + vDSP_Length(cols), vDSP_Length(col.count)) + } + + return dst + } + + set { + precondition(indexIsValidForRow(row.lowerBound, col.lowerBound), "Invalid range") + precondition(indexIsValidForRow(row.upperBound, col.upperBound), "Invalid range") + precondition(newValue.cols == col.count && newValue.rows == row.count, "Matrix dimensions must agree") + + flat.withUnsafeMutableBufferPointer { dstBuf in + let dstPtr = dstBuf.baseAddress! + row.lowerBound * rows + col.lowerBound + vDSP_mmovD(newValue.flat, dstPtr, + vDSP_Length(col.count), vDSP_Length(row.count), + vDSP_Length(newValue.cols), vDSP_Length(cols)) + } + } + } /// Construct new matrix from source using specified extractor. ///