dirichlet WIP
@@ -4,11 +4,15 @@ import io.improbable.keanu.annotation.ExportVertexToPythonBindings;
import io.improbable.keanu.distributions.continuous.Dirichlet;
import io.improbable.keanu.distributions.hyperparam.Diffs;
import io.improbable.keanu.tensor.dbl.DoubleTensor;
+import io.improbable.keanu.vertices.ConstantVertex;
import io.improbable.keanu.vertices.LoadShape;
import io.improbable.keanu.vertices.LoadVertexParam;
+import io.improbable.keanu.vertices.LogProbGraph;
+import io.improbable.keanu.vertices.LogProbGraphSupplier;
import io.improbable.keanu.vertices.SamplableWithManyScalars;
import io.improbable.keanu.vertices.SaveVertexParam;
import io.improbable.keanu.vertices.Vertex;
+import io.improbable.keanu.vertices.bool.BoolVertex;
import io.improbable.keanu.vertices.dbl.Differentiable;
import io.improbable.keanu.vertices.dbl.DoubleVertex;
import io.improbable.keanu.vertices.dbl.KeanuRandom;
@@ -21,9 +25,11 @@ import java.util.Set;
import static io.improbable.keanu.distributions.hyperparam.Diffs.C;
import static io.improbable.keanu.distributions.hyperparam.Diffs.X;
-public class DirichletVertex extends DoubleVertex implements Differentiable, ProbabilisticDouble, SamplableWithManyScalars<DoubleTensor> {
+public class DirichletVertex extends DoubleVertex implements Differentiable, ProbabilisticDouble, SamplableWithManyScalars<DoubleTensor>, LogProbGraphSupplier {
private final DoubleVertex concentration;
+ private static final double EPSILON = 0.00001;
+
private static final String CONCENTRATION_NAME = "concentration";
/**
@@ -83,6 +89,28 @@ public class DirichletVertex extends DoubleVertex implements Differentiable, Pro
return logPdfs.sum();
}
+ @Override
+ public LogProbGraph logProbGraph() {
+ LogProbGraph.DoublePlaceHolderVertex xPlaceHolder = new LogProbGraph.DoublePlaceHolderVertex(this.getShape());
+ LogProbGraph.DoublePlaceHolderVertex concentrationPlaceHolder = new LogProbGraph.DoublePlaceHolderVertex(concentration.getShape());
+
+ final BoolVertex assertXIsGreaterThanEpsilon = xPlaceHolder
+ .sum().minus(1.).abs().greaterThan(ConstantVertex.of(EPSILON))
+ .assertTrue("Sum of values to calculate Dirichlet likelihood for must equal 1");
+
+ final DoubleVertex sumConcentrationLogged = concentrationPlaceHolder.minus(1.).times(xPlaceHolder.log()).sum();
+ final DoubleVertex sumLogGammaConcentration = concentrationPlaceHolder.logGamma().sum();
+ final DoubleVertex logGammaSumConcentration = concentrationPlaceHolder.sum().logGamma();
+
+ final DoubleVertex logProbOutput = sumConcentrationLogged.minus(sumLogGammaConcentration).plus(logGammaSumConcentration);
+
+ return LogProbGraph.builder()
+ .input(this, xPlaceHolder)
+ .input(concentration, concentrationPlaceHolder)
+ .logProbOutput(logProbOutput)
+ .build();
+ }
+
@Override
public Map<Vertex, DoubleTensor> dLogProb(DoubleTensor value, Set<? extends Vertex> withRespectTo) {
Diffs dlnP = Dirichlet.withParameters(concentration.getValue()).dLogProb(value);
