plot.go

// 31 Dec 2021
// This is for doing some plotting. For the moment, this is in the ackwork package,
// but maybe it should be in its own. It definitely goes in its own file, since
// it is tied to one of the chart packages I have tried out.
// What I have learnt about the go-chart package...
//  - The default tick positions are terrible.
//  - If I give it a slice of ticks, it works fine for the main axis, but
//    something terrible happens to the scaling of the axes. I could get
//    it to work by: Make a child type from continuousRange. For this type,
//    define the GetTicks function. go-chart then calls this function, at
//    apparently the right place, and the tick marks come out nicely.

package ackwork

import (
	"fmt"
	"math"
	"os"

	"github.com/wcharczuk/go-chart/v2"
	"github.com/wcharczuk/go-chart/v2/drawing"
	"gitlab.rrz.uni-hamburg.de/Bae5157/axticks"
)

// maketicks gives reasonable default tick locations
func maketicks(axisDscrpt axticks.AxisDscrpt, prcsn int) []chart.Tick {
	xmin, delta, prcsn := axisDscrpt.Xmin, axisDscrpt.Delta, axisDscrpt.Prcsn
	rnge := axisDscrpt.Xmax - axisDscrpt.Xmin
	ntick := int(math.Round((rnge / delta) + 1))
	t := make([]chart.Tick, ntick)
	const fstr = "%.*f"
	t[0].Value, t[0].Label = xmin, fmt.Sprintf(fstr, prcsn, xmin)
	for i := 1; i < ntick; i++ {
		t[i].Value = t[i-1].Value + delta
		t[i].Label = fmt.Sprintf("%.*f", prcsn, t[i].Value)
	}
	return t
}

// range2 just lets us append methods to a range/continuousrange structure
type range2 struct {
	chart.ContinuousRange
}

// GetTicks is of the form wanted by go-chart to return ticks. It has to be
// a method on a range, but then you have to define your own kind of range.
// That is done by defining range2. We don't actually use any of the
// arguments that go-charts offers us.
func (rng *range2) GetTicks(r chart.Renderer, cs chart.Style, vf chart.ValueFormatter) []chart.Tick {
	tmp := []float64{rng.Min, rng.Max}
	if a, err := axticks.Tickpos(tmp); err != nil {
		fmt.Fprintln(os.Stderr, "GetTicks error:", err)
		return nil
	} else {
		rng.Min = a.Xmin
		rng.Max = a.Xmax
		return maketicks(a, a.Prcsn)
	}
}

// plotfWrt writes out a plot of function values to the given
// filename
func plotfWrt(cprm *cprm, fname string) error {
	if !cprm.fplotme {
		return nil
	}
	xdata := cprm.plotnstp // We just unpack for readability
	ydata := cprm.plotf
	tmprtrData := cprm.plotTmprtr

	xaxis := chart.XAxis{
		Name:  "step",
		Range: &range2{},
	}

	fvalAxis := chart.YAxis{
		Name:      "cost (arb units)",
		NameStyle: chart.Style{TextRotationDegrees: 360}, // Zero does not work
		AxisType:  chart.YAxisSecondary,
		Range:     &range2{},
	}
	var tmprtrAxis chart.YAxis
	if cprm.coolme { // If we are not cooling, we do not plot
		tmprtrAxis = chart.YAxis{ // the temperature axis.
			Name:      "temperature",
			NameStyle: chart.Style{TextRotationDegrees: 360},
			Range:     &range2{},
		}
	}

	graph := chart.Chart{
		Width: 800,
		Series: []chart.Series{
			chart.ContinuousSeries{ // The function values
				Name:    "cost",
				YAxis:   chart.YAxisSecondary,
				XValues: xdata,
				YValues: ydata,
				Style: chart.Style{
					StrokeWidth: 5,
					DotWidth:    0,
				},
			},
			chart.ContinuousSeries{
				Name:    "temperature",
				XValues: xdata,
				YValues: tmprtrData,
			},
		},
		YAxis:          tmprtrAxis,
		YAxisSecondary: fvalAxis,
		XAxis:          xaxis,

		Background: chart.Style{
			Padding: chart.Box{
				Left:  50,
				Right: 75,
			},
			FillColor: drawing.ColorTransparent,
		},
	}
	fPlt, err := os.Create(fname)
	if err != nil {
		return err
	}
	defer fPlt.Close()
	if err := graph.Render(chart.PNG, fPlt); err != nil {
		return fmt.Errorf("Render: %w", err)
	}

	return nil
}

// plotxWrt
// For each dimension, allocate space. Copy elements over from the long array.
// Then call the plotter on each series in turn. I think I have to make a
// slice of continuous series and then add them into the chart structure.
func plotxWrt(cprm *cprm, xPltName string, ndim int) error {
	if !cprm.xplotme {
		return nil
	}

	type f64 []float64
	len_used := len(cprm.plotnstp)
	xdata := make([]f64, ndim)
	for i := 0; i < ndim; i++ {
		xdata[i] = make([]float64, len_used)
	}
	var n int
	for i := 0; i < len_used; i++ {
		for j := 0; j < ndim; j++ {
			xdata[j][i] = float64(cprm.plotXtrj[n])
			n++
		}
	}
	series := make([]chart.Series, ndim)
	for i := 0; i < ndim; i++ {
		series[i] = chart.ContinuousSeries{
			Name:    fmt.Sprintf("dimension %d", i),
			XValues: cprm.plotnstp,
			YValues: xdata[i],
		}
	}
	graph := chart.Chart{
		Title: fmt.Sprintf("X trajectories %d dimension", ndim),
		XAxis: chart.XAxis{Name: "step", Range: &range2{}},
		YAxis: chart.YAxis{
			Name:      "x coord",
			Range:     &range2{},
			AxisType:  chart.YAxisSecondary,
			NameStyle: chart.Style{TextRotationDegrees: 360},
		},
		Series: series,
		Background: chart.Style{
			Padding: chart.Box{
				Left: 75,
			},
		},
	}
	xPlt, err := os.Create(xPltName)
	if err != nil {
		return err
	}
	defer xPlt.Close()
	if err := graph.Render(chart.PNG, xPlt); err != nil {
		return fmt.Errorf("plotting X trajectories: %w", err)
	}
	return nil
}