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"}
]}
]
}}

---

local p = {}

local Expression = require("Module:ColinTheCat/Expression")

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

  local options = {
    plots = mw.text.jsonDecode(frame.args.Plots),
    from = parse_xy(frame.args.From or "0, 0"),
    to = parse_xy(frame.args.To),
    origin = parse_xy(frame.args.Origin or "0, 0"),
    grid_step = parse_xy(frame.args.GridStep or "1, 1"),
    axis_labels = {
      x=frame.args.LabelX or "x",
      x_step=parse_expr(frame.args.LabelStepX or 'format("%d", x)'),
      y=frame.args.LabelY or "y",
      y_step=parse_expr(frame.args.LabelStepY or 'format("%d", 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.function_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.function_space.h / options.function_space.w
  }

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

  local svg = mw.html.create()

  draw_grid(svg, options)
  local axes = draw_axes(svg, options)

  local legend = {}
  for _, plot in ipairs(options.plots) do
    draw_plot(svg, options, plot, legend)
  end

  draw_axis_labels(svg, options, axes)

  local fig = mw.html.create("div")
    :cssText("margin: var(--space-xs) auto var(--space-md) auto;")
    :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;",
    }))

  create_legend(fig, legend)

  return tostring(fig)
end

function parse_xy(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=eval_expr(x), y=eval_expr(y) }
end

function create_legend(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

function create_transform(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

function translate_space(space, x, y)
  space.x = space.x + x
  space.y = space.y + y
  space.w = space.w + x
  space.h = space.h + y
end

function draw_grid(svg, options)
  local tx, ty, sx, sy = create_transform(options.function_space, options.content_space)
  local step_x = options.grid_step.x * sx
  local step_y = options.grid_step.y * sy

  local min_x = options.content_space.x
  local min_y = options.content_space.y
  local max_x = options.content_space.w + min_x
  local max_y = options.content_space.h + min_y
  local origin_x, origin_y = tx(options.origin.x), max_y - ty(options.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

function draw_axes(svg, options)
  local tx, ty = create_transform(options.function_space, options.content_space)
  local min_x = options.content_space.x
  local min_y = options.content_space.y
  local max_x = options.content_space.w + min_x
  local max_y = options.content_space.h + min_y
  local origin_x, origin_y = tx(options.origin.x), max_y - ty(options.origin.y)
  local axes = svg:tag("g")
    :attr("stroke-width", 3)
  local axis_x = axes:tag("g")
    :attr("stroke", options.axis_colors.x)
  local axis_y = axes:tag("g")
    :attr("stroke", options.axis_colors.y)

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

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

  return axes
end

function draw_axis_labels(svg, options, axes)
  local labels = 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", options.axis_colors.x)
    :attr("stroke-width", 4)
  local labels_x = labels:tag("g")
    :attr("fill", options.axis_colors.x)
    :attr("text-anchor", "middle")
    :cssText("dominant-baseline: hanging;")
  local notches_y = axes:tag("g")
    :attr("stroke", options.axis_colors.y)
    :attr("stroke-width", 4)
  local labels_y = labels:tag("g")
    :attr("fill", options.axis_colors.y)
    :attr("text-anchor", "end")
    :cssText("dominant-baseline: middle;")

  local tx, ty, sx, sy = create_transform(options.function_space, options.content_space)
  local step_x_f, step_y_f = options.grid_step.x, options.grid_step.y
  local step_x = step_x_f * sx
  local step_y = step_y_f * sy

  local min_x = options.content_space.x
  local min_y = options.content_space.y
  local max_x = options.content_space.w + min_x
  local max_y = options.content_space.h + min_y
  local origin_x_f, origin_y_f = options.origin.x, options.origin.y
  local origin_x, origin_y = tx(origin_x_f), max_y - ty(origin_y_f)

  local notch_length = 10

  local cx = {}
  local k, k_f
  local current_axis
  function cx:lookup(name)
    if name == current_axis then return k_f end
    if name == "format" then return string.format end
    return math[name] or error("Unknown variable: "..name)
  end

  current_axis = "x"
  -- +x
  k = origin_x + step_x
  k_f = origin_x_f + step_x_f
  while(k < max_x) do
    draw_line(notches_x, k, origin_y - notch_length, k, origin_y + notch_length)
    draw_text(labels_x, options.axis_labels.x_step:eval(cx), k, origin_y + notch_length + 6)
    k = k + step_x
    k_f = k_f + step_x_f
  end

  -- -x
  k = origin_x - step_x
  k_f = origin_x_f - step_x_f
  while(k > min_x) do
    draw_line(notches_x, k, origin_y - notch_length, k, origin_y + notch_length)
    draw_text(labels_x, options.axis_labels.x_step:eval(cx), k, origin_y + notch_length + 6)
    k = k - step_x
    k_f = k_f - step_x_f
  end

  current_axis = "y"
  -- +y
  k = origin_y + step_y
  k_f = origin_y_f - step_y_f
  while(k < max_y) do
    draw_line(notches_y, origin_x - notch_length, k, origin_x + notch_length, k)
    draw_text(labels_y, options.axis_labels.y_step:eval(cx), origin_x - notch_length - 6, k)
    k = k + step_x
    k_f = k_f - step_y_f
  end

  -- -y
  k = origin_y - step_y
  k_f = origin_y_f + step_y_f
  while(k > min_y) do
    draw_line(notches_y, origin_x - notch_length, k, origin_x + notch_length, k)
    draw_text(labels_y, options.axis_labels.y_step:eval(cx), origin_x - notch_length - 6, k)
    k = k - step_x
    k_f = k_f + step_y_f
  end

  return labels
end

function draw_arrow(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

function draw_circle(svg, r, x, y)
  return svg:tag("circle")
    :attr("cx", x)
    :attr("cy", y)
    :attr("r", r)
end

function draw_line(svg, x1, y1, x2, y2)
  return svg:tag("line")
    :attr("x1", x1)
    :attr("y1", y1)
    :attr("x2", x2)
    :attr("y2", y2)
end

function draw_text(svg, text, x, y)
  return svg:tag("text")
    :attr("x", x)
    :attr("y", y)
    :wikitext(text)
end

-- Plotting

local plotters = {}
local plot_colors = { "#ff0", "#f0f", "#0ff", "f44" }
local color_idx = 1

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

function draw_plot(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 = get_plot_color()
  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

plotters["function"] = function(svg, options, plot)
  assert(plot.Function, "Function plot missing 'Function'")
  local plot_func = parse_expr(plot.Function)
  local path = ""
  local samples = plot.Samples and tonumber(plot.Samples) or 50
  local tx, ty = create_transform(options.function_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 eval_expr(range.From) or options.from.x
      local rto = range.To and eval_expr(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.function_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

function parse_expr(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

function eval_expr(str, vars)
  local cx = {}
  function cx:lookup(name)
    return vars and vars[name] or math[name] or error("Unknown variable: "..name)
  end

  return parse_expr(str):eval(cx)
end

return p