allow NewAntigenicLikelihood to use FastMatrixParameter

beast-dev · Dec 14, 2023 · 913f15b · 913f15b
1 parent 837ffff
commit 913f15b
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 13 deletions.
diff --git a/src/dr/evolution/coalescent/CoalescentGradient.java b/src/dr/evolution/coalescent/CoalescentGradient.java
@@ -55,7 +55,7 @@ public class CoalescentGradient implements GradientWrtParameterProvider, Reporta
 
     public enum Wrt {
         NODE_HEIGHTS,
-        PARAMETER;
+        PARAMETER
     }
 
     public CoalescentGradient(CoalescentLikelihood likelihood,

diff --git a/src/dr/evomodel/antigenic/NewAntigenicLikelihood.java b/src/dr/evomodel/antigenic/NewAntigenicLikelihood.java
@@ -80,7 +80,7 @@ public NewAntigenicLikelihood(
             Parameter virusDriftParameter,
             Parameter serumDriftParameter,
             MatrixParameter virusSamplingParameter, // TODO Remove
-            MatrixParameter serumLocationsParameter,
+            MatrixParameterInterface serumLocationsParameter,
             CompoundParameter tipTraitsParameter,
             Parameter virusOffsetsParameter,
             Parameter serumOffsetsParameter,
@@ -651,9 +651,20 @@ private Parameter setupSerumBreadths(Parameter serumBreadthsParameter) {
         return serumBreadthsParameter;
     }
 
-    protected void setupLocationsParameter(MatrixParameter locationsParameter, List<String> strains) {
-        locationsParameter.setColumnDimension(mdsDimension);
-        locationsParameter.setRowDimension(strains.size());
+    protected void setupLocationsParameter(MatrixParameterInterface locationsParameter, List<String> strains) {
+        if (locationsParameter instanceof MatrixParameter) {
+            ((MatrixParameter) locationsParameter).setColumnDimension(mdsDimension);
+            ((MatrixParameter) locationsParameter).setRowDimension(strains.size());
+        } else if (locationsParameter instanceof FastMatrixParameter) {
+            FastMatrixParameter fmp = (FastMatrixParameter) locationsParameter;
+            if (fmp.getRowDimension() != mdsDimension) {
+                throw new IllegalArgumentException("Column dim must be " + mdsDimension);
+            }
+            if (fmp.getColumnDimension() != strains.size()) {
+                throw new IllegalArgumentException("Row dim must be " + strains.size());
+            }
+        }
+
         for (int i = 0; i < strains.size(); i++) {
             locationsParameter.getParameter(i).setId(strains.get(i));
         }
@@ -1116,7 +1127,7 @@ private Measurement(final int virus, final int serum, final double virusDate, fi
     private final Parameter serumDriftParameter;
 
     private final MatrixParameter virusSamplingParameter;
-    private final MatrixParameter serumLocationsParameter;
+    private final MatrixParameterInterface serumLocationsParameter;
 
     private final Parameter virusOffsetsParameter;
     private final Parameter serumOffsetsParameter;

diff --git a/src/dr/evomodelxml/antigenic/AntigenicLikelihoodParser.java b/src/dr/evomodelxml/antigenic/AntigenicLikelihoodParser.java
@@ -1,10 +1,7 @@
 package dr.evomodelxml.antigenic;
 
 import dr.evomodel.antigenic.NewAntigenicLikelihood;
-import dr.inference.model.AbstractTransformedCompoundMatrix;
-import dr.inference.model.CompoundParameter;
-import dr.inference.model.MatrixParameter;
-import dr.inference.model.Parameter;
+import dr.inference.model.*;
 import dr.util.Citable;
 import dr.util.DataTable;
 import dr.xml.*;
@@ -78,9 +75,9 @@ public Object parseXMLObject(XMLObject xo) throws XMLParseException {
             // TOD Remove
         }
 
-        MatrixParameter serumLocationsParameter = null;
+        MatrixParameterInterface serumLocationsParameter = null;
         if (xo.hasChildNamed(SERUM_LOCATIONS)) {
-            serumLocationsParameter = (MatrixParameter) xo.getElementFirstChild(SERUM_LOCATIONS);
+            serumLocationsParameter = (MatrixParameterInterface) xo.getElementFirstChild(SERUM_LOCATIONS);
         }
 
         Parameter mdsPrecision = (Parameter) xo.getElementFirstChild(MDS_PRECISION);
@@ -173,7 +170,7 @@ public XMLSyntaxRule[] getSyntaxRules() {
             AttributeRule.newDoubleRule(DRIFT_INITIAL_LOCATIONS, true, "The degree to drift initial virus and serum locations, defaults to 0.0"),
             new ElementRule(TIP_TRAIT, CompoundParameter.class, "Optional parameter of tip locations from the tree", true),
 //                new ElementRule(VIRUS_LOCATIONS, MatrixParameter.class, "Parameter of locations of all virus"),
-            new ElementRule(SERUM_LOCATIONS, MatrixParameter.class, "Parameter of locations of all sera"),
+            new ElementRule(SERUM_LOCATIONS, MatrixParameterInterface.class, "Parameter of locations of all sera"),
             new ElementRule(VIRUS_OFFSETS, Parameter.class, "Optional parameter for virus dates to be stored", true),
             new ElementRule(SERUM_OFFSETS, Parameter.class, "Optional parameter for serum dates to be stored", true),
             new ElementRule(SERUM_POTENCIES, Parameter.class, "Optional parameter for serum potencies", true),

diff --git a/src/dr/inference/model/FastMatrixParameter.java b/src/dr/inference/model/FastMatrixParameter.java
@@ -76,6 +76,14 @@ public FastMatrixParameter(String id, List<Parameter> original, boolean signalCo
         fireParameterChangedEvent(-1, ChangeType.ALL_VALUES_CHANGED);
     }
 
+    @Override
+    public String getDimensionName(int dim) {
+        int pNum = dim / getRowDimension();
+        int index = dim % getRowDimension();
+        String name = getParameter(pNum).getParameterName() + (index + 1);
+        return name;
+    }
+
     private void checkParameterLengths(List<Parameter> parameters) {
         final int length = parameters.get(0).getDimension();
         for (Parameter p : parameters) {
@@ -93,6 +101,17 @@ private void setProxyParameterNames(List<Parameter> original) {
         }
     }
 
+    private void setProxyParameterName(String name, int column) {
+        if (proxyParameterNames == null) {
+            proxyParameterNames = new ArrayList<>();
+            for (int i = 0; i < getColumnDimension(); ++i) {
+                proxyParameterNames.add(null);
+            }
+        }
+
+        proxyParameterNames.set(column, name);
+    }
+
     private List<String> proxyParameterNames;
 
     private String getProxyParameterName(int column) {
@@ -169,6 +188,11 @@ public void setParameterValueNotifyChangedAll(int dim, double value) {
             throw new RuntimeException("Do not call");
         }
 
+        @Override
+        public void setId(String name) {
+            matrix.setProxyParameterName(name, column);
+        }
+
         @Override
         public String getParameterName() {
             String proxyName = matrix.getProxyParameterName(column);

diff --git a/src/dr/inference/model/MatrixParameterInterface.java b/src/dr/inference/model/MatrixParameterInterface.java
@@ -51,6 +51,10 @@ public interface MatrixParameterInterface extends Parameter {
 
     double[] getParameterValues();
 
+    default int getParameterCount() {
+        throw new RuntimeException("Not yet implemented");
+    }
+
     int getUniqueParameterCount();
 
     Parameter getUniqueParameter(int index);