mirror of https://github.com/apache/lucene.git
SOLR-11593: Add support for covariance matrices to the cov Stream Evaluator
This commit is contained in:
parent
44b7bc0177
commit
6d5a7920ae
|
@ -19,13 +19,13 @@ package org.apache.solr.client.solrj.io.eval;
|
|||
import java.io.IOException;
|
||||
import java.math.BigDecimal;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
|
||||
import org.apache.commons.math3.linear.RealMatrix;
|
||||
import org.apache.commons.math3.stat.correlation.Covariance;
|
||||
import org.apache.solr.client.solrj.io.stream.expr.StreamExpression;
|
||||
import org.apache.solr.client.solrj.io.stream.expr.StreamFactory;
|
||||
|
||||
public class CovarianceEvaluator extends RecursiveNumericEvaluator implements TwoValueWorker {
|
||||
public class CovarianceEvaluator extends RecursiveObjectEvaluator implements ManyValueWorker {
|
||||
protected static final long serialVersionUID = 1L;
|
||||
|
||||
public CovarianceEvaluator(StreamExpression expression, StreamFactory factory) throws IOException{
|
||||
|
@ -33,25 +33,26 @@ public class CovarianceEvaluator extends RecursiveNumericEvaluator implements Tw
|
|||
}
|
||||
|
||||
@Override
|
||||
public Object doWork(Object first, Object second) throws IOException{
|
||||
if(null == first){
|
||||
throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - null found for the first value",toExpression(constructingFactory)));
|
||||
}
|
||||
if(null == second){
|
||||
throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - null found for the second value",toExpression(constructingFactory)));
|
||||
}
|
||||
if(!(first instanceof List<?>)){
|
||||
throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - found type %s for the first value, expecting a list of numbers",toExpression(constructingFactory), first.getClass().getSimpleName()));
|
||||
}
|
||||
if(!(second instanceof List<?>)){
|
||||
throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - found type %s for the second value, expecting a list of numbers",toExpression(constructingFactory), first.getClass().getSimpleName()));
|
||||
}
|
||||
public Object doWork(Object ... values) throws IOException{
|
||||
|
||||
if(values.length == 2) {
|
||||
Object first = values[0];
|
||||
Object second = values[1];
|
||||
Covariance covariance = new Covariance();
|
||||
|
||||
return covariance.covariance(
|
||||
((List) first).stream().mapToDouble(value -> ((BigDecimal) value).doubleValue()).toArray(),
|
||||
((List) second).stream().mapToDouble(value -> ((BigDecimal) value).doubleValue()).toArray()
|
||||
);
|
||||
} else if(values.length == 1) {
|
||||
Matrix matrix = (Matrix) values[0];
|
||||
double[][] data = matrix.getData();
|
||||
Covariance covariance = new Covariance(data, true);
|
||||
RealMatrix coMatrix = covariance.getCovarianceMatrix();
|
||||
double[][] coData = coMatrix.getData();
|
||||
return new Matrix(coData);
|
||||
} else {
|
||||
throw new IOException("The cov function expects either two numeric arrays or a matrix as parameters.");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -7056,6 +7056,38 @@ public class StreamExpressionTest extends SolrCloudTestCase {
|
|||
assertEquals(maxa.doubleValue(), 30, .5);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testCovMatrix() throws Exception {
|
||||
String cexpr = "let(a=array(1,2,3), b=array(2,4,6), c=array(4, 8, 12), d=transpose(matrix(a, b, c)), f=cov(d))";
|
||||
ModifiableSolrParams paramsLoc = new ModifiableSolrParams();
|
||||
paramsLoc.set("expr", cexpr);
|
||||
paramsLoc.set("qt", "/stream");
|
||||
String url = cluster.getJettySolrRunners().get(0).getBaseUrl().toString()+"/"+COLLECTIONORALIAS;
|
||||
TupleStream solrStream = new SolrStream(url, paramsLoc);
|
||||
StreamContext context = new StreamContext();
|
||||
solrStream.setStreamContext(context);
|
||||
List<Tuple> tuples = getTuples(solrStream);
|
||||
assertTrue(tuples.size() == 1);
|
||||
List<List<Number>> cm = (List<List<Number>>)tuples.get(0).get("f");
|
||||
assertEquals(cm.size(), 3);
|
||||
List<Number> row1 = cm.get(0);
|
||||
assertEquals(row1.size(), 3);
|
||||
assertEquals(row1.get(0).longValue(), 1);
|
||||
assertEquals(row1.get(1).longValue(), 2);
|
||||
assertEquals(row1.get(2).longValue(), 4);
|
||||
|
||||
List<Number> row2 = cm.get(1);
|
||||
assertEquals(row2.size(), 3);
|
||||
assertEquals(row2.get(0).longValue(), 2);
|
||||
assertEquals(row2.get(1).longValue(), 4);
|
||||
assertEquals(row2.get(2).longValue(), 8);
|
||||
|
||||
List<Number> row3 = cm.get(2);
|
||||
assertEquals(row3.size(), 3);
|
||||
assertEquals(row3.get(0).longValue(), 4);
|
||||
assertEquals(row3.get(1).longValue(), 8);
|
||||
assertEquals(row3.get(2).longValue(), 16);
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
|
Loading…
Reference in New Issue