Module:ColinTheCat/Plotter: Difference between revisions

From Resonite Wiki
No edit summary
No edit summary
 
(94 intermediate revisions by the same user not shown)
Line 1: Line 1:
---This module can generate 2D plots of functions using SVG
local p = {}
local p = {}


---@module "Expression"
local Expression = require("Module:ColinTheCat/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)
function p.Plot2D(frame)
   assert(frame.args.Functions, "Missing 'Functions'")
   assert(frame.args.Plots, "Missing 'Plots'")
   assert(frame.args.To, "Missing 'To'")
   assert(frame.args.To, "Missing 'To'")


  ---@type Plot2DOptions
   local options = {
   local options = {
     functions = mw.text.jsonDecode(frame.args.Functions),
     plots = mw.text.jsonDecode(frame.args.Plots),
     from = parse_xy(frame.args.From or "0 0"),
     from = ParseXY(frame.args.From or "0, 0"),
     to = parse_xy(frame.args.To),
     to = ParseXY(frame.args.To),
     origin = parse_xy(frame.args.Origin or "0 0"),
     origin = ParseXY(frame.args.Origin or "0, 0"),
     grid_step = {
     grid_step = ParseXY(frame.args.GridStep or "1, 1"),
       x=frame.args.GridStepX or "1",
    inline = frame.args.Inline == "true",
       y=frame.args.GridStepY or "1",
    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",
     },
     },
   }
   }


   options.function_space = {
   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,
     x = options.from.x,
     y = options.from.y,
     y = options.from.y,
Line 28: Line 52:
     y = 0,
     y = 0,
     w = 1000,
     w = 1000,
     h = 1000 * options.function_space.h / options.function_space.w
     h = 1000 * options.plot_space.h / options.plot_space.w
   }
   }
  options.svg_size = {w = options.content_space.w, h = options.content_space.h}


   if options.from.x > -20 then options.svg_size.w = options.svg_size.w + 20 end
   -- make space for axis labels
   if options.from.y > -20 then options.svg_size.h = options.svg_size.h + 20 end
  local label_size = { w = 60, h = 40 }
  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
  )
   local x_axis_label_margin = math.max(0, label_size.h - ty(options.origin.y))
  tx, ty, sx, sy = CreateTransform(options.plot_space, options.content_space)


   local svg = mw.html.create()
   local svg = mw.html.create()


   draw_grid(svg, options)
   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")
   local fig = mw.html.create("div")
    :wikitext(frame:callParserFunction("#tag:svg", {
      :cssText(string.format(
      tostring(svg),
        "margin: var(--space-xs) auto var(--space-md) auto;%s",
      viewBox = "0 0 1000 "..(1000 * options.height / options.width),
        options.inline and "" or
      style = "width: 100%; fill: none;",
        " 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 + x_axis_label_margin
        ),
        style = "width: 100%; fill: none; stroke-linecap: butt;",
      }))
 
  CreateLegend(fig, legend)
 
   return tostring(fig)
   return tostring(fig)
end
end


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


function parse_xy(str)
---@param vec Vec2D
  local i = str:find(" ", 1, false)
---@param tx fun(x: number): number
   if i == nil then error("Invalid xy: "..str) end
---@param ty? fun(y: number): number
  local x = tonumber(str:sub(1, i - 1))
---@return Vec2D
  local y = tonumber(str:sub(i + 1))
function TransformVec2D(vec, tx, ty)
   return { x=x, y=y }
   if ty == nil then ty = tx end
   return { x = tx(vec.x), y = ty(vec.y) }
end
end


-- Drawing
---@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


function create_transform(from, to)
---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-from.x) * (to.w/from.w) + to.x
   -- <=> x * (to.w/from.w) - from.x * (to.w/from.w) + to.x
   -- <=> x * (to.w/from.w) - from.x * (to.w/from.w) + to.x
   --        |----sx-----|   |-----------tx--------------|
   --        |----sx-----| |-------------tx--------------|
   local sx, sy = to.w/from.w, to.h/from.h
   local sx, sy = to.w / from.w, to.h / from.h
   local tx = from.x * (to.w/from.w) + to.x
   local tx = to.x - from.x * sx
   local ty = from.y * (to.h/from.h) + to.y
   local ty = to.y - from.y * sy
   local function transform_x(x) return x * sx - tx end
   local function transform_x(x) return x * sx + tx end
   local function transform_y(y) return y * sy - ty end
   local function transform_y(y) return y * sy + ty end
   return transform_x, transform_y, sx, sy
   return transform_x, transform_y, sx, sy
end
end


function draw_grid(svg, options)
---Moves a space around
   local step_x = eval_expr(options.grid_step.x)
---@param space Space2D
   local step_y = eval_expr(options.grid_step.y)
---@param x number
   if step_x <= 0 then error("Invalid x step size") end
---@param y number
   if step_y <= 0 then error("Invalid y step size") end
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


   local tx, ty, sx, sy = create_transform(options.function_space, options.content_space)
---Creates a legend below the plot
   step_x = step_x * sx
---@param container any
   step_y = step_y * sy
---@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


   local min_x = options.content_space.x
-- Drawing
   local min_y = options.content_space.y
 
   local max_x = options.content_space.w + min_x
---Draws a grid in a space
   local max_y = options.content_space.h + min_y
---@param svg any
   local origin_x, origin_y = tx(options.origin.x), ty(options.origin.y)
---@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 = 0 }
   local origin_x, origin_y = origin.x, max_y - origin.y


   local path = ""
   local path = ""
Line 99: Line 219:
   -- +x
   -- +x
   k = origin_x + step_x
   k = origin_x + step_x
   while(k < max_x) do
   while (k < max_x) do
     path = path.."M"..k.." "..min_y.."V"..max_y
     path = path .. "M" .. k .. " " .. min_y .. "V" .. max_y
     k = k + step_x
     k = k + step_x
   end
   end
Line 106: Line 226:
   -- -x
   -- -x
   k = origin_x - step_x
   k = origin_x - step_x
   while(k > min_x) do
   while (k > min_x) do
     path = path.."M"..k.." "..min_y.."V"..max_y
     path = path .. "M" .. k .. " " .. min_y .. "V" .. max_y
     k = k - step_x
     k = k - step_x
   end
   end
Line 113: Line 233:
   -- +y
   -- +y
   k = origin_y + step_y
   k = origin_y + step_y
   while(k < max_y) do
   while (k < max_y) do
     path = path.."M"..min_x.." "..k.."H"..max_x
     path = path .. "M" .. min_x .. " " .. k .. "H" .. max_x
    k = k + step_y
  end
 
  -- -y
  k = origin_y - step_y
  while (k > min_y) do
    path = path .. "M" .. min_x .. " " .. k .. "H" .. max_x
    k = k - step_y
  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(svg, 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"
  -- x0
  k = origin_x
  if max_y <= origin_y then
    DrawLine(notches_x, k, origin_y, k, origin_y + notch_length)
    DrawText(labels_x, labels.x_step:eval(cx), k, origin_y + notch_length + 6)
  end
 
  -- +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
     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"
  -- y0
  k = origin_y
  if min_x >= origin_x then
    DrawLine(notches_y, origin_x - notch_length, k, origin_x, k)
    DrawText(labels_y, labels.y_step:eval(cx), origin_x - notch_length - 6, k)
  end
  -- +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_y
   end
   end


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


   svg:tag("path")
   return labels_g
    :attr("stroke", "#666")
end
    :attr("stroke-width", "1")
 
    :attr("stroke-linecat", "butt")
---Draws an arrow as a <path>
    :attr("d", path)
---@param svg any
    :done()
---@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
end


function new_object(self, o)
---Draws a <line>
   o = o or {}
---@param svg any
  setmetatable(o, self)
---@param x1 number
  self.__index = self
---@param y1 number
  return o
---@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
end


-- Expression Logic
---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


local ExprCtx = {new=new_object}
  return svg:tag("path")
function ExprCtx:lookup(name)
      :attr("stroke", plot.Color)
  return self.vars and self.bars[name] or math[name] or error("Unknown variable: "..name)
      :attr("stroke-width", 3)
      :attr("stroke-linejoin", "bevel")
      :attr("d", path)
end
end


function parse_expr(str)
-- Expression Logic
   local tok = Expression.Tokenizer:new{str=str}
 
---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)
   local expr = Expression.parse_expr(tok)
   if not expr then error("Invalid expression: "..expr) end
   if not expr then error("Invalid expression: " .. expr) end
   if tok:peek() then error("Leftover tokens: "..tok.str) end
   if tok:peek() then error("Leftover tokens: " .. tok.str) end
   return expr
   return expr
end
end


function eval_expr(str, vars)
---Evaluates a simple math expression
   local cx = ExprCtx:new{vars=vars}
---@param str string raw expression
   return parse_expr(str):eval(cx)
---@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
end


return p
return p

Latest revision as of 00:06, 13 February 2024

0 1 2 -1 1
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"}
]}
]
}}
½π π ³⁄₂π -1 1
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 = 40 }
  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
  )
  local x_axis_label_margin = math.max(0, label_size.h - ty(options.origin.y))
  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 + x_axis_label_margin
        ),
        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 = 0 }
  local origin_x, origin_y = origin.x, max_y - 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_y
  end

  -- -y
  k = origin_y - step_y
  while (k > min_y) do
    path = path .. "M" .. min_x .. " " .. k .. "H" .. max_x
    k = k - step_y
  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(svg, 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"
  -- x0
  k = origin_x
  if max_y <= origin_y then
    DrawLine(notches_x, k, origin_y, k, origin_y + notch_length)
    DrawText(labels_x, labels.x_step:eval(cx), k, origin_y + notch_length + 6)
  end

  -- +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"
  -- y0
  k = origin_y
  if min_x >= origin_x then
    DrawLine(notches_y, origin_x - notch_length, k, origin_x, k)
    DrawText(labels_y, labels.y_step:eval(cx), origin_x - notch_length - 6, k)
  end

  -- +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_y
  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_y
  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