Module:ColinTheCat/Plotter

clamp 01

Arguments

Argument	Description	Example	Default
Plots	Plots as JSON	See #Example
From	Lowest x,y to plot	-pi/2,0	0,0
To	Highest x,y to plot	10^2,2
Origin	X/Y Axis origin	0.5,0.5	0,0
GridStep	Coordinate grid step size	pi/2,1	1,1
Inline	Display the plot inline instead of floating	true	(absent)
LabelX, LabelY	Axis label (not implemented)
LabelStepX	Label for grid X-axis increments	format("%sπ", dfrac(x/pi, 1))	format("%d", x)
LabelStepY	Label for grid X-axis increments	format("%.1f", y)	format("%d", x)
ColorX	X-axis color	#44f	#f44
ColorY	Y-axis color	#0f4	#0f0

Underlined Arguments are required.

Available Functions and Constants in Expressions

• printf(format, ...args)
  See Lua string.format.
• dfrac(number, is_factor)
  Formats a float as a fraction, if possible. If it is a factor, then the number 1 or -1 will only produce its sign as an output, useful when formatting as a factor of pi, for example. See #Example below.
• sin,floor,abs,...
  Any constant or function from the Lua math library can be used directly (without the math. prefix).

Example

{{#invoke:ColinTheCat/Plotter|Plot2D
|From=-0.2, -1.2
|To=2 * pi + 0.2, 1.2
|GridStep=0.5 * pi, 1
|LabelStepX=format("%sπ", dfrac(x/pi, 1))
|Plots=
[
{"Type": "function", "Function": "sin(x)", "Samples": 100, "Label": "sin"},
{"Type": "function", "Function": "cos(x)", "Samples": 100, "Label": "cos"},
{"Type": "function", "Function": "tan(x)", "Samples": 35, "Label": "tan", "Ranges": [
  {"To": "(0.5-0.2)*pi"},
  {"From": "(0.5+0.2)*pi", "To": "(1.5-0.2)*pi"},
  {"From": "(1.5+0.2)*pi"}
]}
]
}}

sin

cos

tan

---

---This module can generate 2D plots of functions using SVG
local p = {}

---@module "Expression"
local Expression = require("Module:ColinTheCat/Expression")

---@alias Space2D {x: number, y: number, w: number, h: number}
---@alias Vec2D {x: number, y: number}

---@alias Plot2DOptions {plots: Plot2DItemOptions[], from: Vec2D, to: Vec2D, origin: Vec2D, grid_step: Vec2D, axis_labels: AxisLabels, axis_colors: AxisColors, plot_space: Space2D, content_space: Space2D}
---@alias AxisLabels {x: string, x_step: Expr, y: string, y_step: Expr}
---@alias AxisColors {x: string, y: string}
---@alias Plot2DItemOptions {Type?: string, Color?: string|number, Label?: string}
---@alias FunctionPlotOptions {Type: `function`, Color?: string|number, Label?: string, Function?: string, Samples?: number|string, Ranges?: FunctionPlotRange[]}
---@alias FunctionPlotRange {From?: string, To?: string, Samples?: number|string}

function p.Plot2D(frame)
  assert(frame.args.Plots, "Missing 'Plots'")
  assert(frame.args.To, "Missing 'To'")

  ---@type Plot2DOptions
  local options = {
    plots = mw.text.jsonDecode(frame.args.Plots),
    from = ParseXY(frame.args.From or "0, 0"),
    to = ParseXY(frame.args.To),
    origin = ParseXY(frame.args.Origin or "0, 0"),
    grid_step = ParseXY(frame.args.GridStep or "1, 1"),
    inline = frame.args.Inline == "true",
    axis_labels = {
      x = frame.args.LabelX or "x",
      x_step = ParseExpr(frame.args.LabelStepX or 'format("%.0f", x)'),
      y = frame.args.LabelY or "y",
      y_step = ParseExpr(frame.args.LabelStepY or 'format("%.0f", y)'),
    },
    axis_colors = {
      x = frame.args.ColorX or "#f44",
      y = frame.args.ColorY or "#0f0",
    },
  }

  if options.grid_step.x <= 0 then error("Invalid x step size") end
  if options.grid_step.y <= 0 then error("Invalid y step size") end

  options.plot_space = {
    x = options.from.x,
    y = options.from.y,
    w = options.to.x - options.from.x,
    h = options.to.y - options.from.y,
  }
  options.content_space = {
    x = 0,
    y = 0,
    w = 1000,
    h = 1000 * options.plot_space.h / options.plot_space.w
  }

  -- make space for axis labels
  local label_size = { w = 60, h = 20 }
  local tx, ty, sx, sy = CreateTransform(options.plot_space, options.content_space)
  TranslateSpace(
    options.content_space,
    math.max(0, label_size.w - tx(options.origin.x)),
    0
  )
  tx, ty, sx, sy = CreateTransform(options.plot_space, options.content_space)

  local svg = mw.html.create()

  local abs_origin = TransformVec2D(options.origin, tx, ty)
  local abs_step = ScaleVec2D(options.grid_step, sx, sy)

  DrawGrid(svg, options.content_space, abs_step, abs_origin)
  local axes = DrawAxes(svg, options.content_space, abs_origin, options.axis_colors)

  ---@alias Legend {color: string, label: string}[]
  ---@type Legend
  local legend = {}
  for _, plot in ipairs(options.plots) do
    DrawPlot(svg, options, plot, legend)
  end

  DrawAxisLabels(
    svg, axes, options.plot_space,
    options.content_space, abs_step, abs_origin,
    options.axis_labels, options.axis_colors
  )

  local fig = mw.html.create("div")
      :cssText(string.format(
        "margin: var(--space-xs) auto var(--space-md) auto;%s",
        options.inline and "" or
        " margin-left: var(--space-lg); float: right; clear: right; max-width: 500px; width: 100%;"
      ))
      :wikitext(frame:callParserFunction("#tag:svg", {
        tostring(svg),
        viewBox = string.format(
          "0 0 %s %s",
          options.content_space.w + options.content_space.x,
          options.content_space.h + math.max(0, label_size.h - ty(options.origin.y))
        ),
        style = "width: 100%; fill: none; stroke-linecap: butt;",
      }))

  CreateLegend(fig, legend)

  return tostring(fig)
end

---@param str string
---@return Vec2D
function ParseXY(str)
  local i = str:find(",", 1, false)
  if i == nil then error("Invalid xy: " .. str) end
  local x = str:sub(1, i - 1)
  local y = str:sub(i + 1)
  return { x = tonumber(EvalExpr(x)), y = tonumber(EvalExpr(y)) }
end

---@param vec Vec2D
---@param tx fun(x: number): number
---@param ty? fun(y: number): number
---@return Vec2D
function TransformVec2D(vec, tx, ty)
  if ty == nil then ty = tx end
  return { x = tx(vec.x), y = ty(vec.y) }
end

---@param vec Vec2D
---@param sx number
---@param sy? number
---@return Vec2D
function ScaleVec2D(vec, sx, sy)
  if sy == nil then sy = sx end
  return { x = vec.x * sx, y = vec.y * sy }
end

---Creates a transform from one space to another
---@param from Space2D
---@param to Space2D
---@return fun(x: number): number tx Transform X
---@return fun(y: number): number ty Transform Y
---@return number sx Scale X
---@return number xy Scale Y
function CreateTransform(from, to)
  --     (x-from.x) * (to.w/from.w) + to.x
  -- <=> x * (to.w/from.w) - from.x * (to.w/from.w) + to.x
  --         |----sx-----| |-------------tx--------------|
  local sx, sy = to.w / from.w, to.h / from.h
  local tx = to.x - from.x * sx
  local ty = to.y - from.y * sy
  local function transform_x(x) return x * sx + tx end
  local function transform_y(y) return y * sy + ty end
  return transform_x, transform_y, sx, sy
end

---Moves a space around
---@param space Space2D
---@param x number
---@param y number
function TranslateSpace(space, x, y)
  space.x = space.x + x
  space.y = space.y + y
  --space.w = space.w + x
  --space.h = space.h + y
end

---Gets the edge bounds of a space
---@param space Space2D
---@return number min_x
---@return number min_y
---@return number max_x
---@return number max_y
function SpaceBounds(space)
  return space.x, space.y, space.w + space.x, space.h + space.y
end

---Creates a legend below the plot
---@param container any
---@param legend Legend
---@return any
function CreateLegend(container, legend)
  local foot = container:tag("div")
      :cssText(
        "font-size: 0.8125rem; text-align: center; color: var(--color-base--subtle); "
        .. "margin-top: var(--space-xs); padding-inline: var(--border-radius--small);"
      )
  for _, item in ipairs(legend) do
    foot
        :tag("div")
        :cssText(
          "display: inline-block; height: 0.9em; aspect-ratio: 3/2; margin-inline: 0.3em; vertical-align: middle; "
          .. "background-color: " .. item.color .. ";"
        )
        :done()
        :wikitext(item.label)
  end
  return legend
end

-- Drawing

---Draws a grid in a space
---@param svg any
---@param space Space2D
---@param step? Vec2D Grid steps, default: (1,1)
---@param origin? Vec2D Alignment point in the space, default: (0,0)
---@return any path
function DrawGrid(svg, space, step, origin)
  step = step or { x = 1, y = 1 }
  local step_x, step_y = step.x, step.y
  local min_x, min_y, max_x, max_y = SpaceBounds(space)
  origin = origin or { x = 0, y = max_y }
  local origin_x, origin_y = origin.x, origin.y

  local path = ""
  local k

  -- +x
  k = origin_x + step_x
  while (k < max_x) do
    path = path .. "M" .. k .. " " .. min_y .. "V" .. max_y
    k = k + step_x
  end

  -- -x
  k = origin_x - step_x
  while (k > min_x) do
    path = path .. "M" .. k .. " " .. min_y .. "V" .. max_y
    k = k - step_x
  end

  -- +y
  k = origin_y + step_y
  while (k < max_y) do
    path = path .. "M" .. min_x .. " " .. k .. "H" .. max_x
    k = k + step_x
  end

  -- -y
  k = origin_y - step_y
  while (k > min_y) do
    path = path .. "M" .. min_x .. " " .. k .. "H" .. max_x
    k = k - step_x
  end

  return svg:tag("path")
      :attr("stroke", "#666")
      :attr("stroke-width", 2)
      :attr("d", path)
end

---Draws axis arrows in a space
---@param svg any
---@param space Space2D
---@param origin Vec2D
---@param colors AxisColors
---@return unknown
function DrawAxes(svg, space, origin, colors)
  local min_x, min_y, max_x, max_y = SpaceBounds(space)
  origin = origin or { x = 0, y = 0 }
  local origin_x, origin_y = origin.x, max_y - origin.y

  local axes = svg:tag("g")
      :attr("stroke-width", 3)
  local axis_x = axes:tag("g")
      :attr("stroke", colors.x)
  local axis_y = axes:tag("g")
      :attr("stroke", colors.y)

  DrawArrow(axis_x, origin_x, origin_y, max_x - 5, origin_y)
  DrawArrow(axis_y, origin_x, origin_y, origin_x, min_y + 5)

  if min_x < origin_x then
    DrawLine(axis_x, origin_x - 10, origin_y, min_x, origin_y)
        :attr("stroke-dasharray", "10 10")
  end
  if max_y > origin_y then
    DrawLine(axis_y, origin_x, origin_y + 10, origin_x, max_y)
        :attr("stroke-dasharray", "10 10")
  end
  DrawCircle(axes, 4, origin_x, origin_y)
      :attr("fill", "#fff")

  return axes
end

---comment
---@param svg any
---@param axes any
---@param space Space2D
---@param step? Vec2D
---@param origin? Vec2D
---@param labels AxisLabels
---@param colors AxisColors
---@return unknown
function DrawAxisLabels(svg, axes, pspace, space, step, origin, labels, colors)
  local labels_g = svg:tag("g")
      :attr("stroke", "none")
      :cssText("font-size: 28px; text-shadow: 0 0 3px #000; font-weight: 600;")
  local notches_x = axes:tag("g")
      :attr("stroke", colors.x)
      :attr("stroke-width", 4)
  local labels_x = labels_g:tag("g")
      :attr("fill", colors.x)
      :attr("text-anchor", "middle")
      :cssText("dominant-baseline: hanging;")
  local notches_y = axes:tag("g")
      :attr("stroke", colors.y)
      :attr("stroke-width", 4)
  local labels_y = labels_g:tag("g")
      :attr("fill", colors.y)
      :attr("text-anchor", "end")
      :cssText("dominant-baseline: middle;")

  step = step or { x = 1, y = 1 }
  local step_x, step_y = step.x, step.y
  local min_x, min_y, max_x, max_y = SpaceBounds(space)
  origin = origin or { x = 0, y = 0 }
  local origin_x, origin_y = origin.x, max_y - origin.y

  ---Transform to plot space
  local tx, ty = CreateTransform(space, pspace)

  local notch_length = 10

  local cx = {}
  local k
  local current_axis
  function cx:lookup(name)
    if name == current_axis then return name == "x" and tx(k) or ty(max_y-k) end
    if name == "format" then return string.format end
    if name == "dfrac" then return DisplayFrac end
    return math[name] or error("Unknown variable: " .. name)
  end

  current_axis = "x"
  -- +x
  k = origin_x + step_x
  while (k < max_x) do
    DrawLine(notches_x, k, origin_y - notch_length, k, origin_y + notch_length)
    DrawText(labels_x, labels.x_step:eval(cx), k, origin_y + notch_length + 6)
    k = k + step_x
  end

  -- -x
  k = origin_x - step_x
  while (k > min_x) do
    DrawLine(notches_x, k, origin_y - notch_length, k, origin_y + notch_length)
    DrawText(labels_x, labels.x_step:eval(cx), k, origin_y + notch_length + 6)
    k = k - step_x
  end

  current_axis = "y"
  -- +y
  k = origin_y + step_y
  while (k < max_y) do
    DrawLine(notches_y, origin_x - notch_length, k, origin_x + notch_length, k)
    DrawText(labels_y, labels.y_step:eval(cx), origin_x - notch_length - 6, k)
    k = k + step_x
  end

  -- -y
  k = origin_y - step_y
  while (k > min_y) do
    DrawLine(notches_y, origin_x - notch_length, k, origin_x + notch_length, k)
    DrawText(labels_y, labels.y_step:eval(cx), origin_x - notch_length - 6, k)
    k = k - step_x
  end

  return labels_g
end

---Draws an arrow as a <path>
---@param svg any
---@param x1 number
---@param y1 number
---@param x2 number
---@param y2 number
---@return any path
function DrawArrow(svg, x1, y1, x2, y2)
  local dx, dy = x2 - x1, y2 - y1
  local len = math.sqrt(dx * dx + dy * dy)
  dx, dy = dx / len, dy / len
  local leftx, lefty = dy, -dx
  local arrow_size = 6
  return svg:tag("path")
      :attr("stroke-linejoin", "round")
      :attr("stroke-linecap", "round")
      :attr("d", string.format(
        "M%s %sL%s %sL%s %sL%s %sL%s %s",
        x1, y1,
        x2, y2,
        x2 + arrow_size * (leftx - dx), y2 + arrow_size * (lefty - dy),
        x2, y2,
        x2 + arrow_size * (-leftx - dx), y2 + arrow_size * (-lefty - dy)
      ))
end

---Draws a <circle>
---@param svg any
---@param r number
---@param x number
---@param y number
---@return any circle
function DrawCircle(svg, r, x, y)
  return svg:tag("circle")
      :attr("cx", x)
      :attr("cy", y)
      :attr("r", r)
end

---Draws a <line>
---@param svg any
---@param x1 number
---@param y1 number
---@param x2 number
---@param y2 number
---@return any line
function DrawLine(svg, x1, y1, x2, y2)
  return svg:tag("line")
      :attr("x1", x1)
      :attr("y1", y1)
      :attr("x2", x2)
      :attr("y2", y2)
end

---Draws a <text>
---@param svg any
---@param text string
---@param x number
---@param y number
---@return any text
function DrawText(svg, text, x, y)
  return svg:tag("text")
      :attr("x", x)
      :attr("y", y)
      :wikitext(text)
end

-- Plotting

---@alias Plot2DHandler fun(svg, options: Plot2DOptions, plot: Plot2DItemOptions): any

---@type table<string, Plot2DHandler>
local plotters = {}
local plot_colors = { "#ff0", "#f0f", "#0ff", "f44" }
local color_idx = 1

function GetPlotColor()
  local color = plot_colors[color_idx]
  color_idx = color_idx == #plot_colors and 1 or color_idx + 1
  return color
end

---Draws a plot
---@param svg any
---@param options any
---@param plot Plot2DItemOptions
---@param legend {color: string, label: string}[] Append legend data to this table
---@return unknown
function DrawPlot(svg, options, plot, legend)
  assert(plot.Type, "Missing plot 'Type'")
  local plotter = plotters[plot.Type]
  if not plotter then error("Unknown plot type '" .. plot.Type .. "'") end
  if not plot.Color then
    plot.Color = GetPlotColor()
  elseif type(plot.Color) == "number" then
    plot.Color = plot_colors[plot.Color]
  end
  if plot.Label then
    table.insert(legend, { color = plot.Color, label = plot.Label })
  end
  return plotter(svg, options, plot)
end

---@type Plot2DHandler
---@param plot FunctionPlotOptions
plotters["function"] = function(svg, options, plot)
  assert(plot.Function, "Function plot missing 'Function'")
  local plot_func = ParseExpr(plot.Function)
  local path = ""
  local samples = plot.Samples and tonumber(plot.Samples) or 50
  local tx, ty = CreateTransform(options.plot_space, options.content_space)

  local ranges = {}
  if plot.Ranges then
    assert(#plot.Ranges > 0, "Empty 'Ranges'")
    for _, range in ipairs(plot.Ranges) do
      local rfrom = range.From and EvalExpr(range.From) or options.from.x
      local rto = range.To and EvalExpr(range.To) or options.to.x
      assert(rfrom < rto, "Invalid range")
      local rsamples = range.Samples and tonumber(range.Samples) or samples
      table.insert(
        ranges,
        {
          from = rfrom,
          to = rto,
          samples = rsamples,
          step = (rto - rfrom) / (rsamples - 1),
        }
      )
    end
  else
    table.insert(
      ranges,
      {
        from = options.from.x,
        to = options.to.x,
        samples = samples,
        step = options.plot_space.w / (samples - 1)
      }
    )
  end

  local max_y = options.content_space.h + options.content_space.y

  local cx = {}
  local x = 0
  function cx:lookup(name)
    if name == "x" then
      return x
    else
      return math[name] or error("Unknown variable: " .. name)
    end
  end

  for _, range in ipairs(ranges) do
    local path_op = "M"
    for i = 0, range.samples - 1 do
      x = range.from + i * range.step
      local y = plot_func:eval(cx)
      path = path .. path_op .. tx(x) .. " " .. (max_y - ty(y))
      path_op = "L"
    end
  end

  return svg:tag("path")
      :attr("stroke", plot.Color)
      :attr("stroke-width", 3)
      :attr("stroke-linejoin", "bevel")
      :attr("d", path)
end

-- Expression Logic

---Parses any expression and validates the result
---@param str string raw expression
---@return Expr expression
function ParseExpr(str)
  local tok = Expression.Tokenizer:new { str = str }
  local expr = Expression.parse_expr(tok)
  if not expr then error("Invalid expression: " .. expr) end
  if tok:peek() then error("Leftover tokens: " .. tok.str) end
  return expr
end

---Evaluates a simple math expression
---@param str string raw expression
---@param vars? table<string, number> variables
---@return number
function EvalExpr(str, vars)
  local cx = {}
  function cx:lookup(name)
    return vars and vars[name] or math[name] or error("Unknown variable: " .. name)
  end

  return ParseExpr(str):eval(cx)
end

local nice_frac_table = {
  -- 0.
  ["5000"] = "½",
  ["3333"] = "⅓",
  ["6667"] = "⅔",
  ["2500"] = "¼",
  ["7500"] = "¾",
  ["2000"] = "⅕",
  ["4000"] = "⅖",
  ["6000"] = "⅗",
  ["8000"] = "⅘",
  ["1667"] = "⅙",
  ["8333"] = "⅚",
  ["1429"] = "⅐",
  ["1250"] = "⅛",
  ["3750"] = "⅜",
  ["6250"] = "⅝",
  ["8750"] = "⅞",
  ["1111"] = "⅑",
  ["1000"] = "⅒",
  -- 1.
  ["15000"] = "³⁄₂",
  ["13333"] = "⁴⁄₃",
  ["12500"] = "⁵⁄₄",
  ["12000"] = "⁶⁄₅",
  ["11667"] = "⁷⁄₆",
  ["11429"] = "⁸⁄₇",
  ["11250"] = "⁹⁄₈",
  ["11111"] = "¹⁰⁄₉",
  -- 2.
  ["25000"] = "⁵⁄₂",
}

function DisplayFrac(num, is_factor)
  is_factor = is_factor and is_factor ~= 0
  local sign = "";
  local pnum = num < 0 and -num or num
  local index = tostring(math.floor(pnum * 10000 + 0.5))
  local nice = nice_frac_table[index]

  if nice then return sign .. nice end
  if is_factor and index == "10000" then return sign end
  if index:sub(2) == "0000" then return tostring(math.floor(num)) end
  return string.format("%.2f", num)
end

return p