keep improving SVS related files #64

bioDS · Jul 18, 2024 · 44313af · 44313af
1 parent 1055d9c
commit 44313af
Showing 1 changed file with 25 additions and 10 deletions.
diff --git a/...t/src/main/java/phylonco/beast/evolution/populationmodel/StochasticVariableSelection.java b/...t/src/main/java/phylonco/beast/evolution/populationmodel/StochasticVariableSelection.java
@@ -1,34 +1,45 @@
 package phylonco.beast.evolution.populationmodel;
 
+import beast.base.core.BEASTInterface;
 import beast.base.core.Description;
 import beast.base.core.Input;
 import beast.base.core.Loggable;
 import beast.base.evolution.tree.coalescent.PopulationFunction;
-import beast.base.inference.parameter.RealParameter;
+import beast.base.inference.parameter.IntegerParameter;
 
 import java.io.PrintStream;
 import java.util.ArrayList;
 import java.util.List;
 
 @Description("Stochastic variable selection for different population growth models.")
 public class StochasticVariableSelection extends PopulationFunction.Abstract implements Loggable {
-    public final Input<RealParameter> indicatorInput = new Input<>("indicator",
+    public final Input<IntegerParameter> indicatorInput = new Input<>("indicator",
             "The indicator for selecting the population model.", Input.Validate.REQUIRED);
-    public final Input<PopulationFunction[]> modelsInput = new Input<>("models",
-            "The array of population models.", new PopulationFunction[0]);
+    public final Input<List<PopulationFunction>> modelsInput = new Input<>("models",
+            "The list of population models.", new ArrayList<>());
 
     private PopulationFunction selectedModel;
 
     @Override
     public void initAndValidate() {
+        if (indicatorInput.get() != null && indicatorInput.get() instanceof IntegerParameter) {
+            IntegerParameter IParam = indicatorInput.get();
+            IParam.setBounds(Math.max(0, IParam.getLower()), Math.min(3, IParam.getUpper()));
+        }
+
         int indicator = (int) indicatorInput.get().getArrayValue();
-        PopulationFunction[] models = modelsInput.get();
 
-        if (indicator < 0 || indicator >= models.length) {
+        List<PopulationFunction> modelsList = modelsInput.get();
+
+        if (indicator < 0 || indicator >= modelsList.size()) {
             throw new IllegalArgumentException("Invalid indicator value: " + indicator);
         }
 
-        selectedModel = models[indicator];
+        if (modelsList.size() != 4) {
+            throw new IllegalArgumentException("There must be exactly 4 population models.");
+        }
+
+        selectedModel = modelsList.get(indicator);
 
         if (selectedModel == null) {
             throw new IllegalArgumentException("Selected model is null. Indicator: " + indicator);
@@ -53,9 +64,13 @@ public double getInverseIntensity(double x) {
     @Override
     public List<String> getParameterIds() {
         List<String> ids = new ArrayList<>();
-        ids.add(indicatorInput.get().getID());
+        if (indicatorInput.get() instanceof BEASTInterface) {
+            ids.add(((BEASTInterface) indicatorInput.get()).getID());
+        }
         for (PopulationFunction model : modelsInput.get()) {
-            ids.addAll(model.getParameterIds());
+            if (model instanceof BEASTInterface) {
+                ids.add(((BEASTInterface) model).getID());
+            }
         }
         return ids;
     }
@@ -68,4 +83,4 @@ public void log(long sample, PrintStream out) {}
 
     @Override
     public void close(PrintStream out) {}
-}
+}