Merge pull request #152 from mumax/mumax3-convert/arrows-vs-scalar-color

allow to combine arrows with scalar colormaps
mumax · Jan 25, 2018 · e2c09f4 · e2c09f4
2 parents fd1a693 + 530ef1c
commit e2c09f4
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 27 deletions.
diff --git a/cmd/mumax3-convert/main.go b/cmd/mumax3-convert/main.go
@@ -82,7 +82,7 @@ var (
 )
 
 var (
-	colormap []color.RGBA
+	colormap []draw.ColorMapSpec
 )
 
 type task struct {
@@ -98,7 +98,8 @@ func main() {
 		log.Fatal("no input files")
 	}
 
-	colormap = parseColors(*flag_color)
+	colormap = make([]draw.ColorMapSpec, 1, 1)
+	colormap[0].Cmap = parseColors(*flag_color)
 
 	// politely try to make the output directory
 	if *flag_dir != "" {
@@ -324,8 +325,13 @@ func preprocess(f *data.Slice) {
 	if *flag_normpeak {
 		normpeak(f)
 	}
+	colormap[0].Ccomp = -1
 	if *flag_comp != "" {
-		*f = *f.Comp(parseComp(*flag_comp))
+		c := parseComp(*flag_comp)
+		colormap[0].Ccomp = c
+		if *flag_arrows == 0 {
+			*f = *f.Comp(c)
+		}
 	}
 	crop(f)
 	if *flag_resize != "" {

diff --git a/draw/colorscale.go b/draw/colorscale.go
@@ -3,6 +3,11 @@ package draw
 import "image/color"
 import "fmt"
 
+type ColorMapSpec struct {
+	Cmap  []color.RGBA
+	Ccomp int
+}
+
 func ColorMap(min, max, value float32, colormap ...color.RGBA) color.RGBA {
 	// default colormap: black-white
 	if len(colormap) < 1 {

diff --git a/draw/encode.go b/draw/encode.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"github.com/mumax/3/data"
 	"image"
-	"image/color"
 	"image/gif"
 	"image/jpeg"
 	"image/png"
@@ -15,7 +14,7 @@ import (
 	"strings"
 )
 
-func RenderFile(fname string, f *data.Slice, min, max string, arrowSize int, colormap ...color.RGBA) error {
+func RenderFile(fname string, f *data.Slice, min, max string, arrowSize int, colormap ...ColorMapSpec) error {
 	out, err := os.Create(fname)
 	if err != nil {
 		return err
@@ -24,7 +23,7 @@ func RenderFile(fname string, f *data.Slice, min, max string, arrowSize int, col
 	return RenderFormat(out, f, min, max, arrowSize, fname, colormap...)
 }
 
-func RenderFormat(out io.Writer, f *data.Slice, min, max string, arrowSize int, format string, colormap ...color.RGBA) error {
+func RenderFormat(out io.Writer, f *data.Slice, min, max string, arrowSize int, format string, colormap ...ColorMapSpec) error {
 	var codecs = map[string]codec{".png": PNG, ".jpg": JPEG100, ".gif": GIF256}
 	ext := strings.ToLower(path.Ext(format))
 	enc := codecs[ext]
@@ -38,7 +37,7 @@ func RenderFormat(out io.Writer, f *data.Slice, min, max string, arrowSize int,
 type codec func(io.Writer, image.Image) error
 
 // Render data and encode with arbitrary codec.
-func Render(out io.Writer, f *data.Slice, min, max string, arrowSize int, encode codec, colormap ...color.RGBA) error {
+func Render(out io.Writer, f *data.Slice, min, max string, arrowSize int, encode codec, colormap ...ColorMapSpec) error {
 	img := Image(f, min, max, arrowSize, colormap...)
 	buf := bufio.NewWriter(out)
 	defer buf.Flush()

diff --git a/draw/image.go b/draw/image.go
@@ -10,42 +10,57 @@ import (
 )
 
 // Renders an image of slice. fmin, fmax = "auto" or a number to set the min/max color scale.
-func Image(f *data.Slice, fmin, fmax string, arrowSize int, colormap ...color.RGBA) *image.RGBA {
+func Image(f *data.Slice, fmin, fmax string, arrowSize int, colormap ...ColorMapSpec) *image.RGBA {
 	img := new(image.RGBA)
 	On(img, f, fmin, fmax, arrowSize, colormap...)
 	return img
 }
 
 // Render on existing image buffer. Resize it if needed
-func On(img *image.RGBA, f *data.Slice, fmin, fmax string, arrowSize int, colormap ...color.RGBA) {
+func On(img *image.RGBA, f *data.Slice, fmin, fmax string, arrowSize int, colormap ...ColorMapSpec) {
+	cm := colormap[0]
 	dim := f.NComp()
 	switch dim {
 	default:
 		log.Fatalf("unsupported number of components: %v", dim)
 	case 3:
-		drawVectors(img, f.Vectors(), arrowSize)
-	case 1:
-		min, max := extrema(f.Host()[0])
-		if fmin != "auto" {
-			m, err := strconv.ParseFloat(fmin, 32)
-			if err != nil {
-				util.Fatal("draw: scale:", err)
+		if cm.Ccomp >= 0 {
+			ff := f.Comp(cm.Ccomp)
+			min, max := parseMinMax(ff, fmin, fmax)
+			drawFloats(img, ff.Scalars(), min, max, cm.Cmap...)
+			if arrowSize > 0 {
+				drawArrows(img, f.Vectors(), arrowSize)
 			}
-			min = float32(m)
+		} else {
+			drawVectors(img, f.Vectors(), arrowSize)
 		}
-		if fmax != "auto" {
-			m, err := strconv.ParseFloat(fmax, 32)
-			if err != nil {
-				util.Fatal("draw: scale:", err)
-			}
-			max = float32(m)
+	case 1:
+		min, max := parseMinMax(f, fmin, fmax)
+		drawFloats(img, f.Scalars(), min, max, cm.Cmap...)
+	}
+}
+
+func parseMinMax(f *data.Slice, fmin, fmax string) (min, max float32) {
+	min, max = extrema(f.Host()[0])
+	if fmin != "auto" {
+		m, err := strconv.ParseFloat(fmin, 32)
+		if err != nil {
+			util.Fatal("draw: scale:", err)
 		}
-		if min == max {
-			min -= 1
-			max += 1 // make it gray instead of black
+		min = float32(m)
+	}
+	if fmax != "auto" {
+		m, err := strconv.ParseFloat(fmax, 32)
+		if err != nil {
+			util.Fatal("draw: scale:", err)
 		}
-		drawFloats(img, f.Scalars(), min, max, colormap...)
+		max = float32(m)
 	}
+	if min == max {
+		min -= 1
+		max += 1 // make it gray instead of black
+	}
+	return
 }
 
 // Draws rank 4 tensor (3D vector field) as image