A V0 implementation tested against matlab
Supports: - HDF5 file load and store - Precision timing of the process to stdout - logging (verbose mode) to stdout - Command line arguments and help
This commit is contained in:
+141
-87
@@ -1,107 +1,161 @@
|
||||
/*!
|
||||
* \file main.cpp
|
||||
* \brief Main application file
|
||||
*
|
||||
* \author
|
||||
* Christos Choutouridis AEM:8997
|
||||
* <cchoutou@ece.auth.gr>
|
||||
*/
|
||||
|
||||
#include <iostream>
|
||||
#include <cblas.h>
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <exception>
|
||||
#include <unistd.h>
|
||||
#include <cstdio>
|
||||
|
||||
#include <v0.hpp>
|
||||
#include <v1.hpp>
|
||||
#include <matrix.hpp>
|
||||
#include <utils.hpp>
|
||||
#include <config.h>
|
||||
|
||||
// Global session data
|
||||
session_t session;
|
||||
Log logger;
|
||||
Timing timer;
|
||||
|
||||
/*!
|
||||
* Function to compute squared Euclidean distances
|
||||
*
|
||||
* \fn void pdist2(const double*, const double*, double*, int, int, int)
|
||||
* \param X m x d matrix
|
||||
* \param Y n x d matrix
|
||||
* \param D2 m x n matrix to store distances
|
||||
* \param m number of rows in X
|
||||
* \param n number of rows in Y
|
||||
* \param d number of columns in both X and Y
|
||||
* A small command line argument parser
|
||||
* \return The status of the operation
|
||||
*/
|
||||
void pdist2(const double* X, const double* Y, double* D2, int m, int n, int d){
|
||||
// Compute the squared norms of each row in X and Y
|
||||
std::vector<double> X_norms(m), Y_norms(n);
|
||||
for (int i = 0; i < m; ++i) {
|
||||
X_norms[i] = cblas_ddot(d, X + i * d, 1, X + i * d, 1);
|
||||
}
|
||||
for (int j = 0; j < n; ++j) {
|
||||
Y_norms[j] = cblas_ddot(d, Y + j * d, 1, Y + j * d, 1);
|
||||
}
|
||||
bool get_options(int argc, char* argv[]){
|
||||
bool status =true;
|
||||
|
||||
// Compute -2 * X * Y'
|
||||
cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, n, d, -2.0, X, d, Y, d, 0.0, D2, n);
|
||||
// iterate over the passed arguments
|
||||
for (int i=1 ; i<argc ; ++i) {
|
||||
std::string arg(argv[i]); // get current argument
|
||||
|
||||
// Step 3: Add the squared norms to each entry in D2
|
||||
for (int i = 0; i < m; ++i) {
|
||||
for (int j = 0; j < n; ++j) {
|
||||
D2[i * n + j] += X_norms[i] + Y_norms[j];
|
||||
D2[i * n + j] = std::max(D2[i * n + j], 0.0); // Ensure non-negative
|
||||
D2[i * n + j] = std::sqrt(D2[i * n + j]); // Take the square root of each
|
||||
if (arg == "-c" || arg == "--corpus") {
|
||||
if (i+2 < argc) {
|
||||
session.corpusMtxFile = std::string(argv[++i]);
|
||||
session.corpusDataSet = std::string(argv[++i]);
|
||||
}
|
||||
else
|
||||
status = false;
|
||||
}
|
||||
else if (arg == "-o" || arg == "--output") {
|
||||
if (i+3 < argc) {
|
||||
session.outMtxFile = std::string(argv[++i]);
|
||||
session.outMtxIdxDataSet = std::string(argv[++i]);
|
||||
session.outMtxDstDataSet = std::string(argv[++i]);
|
||||
}
|
||||
else
|
||||
status = false;
|
||||
}
|
||||
else if (arg == "-q" || arg == "--query") {
|
||||
if (i+2 < argc) {
|
||||
session.queryMtxFile = std::string(argv[++i]);
|
||||
session.queryDataSet = std::string(argv[++i]);
|
||||
session.queryMtx = true;
|
||||
}
|
||||
else
|
||||
status = false;
|
||||
}
|
||||
else if (arg == "-k") {
|
||||
session.k = (i+1 < argc) ? std::atoi(argv[++i]) : session.k;
|
||||
}
|
||||
else if (arg == "-n" || arg == "--max_trheads") {
|
||||
session.max_threads = (i+1 < argc) ? std::atoi(argv[++i]) : session.max_threads;
|
||||
}
|
||||
else if (arg == "-t" || arg == "--timing")
|
||||
session.timing = true;
|
||||
else if (arg == "-v" || arg == "--verbose")
|
||||
session.verbose = true;
|
||||
else if (arg == "-h" || arg == "--help") {
|
||||
std::cout << "annsearch - an aproximation knnsearch utility\n\n";
|
||||
std::cout << "annsearch -c <file> [-k <N>] [-o <file>] [-q <file>] [-n <threads>] [-t] [-v]\n";
|
||||
std::cout << '\n';
|
||||
std::cout << "Options:\n\n";
|
||||
std::cout << " -c | --corpus <file> <dataset>\n";
|
||||
std::cout << " Path to hdf5 file to open and name of the dataset to load\n\n";
|
||||
std::cout << " -o | --output <file> <idx-dataset> <dst-dataset> \n";
|
||||
std::cout << " Path to <file> to store the data and the names of the datasets.\n\n";
|
||||
std::cout << " -q | --query <file> <dataset>\n";
|
||||
std::cout << " Path to hdf5 file to open and name of the dataset to load\n";
|
||||
std::cout << " If not defined, the corpus is used\n\n";
|
||||
std::cout << " -k <number>\n";
|
||||
std::cout << " Set the number of closest neighbors to find. \n\n";
|
||||
std::cout << " -n | --max_trheads <threads>\n";
|
||||
std::cout << " Reduce the thread number for the execution to <threads>. <threads> must be less or equal to available CPUs.\n\n";
|
||||
std::cout << " -t | --timing\n";
|
||||
std::cout << " Request timing measurements output to stdout.\n\n";
|
||||
std::cout << " -v | --verbose\n";
|
||||
std::cout << " Request a more verbose output to stdout.\n\n";
|
||||
std::cout << " -h | --help <size>\n";
|
||||
std::cout << " Prints this and exit.\n\n";
|
||||
std::cout << "Examples:\n\n";
|
||||
std::cout << " ...Example case...:\n";
|
||||
std::cout << " > ./annsearch -i <MFILE> ... \n\n";
|
||||
|
||||
exit(0);
|
||||
}
|
||||
else { // parse error
|
||||
std::cout << "Invocation error. Try -h for details.\n";
|
||||
status = false;
|
||||
}
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
void quickselect(std::vector<std::pair<double, int>>& vec, int k) {
|
||||
std::nth_element(
|
||||
vec.begin(),
|
||||
vec.begin() + k,
|
||||
vec.end(),
|
||||
[](const std::pair<double, int>& a, const std::pair<double, int>& b) {
|
||||
return a.first < b.first;
|
||||
});
|
||||
vec.resize(k); // Keep only the k smallest elements
|
||||
}
|
||||
NAMESPACE_VERSION;
|
||||
|
||||
// K-nearest neighbor search function
|
||||
void knnsearch(const double* C, const double* Q, int m, int n, int d, int k,
|
||||
std::vector<std::vector<int>>& idx, std::vector<std::vector<double>>& dst) {
|
||||
std::vector<double> D(m * n);
|
||||
pdist2(C, Q, D.data(), m, n, d);
|
||||
int main(int argc, char* argv[]) try {
|
||||
// Instantiate matrixes
|
||||
MatrixDst Corpus;
|
||||
MatrixDst Query;
|
||||
MatrixIdx Idx;
|
||||
MatrixDst Dst;
|
||||
|
||||
idx.resize(n, std::vector<int>(k));
|
||||
dst.resize(n, std::vector<double>(k));
|
||||
// try to read command line
|
||||
if (!get_options(argc, argv))
|
||||
exit(1);
|
||||
|
||||
for (int j = 0; j < n; ++j) {
|
||||
// Create a vector of pairs (distance, index) for the j-th query
|
||||
std::vector<std::pair<double, int>> dst_idx(m);
|
||||
for (int i = 0; i < m; ++i) {
|
||||
dst_idx[i] = {D[i * n + j], i};
|
||||
}
|
||||
if (access(session.outMtxFile.c_str(), F_OK) == 0)
|
||||
std::remove(session.outMtxFile.c_str());
|
||||
|
||||
// Find the k smallest distances using quickSelectKSmallest
|
||||
quickselect(dst_idx, k);
|
||||
// Load data
|
||||
timer.start();
|
||||
Mtx::load<MatrixDst, DstHDF5Type>(session.corpusMtxFile, session.corpusDataSet, Corpus);
|
||||
if (session.queryMtx)
|
||||
Mtx::load<MatrixDst, DstHDF5Type>(session.corpusMtxFile, session.corpusDataSet, Query);
|
||||
timer.stop();
|
||||
timer.print_dt("Load hdf5 files");
|
||||
|
||||
// Sort the k smallest results by distance for consistency
|
||||
std::sort(dst_idx.begin(), dst_idx.end());
|
||||
logger << "Start knnsearch ...";
|
||||
timer.start();
|
||||
if (session.queryMtx)
|
||||
knnsearch(Corpus, Query, session.k, Idx, Dst);
|
||||
else
|
||||
knnsearch(Corpus, Corpus, session.k, Idx, Dst);
|
||||
timer.stop();
|
||||
logger << " Done" << logger.endl;
|
||||
timer.print_dt("knnsearch");
|
||||
|
||||
// Store the indices and distances
|
||||
for (int i = 0; i < k; ++i) {
|
||||
idx[j][i] = dst_idx[i].second;
|
||||
dst[j][i] = dst_idx[i].first;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int main(){
|
||||
|
||||
int m = 5; // Number of points in C (corpus)
|
||||
int n = 3; // Number of points in Q (query)
|
||||
int d = 2; // Dimensions
|
||||
int k = 2; // Number of nearest neighbors to find
|
||||
|
||||
double C[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0}; // m x d matrix
|
||||
double Q[] = {1.5, 2.5, 3.5, 4.5, 5.5, 6.5}; // n x d matrix
|
||||
|
||||
std::vector<std::vector<int>> idx;
|
||||
std::vector<std::vector<double>> dst;
|
||||
|
||||
knnsearch(C, Q, m, n, d, k, idx, dst);
|
||||
|
||||
// Print results
|
||||
for (int i = 0; i < n; ++i) {
|
||||
std::cout << "Query point " << i << ":\n";
|
||||
for (int j = 0; j < k; ++j) {
|
||||
std::cout << " Neighbor " << j <<": Index = " << idx[i][j] <<", Distance = " << dst[i][j] << '\n';
|
||||
}
|
||||
}
|
||||
// Store data
|
||||
timer.start();
|
||||
Mtx::store<MatrixIdx, IdxHDF5Type>(session.outMtxFile, session.outMtxIdxDataSet, Idx);
|
||||
Mtx::store<MatrixDst, DstHDF5Type>(session.outMtxFile, session.outMtxDstDataSet, Dst);
|
||||
timer.stop();
|
||||
timer.print_dt("Store hdf5 files");
|
||||
|
||||
return 0;
|
||||
}
|
||||
catch (std::exception& e) {
|
||||
//we probably pollute the user's screen. Comment `cerr << ...` if you don't like it.
|
||||
std::cerr << "Error: " << e.what() << '\n';
|
||||
exit(1);
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user