Generalize aggregation stuff to simplify newer aggregations based in the evaluated linear model

Author: mbok

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

@@ -16,186 +16,99 @@
 
 package org.scaleborn.elasticsearch.linreg.aggregation.stats;
 
-import static java.util.Collections.emptyMap;
-
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.common.io.stream.StreamInput;
-import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.logging.Loggers;
-import org.elasticsearch.common.xcontent.XContentBuilder;
-import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
-import org.scaleborn.elasticsearch.linreg.aggregation.support.StateInputStreamAdapter;
-import org.scaleborn.elasticsearch.linreg.aggregation.support.StateOutputStreamAdapter;
-import org.scaleborn.linereg.evaluation.DerivationEquation;
-import org.scaleborn.linereg.evaluation.DerivationEquationBuilder;
+import org.scaleborn.elasticsearch.linreg.aggregation.support.BaseInternalAggregation;
+import org.scaleborn.linereg.calculation.statistics.Statistics;
+import org.scaleborn.linereg.calculation.statistics.StatsCalculator;
+import org.scaleborn.linereg.calculation.statistics.StatsModel;
 import org.scaleborn.linereg.evaluation.SlopeCoefficients;
-import org.scaleborn.linereg.evaluation.commons.CommonsMathSolver;
-import org.scaleborn.linereg.statistics.Statistics;
-import org.scaleborn.linereg.statistics.StatsBuilder;
-import org.scaleborn.linereg.statistics.StatsModel;
-import org.scaleborn.linereg.statistics.StatsSampling;
 
 /**
  * Created by mbok on 21.03.17.
  */
-public class InternalStats extends InternalAggregation implements Stats {
+public class InternalStats extends
+    BaseInternalAggregation<StatsAggregationSampling, StatsResults, InternalStats> implements
+    Stats {
 
   private static final Logger LOGGER = Loggers.getLogger(InternalStats.class);
-  /**
-   * per shard sampling needed to compute stats
-   */
-  private StatsSampling<?> sampling;
-  /**
-   * final result
-   */
-  private Statistics results;
 
-  /**
-   * Features count
-   */
-  private int featuresCount;
+  private static final StatsCalculator statsCalculator = new StatsCalculator();
 
   /**
    * per shard ctor
    */
-  protected InternalStats(String name, int featuresCount, StatsSampling<?> linRegSampling,
-      Statistics results,
-      List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) {
-    super(name, pipelineAggregators, metaData);
-    this.featuresCount = featuresCount;
-    this.sampling = linRegSampling;
-    this.results = results;
+  protected InternalStats(final String name, final int featuresCount,
+      final StatsAggregationSampling linRegSampling,
+      final StatsResults results,
+      final List<PipelineAggregator> pipelineAggregators, final Map<String, Object> metaData) {
+    super(name, featuresCount, linRegSampling, results, pipelineAggregators, metaData);
   }
 
   /**
    * Read from a stream.
    */
-  public InternalStats(StreamInput in) throws IOException {
-    super(in);
-    this.featuresCount = in.readInt();
-    if (in.readBoolean()) {
-      this.sampling = StatsAggregationBuilder.buildSampling(this.featuresCount);
-      StateInputStreamAdapter streamAdapter = new StateInputStreamAdapter(in);
-      this.sampling.loadState(streamAdapter);
-    }
-    if (in.readBoolean()) {
-      this.results = new DefaultStatistics(in.readDouble(), in.readDouble());
-    }
+  public InternalStats(final StreamInput in) throws IOException {
+    super(in, StatsResults::new);
   }
 
   @Override
-  protected void doWriteTo(StreamOutput out) throws IOException {
-    out.writeInt(this.featuresCount);
-    out.writeBoolean(this.sampling != null);
-    if (this.sampling != null) {
-      StateOutputStreamAdapter outAdapter = new StateOutputStreamAdapter(out);
-      this.sampling.saveState(outAdapter);
-    }
-    out.writeBoolean(this.results != null);
-    if (this.results != null) {
-      out.writeDouble(this.results.getRss());
-      out.writeDouble(this.results.getMse());
-    }
+  protected StatsAggregationSampling buildSampling(final int featuresCount) {
+    return StatsAggregationBuilder.buildSampling(featuresCount);
   }
 
+
   @Override
   public String getWriteableName() {
     return StatsAggregationBuilder.NAME;
   }
 
-  static class Fields {
-
-    public static final String RSS = "rss";
-    public static final String MSE = "mse";
-  }
-
   @Override
   public double getRss() {
-    if (results == null) {
+    if (this.results == null) {
       return Double.NaN;
     }
-    return results.getRss();
+    return this.results.statistics.getRss();
   }
 
   @Override
   public double getMse() {
-    if (results == null) {
+    if (this.results == null) {
       return Double.NaN;
     }
-    return results.getMse();
+    return this.results.statistics.getMse();
   }
 
   @Override
-  public XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
-    if (results != null) {
-      // RSS
-      builder.field(Fields.RSS, results.getRss());
-      // MSE
-      builder.field(Fields.MSE, results.getMse());
+  public Object getDoProperty(final String element) {
+    switch (element) {
+      case "rss":
+        return getRss();
+      case "mse":
+        return getMse();
     }
-    return builder;
+    return null;
   }
 
   @Override
-  public Object getProperty(List<String> path) {
-    if (path.isEmpty()) {
-      return this;
-    } else if (path.size() == 1) {
-      String element = path.get(0);
-      if (results == null) {
-        return emptyMap();
-      }
-      switch (element) {
-        case "rss":
-          return results.getRss();
-        case "mse":
-          return results.getMse();
-        default:
-          throw new IllegalArgumentException(
-              "Found unknown path element [" + element + "] in [" + getName() + "]");
-      }
-    } else {
-      throw new IllegalArgumentException("path not supported for [" + getName() + "]: " + path);
-    }
+  protected InternalStats buildInternalAggregation(final String name, final int featuresCount,
+      final StatsAggregationSampling linRegSampling, final StatsResults results,
+      final List<PipelineAggregator> pipelineAggregators, final Map<String, Object> metaData) {
+    return new InternalStats(name, featuresCount, linRegSampling, results, pipelineAggregators,
+        metaData);
   }
 
-  @SuppressWarnings("unchecked")
   @Override
-  public InternalAggregation doReduce(List<InternalAggregation> aggregations,
-      ReduceContext reduceContext) {
-    // merge samples across all shards
-    List<InternalAggregation> aggs = new ArrayList<>(aggregations);
-    aggs.removeIf(p -> ((InternalStats) p).sampling == null);
-
-    // return empty result iff all samples are null
-    if (aggs.isEmpty()) {
-      return new InternalStats(name, featuresCount, null, new DefaultStatistics(0, 0),
-          pipelineAggregators(),
-          getMetaData());
-    }
-
-    StatsSampling composedSampling = StatsAggregationBuilder
-        .buildSampling(featuresCount);
-    for (int i = 0; i < aggs.size(); ++i) {
-      LOGGER.info("Merging sampling={}: {}", i, ((InternalStats) aggs.get(i)).sampling);
-      composedSampling.merge(((InternalStats) aggs.get(i)).sampling);
-    }
-
-    // Linear regression evaluation
-    DerivationEquation derivationEquation = new DerivationEquationBuilder()
-        .buildDerivationEquation(composedSampling);
-    CommonsMathSolver commonsMathSolver = new CommonsMathSolver();
-    SlopeCoefficients slopeCoefficients = commonsMathSolver.solveCoefficients(derivationEquation);
-    StatsBuilder statsBuilder = new StatsBuilder();
-    Statistics statsResult = statsBuilder
-        .buildStats(new StatsModel(composedSampling, slopeCoefficients));
-    LOGGER.info("Evaluated linear with {} and stats {}", slopeCoefficients, statsResult);
-    return new InternalStats(name, featuresCount, composedSampling, statsResult,
-        pipelineAggregators(), getMetaData());
+  protected StatsResults buildResults(final StatsAggregationSampling composedSampling,
+      final SlopeCoefficients slopeCoefficients) {
+    final Statistics stats = statsCalculator
+        .calculate(new StatsModel(composedSampling, slopeCoefficients));
+    return new StatsResults(slopeCoefficients, stats);
   }
+
 }

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

@@ -16,7 +16,7 @@
 
 package org.scaleborn.elasticsearch.linreg.aggregation.stats;
 
-import org.scaleborn.linereg.statistics.Statistics;
+import org.scaleborn.linereg.calculation.statistics.Statistics;
 
 /**
  * Created by mbok on 21.03.17.

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

@@ -28,8 +28,6 @@
 import org.scaleborn.elasticsearch.linreg.aggregation.support.BaseAggregationBuilder;
 import org.scaleborn.linereg.sampling.exact.ExactModelSamplingFactory;
 import org.scaleborn.linereg.sampling.exact.ExactSamplingContext;
-import org.scaleborn.linereg.statistics.StatsSampling;
-import org.scaleborn.linereg.statistics.StatsSampling.StatsSamplingProxy;
 
 /**
  * Created by mbok on 21.03.17.
@@ -42,21 +40,21 @@ public class StatsAggregationBuilder extends
   private static final ExactModelSamplingFactory MODEL_SAMPLING_FACTORY = new ExactModelSamplingFactory();
 
 
-  public StatsAggregationBuilder(String name) {
+  public StatsAggregationBuilder(final String name) {
     super(name);
   }
 
-  public StatsAggregationBuilder(StreamInput in) throws IOException {
+  public StatsAggregationBuilder(final StreamInput in) throws IOException {
     super(in);
   }
 
   @Override
-  protected StatsAggregatorFactory innerInnerBuild(SearchContext context,
-      List<NamedValuesSourceConfigSpec<Numeric>> configs, MultiValueMode multiValueMode,
-      AggregatorFactory<?> parent, AggregatorFactories.Builder subFactoriesBuilder)
+  protected StatsAggregatorFactory innerInnerBuild(final SearchContext context,
+      final List<NamedValuesSourceConfigSpec<Numeric>> configs, final MultiValueMode multiValueMode,
+      final AggregatorFactory<?> parent, final AggregatorFactories.Builder subFactoriesBuilder)
       throws IOException {
-    return new StatsAggregatorFactory(name, configs, multiValueMode, context, parent,
-        subFactoriesBuilder, metaData);
+    return new StatsAggregatorFactory(this.name, configs, multiValueMode, context, parent,
+        subFactoriesBuilder, this.metaData);
   }
 
 
@@ -65,13 +63,14 @@ public String getType() {
     return NAME;
   }
 
-  static StatsSampling<?> buildSampling(final int featuresCount) {
-    ExactSamplingContext context = MODEL_SAMPLING_FACTORY
+  static StatsAggregationSampling buildSampling(final int featuresCount) {
+    final ExactSamplingContext context = MODEL_SAMPLING_FACTORY
         .createContext(featuresCount);
-    StatsSamplingProxy<?> statsSampling = new StatsSamplingProxy<>(context,
-        MODEL_SAMPLING_FACTORY.createResponseVarianceTermSampling(context),
+    final StatsAggregationSampling statsSampling = new StatsAggregationSampling(context,
         MODEL_SAMPLING_FACTORY.createCoefficientLinearTermSampling(context),
-        MODEL_SAMPLING_FACTORY.createCoefficientSquareTermSampling(context));
+        MODEL_SAMPLING_FACTORY.createCoefficientSquareTermSampling(context),
+        MODEL_SAMPLING_FACTORY.createInterceptSampling(context),
+        MODEL_SAMPLING_FACTORY.createResponseVarianceTermSampling(context));
     return statsSampling;
   }
 }

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

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 2017 Scaleborn UG, www.scaleborn.com
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.scaleborn.elasticsearch.linreg.aggregation.stats;
+
+import java.io.IOException;
+import org.scaleborn.elasticsearch.linreg.aggregation.support.BaseSampling;
+import org.scaleborn.linereg.calculation.statistics.StatsSampling;
+import org.scaleborn.linereg.sampling.io.StateInputStream;
+import org.scaleborn.linereg.sampling.io.StateOutputStream;
+
+/**
+ * Created by mbok on 08.04.17.
+ */
+public class StatsAggregationSampling extends BaseSampling<StatsAggregationSampling> implements
+    StatsSampling<StatsAggregationSampling> {
+
+  private final ResponseVarianceTermSampling responseVarianceTermSampling;
+
+  public StatsAggregationSampling(
+      final SamplingContext<?> samplingContext,
+      final CoefficientLinearTermSampling<?> coefficientLinearTermSampling,
+      final CoefficientSquareTermSampling<?> coefficientSquareTermSampling,
+      final InterceptSampling<?> interceptSampling,
+      final ResponseVarianceTermSampling<?> responseVarianceTermSampling) {
+    super(samplingContext, coefficientLinearTermSampling, coefficientSquareTermSampling,
+        interceptSampling);
+    this.responseVarianceTermSampling = responseVarianceTermSampling;
+  }
+
+  @Override
+  public void merge(final StatsAggregationSampling fromSample) {
+    super.merge(fromSample);
+    this.responseVarianceTermSampling.merge(fromSample.responseVarianceTermSampling);
+  }
+
+  @Override
+  public void sample(final double[] featureValues, final double responseValue) {
+    super.sample(featureValues, responseValue);
+    this.responseVarianceTermSampling.sample(featureValues, responseValue);
+  }
+
+  @Override
+  public void saveState(final StateOutputStream destination) throws IOException {
+    super.saveState(destination);
+    this.responseVarianceTermSampling.saveState(destination);
+  }
+
+  @Override
+  public void loadState(final StateInputStream source) throws IOException {
+    super.loadState(source);
+    this.responseVarianceTermSampling.loadState(source);
+  }
+
+  @Override
+  public double getResponseVariance() {
+    return this.responseVarianceTermSampling.getResponseVariance();
+  }
+}

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

@@ -36,7 +36,6 @@
 import org.elasticsearch.search.aggregations.support.NamedValuesSourceSpec;
 import org.elasticsearch.search.aggregations.support.ValuesSource.Numeric;
 import org.elasticsearch.search.internal.SearchContext;
-import org.scaleborn.linereg.statistics.StatsSampling;
 
 /**
  * Created by mbok on 21.03.17.
@@ -50,7 +49,7 @@ public class StatsAggregator extends MetricsAggregator {
    */
   final NumericMultiValuesSource valuesSources;
 
-  protected ObjectArray<StatsSampling<?>> samplings;
+  protected ObjectArray<StatsAggregationSampling> samplings;
 
   public StatsAggregator(final String name,
       final List<NamedValuesSourceSpec<Numeric>> valuesSources,
@@ -95,7 +94,7 @@ public void collect(final int doc, final long bucket) throws IOException {
         if (includeDocument(doc) == true) {
           StatsAggregator.this.samplings = bigArrays
               .grow(StatsAggregator.this.samplings, bucket + 1);
-          StatsSampling<?> sampling = StatsAggregator.this.samplings.get(bucket);
+          StatsAggregationSampling sampling = StatsAggregator.this.samplings.get(bucket);
           // add document fields to correlation stats
           if (sampling == null) {
             sampling = StatsAggregationBuilder.buildSampling(this.fieldNames.length - 1);