dorun.go

// Aug 2021
package mcwork

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"math"
	"math/rand"
	"os"
	"sync"

	"example.com/ackley_mc/ackley"
)

var seedLocker sync.Mutex

// getSeed returns the random number seed, but we have to put a lock
// around it so multiple runs get consistent successive seeds
func getSeed() int {
	seedLocker.Lock()
	r := Seed
	Seed++
	seedLocker.Unlock()
	return r
}

type withBytes interface {
	Bytes() []byte
}

type cprm struct { // parameters calculated from input
	rand       *rand.Rand
	coolMult   float64        // multiplier for temperature
	fOutRaw    io.WriteCloser // the raw file handle for csv files
	fOut       *bufio.Writer  // buffered version of fOutRaw
	doTxtFval  bool           // Are we writing text files (csv) ?
	plotnstp   []float64      // Number of steps. float64 for the plot library
	plotf      []float64      // Function values for plotting
	plotTmprtr []float64      // Temperature values for plotting
	plotXtrj   []float32      // for plotting trajectories
	coolme     bool           // Cooling ? or just constant temperature ?
	fplotWrt   io.Writer      // Where to write the plot of function values
	xplotWrt   io.Writer      // Where to write the plot of x trajectories
}

// isNotSane checks for obviously silly parameters
func isNotSane(mcPrm *McPrm) error {
	if mcPrm.fOutName == "" && mcPrm.fPltName == "" && mcPrm.xPltName == "" {
		if false { // if there are no graphics compiled in, ..
			return fmt.Errorf("All output files are empty. You would not see anything")
		}
	}
	if mcPrm.FnlTmp > mcPrm.IniTmp {
		const finalTempErr = "Final temp %g higher than initial temp %g"
		return fmt.Errorf(finalTempErr, mcPrm.FnlTmp, mcPrm.IniTmp)
	}
	if mcPrm.IniTmp < 0 {
		return fmt.Errorf("negative initial temperature %g", mcPrm.IniTmp)
	}
	return nil
}

// setupRun does things like get the cooling rate, seed the random numbers.
func setupRun(mcPrm *McPrm, cprm *cprm) error {
	var err error
	cprm.rand = rand.New(rand.NewSource(int64(getSeed())))

	if mcPrm.IniTmp != mcPrm.FnlTmp { // cooling or constant temperature ?
		var coolrate float64
		cprm.coolme = true
		nStep := float64(mcPrm.NStep)
		fnlTmp, iniTmp := float64(mcPrm.FnlTmp), float64(mcPrm.IniTmp)
		if fnlTmp == 0 { // avoid divide by zeroes
			fnlTmp = math.SmallestNonzeroFloat32
		}
		coolrate = -1 / nStep
		coolrate *= math.Log(fnlTmp / iniTmp)
		cprm.coolMult = math.Exp(-coolrate)
	}
	if mcPrm.fOutName != "" { // We are going to write text, .csv
		if mcPrm.fOutName, err = setSuffix(mcPrm.fOutName, ".csv"); err != nil {
			return err
		}
		if cprm.fOutRaw, err = os.Create(mcPrm.fOutName); err != nil {
			return err
		}
		cprm.fOut = bufio.NewWriter(cprm.fOutRaw)
		cprm.doTxtFval = true // other parts of code can see that we are writing to file
	}

	if mcPrm.fPltName != "" { // Function value plot to file
		if mcPrm.fPltName, err = setSuffix(mcPrm.fPltName, ".png"); err != nil {
			return fmt.Errorf("broke setting plot filename: %w", err)
		}
		if cprm.fplotWrt, err = os.Create(mcPrm.fPltName); err != nil {
			return err
		}
	} else { // Must be writing to a screen
		var w bytes.Buffer // could be fancier. Preallocate and use bytes.Newbuffer
		cprm.fplotWrt = &w
	}
	n_alloc := mcPrm.NStep / 5 // About a fifth of the number of steps
	cprm.plotnstp = make([]float64, 0, n_alloc)
	cprm.plotf = make([]float64, 0, n_alloc)
	if cprm.coolme {
		cprm.plotTmprtr = make([]float64, 0, n_alloc)
	}

	if mcPrm.xPltName != "" {
		if mcPrm.xPltName, err = setSuffix(mcPrm.xPltName, ".png"); err != nil {
			return fmt.Errorf("broke setting plot filename: %w", err)
		}
		if cprm.xplotWrt, err = os.Create(mcPrm.xPltName); err != nil {
			return err
		}
	} else { // Write to buffer, to plot on screen
		var w bytes.Buffer
		cprm.xplotWrt = &w // I think we can remove the preceding line
	}
	cprm.plotXtrj = make([]float32, 0, int(n_alloc)*len(mcPrm.XIni))
	if err = isNotSane(mcPrm); err != nil {
		return err
	}
	return nil
}

// newx provides a random step to give us new coordinates.
// An earlier version moved all dimensions at once, but required
// teeny moves.
func newx(xold []float32, xT []float32, rand *rand.Rand, xDlta float64) {
	var iDim int
	if (len(xold)) > 1 {
		iDim = rand.Intn(len(xold)) // pick one dimension to change
	}
	t := 2.*rand.Float64() - 1
	t *= xDlta
	copy(xT, xold)
	xT[iDim] = xT[iDim] + float32(t)
}

// printfVal is the loop to print out the function value and coordinates
// for later plotting.
func printfVal(fOut io.Writer, x []float32, n int, tmprtr float64, fOld float64) {
	fmt.Fprintf(fOut, "%d,%.4g,%.5g", n, tmprtr, fOld)
	for _, xx := range x {
		fmt.Fprintf(fOut, ",%.3g", xx)
	}
	fmt.Fprintln(fOut)
}

// saveBest checks if we have found a best value for the function. If so,
// update the location where we store it.
var bestfval float64
var bestX []float32
var bestN int

func saveBest(x []float32, fTrial float64, n int) {
	if bestX == nil {
		bestX = make([]float32, len(x))
	}
	if fTrial < bestfval || n == 0 {
		bestfval = fTrial
		bestN = n
		copy(bestX, x)
	}
}

// saveStep is called when we accept a change. It writes or saves for
// later plotting. Every time we append values, we have to do it twice
// so we get a proper stepped appearance.
func saveStep(cprm *cprm, n int, tmprtr float64, x []float32, fTrial float64) {
	if cprm.coolme { // Only keep the best values found if we are annealing
		saveBest(x, fTrial, n)
	}
	fprev := fTrial
	xprev := x
	if n != 0 {
		fprev = cprm.plotf[len(cprm.plotf)-1]
		xprev = cprm.plotXtrj[len(cprm.plotXtrj)-len(x):]
	}
	cprm.plotnstp = append(cprm.plotnstp, float64(n), float64(n)) // num steps
	if cprm.fplotWrt != nil {                                     // Function values
		cprm.plotf = append(cprm.plotf, fprev)            // copy the last coordinates
		cprm.plotf = append(cprm.plotf, fTrial)           // and append them
		cprm.plotTmprtr = append(cprm.plotTmprtr, tmprtr) // then the new ones
		cprm.plotTmprtr = append(cprm.plotTmprtr, tmprtr) // and again
	}
	if cprm.xplotWrt != nil { // The trajectory
		cprm.plotXtrj = append(cprm.plotXtrj, xprev...)
		cprm.plotXtrj = append(cprm.plotXtrj, x...)
	}
	if cprm.doTxtFval {
		printfVal(cprm.fOut, x, n, tmprtr, fTrial)
	}
}

// doRun does a Monte Carlo run.
// We only need single precision for most things, except for function
// values, but the interface to the plot and interface packages mostly
// want double precision, so we have a lot of float64 that shoudl be
// float32.
func singleRun(mcPrm *McPrm) (MCresult, error) {
	var cprm cprm
	broken := MCresult{}
	if err := setupRun(mcPrm, &cprm); err != nil {
		return broken, err
	}
	if cprm.doTxtFval { // The text file documenting the run
		defer cprm.fOutRaw.Close()
		defer cprm.fOut.Flush()
	}

	if c, ok := cprm.fplotWrt.(io.Closer); ok {
		defer c.Close() // file handle for function values
	}
	if c, ok := cprm.xplotWrt.(io.Closer); ok {
		defer c.Close() // file handle for X trajectories
	}

	var nAcc int                           // Counter, number of accepted moves
	x := make([]float32, len(mcPrm.XIni))  // current position
	xT := make([]float32, len(mcPrm.XIni)) // trial position

	for i := range mcPrm.XIni { // Initial coordinates from
		x[i] = float32(mcPrm.XIni[i]) // the control structure
	}

	fOld := ackley.Ackley(x) // Initial function value
	fTrial := fOld
	tmprtr := float64(mcPrm.IniTmp)
	xDlta := mcPrm.XDlta // Step size which might be adjusted on the fly
	saveStep(&cprm, 0, tmprtr, x, fOld)

	for n := 0; n < mcPrm.NStep; n++ {
		var accept bool
		newx(x, xT, cprm.rand, xDlta)
		fTrial = ackley.Ackley(xT)
		if fTrial <= fOld {
			accept = true
		} else {
			delta := fTrial - fOld         // Decision as to whether to
			t := math.Exp(-delta / tmprtr) // accept or reject the trial
			if cprm.rand.Float64() < t {
				accept = true
			}
		}
		if accept {
			nAcc++
			copy(x, xT)
			fOld = fTrial
			saveStep(&cprm, n+1, tmprtr, x, fTrial)
		}
		if cprm.coolme {
			tmprtr *= cprm.coolMult
		}
	}
	// On plots, we want the last step, even if nothing changed.
	if fTrial != fOld { // last move was not accepted, so we have to add last step
		saveStep(&cprm, mcPrm.NStep, tmprtr, x, fOld)
	}

	defer fmt.Println("n accepted:", nAcc, "of", mcPrm.NStep)
	if bestX != nil {
		s := "Best function value: %.2f at %.2v at step %d\n"
		defer fmt.Printf(s, bestfval, bestX, bestN)
	}
	if err := plotfWrt(&cprm); err != nil {
		return broken, err
	}
	if err := plotxWrt(&cprm, len(mcPrm.XIni)); err != nil {
		return broken, err
	}

	var fdata, xdata []byte
	if funcBuf, ok := cprm.fplotWrt.(withBytes); ok {
		fdata = funcBuf.Bytes()
	}
	if xBuf, ok := cprm.xplotWrt.(withBytes); ok {
		xdata = xBuf.Bytes()
	}

	return MCresult{
		Fdata: fdata, Xdata: xdata,
		NStep: mcPrm.NStep, NAcc: nAcc,
		Bestfval: bestfval, BestX: bestX,
		LastX: x,
	}, nil
}

// DoRun will do a run with equilibration. It will call singleRun twice.
// On the first run, we make sure there is no cooling and output is
// not written or goes to the garbage (io.Discard).
func DoRun(mcPrm *McPrm) (MCresult, error) {
	var tmpX []float32
	if mcPrm.NEquil > 0 {
		tmpX = make([]float32, len(mcPrm.XIni))
		mcEquil := *mcPrm
		mcEquil.FnlTmp = mcEquil.IniTmp
		mcEquil.NStep = mcEquil.NEquil
		mcEquil.NEquil = 0
		mcEquil.fOutName = ""
		mcEquil.fPltName, mcEquil.xPltName = "", ""
		if eqResult, err := singleRun(&mcEquil); err != nil {
			return eqResult, err
		} else {
			copy(tmpX, mcPrm.XIni)
			copy(mcPrm.XIni, eqResult.LastX)
		}
	}
	mcResult, err := singleRun(mcPrm)
	if tmpX != nil {
		copy(mcPrm.XIni, tmpX)
	}
	return mcResult, err
}