fermi/f.go

388 lines
9.8 KiB
Go

package main
import (
"bufio"
"errors"
"fmt"
"git.nunosempere.com/NunoSempere/fermi/sample"
"math"
"os"
"sort"
"strconv"
"strings"
)
const NORMAL90CONFIDENCE = 1.6448536269514727
const GENERAL_ERR_MSG = "Valid inputs: 2 || * 2 || / 2 || 2 20 || * 2 20 || / 2 20 || clean || =: var || op var || clean || help || debug || exit"
const N_SAMPLES = 1_000_000
// Distribution interface
// https://go.dev/tour/methods/9
type Dist interface {
Samples() []float64
}
type Scalar float64
type Lognormal struct {
low float64
high float64
}
type Beta struct {
a float64
b float64
}
type FilledSamples struct {
xs []float64
}
func (p Scalar) Samples() []float64 {
xs := make([]float64, N_SAMPLES)
for i := 0; i < N_SAMPLES; i++ {
xs[i] = float64(p)
}
return xs
}
func (ln Lognormal) Samples() []float64 {
sampler := func(r sample.Src) float64 { return sample.Sample_to(ln.low, ln.high, r) }
// return sample.Sample_parallel(sampler, N_SAMPLES)
// Can't do parallel because then I'd have to await throughout the code
return sample.Sample_serially(sampler, N_SAMPLES)
}
func (beta Beta) Samples() []float64 {
sampler := func(r sample.Src) float64 { return sample.Sample_beta(beta.a, beta.b, r) }
// return sample.Sample_parallel(sampler, N_SAMPLES)
return sample.Sample_serially(sampler, N_SAMPLES)
}
func (fs FilledSamples) Samples() []float64 {
return fs.xs
}
// Parse line into Distribution
func parseLineErr(err_msg string) (string, Dist, error) {
fmt.Println(GENERAL_ERR_MSG)
fmt.Println(err_msg)
var errorDist Dist
return "", errorDist, errors.New(err_msg)
}
func parseLine(line string, vars map[string]Dist) (string, Dist, error) {
words := strings.Split(strings.TrimSpace(line), " ")
op := ""
var dist Dist
switch words[0] {
case "*", "/", "+", "-":
op = words[0]
words = words[1:]
default:
op = "*" // later, change the below to
}
switch len(words) {
case 0:
return parseLineErr("Operator must have operand; can't operate on nothing")
case 1:
var_word, var_word_exists := vars[words[0]]
single_float, err1 := strconv.ParseFloat(words[0], 64) // abstract this away to search for K/M/B/T/etc.
switch {
case var_word_exists:
dist = var_word
case err1 == nil:
dist = Scalar(single_float)
case err1 != nil && !var_word_exists:
return parseLineErr("Trying to operate on a scalar, but scalar is neither a float nor an assigned variable")
}
case 2:
new_low, err1 := strconv.ParseFloat(words[0], 64)
new_high, err2 := strconv.ParseFloat(words[1], 64)
if err1 != nil || err2 != nil {
return parseLineErr("Trying to operate by a distribution, but distribution is not specified as two floats")
}
dist = Lognormal{low: new_low, high: new_high}
case 3:
if words[0] == "beta" || words[0] == "b" {
a, err1 := strconv.ParseFloat(words[1], 64)
b, err2 := strconv.ParseFloat(words[2], 64)
if err1 != nil || err2 != nil {
return parseLineErr("Trying to specify a beta distribution? Try beta 1 2")
}
dist = Beta{a: a, b: b}
} else {
return parseLineErr("Input not understood or not implemented yet")
}
default:
return parseLineErr("Input not understood or not implemented yet")
}
return op, dist, nil
}
func multiplyLogDists(l1 Lognormal, l2 Lognormal) Lognormal {
logmean1 := (math.Log(l1.high) + math.Log(l1.low)) / 2.0
logstd1 := (math.Log(l1.high) - math.Log(l1.low)) / (2.0 * NORMAL90CONFIDENCE)
logmean2 := (math.Log(l2.high) + math.Log(l2.low)) / 2.0
logstd2 := (math.Log(l2.high) - math.Log(l2.low)) / (2.0 * NORMAL90CONFIDENCE)
logmean_product := logmean1 + logmean2
logstd_product := math.Sqrt(logstd1*logstd1 + logstd2*logstd2)
h := logstd_product * NORMAL90CONFIDENCE
loglow := logmean_product - h
loghigh := logmean_product + h
return Lognormal{low: math.Exp(loglow), high: math.Exp(loghigh)}
}
func multiplyBetaDists(beta1 Beta, beta2 Beta) Beta {
return Beta{a: beta1.a + beta2.a, b: beta1.b + beta2.b}
}
func operateAsSamples(dist1 Dist, dist2 Dist, op string) (Dist, error) {
xs := dist1.Samples()
ys := dist2.Samples()
// fmt.Printf("xs: %v\n", xs)
// fmt.Printf("ys: %v\n", ys)
zs := make([]float64, N_SAMPLES)
for i := 0; i < N_SAMPLES; i++ {
switch op {
case "*":
zs[i] = xs[i] * ys[i]
case "/":
if ys[0] != 0 {
zs[i] = xs[i] / ys[i]
} else {
fmt.Println("Error: When dividing as samples, division by zero")
return nil, errors.New("Division by zero")
}
case "+":
zs[i] = xs[i] + ys[i]
case "-":
zs[i] = xs[i] - ys[i]
}
}
// fmt.Printf("%v\n", zs)
return FilledSamples{xs: zs}, nil
}
func multiplyDists(old_dist Dist, new_dist Dist) (Dist, error) {
switch o := old_dist.(type) {
case Lognormal:
{
switch n := new_dist.(type) {
case Lognormal:
return multiplyLogDists(o, n), nil
case Scalar:
return multiplyLogDists(o, Lognormal{low: float64(n), high: float64(n)}), nil
default:
return operateAsSamples(old_dist, new_dist, "*")
}
}
case Scalar:
{
if o == 1 {
return new_dist, nil
}
switch n := new_dist.(type) {
case Lognormal:
return multiplyLogDists(Lognormal{low: float64(o), high: float64(o)}, n), nil
case Scalar:
return Scalar(float64(o) * float64(n)), nil
default:
return operateAsSamples(old_dist, new_dist, "*")
}
}
case Beta:
switch n := new_dist.(type) {
case Beta:
return multiplyBetaDists(o, n), nil
default:
return operateAsSamples(old_dist, new_dist, "*")
}
default:
return operateAsSamples(old_dist, new_dist, "*")
}
}
func divideDists(old_dist Dist, new_dist Dist) (Dist, error) {
switch o := old_dist.(type) {
case Lognormal:
{
switch n := new_dist.(type) {
case Lognormal:
return multiplyLogDists(o, Lognormal{low: 1.0 / n.high, high: 1.0 / n.low}), nil
case Scalar:
return multiplyLogDists(o, Lognormal{low: 1.0 / float64(n), high: 1.0 / float64(n)}), nil
default:
return operateAsSamples(old_dist, new_dist, "/")
}
}
case Scalar:
{
switch n := new_dist.(type) {
case Lognormal:
return multiplyLogDists(Lognormal{low: float64(o), high: float64(o)}, Lognormal{low: 1.0 / n.high, high: 1.0 / n.low}), nil
case Scalar:
return Scalar(float64(o) / float64(n)), nil
default:
return operateAsSamples(old_dist, new_dist, "/")
}
}
default:
return operateAsSamples(old_dist, new_dist, "/")
}
}
func joinDists(old_dist Dist, new_dist Dist, op string) (Dist, error) {
switch op {
case "*":
return multiplyDists(old_dist, new_dist)
case "/":
return divideDists(old_dist, new_dist)
case "+":
return operateAsSamples(old_dist, new_dist, "+")
case "-":
return operateAsSamples(old_dist, new_dist, "-")
default:
return old_dist, errors.New("Can't combine distributions in this way")
}
}
/* Pretty print distributions */
func prettyPrintFloat(f float64) {
switch {
case math.Abs(f) >= 1_000_000_000_000:
fmt.Printf("%.1fT", f/1_000_000_000_000)
case math.Abs(f) >= 1_000_000_000:
fmt.Printf("%.1fB", f/1_000_000_000)
case math.Abs(f) >= 1_000_000:
fmt.Printf("%.1fM", f/1_000_000)
case math.Abs(f) >= 1_000:
fmt.Printf("%.1fK", f/1_000)
case math.Abs(f) <= 0.0001:
fmt.Printf("%.5f", f)
case math.Abs(f) <= 0.001:
fmt.Printf("%.4f", f)
case math.Abs(f) <= 0.01:
fmt.Printf("%.3f", f)
case math.Abs(f) <= 0.1:
fmt.Printf("%.2f", f)
default:
fmt.Printf("%.1f", f)
}
}
func prettyPrint2Floats(low float64, high float64) {
prettyPrintFloat(low)
fmt.Printf(" ")
prettyPrintFloat(high)
fmt.Printf("\n")
}
func prettyPrintDist(dist Dist) {
switch v := dist.(type) {
case Lognormal:
fmt.Printf("=> ")
prettyPrint2Floats(v.low, v.high)
case FilledSamples:
tmp_xs := make([]float64, N_SAMPLES)
copy(tmp_xs, v.xs)
sort.Slice(tmp_xs, func(i, j int) bool {
return tmp_xs[i] < tmp_xs[j]
})
low_int := N_SAMPLES / 20
low := tmp_xs[low_int]
high_int := N_SAMPLES * 19 / 20
high := tmp_xs[high_int]
prettyPrint2Floats(low, high)
case Beta:
fmt.Printf("=> beta ")
prettyPrint2Floats(v.a, v.b)
case Scalar:
fmt.Printf("=> scalar ")
w := float64(v)
prettyPrintFloat(w)
fmt.Println()
default:
fmt.Printf("%v", v)
}
}
/* Main event loop */
func main() {
reader := bufio.NewReader(os.Stdin)
var init_dist Dist
init_dist = Scalar(1) // Lognormal{low: 1, high: 1}
old_dist := init_dist
vars := make(map[string]Dist)
// Could eventually be a more complex struct with:
// { Dist, VariableMaps, ConfigParams } or smth
EventForLoop:
for {
input, _ := reader.ReadString('\n')
if strings.TrimSpace(input) == "" {
continue EventForLoop
}
{
words := strings.Split(strings.TrimSpace(input), " ")
switch {
case words[0] == "exit" || words[0] == "e":
break EventForLoop
case words[0] == "help" || words[0] == "h":
fmt.Println(GENERAL_ERR_MSG)
continue EventForLoop
case words[0] == "debug" || words[0] == "d":
fmt.Printf("Old dist: %v\n", old_dist)
fmt.Printf("Vars: %v\n", vars)
continue EventForLoop
case words[0] == "=:" && len(words) == 2:
vars[words[1]] = old_dist
fmt.Printf("%s ", words[1])
prettyPrintDist(old_dist)
continue EventForLoop
case words[0] == "." || words[0] == "clean" || words[0] == "c":
old_dist = init_dist
fmt.Println()
continue EventForLoop
case words[0] == "=." && len(words) == 2:
vars[words[1]] = old_dist
fmt.Printf("%s ", words[1])
prettyPrintDist(old_dist)
old_dist = init_dist
fmt.Println()
continue EventForLoop
// Other possible cases:
// Save to file
// Sample n samples
// Save stack to a variable?
// clean stack
// Define a function? No, too much of a nerdsnipea
}
}
op, new_dist, err := parseLine(input, vars)
if err != nil {
continue EventForLoop
}
joint_dist, err := joinDists(old_dist, new_dist, op)
if err != nil {
fmt.Printf("%v\n", err)
fmt.Printf("Dist on stack: ")
prettyPrintDist(old_dist)
continue EventForLoop
}
old_dist = joint_dist
prettyPrintDist(old_dist)
}
}