add flag for reordering matrix

cmd/root.go

@@ -17,6 +17,7 @@ along with this program. If not, see <http://www.gnu.org/licenses/>.
 package cmd
 
 import (
+	"errors"
 	"fmt"
 	"os"
 
@@ -34,6 +35,8 @@ var pythonFmt bool
 
 var oneline bool
 
+var nodeOrder []string
+
 // rootCmd represents the base command when called without any subcommands
 var rootCmd = &cobra.Command{
 	Use:   "dptdist",
@@ -77,7 +80,11 @@ func parse(cmd *cobra.Command, args []string) {
 	}
 	fmtOptions = append(fmtOptions, mat.Squeeze())
 
-	a := graph.AdjacencyMatrix()
+	a, err := orderedAdjMatrix(graph)
+	if err != nil {
+		panic(err)
+	}
+
 	out := mat.Formatted(a, fmtOptions...)
 
 	// for matlab and python formats, %#v outputs as matrix and %v is
@@ -89,6 +96,36 @@ func parse(cmd *cobra.Command, args []string) {
 	}
 }
 
+func orderedAdjMatrix(g *markov.AbsorbingMarkovChain) (*mat.Dense, error) {
+	matrix := g.AdjacencyMatrix()
+	if len(nodeOrder) == 0 {
+		return matrix, nil
+	}
+
+	nodeIndexes := make(map[string]int)
+	for i := 0; i < len(nodeOrder); i++ {
+		nodeIndexes[nodeOrder[i]] = i
+	}
+
+	nodes := g.Nodes()
+	newOrder := make([]int, nodes.Len())
+	for nodes.Next() {
+		node := nodes.Node().(*markov.Node)
+		id := node.DOTID()
+
+		var ok bool
+		newOrder[node.ID()], ok = nodeIndexes[id]
+		if !ok {
+			return nil, errors.New("node '" + id + "' not in given order")
+		}
+	}
+
+	matrix.PermuteRows(newOrder, true)
+	matrix.PermuteCols(newOrder, true)
+
+	return matrix, nil
+}
+
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
@@ -115,4 +152,6 @@ func init() {
 	rootCmd.MarkFlagsMutuallyExclusive("matlab", "python")
 
 	rootCmd.Flags().BoolVar(&oneline, "oneline", false, "output on one line")
+
+	rootCmd.Flags().StringSliceVarP(&nodeOrder, "order", "o", nil, "order of nodes in rows/columns of output")
 }