diff --git a/prism/src/pta/PTAAbstractRefine.java b/prism/src/pta/PTAAbstractRefine.java
index 92463e74..07bc196d 100644
--- a/prism/src/pta/PTAAbstractRefine.java
+++ b/prism/src/pta/PTAAbstractRefine.java
@@ -166,7 +166,7 @@ public class PTAAbstractRefine extends STPGAbstractRefine
 
 		// Check for deadlocks in the constructed STPG
 		// (should never occur because of the restrictions we impose on PTAs)
-		if (((STPG) abstraction).steps.get(src).size() == 0) {
+		if (abstraction.getNumChoices(src) == 0) {
 			throw new PrismException("STPG has deadlock in state #" + src + ":" + graph.states.get(src));
 		}
 	}
@@ -298,7 +298,7 @@ public class PTAAbstractRefine extends STPGAbstractRefine
 			for (int in : choiceList) {
 				// If using BitSets for action labels (i.e. have not cached zones for reuse)
 				// need to build the corresponding zone
-				action = ((STPG) abstraction).steps.get(splitState).get(in).action;
+				action = ((STPG) abstraction).getChoice(splitState, in).action;
 				if (!storeValidZones) {
 					actionBitSet = (BitSet) action;
 					List<SymbolicTransition> sts = graph.trans.get(splitState);
diff --git a/prism/src/pta/PTAExpected.java b/prism/src/pta/PTAExpected.java
index ac51a9a0..e6a9c450 100644
--- a/prism/src/pta/PTAExpected.java
+++ b/prism/src/pta/PTAExpected.java
@@ -332,7 +332,7 @@ public class PTAExpected
 				if (distr.isEmpty())
 					continue;
 				// Add distribution
-				choice = mdp.addDistribution(src, distr);
+				choice = mdp.addChoice(src, distr);
 				// Compute reward
 				if (min) {
 					rewSum = 0.0;
@@ -406,7 +406,7 @@ public class PTAExpected
 		// Add transition + reward
 		distr = new Distribution();
 		distr.add(initialStates.get(0), 1.0);
-		mdp.addDistribution(src, distr);
+		mdp.addChoice(src, distr);
 		if (!min)
 			mdp.setTransitionReward(src, 0, getMinMaxForZone(graph.states.get(initialStates.get(0)).zone, someClock,
 					min));
diff --git a/prism/src/pta/ReachabilityGraph.java b/prism/src/pta/ReachabilityGraph.java
index 54c24843..070c4208 100644
--- a/prism/src/pta/ReachabilityGraph.java
+++ b/prism/src/pta/ReachabilityGraph.java
@@ -198,7 +198,7 @@ public class ReachabilityGraph
 					distr.add(dest, edge.getProbability());
 				}
 				if (!distr.isEmpty())
-					mdp.addDistribution(src, distr);
+					mdp.addChoice(src, distr);
 			}
 		}