knn.hpp

Go to the documentation of this file.

/*
 *   Copyright 2021 Huawei Technologies Co., Ltd.
 *
 * 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.
 */

#ifndef _H_GRB_KNN
#define _H_GRB_KNN

#include "graphblas/algorithms/mpv.hpp"

#include <graphblas.hpp>


namespace grb {

    namespace algorithms {

        template< Descriptor descr, typename OutputType, typename InputType >
        RC knn(
            Vector< OutputType > &u, const Matrix< InputType > &A,
            const size_t source, const size_t k,
            Vector< bool > &buf1
        ) {
            // the nearest-neighbourhood ring
            Semiring<
                operators::logical_or< bool >, operators::logical_and< bool >,
                identities::logical_false, identities::logical_true
            > ring;

            // check input
            const size_t n = nrows( A );
            if( n != ncols( A ) ) {
                return MISMATCH;
            }
            if( size( buf1 ) != n ) {
                return MISMATCH;
            }
            if( size( u ) != n ) {
                return MISMATCH;
            }
            if( capacity( u ) != n ) {
                return ILLEGAL;
            }
            if( capacity( buf1 ) != n ) {
                return ILLEGAL;
            }

            // prepare
            RC ret = SUCCESS;
            if( nnz( u ) != 0 ) {
                ret = clear( u );
            }
            if( nnz( buf1 ) != 0 ) {
                ret = ret ? ret : clear( buf1 );
            }
#ifdef _DEBUG
            std::cout << "grb::algorithms::knn called with source " << source << " "
                << "and k " << k << ".\n";
#endif
            ret = ret ? ret : setElement( buf1, true, source );

            // do sparse matrix powers on the given ring
            if( ret == SUCCESS ) {
                if( descr & descriptors::transpose_matrix ) {
                    ret = mpv< (descr | descriptors::add_identity) &
                        ~( descriptors::transpose_matrix )
                    >( u, A, k, buf1, buf1, ring );
                } else {
                    ret = mpv< descr | descriptors::add_identity |
                        descriptors::transpose_matrix
                    >( u, A, k, buf1, buf1, ring );
                }
            }

            // done
            return ret;
        }

    } // namespace algorithms

} // namespace grb

#endif