diff --git a/src/main/java/gov/usgs/earthquake/nshmp/model/GmmTree.java b/src/main/java/gov/usgs/earthquake/nshmp/model/GmmTree.java index 44d3dc948f1a5d522900356192c6f406e97fafa6..b92ad9407d8ab2d7ba2b7b5e80396853f511c771 100644 --- a/src/main/java/gov/usgs/earthquake/nshmp/model/GmmTree.java +++ b/src/main/java/gov/usgs/earthquake/nshmp/model/GmmTree.java @@ -1,9 +1,11 @@ package gov.usgs.earthquake.nshmp.model; +import static java.util.stream.Collectors.toUnmodifiableMap; + +import java.util.Map; import java.util.Optional; import java.util.OptionalDouble; import java.util.Set; -import java.util.stream.Collectors; import com.google.gson.JsonElement; @@ -37,7 +39,7 @@ public class GmmTree { */ private final LogicTree<Gmm> tree; - private final Set<Gmm> gmms; + private final Map<Gmm, Double> gmms; /* GmmConfig fields. */ private final OptionalDouble maxDistance; @@ -45,8 +47,9 @@ public class GmmTree { private GmmTree(LogicTree<Gmm> tree, Optional<GmmConfig> config) { this.tree = tree; this.gmms = tree.stream() - .map(Branch::value) - .collect(Collectors.toUnmodifiableSet()); + .collect(toUnmodifiableMap( + Branch::value, + Branch::weight)); this.maxDistance = config.isPresent() ? OptionalDouble.of(config.orElseThrow().maxDistance) : OptionalDouble.empty(); @@ -64,7 +67,15 @@ public class GmmTree { /** The set of GMMs in the logic tree. */ public Set<Gmm> gmms() { - return gmms; + return gmms.keySet(); + } + + /** + * The logic tree weight of the requested GMM. + * @throws NullPointerException if the supplied GMM is not in this tree + */ + public double gmmWeight(Gmm gmm) { + return gmms.get(gmm); } /** The maximum applicable distance for the GMMs in the logic tree. */