#5 First version for R² as additional value of the statistics aggregation response

Author: mbok

gradle.properties

@@ -4,4 +4,4 @@ wagon-ssh-external.version=2.10
 commons-math3.version=3.6.1
 group=org.scaleborn.elasticsearch.plugin
 name=elasticsearch-linear-regression
-version=5.5.1.1
+version=5.5.1.2-SNAPSHOT

src/main/java/org/scaleborn/elasticsearch/linreg/aggregation/stats/InternalStats.java

@@ -84,13 +84,23 @@ public double getMse() {
     return this.results.statistics.getMse();
   }
 
+  @Override
+  public double getR2() {
+    if (this.results == null) {
+      return Double.NaN;
+    }
+    return this.results.statistics.getR2();
+  }
+
   @Override
   public Object getDoProperty(final String element) {
     switch (element) {
       case "rss":
         return getRss();
       case "mse":
         return getMse();
+      case "r2":
+        return getR2();
     }
     return null;
   }

src/main/java/org/scaleborn/elasticsearch/linreg/aggregation/stats/StatsResults.java

@@ -34,6 +34,7 @@ static class Fields {
 
     public static final String RSS = "rss";
     public static final String MSE = "mse";
+    public static final String R2 = "r2";
   }
 
   final Statistics statistics;
@@ -46,14 +47,15 @@ public StatsResults(final SlopeCoefficients slopeCoefficients, final double inte
 
   public StatsResults(final StreamInput in) throws IOException {
     super(in);
-    this.statistics = new DefaultStatistics(in.readDouble(), in.readDouble());
+    this.statistics = new DefaultStatistics(in.readDouble(), in.readDouble(), in.readDouble());
   }
 
   @Override
   public void writeTo(final StreamOutput out) throws IOException {
     super.writeTo(out);
     out.writeDouble(this.statistics.getRss());
     out.writeDouble(this.statistics.getMse());
+    out.writeDouble(this.statistics.getR2());
   }
 
   @Override
@@ -63,6 +65,8 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa
     builder.field(Fields.RSS, this.statistics.getRss());
     // MSE
     builder.field(Fields.MSE, this.statistics.getMse());
+    // R2
+    builder.field(Fields.R2, this.statistics.getR2());
     return super.toXContent(builder, params);
   }
 

src/main/java/org/scaleborn/linereg/calculation/statistics/Statistics.java

@@ -33,17 +33,24 @@ public interface Statistics {
    */
   double getMse();
 
+  /**
+   * @return R², coefficient of determination
+   */
+  double getR2();
+
   /**
    * Default statistics bean.
    */
   public class DefaultStatistics implements Statistics {
 
     private final double rss;
     private final double mse;
+    private final double r2;
 
-    public DefaultStatistics(final double rss, final double mse) {
+    public DefaultStatistics(final double rss, final double mse, final double r2) {
       this.rss = rss;
       this.mse = mse;
+      this.r2 = r2;
     }
 
     @Override
@@ -56,11 +63,17 @@ public double getMse() {
       return this.mse;
     }
 
+    @Override
+    public double getR2() {
+      return this.r2;
+    }
+
     @Override
     public String toString() {
       return "DefaultStatistics{" +
           "rss=" + this.rss +
           ", mse=" + this.mse +
+          ", r2=" + this.r2 +
           '}';
     }
   }

src/main/java/org/scaleborn/linereg/calculation/statistics/StatsCalculator.java

@@ -29,7 +29,8 @@ public Statistics calculate(final StatsModel model) {
         .getCovarianceLowerTriangularMatrix();
     final double[] slopeCoefficients = model.getSlopeCoefficients().getCoefficients();
 
-    double squaredError = model.getStatsSampling().getResponseVariance();
+    final double responseVariance = model.getStatsSampling().getResponseVariance();
+    double squaredError = responseVariance;
 
     for (int i = 0; i < featuresCount; i++) {
       final double c = slopeCoefficients[i];
@@ -59,6 +60,11 @@ public double getRss() {
       public double getMse() {
         return rss / model.getStatsSampling().getCount();
       }
+
+      @Override
+      public double getR2() {
+        return 1 - (rss / responseVariance);
+      }
     };
   }
 }

src/test/java/org/scaleborn/linereg/TestModels.java

@@ -124,6 +124,7 @@ public void assertCoefficients(final double[] givenCoefficients, final double de
     public void assertStatistics(final Statistics statistics) {
       assertEquals("RSS not equal", this.expectedStatistics.getRss(), statistics.getRss(), 0.0001d);
       assertEquals("MSE not equal", this.expectedStatistics.getMse(), statistics.getMse(), 0.0001d);
+      assertEquals("R² not equal", this.expectedStatistics.getR2(), statistics.getR2(), 0.0001d);
     }
   }
 
@@ -141,7 +142,7 @@ public void assertStatistics(final Statistics statistics) {
        * Residual Sum of Squares: rss = 5.459016393
        * Coefficient of Determination: R2 = 0.787704918
        */
-      new DefaultStatistics(5.459016393, 5.459016393 / 7));
+      new DefaultStatistics(5.459016393, 5.459016393 / 7, 0.787704918));
 
   public static TestModel MULTI_FEATURES_2_MODEL_1 = new TestModel(3, 2, new double[][]{
       new double[]{-2, 3, 5},
@@ -157,7 +158,7 @@ public void assertStatistics(final Statistics statistics) {
        * Coefficient of Determination: R2 = 8.835223808·10-1
       */
   }, new double[]{-0.5496314882d, 0.3070409283d},
-      new DefaultStatistics(2.99513878d, 2.99513878d / 7));
+      new DefaultStatistics(2.99513878d, 2.99513878d / 7, 0.8835223808));
 
   public static TestModel MULTI_FEATURES_3_MODEL_1 = new TestModel(4, 3, new double[][]{
       new double[]{4, -2, 3, 5},
@@ -173,7 +174,7 @@ public void assertStatistics(final Statistics statistics) {
        * Coefficient of Determination: R2 = 8.947697777·10-1
       */
   }, new double[]{-0.03116979852d, -0.6272993725d, 0.3079647314d},
-      new DefaultStatistics(2.705920002d, 2.705920002d / 7));
+      new DefaultStatistics(2.705920002d, 2.705920002d / 7, 0.8947697777));
 
   /**
    * Reference data set Longley: from http://www.itl.nist.gov/div898/strd/lls/data/Longley.shtml
@@ -211,6 +212,6 @@ public void assertStatistics(final Statistics statistics) {
        */
       , new double[]{15.0618722713733d, -0.358191792925910E-01d, -2.02022980381683d,
       -1.03322686717359d, -0.511041056535807E-01d, 1829.15146461355d},
-      new DefaultStatistics(836424.055505915d, 836424.055505915d / 16));
+      new DefaultStatistics(836424.055505915d, 836424.055505915d / 16, 0.995479004577296));
 
 }