diff --git a/cmd/root.go b/cmd/root.go index a43fd45..a404c75 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -17,6 +17,7 @@ along with this program. If not, see . 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") }