Improve printing of complex numbers and Square roots.

corywalker · Nov 1, 2018 · cfe2dd4 · cfe2dd4
1 parent 5f2b9f5
commit cfe2dd4
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/expreduce/atoms/ex_complex.go b/expreduce/atoms/ex_complex.go
@@ -18,6 +18,15 @@ func (cmplx *Complex) StringForm(p expreduceapi.ToStringParams) string {
 	if p.Form == "FullForm" {
 		return fmt.Sprintf("Complex[%v, %v]", cmplx.Re, cmplx.Im)
 	}
+	reInt, reIsInt := cmplx.Re.(*Integer)
+	imInt, imIsInt := cmplx.Im.(*Integer)
+	if reIsInt && reInt.Val.Sign() == 0 {
+		if imIsInt && imInt.Val.Int64() == 1 {
+			return "I"
+		}
+		p.PreviousHead = "System`Times"
+		return fmt.Sprintf("(%v*I)", cmplx.Im.StringForm(p))
+	}
 	p.PreviousHead = "System`Plus"
 	return fmt.Sprintf("(%v + %v*I)", cmplx.Re.StringForm(p), cmplx.Im.StringForm(p))
 }

diff --git a/expreduce/builtin_power.go b/expreduce/builtin_power.go
@@ -1,6 +1,7 @@
 package expreduce
 
 import (
+	"fmt"
 	"math"
 	"math/big"
 
@@ -122,6 +123,18 @@ func getPowerDefinitions() (defs []Definition) {
 		Name:    "Power",
 		Default: "1",
 		toString: func(this expreduceapi.ExpressionInterface, params expreduceapi.ToStringParams) (bool, string) {
+			if this.Len() == 2 {
+				if atoms.IsSameQ(this.GetPart(2), atoms.NewRational(big.NewInt(1), big.NewInt(2))) {
+					nextParams := params
+					nextParams.PreviousHead = "<TOPLEVEL>"
+					if params.Form == "TeXForm" {
+						return true, fmt.Sprintf("\\sqrt{%v}", this.GetPart(1).StringForm(nextParams))
+					}
+					if params.Form == "InputForm" {
+						return true, fmt.Sprintf("Sqrt[%v]", this.GetPart(1).StringForm(nextParams))
+					}
+				}
+			}
 			return toStringInfixAdvanced(this.GetParts()[1:], "^", "System`Power", false, "", "", params)
 		},
 		legacyEvalFn: func(this expreduceapi.ExpressionInterface, es expreduceapi.EvalStateInterface) expreduceapi.Ex {