Merge pull request #5231 from FlorentinD/config-trainer-method-cleanup

Simplify TrainerConfig::methodName

Author: Mats-SX

ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/ClassifierFactory.java

@@ -75,7 +75,7 @@ public static MemoryEstimation dataMemoryEstimation(
         int featureDimension,
         boolean isReduced
     ) {
-        switch (TrainingMethod.valueOf(trainerConfig.methodName())) {
+        switch (trainerConfig.method()) {
             case LogisticRegression:
                 return LogisticRegressionData.memoryEstimation(isReduced, numberOfClasses, MemoryRange.of(featureDimension));
             case RandomForest:

ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/ClassifierTrainerFactory.java

@@ -32,7 +32,7 @@
 import java.util.Optional;
 import java.util.function.LongUnaryOperator;
 
-public class ClassifierTrainerFactory {
+public final class ClassifierTrainerFactory {
 
     private ClassifierTrainerFactory() {}
 
@@ -45,7 +45,7 @@ public static ClassifierTrainer create(
         Optional<Long> randomSeed,
         boolean reduceClassCount
     ) {
-        switch (TrainingMethod.valueOf(config.methodName())) {
+        switch (config.method()) {
             case LogisticRegression: {
                 return new LogisticRegressionTrainer(
                     concurrency,
@@ -79,7 +79,7 @@ public static MemoryEstimation memoryEstimation(
         MemoryRange featureDimension,
         boolean isReduced
     ) {
-        switch (TrainingMethod.valueOf(config.methodName())) {
+        switch (config.method()) {
             case LogisticRegression:
                 return LogisticRegressionTrainer.memoryEstimation(
                     isReduced,

ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/RegressionTrainerFactory.java

@@ -41,7 +41,7 @@ public static RegressorTrainer create(
         int concurrency,
         Optional<Long> randomSeed
     ) {
-        switch (TrainingMethod.valueOf(config.methodName())) {
+        switch (config.method()) {
             case LinearRegression: {
                 return new LinearRegressionTrainer(
                     concurrency,
@@ -59,7 +59,7 @@ public static RegressorTrainer create(
                 );
             }
             default:
-                throw new IllegalStateException(formatWithLocale("Method %s is not a regression method", config.methodName()));
+                throw new IllegalStateException(formatWithLocale("Method %s is not a regression method", config.method()));
         }
     }
 }

ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/TrainerConfig.java

@@ -30,18 +30,18 @@
 public interface TrainerConfig extends ToMapConvertible {
 
     @Configuration.Ignore
-    String methodName();
+    TrainingMethod method();
 
     @Value.Derived
     @Configuration.Ignore
     default TunableTrainerConfig toTunableConfig() {
-        return TunableTrainerConfig.of(toMap(), TrainingMethod.valueOf(methodName()));
+        return TunableTrainerConfig.of(toMap(), method());
     }
 
     @Configuration.Ignore
     default Map<String, Object> toMapWithTrainerMethod() {
         var mapWithTrainerMethod = new HashMap<>(toMap());
-        mapWithTrainerMethod.put("methodName", methodName());
+        mapWithTrainerMethod.put("methodName", method().name());
 
         return mapWithTrainerMethod;
     }

ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/linearregression/LinearRegressionTrainConfig.java

@@ -52,7 +52,7 @@ static LinearRegressionTrainConfig of(Map<String, Object> params) {
 
     @Override
     @Configuration.Ignore
-    default String methodName() {
-        return TrainingMethod.LinearRegression.name();
+    default TrainingMethod method() {
+        return TrainingMethod.LinearRegression;
     }
 }

ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/logisticregression/LogisticRegressionTrainConfig.java

@@ -52,7 +52,7 @@ static LogisticRegressionTrainConfig of(Map<String, Object> params) {
 
     @Override
     @Configuration.Ignore
-    default String methodName() {
-        return TrainingMethod.LogisticRegression.name();
+    default TrainingMethod method() {
+        return TrainingMethod.LogisticRegression;
     }
 }

ml/ml-algo/src/main/java/org/neo4j/gds/ml/models/randomforest/RandomForestTrainerConfig.java

@@ -55,8 +55,8 @@ default int numberOfDecisionTrees() {
 
     @Override
     @Configuration.Ignore
-    default String methodName() {
-        return TrainingMethod.RandomForest.name();
+    default TrainingMethod method() {
+        return TrainingMethod.RandomForest;
     }
 
     static RandomForestTrainerConfig of(Map<String, Object> params) {

ml/ml-algo/src/main/java/org/neo4j/gds/ml/nodeClassification/ClassificationMetricComputer.java

@@ -29,7 +29,6 @@
 import org.neo4j.gds.ml.models.ClassifierFactory;
 import org.neo4j.gds.ml.models.Features;
 import org.neo4j.gds.ml.models.TrainerConfig;
-import org.neo4j.gds.ml.models.TrainingMethod;
 import org.openjdk.jol.util.Multiset;
 
 import java.util.function.LongUnaryOperator;
@@ -120,7 +119,7 @@ public static MemoryEstimation estimateEvaluation(
             .rangePerNode(
                 "classifier runtime",
                 nodeCount -> ClassifierFactory.runtimeOverheadMemoryEstimation(
-                    TrainingMethod.valueOf(config.methodName()),
+                    config.method(),
                     batchSize,
                     fudgedClassCount,
                     fudgedFeatureCount,

ml/ml-algo/src/test/java/org/neo4j/gds/ml/metrics/StatsMapTest.java

@@ -22,6 +22,7 @@
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.core.GraphDimensions;
 import org.neo4j.gds.ml.models.TrainerConfig;
+import org.neo4j.gds.ml.models.TrainingMethod;
 
 import java.util.List;
 import java.util.Map;
@@ -79,18 +80,18 @@ void toMap() {
 
     private static final class TestTrainerConfig implements TrainerConfig {
 
-        private final String method;
+        private final String name;
 
-        private TestTrainerConfig(String method) {this.method = method;}
+        private TestTrainerConfig(String name) {this.name = name;}
 
         @Override
         public Map<String, Object> toMap() {
-            return Map.of("method", method);
+            return Map.of("name", name);
         }
 
         @Override
-        public String methodName() {
-            return method;
+        public TrainingMethod method() {
+            return TrainingMethod.RandomForest;
         }
     }
 

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/TrainingPipeline.java

@@ -129,7 +129,7 @@ public void addTrainerConfig(TunableTrainerConfig trainingConfig) {
     }
 
     public void addTrainerConfig(TrainerConfig trainingConfig) {
-        this.trainingParameterSpace.get(TrainingMethod.valueOf(trainingConfig.methodName())).add(trainingConfig.toTunableConfig());
+        this.trainingParameterSpace.get(trainingConfig.method()).add(trainingConfig.toTunableConfig());
     }
 
     public AutoTuningConfig autoTuningConfig() {