tf-kmeans/dist/index.js

A Library for Calculating K-Means using Tensorflow

"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : new P(function (resolve) { resolve(result.value); }).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (Object.hasOwnProperty.call(mod, k)) result[k] = mod[k];
    result["default"] = mod;
    return result;
};
Object.defineProperty(exports, "__esModule", { value: true });
const tf = __importStar(require("@tensorflow/tfjs-core"));
class KMeans {
    constructor({ k = 2, maxIter = 10, distanceFunction = KMeans.EuclideanDistance } = {}) {
        this.k = 2;
        this.maxIter = 200;
        this.distanceFunction = KMeans.EuclideanDistance;
        this.k = k;
        this.maxIter = maxIter;
        this.distanceFunction = distanceFunction;
    }
    static EuclideanDistance(values, centroids) {
        return tf.tidy(() => values.squaredDifference(centroids).sum(1).sqrt());
    }
    GenerateIndices(rows) {
        const indices = [];
        indices.length = rows;
        for (let i = 0; i < indices.length; ++i)
            indices[i] = i;
        return indices;
    }
    NewCentroid(values, assignments, cluster, rows) {
        return tf.tidy(() => {
            let selectedIndices = [];
            selectedIndices.length = rows;
            selectedIndices = selectedIndices.fill(cluster);
            const selectedIndicesT = tf.tensor(selectedIndices);
            let where = tf.equal(assignments, selectedIndicesT).asType("int32");
            where = where.reshape([where.shape[0], 1]);
            const count = where.sum();
            const newCentroid = values.mul(where).sum(0).div(count);
            return newCentroid;
        });
    }
    NewCentroids(values, assignments) {
        return tf.tidy(() => {
            const rows = values.shape[0];
            const centroids = [];
            for (let cluster = 0; cluster < this.k; ++cluster) {
                centroids.push(this.NewCentroid(values, assignments, cluster, rows));
            }
            return tf.stack(centroids);
        });
    }
    AssignCluster(value, centroids) {
        return tf.tidy(() => this.distanceFunction(value, centroids).argMin(0));
    }
    AssignClusters(values, centroids) {
        return tf.tidy(() => {
            const rows = values.shape[0];
            const minIndexes = [];
            for (const index of this.GenerateIndices(rows)) {
                const value = values.gather(index);
                minIndexes.push(this.AssignCluster(value, centroids));
                value.dispose();
            }
            return tf.stack(minIndexes);
        });
    }
    RandomSample(vals) {
        return tf.tidy(() => {
            const rows = vals.shape[0];
            if (rows < this.k)
                throw new Error("Rows are Less than K");
            const indicesRaw = tf.util.createShuffledIndices(rows).slice(0, this.k);
            const indices = [];
            indicesRaw.forEach((index) => indices.push(index));
            return tf.gatherND(vals, tf.tensor(indices, [this.k, 1], "int32"));
        });
    }
    CheckCentroidSimmilarity(newCentroids, centroids, vals) {
        return tf.tidy(() => newCentroids
            .equal(centroids)
            .asType("int32")
            .sum(1)
            .div(vals.shape[1])
            .sum()
            .equal(this.k)
            .dataSync()[0]);
    }
    TrainSingleStep(values) {
        return tf.tidy(() => {
            const predictions = this.Predict(values);
            const newCentroids = this.NewCentroids(values, predictions);
            return [newCentroids, predictions];
        });
    }
    Train(values, callback = (_, __) => { }) {
        this.centroids = this.RandomSample(values);
        let iter = 0;
        while (true) {
            let [newCentroids, predictions] = this.TrainSingleStep(values);
            const same = this.CheckCentroidSimmilarity(newCentroids, this.centroids, values);
            if (same || iter >= this.maxIter) {
                newCentroids.dispose();
                return predictions;
            }
            this.centroids.dispose();
            this.centroids = newCentroids;
            ++iter;
            callback(this.centroids, predictions);
        }
    }
    TrainAsync(values, callback = (_iter, _centroid, _predictions) => __awaiter(this, void 0, void 0, function* () { })) {
        return __awaiter(this, void 0, void 0, function* () {
            this.centroids = this.RandomSample(values);
            let iter = 0;
            while (true) {
                let [newCentroids, predictions] = this.TrainSingleStep(values);
                const same = this.CheckCentroidSimmilarity(newCentroids, this.centroids, values);
                if (same || iter >= this.maxIter) {
                    newCentroids.dispose();
                    return predictions;
                }
                this.centroids.dispose();
                this.centroids = newCentroids;
                yield callback(iter, this.centroids, predictions);
                ++iter;
            }
        });
    }
    Predict(y) {
        return tf.tidy(() => {
            if (y.shape[1] == null)
                y = y.reshape([1, y.shape[0]]);
            return this.AssignClusters(y, this.centroids);
        });
    }
    Centroids() {
        return this.centroids;
    }
    Dispose() {
        this.centroids.dispose();
    }
}
exports.default = KMeans;
//# sourceMappingURL=index.js.map