Improving GUI robustness wrt deme count.

Author: tgvaughan

src/beast/app/multitypetree/beauti/InitMigrationModelConnector.java

@@ -23,8 +23,11 @@
 import beast.evolution.likelihood.TreeLikelihood;
 import beast.evolution.tree.SCMigrationModel;
 import beast.evolution.tree.StructuredCoalescentMultiTypeTree;
+import beast.evolution.tree.TraitSet;
 
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 
 /**
@@ -37,6 +40,15 @@
  * @author Tim Vaughan ([email protected])
  */
 public class InitMigrationModelConnector {
+
+    public static int uniqueTraitsInData(StructuredCoalescentMultiTypeTree scTree) {
+        Set<String> uniqueTypes = new HashSet<>();
+        TraitSet typeTraitSet = scTree.typeTraitInput.get();
+        for (String taxonName : typeTraitSet.taxaInput.get().getTaxaNames())
+            uniqueTypes.add(typeTraitSet.getStringValue(taxonName));
+
+        return uniqueTypes.size();
+    }
 
     public static boolean customConnector(BeautiDoc doc) {
 
@@ -56,6 +68,37 @@ public static boolean customConnector(BeautiDoc doc) {
             String rateMatrixStr = getParameterString((RealParameter)migModel.rateMatrixInput.get());
             String popSizesStr = getParameterString((RealParameter)migModel.popSizesInput.get());
 
+            // Ensure model has minimum number of demes
+            int uniqueTraitCount = uniqueTraitsInData(tree);
+            StringBuilder rateMatrixStrBuilder = new StringBuilder();
+            StringBuilder popSizesStrBuilder = new StringBuilder();
+            if (migModel.getNTypes()<uniqueTraitCount) {
+                for (int i=0; i<uniqueTraitCount; i++) {
+                    popSizesStrBuilder.append(" 1.0");
+                    for (int j=0; j<uniqueTraitCount; j++) {
+                        if (j == i)
+                            continue;
+
+                        rateMatrixStrBuilder.append(" 1.0");
+                    }
+                }
+
+                popSizesStr = popSizesStrBuilder.toString();
+                rateMatrixStr = rateMatrixStrBuilder.toString();
+
+                ((RealParameter)migModel.popSizesInput.get()).setDimension(uniqueTraitCount);
+                ((RealParameter)migModel.popSizesInput.get()).valuesInput.setValue(popSizesStr,
+                        (RealParameter)migModel.popSizesInput.get());
+
+                ((RealParameter)migModel.rateMatrixInput.get()).setDimension(uniqueTraitCount*(uniqueTraitCount-1));
+                ((RealParameter)migModel.rateMatrixInput.get()).valuesInput.setValue(rateMatrixStr,
+                        (RealParameter)migModel.rateMatrixInput.get());
+
+                ((RealParameter)migModel.popSizesInput.get()).initAndValidate();
+                ((RealParameter)migModel.rateMatrixInput.get()).initAndValidate();
+                migModel.initAndValidate();
+            }
+
             ((RealParameter)migModelInit.popSizesInput.get()).setDimension(migModel.getNTypes());
 
             ((RealParameter)migModelInit.popSizesInput.get()).valuesInput.setValue(

src/beast/evolution/tree/StructuredCoalescentMultiTypeTree.java

@@ -29,9 +29,7 @@
 
 import java.io.FileNotFoundException;
 import java.io.PrintStream;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 
 /**
  * A multi-type tree generated randomly from leaf types and a migration matrix
@@ -118,7 +116,7 @@ public void initAndValidate() {
         // Obtain leaf colours from explicit input or alignment:
         leafTypes = Lists.newArrayList();
         leafNames = Lists.newArrayList();
-        if (leafTypesInput.get() != null) {            
+        if (leafTypesInput.get() != null) {
             for (int i=0; i<leafTypesInput.get().getDimension(); i++) {
                 leafTypes.add(leafTypesInput.get().getValue(i));
                 leafNames.add(String.valueOf(i));
@@ -137,7 +135,14 @@ public void initAndValidate() {
             }
 
         }
-        
+
+        // Count unique leaf types:
+        int nUniqueLeafTypes = new HashSet<>(leafTypes).size();
+        if (nUniqueLeafTypes > migModel.getNTypes())
+            throw new IllegalArgumentException("There are " + nUniqueLeafTypes
+                    + " unique leaf types but the model only includes "
+                    + migModel.getNTypes() + " unique types!");
+
         nLeaves = leafTypes.size();
 
         // Set leaf times if specified: