Refactored the encoding of Segmentation elements. Fixes jpmml/jpmml-s…

…klearn#171
jpmml · Aug 20, 2022 · f354c19 · f354c19
1 parent 9933ebc
commit f354c19
Showing 1 changed file with 9 additions and 12 deletions.
diff --git a/pmml-converter/src/main/java/org/jpmml/converter/mining/MiningModelUtil.java b/pmml-converter/src/main/java/org/jpmml/converter/mining/MiningModelUtil.java
@@ -175,11 +175,8 @@ public MiningModel createModelChain(List<? extends Model> models, Segmentation.M
 
 		MiningFunction miningFunction = lastModel.requireMiningFunction();
 
-		Segmentation segmentation = createSegmentation(Segmentation.MultipleModelMethod.MODEL_CHAIN, models)
-			.setMissingPredictionTreatment(missingPredictionTreatment);
-
 		MiningModel miningModel = new MiningModel(miningFunction, createMiningSchema(models))
-			.setSegmentation(segmentation);
+			.setSegmentation(createSegmentation(Segmentation.MultipleModelMethod.MODEL_CHAIN, missingPredictionTreatment, models));
 
 		return miningModel;
 	}
@@ -211,11 +208,8 @@ public MiningModel createMultiModelChain(List<? extends Model> models, Segmentat
 			}
 		}
 
-		Segmentation segmentation = createSegmentation(Segmentation.MultipleModelMethod.MULTI_MODEL_CHAIN, models)
-			.setMissingPredictionTreatment(missingPredictionTreatment);
-
 		MiningModel miningModel = new MiningModel(miningFunction, createMiningSchema(models))
-			.setSegmentation(segmentation);
+			.setSegmentation(createSegmentation(Segmentation.MultipleModelMethod.MULTI_MODEL_CHAIN, missingPredictionTreatment, models));
 
 		return miningModel;
 	}
@@ -248,12 +242,12 @@ public MiningSchema createMiningSchema(List<? extends Model> models){
 	}
 
 	static
-	public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, List<? extends Model> models){
-		return createSegmentation(multipleModelMethod, models, null);
+	public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, Segmentation.MissingPredictionTreatment missingPredictionTreatment, List<? extends Model> models){
+		return createSegmentation(multipleModelMethod, missingPredictionTreatment, models, null);
 	}
 
 	static
-	public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, List<? extends Model> models, List<? extends Number> weights){
+	public Segmentation createSegmentation(Segmentation.MultipleModelMethod multipleModelMethod, Segmentation.MissingPredictionTreatment missingPredictionTreatment, List<? extends Model> models, List<? extends Number> weights){
 
 		if((weights != null) && (models.size() != weights.size())){
 			throw new IllegalArgumentException();
@@ -275,7 +269,10 @@ public Segmentation createSegmentation(Segmentation.MultipleModelMethod multiple
 			segments.add(segment);
 		}
 
-		return new Segmentation(multipleModelMethod, segments);
+		Segmentation segmentation = new Segmentation(multipleModelMethod, segments)
+			.setMissingPredictionTreatment(missingPredictionTreatment);
+
+		return segmentation;
 	}
 
 	static