plot: Contour Fixes (#271)

Fixes #270.

Changes:
- Fixes interpolation at domain edges.
- Adds the option to specify a comparison operator.

src/lib/plot/contour.typ

@@ -4,29 +4,54 @@
 
 // Find contours of a 2D array by using marching squares algorithm
 //
-// - data (array): 2D float array (rows => columns)
-// - offset (float): Value (offset) a cell must be greater or equal to to count as true
+// - data (array): A 2D array of floats where the first index is the row and the second index is the column
+// - offset (float): Z value threshold of a cell compare with `op` to, to count as true
+// - op (auto,string,funciton): Z value comparison oparator:
+//     / `">", ">=", "<", "<=", "!=", "=="`: Use the passed operator to compare z.
+//     / `auto`: Use ">=" for positive z values, "<=" for negative z values.
+//     / `<function>`: If set to a function, that function gets called
+//       with two arguments, the z value `z1` to compare against and
+//       the z value `z2` of the data and must return a boolean: `(z1, z2) => boolean`.
 // - interpolate (bool): Enable cell interpolation for smoother lines
 // - contour-limit (int): Contour limit after which the algorithm panics
 // -> array: Array of contour point arrays
-#let find-contours(data, offset, interpolate: true, contour-limit: 50) = {
-  assert(data != none)
-  assert(type(data) == array)
-  assert(type(offset) in (int, float))
+#let find-contours(data, offset, op: auto, interpolate: true, contour-limit: 50) = {
+  assert(data != none and type(data) == array,
+    message: "Data must be of type array")
+  assert(type(offset) in (int, float),
+    message: "Offset must be numeric")
 
   let n-rows = data.len()
   let n-cols = data.at(0).len()
   if n-rows < 2 or n-cols < 2 {
     return ()
   }
 
+  assert(op == auto or type(op) in (str, function),
+    message: "Operator must be of type auto, string or function")
+  if op == auto {
+    op = if offset < 0 { "<=" } else { ">=" }
+  }
+  if type(op) == str {
+    assert(op in ("<", "<=", ">", ">=", "==", "!="),
+      message: "Operator must be one of: <, <=, >, >=, != or ==")
+  }
+
   // Return if data is set
-  let is-set(v) = {
-    return if offset < 0 {
-      v <= offset
-    } else {
-      v >= offset
-    }
+  let is-set = if type(op) == function {
+    v => op(offset, v)
+  } else if op == "==" {
+    v => v == offset
+  } else if op == "!=" {
+    v => v != offset
+  } else if op == "<" {
+    v => v < offset
+  } else if op == "<=" {
+    v => v <= offset
+  } else if op == ">" {
+    v => v > offset
+  } else if op == ">=" {
+    v => v >= offset
   }
 
   // Build a binary map that has 0 for unset and 1 for set cells
@@ -45,7 +70,7 @@
     if x >= 0 and x < n-cols and y >= 0 and y < n-rows {
       return float(data.at(y).at(x))
     }
-    return 0
+    return none
   }
 
   // Get case (0 to 15)
@@ -55,6 +80,8 @@
 
   let lerp(a, b) = {
     if a == b { return a }
+    else if a == none { return 1 }
+    else if b == none { return 0 }
     return (offset - a) / (b - a)
   }
 
@@ -198,7 +225,7 @@
 #let _stroke(self, ctx) = {
   for c in self.contours {
     for p in c.stroke-paths {
-      draw.line(..p, fill: none)
+      draw.line(..p, fill: none, close: p.first() == p.last())
     }
   }
 }
@@ -208,14 +235,16 @@
   if not self.fill { return }
   for c in self.contours {
     for p in c.fill-paths {
-      draw.line(..p, stroke: none)
+      draw.line(..p, stroke: none, close: p.first() == p.last())
     }
   }
 }
 
 /// Add a contour plot of a sampled function or a matrix.
 ///
-/// - data (array, function): Matrix or `(x, y) => z` function
+/// - data (array, function): A function of the signature `(x, y) => z`
+///                           or an array of floats where the first
+///                           index is the row and the second index is the column.
 ///
 ///                           *Examples:*
 ///                           - `(x, y) => x > 0`
@@ -228,6 +257,12 @@
 /// - x-samples (int): X axis domain samples (2 < n)
 /// - y-samples (int): Y axis domain samples (2 < n)
 /// - interpolate (bool): Use linear interpolation between sample values
+/// - op (auto,string,funciton): Z value comparison oparator:
+///   / `">", ">=", "<", "<=", "!=", "=="`: Use the operator for comparison.
+///   / `auto`: Use ">=" for positive z values, "<=" for negative z values.
+///   / `<function>`: Call comparison function of the format `(plot-z, data-z) => boolean`,
+///                   where `plot-z` is the z-value from the plots `z` argument and `data-z`
+///                   is the z-value of the data getting plotted.
 /// - fill (bool): Fill each contour
 /// - style (style): Style to use, can be used with a palette function
 /// - axes (array): Name of the axes to use ("x", "y"), note that not all
@@ -240,6 +275,7 @@
                  x-samples: 25,
                  y-samples: 25,
                  interpolate: true,
+                 op: auto,
                  axes: ("x", "y"),
                  style: (:),
                  fill: false,
@@ -262,7 +298,7 @@
   let contours = ()
   let z = if type(z) == array { z } else { (z,) }
   for z in z {
-    for contour in find-contours(data, z, interpolate: interpolate) {
+    for contour in find-contours(data, z, op: op, interpolate: interpolate) {
       let line-data = contour.map(pt => {
         (pt.at(0) * dx + x-min,
          pt.at(1) * dy + y-min)

tests/plot/contour/test.typ

@@ -73,3 +73,56 @@
     )
   })
 }))
+
+/* Complex contour #270 */
+#box(stroke: 2pt + red, canvas({
+  plot.plot(size: (8, 8), {
+    // x >= 0
+    plot.add-contour(
+      (x, y) => x,
+      z: 0,
+      y-samples: 2,
+      x-samples: 2,
+      x-domain: (0, 10),
+      y-domain: (-10, 10),
+      fill: true,
+    )
+
+    // y >= 0
+    plot.add-contour(
+      (x, y) => y,
+      z: 0,
+      y-samples: 2,
+      x-samples: 2,
+      x-domain: (-10, 10),
+      y-domain: (0, 10),
+      fill: true,
+    )
+
+    // hyperbola
+    plot.add-contour(
+      (x, y) => (x - 1) * (y - 1),
+      x-domain: (-10, 10),
+      y-domain: (-10, 10),
+      fill: true,
+      z: 1,
+    )
+
+    // circle
+    plot.add-contour(
+      (x, y) => (calc.pow((x - 1), 2) + calc.pow((y - 1), 2)),
+      x-domain: (-10, 10),
+      y-domain: (-10, 10),
+      z: 9,
+      op: "<=",
+      fill: true,
+    )
+
+    // line
+    plot.add-contour(
+      (x, y) => x + 1 - y,
+      x-domain: (-10, 10),
+      y-domain: (-10, 10),
+    )
+  })
+}))