- Dropped the Comparable requirement down in to a comparator function on Interpolator. This allows us to Heatmap sequences of arbitrary types (for example, using a KeyPath to an int or float to grade them).

karwa · karwa · commit 3cdffbb7e4c4 · 2020-02-03T04:32:08.000+01:00
- Moved Interpolator to its own file. It should possibly be given a new name, too.
- Added convenience functions for generating a Heatmap from arbitrary RACs or 2D Sequences. This allows for some really cool things.
diff --git a/Sources/SwiftPlot/Heatmap.swift b/Sources/SwiftPlot/Heatmap.swift
@@ -1,11 +1,23 @@
 
-public struct Heatmap<SeriesType>
-where SeriesType: Sequence, SeriesType.Element: Sequence,
-      SeriesType.Element.Element: Comparable {
+// Todo list for Heatmap:
+// - Colormaps
+// - Shift grid to block bounds
+// - Draw grid over blocks?
+// - Spacing between blocks
+// - Setting X/Y axis labels
+// - Displaying colormap next to plot
+// - Collection slicing by filter closure
+
+/// A heatmap is a plot of 2-dimensional data, where each value is assigned a colour value along a gradient.
+///
+/// Use the `interpolator` property to control how values are graded. For example, if your data structure has
+/// a salient integer or floating-point property, `Interpolator.linearByKeyPath` will allow you to grade values by that property.
+public struct Heatmap<SeriesType> where SeriesType: Sequence, SeriesType.Element: Sequence {
   
   public typealias Element = SeriesType.Element.Element
 
   public var layout = GraphLayout()
+  
   public var values: SeriesType
   public var interpolator: Interpolator<Element>
   
@@ -17,50 +29,13 @@ where SeriesType: Sequence, SeriesType.Element: Sequence,
   }
 }
 
-// Initialisers with default arguments.
-
-extension Heatmap
-  where SeriesType: ExpressibleByArrayLiteral, SeriesType.Element: ExpressibleByArrayLiteral,
-        SeriesType.ArrayLiteralElement == SeriesType.Element {
-  
-  public init(interpolator: Interpolator<Element>) {
-    self.init(values: [[]], interpolator: interpolator)
-  }
-}
-
-extension Heatmap
-  where SeriesType: ExpressibleByArrayLiteral, SeriesType.Element: ExpressibleByArrayLiteral,
-        SeriesType.ArrayLiteralElement == SeriesType.Element, Element: FloatConvertible {
-  
-  public init(values: SeriesType) {
-    self.init(values: values, interpolator: .linear)
-  }
-  
-  public init() {
-    self.init(interpolator: .linear)
-  }
-}
-
-extension Heatmap
-  where SeriesType: ExpressibleByArrayLiteral, SeriesType.Element: ExpressibleByArrayLiteral,
-        SeriesType.ArrayLiteralElement == SeriesType.Element, Element: FixedWidthInteger {
-  
-  public init(values: SeriesType) {
-    self.init(values: values, interpolator: .linear)
-  }
-  
-  public init() {
-    self.init(interpolator: .linear)
-  }
-}
-
 // Layout and drawing.
 
 extension Heatmap: HasGraphLayout, Plot {
   
   public struct DrawingData {
     var values: SeriesType?
-    var range: ClosedRange<Element>?
+    var range: (min: Element, max: Element)?
     var itemSize = Size.zero
     var rows = 0
     var columns = 0
@@ -70,7 +45,6 @@ extension Heatmap: HasGraphLayout, Plot {
     
     var results = DrawingData()
     var markers = PlotMarkers()
-    
     // Extract the first (inner) element as a starting point.
     guard let firstElem = values.first(where: { _ in true })?.first(where: { _ in true }) else {
       return (results, nil)
@@ -83,16 +57,16 @@ extension Heatmap: HasGraphLayout, Plot {
     for row in values {
       var columnsInRow = 0
       for column in row {
-        maxValue = max(maxValue, column)
-        minValue = min(minValue, column)
+        maxValue = interpolator.compare(maxValue, column) ? column : maxValue
+        minValue = interpolator.compare(minValue, column) ? minValue : column
         columnsInRow += 1
       }
       maxColumns = max(maxColumns, columnsInRow)
       totalRows += 1
     }
     // Update results.
     results.values = values
-    results.range = minValue...maxValue
+    results.range = (minValue, maxValue)
     results.rows = totalRows
     results.columns = maxColumns
     results.itemSize = Size(
@@ -116,20 +90,18 @@ extension Heatmap: HasGraphLayout, Plot {
   }
   
   public func drawData(_ data: DrawingData, size: Size, renderer: Renderer) {
-    
-    
     guard let values = data.values, let range = data.range else { return }
     
     for (rowIdx, row) in values.enumerated() {
-      for (columnIdx, column) in row.enumerated() {
+      for (columnIdx, element) in row.enumerated() {
         let rect = Rect(
           origin: Point(Float(columnIdx) * data.itemSize.width,
                         Float(rowIdx) * data.itemSize.height),
-          size: data.itemSize)
-        renderer.drawSolidRect(rect,
-                               fillColor: getColor(of: column, range: range),
-                               hatchPattern: .none)
-//        renderer.drawText(text: String(describing: column),
+          size: data.itemSize
+        )
+        let color = getColor(of: element, min: range.min, max: range.max)
+        renderer.drawSolidRect(rect, fillColor: color, hatchPattern: .none)
+//        renderer.drawText(text: String(describing: element),
 //                          location: rect.origin + Point(50,50),
 //                          textSize: 20,
 //                          color: .white,
@@ -139,70 +111,82 @@ extension Heatmap: HasGraphLayout, Plot {
     }
   }
   
-  func getColor(of value: Element, range: ClosedRange<Element>) -> Color {
+  private func getColor(of value: Element, min: Element, max: Element) -> Color {
     let startColor = Color.orange
     let endColor = Color.purple
-    let interp = interpolator.callAsFunction(value, in: range)
-    
-    return lerp(startColor: startColor, endColor: endColor, interp)
+    let offset = interpolator.interpolate(value, min, max)
+    return lerp(startColor: startColor, endColor: endColor, offset)
   }
 }
 
+// Initialisers with default arguments.
 
-// Interpolator.
+extension Heatmap
+  where SeriesType: ExpressibleByArrayLiteral, SeriesType.Element: ExpressibleByArrayLiteral,
+        SeriesType.ArrayLiteralElement == SeriesType.Element {
+  
+  public init(interpolator: Interpolator<Element>) {
+    self.init(values: [[]], interpolator: interpolator)
+  }
+}
 
-public struct Interpolator<Element> where Element: Comparable {
-  public var interpolate: (Element, ClosedRange<Element>) -> Float
+extension Heatmap
+  where SeriesType: ExpressibleByArrayLiteral, SeriesType.Element: ExpressibleByArrayLiteral,
+        SeriesType.ArrayLiteralElement == SeriesType.Element, Element: FloatConvertible {
   
-  public init(_ block: @escaping (Element, ClosedRange<Element>)->Float) {
-    self.interpolate = block
+  public init(values: SeriesType) {
+    self.init(values: values, interpolator: .linear)
   }
-  public func callAsFunction(_ item: Element, in range: ClosedRange<Element>) -> Float {
-    interpolate(item, range)
+  
+  public init() {
+    self.init(interpolator: .linear)
   }
 }
 
-extension Interpolator where Element: FloatConvertible {
-  public static var linear: Interpolator {
-    Interpolator { value, range in
-      let value = Float(value)
-      let range = Float(range.lowerBound)...Float(range.upperBound)
-      let totalDistance = range.lowerBound.distance(to: range.upperBound)
-      let valueOffset   = range.lowerBound.distance(to: value)
-      return valueOffset/totalDistance
-    }
+extension Heatmap
+  where SeriesType: ExpressibleByArrayLiteral, SeriesType.Element: ExpressibleByArrayLiteral,
+        SeriesType.ArrayLiteralElement == SeriesType.Element, Element: FixedWidthInteger {
+  
+  public init(values: SeriesType) {
+    self.init(values: values, interpolator: .linear)
   }
-}
-extension Interpolator where Element: FixedWidthInteger {
-  public static var linear: Interpolator {
-    Interpolator { value, range in
-      let distance = range.lowerBound.distance(to: range.upperBound)
-      let valDist = range.lowerBound.distance(to: value)
-      return Float(valDist)/Float(distance)
-    }
+  
+  public init() {
+    self.init(interpolator: .linear)
   }
 }
 
-extension Interpolator {
+// Collection construction shorthand.
+
+extension Sequence where Element: Sequence {
   
-  public static func linearByKeyPath<T>(_ kp: KeyPath<Element, T>) -> Interpolator<Element>
-    where T: FloatConvertible {
-      let i = Interpolator<T>.linear
-      return Interpolator { value, range in
-        let value = value[keyPath: kp]
-        let range = range.lowerBound[keyPath: kp]...range.upperBound[keyPath: kp]
-        return i.interpolate(value, range)
-      }
+  /// Returns a heatmap of values from this 2-dimensional sequence.
+  /// - parameters:
+  ///   - interpolator: A function or `KeyPath` which maps values to a continuum between 0 and 1.
+  /// - returns: A heatmap plot of the sequence's inner items.
+  public func heatmap(interpolator: Interpolator<Element.Element>) -> Heatmap<Self> {
+    return Heatmap(values: self, interpolator: interpolator)
   }
+}
+
+extension RandomAccessCollection {
   
-  public static func linearByKeyPath<T>(_ kp: KeyPath<Element, T>) -> Interpolator<Element>
-    where T: FixedWidthInteger {
-      let i = Interpolator<T>.linear
-      return Interpolator { value, range in
-        let value = value[keyPath: kp]
-        let range = range.lowerBound[keyPath: kp]...range.upperBound[keyPath: kp]
-        return i.interpolate(value, range)
+  /// Returns a heatmap of this collection's values, generated by slicing rows with the given width.
+  /// - parameters:
+  ///   - width:        The width of the heatmap to generate. Must be greater than 0.
+  ///   - interpolator: A function or `KeyPath` which maps values to a continuum between 0 and 1.
+  /// - returns: A heatmap plot of the collection's values.
+  public func heatmap(width: Int, interpolator: Interpolator<Element>) -> Heatmap<[SubSequence]> {
+    precondition(width > 0, "Cannot build a histogram with zero or negative width")
+    let height = Int((Float(count) / Float(width)).rounded(.up))
+    return (0..<height).map { row -> SubSequence in
+      guard let start = index(startIndex, offsetBy: row * width, limitedBy: endIndex) else {
+        return self[startIndex..<startIndex]
       }
+      guard let end = index(start, offsetBy: width, limitedBy: endIndex) else {
+        return self[start..<endIndex]
+      }
+      return self[start..<end]
+    }.heatmap(interpolator: interpolator)
   }
-  
 }
diff --git a/Sources/SwiftPlot/Interpolator.swift b/Sources/SwiftPlot/Interpolator.swift
@@ -0,0 +1,79 @@
+
+/// An `Interpolator` maps values to a continuum between 0 and 1
+public struct Interpolator<Element> {
+  public var compare: (Element, Element) -> Bool
+  public var interpolate: (Element, Element, Element) -> Float
+  
+  public init(
+    compare areInIncreasingOrder: @escaping (Element, Element) -> Bool,
+    interpolate: @escaping (Element, Element, Element)->Float
+  ) {
+    self.compare = areInIncreasingOrder
+    self.interpolate = interpolate
+  }
+}
+
+extension Interpolator {
+  
+  public var inverted: Interpolator<Element> {
+    Interpolator(
+      compare: { self.compare($0, $1) },
+      interpolate: { 1 - self.interpolate($0, $1, $2) }
+    )
+  }
+}
+
+extension Interpolator where Element: Comparable {
+  public init(interpolate: @escaping (Element, Element, Element)->Float) {
+    self.init(compare: <, interpolate: interpolate)
+  }
+}
+
+// Linear mapping for numeric types.
+
+extension Interpolator where Element: FloatConvertible {
+  public static var linear: Interpolator {
+    Interpolator { value, min, max in
+      let value = Float(value)
+      let range = Float(min)...Float(max)
+      let totalDistance = range.lowerBound.distance(to: range.upperBound)
+      let valueOffset   = range.lowerBound.distance(to: value)
+      return valueOffset/totalDistance
+    }
+  }
+}
+extension Interpolator where Element: FixedWidthInteger {
+  public static var linear: Interpolator {
+    Interpolator { value, min, max in
+      let distance = min.distance(to: max)
+      let valDist = min.distance(to: value)
+      return Float(valDist)/Float(distance)
+    }
+  }
+}
+
+// Mapping by key-paths.
+
+extension Interpolator {
+  
+  public static func linearByKeyPath<T>(_ kp: KeyPath<Element, T>) -> Interpolator<Element>
+    where T: FloatConvertible {
+      let i = Interpolator<T>.linear
+      return Interpolator(
+        compare: { $0[keyPath: kp] < $1[keyPath: kp] },
+        interpolate: { value, min, max in
+          return i.interpolate(value[keyPath: kp], min[keyPath: kp], max[keyPath: kp])
+      })
+  }
+  
+  public static func linearByKeyPath<T>(_ kp: KeyPath<Element, T>) -> Interpolator<Element>
+    where T: FixedWidthInteger {
+      let i = Interpolator<T>.linear
+      return Interpolator(
+        compare: { $0[keyPath: kp] < $1[keyPath: kp] },
+        interpolate: { value, min, max in
+          return i.interpolate(value[keyPath: kp], min[keyPath: kp], max[keyPath: kp])
+      })
+  }
+  
+}
diff --git a/Tests/SwiftPlotTests/Heatmap/heatmap.swift b/Tests/SwiftPlotTests/Heatmap/heatmap.swift
