fnplot.go

package fnplot

import (
	"fmt"
	"log"
	"math/big"
	"math/rand"
	"reflect"
	"sort"
	"sync"
	"time"

	"github.com/leanovate/gopter"
	"github.com/pkg/errors"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/plotutil"
	"gonum.org/v1/plot/vg"
)

type generatorResult interface {
	RetrieveAsValue() (reflect.Value, bool)
}

type sortablePoints plotter.XYs

func (sp sortablePoints) Len() int           { return len(sp) }
func (sp sortablePoints) Swap(i, j int)      { sp[i], sp[j] = sp[j], sp[i] }
func (sp sortablePoints) Less(i, j int) bool { return sp[i].X < sp[j].X }

type ioPair struct {
	input  Values
	output Values
}

type ValuesSet struct {
	pairs     []ioPair
	minInput  *big.Float
	maxInput  *big.Float
	minOutput *big.Float
	maxOutput *big.Float
	mu        sync.RWMutex
}

// TODO: Consider using a channel instead of a synchronized slice.
func (set *ValuesSet) insert(input, output Values) error {
	set.mu.Lock()
	defer set.mu.Unlock()

	set.pairs = append(set.pairs, ioPair{input: input, output: output})
	in, err := input.Scalar()
	if err != nil {
		return errors.WithMessage(err, "error converting input to int")
	}
	if set.minInput == nil || set.minInput.Cmp(in) == 1 {
		set.minInput = in
	}
	if set.maxInput == nil || set.maxInput.Cmp(in) == -1 {
		set.maxInput = in
	}
	out, err := output.Scalar()
	if err != nil {
		return errors.WithMessage(err, "error converting output to int")
	}
	if set.minOutput == nil || set.minOutput.Cmp(out) == 1 {
		set.minOutput = out
	}
	if set.maxOutput == nil || set.maxOutput.Cmp(out) == -1 {
		set.maxOutput = out
	}
	return nil
}

func (set *ValuesSet) PointsOn(xAxis, yAxis Axis) (plotter.XYs, error) {
	set.mu.RLock()
	defer set.mu.RUnlock()

	xAxis.SetMaxValue(set.maxInput)
	yAxis.SetMaxValue(set.maxOutput)

	points := make(plotter.XYs, len(set.pairs))
	for i := range set.pairs {
		inputScalar, err := set.pairs[i].input.Scalar()
		if err != nil {
			return nil, errors.WithMessage(err, fmt.Sprintf("error converting input %d to int", i))
		}
		points[i].X = xAxis.Point(inputScalar)

		outputScalar, err := set.pairs[i].output.Scalar()
		if err != nil {
			return nil, errors.WithMessage(err, fmt.Sprintf("error converting output %d to int", i))
		}
		points[i].Y = yAxis.Point(outputScalar)
	}
	sort.Sort(sortablePoints(points))
	return points, nil
}

// A Fn is a plottable function that holds the function to plot, the input
// generators, and the inputs and outputs as scalars.
type Fn struct {
	p   gopter.Prop
	set *ValuesSet
}

// errorProp creates a property that will always fail with an error.
// Mostly used as a fallback when setup/initialization fails
func errorProp(err error) gopter.Prop {
	return func(genParams *gopter.GenParameters) *gopter.PropResult {
		return &gopter.PropResult{
			Status: gopter.PropError,
			Error:  err,
		}
	}
}

// forAllGens returns a gopter.Prop that will run the provided function with
// inputs generated by the provided generators. The input/output pairs are
// inserted into the given ValuesSet.
// Based on "github.com/leanovate/gopter/prop".ForAllNoShrink:
// https://github.com/leanovate/gopter/blob/293686f39f478c1a469f003eaf0518d15c7c4509/prop/forall_no_shrink.go#L18
func forAllGens(vs *ValuesSet, fn interface{}, gens ...gopter.Gen) gopter.Prop {
	fnVal := reflect.ValueOf(fn)
	fnType := fnVal.Type()
	if fnType.Kind() != reflect.Func {
		return errorProp(fmt.Errorf("Expecting fn to be a func, got %s", fnVal.Kind().String()))
	}
	numArgs := len(gens)
	if fnType.NumIn() != numArgs {
		return errorProp(fmt.Errorf("Number of parameters does not match number of generators: %d != %d", fnType.NumIn(), numArgs))
	}

	return gopter.SaveProp(func(genParams *gopter.GenParameters) *gopter.PropResult {
		genResults := make([]*gopter.GenResult, len(gens))
		args := make([]reflect.Value, len(gens))
		var ok bool
		for i, gen := range gens {
			result := gen(genParams)
			genResults[i] = result
			args[i], ok = result.RetrieveAsValue()
			if !ok {
				return &gopter.PropResult{
					Status: gopter.PropUndecided,
				}
			}
		}

		results := fnVal.Call(args)
		vs.insert(args, results)

		// TODO: Is this necessary?
		result := &gopter.PropResult{Status: gopter.PropTrue}
		for i, genResult := range genResults {
			result = result.AddArgs(gopter.NewPropArg(genResult, 0, args[i].Interface(), args[i].Interface()))
		}
		return result
	})
}

func NewFn(fn interface{}, samples int, gens ...Generator) Fn {
	gopterGens := make([]gopter.Gen, len(gens))
	for i := range gens {
		gopterGens[i] = gopter.Gen(gens[i])
	}
	vs := &ValuesSet{
		pairs: make([]ioPair, 10),
	}
	f := Fn{
		p:   forAllGens(vs, fn, gopterGens...),
		set: vs,
	}
	f.run(samples)
	return f
}

// run runs the function with the set of input generators.
func (fn Fn) run(samples int) error {
	res := fn.p.Check(&gopter.TestParameters{
		MinSuccessfulTests: samples,
		MaxSize:            samples,
		Seed:               time.Now().UnixNano(),
		Rng:                rand.New(gopter.NewLockedSource(time.Now().UnixNano())),
		Workers:            10, // TODO: Make configurable

		// The following values are irrelevant because we're not discarding any
		// samples.
		MaxDiscardRatio: 0,
		MaxShrinkCount:  0,
		MinSize:         0,
	})
	return res.Error
}

func (fn Fn) ValuesSet() *ValuesSet {
	return fn.set
}

type Plot struct {
	Title string
	Fn    Fn
	X, Y  Axis
}

// Save writes the plot as an image to the given filename. The image format is
// determined by the file extension.
func (pl Plot) Save(filename string) error {
	p, err := plot.New()
	if err != nil {
		return errors.WithMessage(err, "error creating plot")
	}
	p.Title.Text = pl.Title
	p.X.Label.Text = " "
	p.Y.Label.Text = " "

	points, err := pl.Fn.ValuesSet().PointsOn(pl.X, pl.Y)
	if err != nil {
		log.Fatalf("Error generating X,Y points: %s", err)
	}
	err = plotutil.AddLinePoints(p, "Fn", points)
	if err == plotter.ErrInfinity {
		return errors.New("infinity value found, consider using an axis that supports scaling")
	} else if err != nil {
		return err
	}

	// Save the plot to a file. The format is determined by the file extension.
	err = p.Save(20*vg.Inch, 4*vg.Inch, filename)
	return errors.WithMessage(err, "error writing plot image")
}