412 lines
11 KiB
Go
412 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"fmt"
|
|
"git.nunosempere.com/NunoSempere/fermi/sample"
|
|
"git.nunosempere.com/NunoSempere/fermi/pretty"
|
|
"math"
|
|
"os"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
/* Types and interfaces */
|
|
|
|
type Stack struct {
|
|
old_dist Dist
|
|
vars map[string]Dist
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
/* Dist interface functions */
|
|
// https://go.dev/tour/methods/9
|
|
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
|
|
}
|
|
|
|
/* Constants */
|
|
const GENERAL_ERR_MSG = " Operation | Variable assignment | Special\n" +
|
|
" Operation: operator operand\n" +
|
|
" operator: (empty) | * | / | + | -\n" +
|
|
" operand: scalar | lognormal | beta | variable\n" +
|
|
" lognormal: low high\n" +
|
|
" beta: beta alpha beta\n" +
|
|
" Variable assignment: =: variable_name\n" +
|
|
" Clear stack: . | c | clear\n" +
|
|
" Variable assignment and clear stack: =. variable_name\n" +
|
|
" Other special operations: help | debug | exit\n" +
|
|
" Examples: \n" +
|
|
" + 2\n" +
|
|
" / 2.5\n" +
|
|
" * 1 10 (interpreted as lognormal)\n" +
|
|
" + 1 10\n" +
|
|
" * beta 1 10\n" +
|
|
" 1 10 (multiplication taken as default operation)\n" +
|
|
" =: x\n" +
|
|
" .\n" +
|
|
" 1 100\n" +
|
|
" + x\n" +
|
|
" exit\n"
|
|
const NORMAL90CONFIDENCE = 1.6448536269514727
|
|
const INIT_DIST Scalar = Scalar(1)
|
|
const N_SAMPLES = 100_000
|
|
|
|
/* Pretty print for distributions */
|
|
// Needs types
|
|
func prettyPrintDist(dist Dist) {
|
|
switch v := dist.(type) {
|
|
case Lognormal:
|
|
fmt.Printf("=> ")
|
|
pretty.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]
|
|
fmt.Printf("=> ")
|
|
pretty.PrettyPrintFloat(low)
|
|
fmt.Printf(" ")
|
|
pretty.PrettyPrintFloat(high)
|
|
fmt.Printf(" (")
|
|
pretty.PrettyPrintInt(N_SAMPLES)
|
|
fmt.Printf(" samples)\n")
|
|
case Beta:
|
|
fmt.Printf("=> beta ")
|
|
pretty.PrettyPrint2Floats(v.a, v.b)
|
|
case Scalar:
|
|
fmt.Printf("=> scalar ")
|
|
w := float64(v)
|
|
pretty.PrettyPrintFloat(w)
|
|
fmt.Println()
|
|
default:
|
|
fmt.Printf("%v", v)
|
|
}
|
|
}
|
|
|
|
// 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 parseLineIntoOpAndDist(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 operateDistsAsSamples(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 operateDistsAsSamples(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 operateDistsAsSamples(old_dist, new_dist, "*")
|
|
}
|
|
}
|
|
case Beta:
|
|
switch n := new_dist.(type) {
|
|
case Beta:
|
|
return multiplyBetaDists(o, n), nil
|
|
default:
|
|
return operateDistsAsSamples(old_dist, new_dist, "*")
|
|
}
|
|
default:
|
|
return operateDistsAsSamples(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:
|
|
// to do: check division by zero
|
|
return multiplyLogDists(o, Lognormal{low: 1.0 / n.high, high: 1.0 / n.low}), nil
|
|
case Scalar:
|
|
// to do: check division by zero
|
|
return multiplyLogDists(o, Lognormal{low: 1.0 / float64(n), high: 1.0 / float64(n)}), nil
|
|
default:
|
|
return operateDistsAsSamples(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:
|
|
// to do: check division by zero
|
|
return Scalar(float64(o) / float64(n)), nil
|
|
default:
|
|
return operateDistsAsSamples(old_dist, new_dist, "/")
|
|
}
|
|
}
|
|
default:
|
|
return operateDistsAsSamples(old_dist, new_dist, "/")
|
|
}
|
|
}
|
|
|
|
/* Combine old dist and new line */
|
|
// We want this as a function to be able to have parenthesis/recusion, possibly functions
|
|
func operateStackWithDist(stack Stack, new_dist Dist, op string) Stack {
|
|
|
|
var combined_dist Dist
|
|
var err error
|
|
switch op {
|
|
case "*":
|
|
if combined_dist, err = multiplyDists(stack.old_dist, new_dist); err == nil {
|
|
stack.old_dist = combined_dist
|
|
}
|
|
case "/":
|
|
if combined_dist, err = divideDists(stack.old_dist, new_dist); err == nil {
|
|
stack.old_dist = combined_dist
|
|
}
|
|
case "+":
|
|
if combined_dist, err = operateDistsAsSamples(stack.old_dist, new_dist, "+"); err == nil {
|
|
stack.old_dist = combined_dist
|
|
}
|
|
case "-":
|
|
if combined_dist, err = operateDistsAsSamples(stack.old_dist, new_dist, "-"); err == nil {
|
|
stack.old_dist = combined_dist
|
|
}
|
|
default:
|
|
fmt.Println("Can't combine distributions in this way")
|
|
}
|
|
return stack
|
|
|
|
}
|
|
|
|
func runRepl(stack Stack, reader *bufio.Reader) Stack {
|
|
|
|
|
|
replForLoop:
|
|
for {
|
|
new_line, _ := reader.ReadString('\n')
|
|
words := strings.Split(strings.TrimSpace(new_line), " ")
|
|
switch {
|
|
/* Empty line case */
|
|
case strings.TrimSpace(new_line) == "":
|
|
continue replForLoop
|
|
/* Parenthesis */
|
|
case len(words) == 2 && (words[0] == "*" || words[0] == "+" || words[0] == "-" || words[0] == "/") && words[1] == "(":
|
|
new_stack := runRepl(Stack{old_dist: INIT_DIST, vars: stack.vars}, reader) // to do: think if I want to shadow variables or not (right now, variables persist, since I'm copying the map(.
|
|
stack = operateStackWithDist(stack, new_stack.old_dist, words[0])
|
|
prettyPrintDist(stack.old_dist)
|
|
case len(words) == 1 && words[0] == ")":
|
|
return stack
|
|
/* Special operations */
|
|
case words[0] == "exit" || words[0] == "e":
|
|
os.Exit(0)
|
|
case words[0] == "help" || words[0] == "h":
|
|
fmt.Println(GENERAL_ERR_MSG)
|
|
continue replForLoop
|
|
case words[0] == "debug" || words[0] == "d":
|
|
fmt.Printf("Old dist: %v\n", stack.old_dist)
|
|
fmt.Printf("Vars: %v\n", stack.vars)
|
|
continue replForLoop
|
|
case words[0] == "clear" || words[0] == "c" || words[0] == ".":
|
|
stack.old_dist = INIT_DIST
|
|
fmt.Println()
|
|
continue replForLoop
|
|
/* Variable assignment */
|
|
case words[0] == "=:" && len(words) == 2:
|
|
stack.vars[words[1]] = stack.old_dist
|
|
fmt.Printf("%s ", words[1])
|
|
prettyPrintDist(stack.old_dist)
|
|
continue replForLoop
|
|
case words[0] == "=." && len(words) == 2:
|
|
stack.vars[words[1]] = stack.old_dist
|
|
fmt.Printf("%s ", words[1])
|
|
prettyPrintDist(stack.old_dist)
|
|
stack.old_dist = INIT_DIST
|
|
fmt.Println()
|
|
continue replForLoop
|
|
default:
|
|
op, new_dist, err := parseLineIntoOpAndDist(new_line, stack.vars)
|
|
if err != nil {
|
|
continue replForLoop
|
|
}
|
|
stack = operateStackWithDist(stack, new_dist, op)
|
|
prettyPrintDist(stack.old_dist)
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Main event loop */
|
|
func main() {
|
|
|
|
reader := bufio.NewReader(os.Stdin)
|
|
stack := Stack{old_dist: INIT_DIST, vars: make(map[string]Dist)}
|
|
runRepl(stack, reader)
|
|
}
|