aiom_pack/src/models/gatekeeper.js

Framework for interdependent (mcmc-like) behavioral experiments

// different kinds of gatekeepers here
const sampling = require('./sampling');
const transformer = require('./transformer');
const { pool } = require('../core/database');
const lower_bound = Number(process.env.lower_bound);
const upper_bound = Number(process.env.upper_bound);
var n = require("numeric");
var sqrt2PI = Math.sqrt(Math.PI * 2);

/**
 * Numerically stable log-sum-exp.
 * @param {number[]} arr Array of numbers (logarithms).
 * @returns {number} log(sum(exp(arr_i))).
 */
function logSumExp(arr) {
    if (arr.length === 0) return -Infinity;
    const maxVal = Math.max(...arr);
    if (maxVal === -Infinity) return -Infinity;

    let sumExp = 0;
    for (let i = 0; i < arr.length; i++) {
        sumExp += Math.exp(arr[i] - maxVal);
    }
    return maxVal + Math.log(sumExp);
}


function GaussianKDE(parameters) {
    this.tree_data = parameters.tree_data;
    this.bandwidth = parameters.bandwidth;
    this.dimensionality = parameters.dimensionality;
    this.n_samples = parameters.n_samples;
    let mean = new Array(this.dimensionality).fill(0);

    for (let i = 0; i < this.tree_data.length; i++) {
        for (let j = 0; j < this.dimensionality; j++) {
            mean[j] += this.tree_data[i][j];
        }
    }
    for (let j = 0; j < this.dimensionality; j++) {
        mean[j] /= this.tree_data.length;
    }
    this.mean = mean;
}

/**
 * Evaluates the density function of the gaussian at the given point
 */
GaussianKDE.prototype.density = function(x) {
    if (x.length !== this.dimensionality) {
        throw new Error(`Input point dimensionality (${x.length}) does not match model dimensionality (${this.dimensionality}).`);
    }
    const logKernelConstant = -this.dimensionality * Math.log(this.bandwidth) - (this.dimensionality / 2) * Math.log(2 * Math.PI);

    const logKernelTerms = this.tree_data.map(trainSample => {
        let squaredDistance = 0;
        for (let j = 0; j < this.dimensionality; j++) {
            const diff = x[j] - trainSample[j];
            squaredDistance += diff * diff;
        }
        const expArgument = -0.5 * squaredDistance / (this.bandwidth * this.bandwidth);
        return expArgument + logKernelConstant;
    });

    const logLikelihoodSum = logSumExp(logKernelTerms);

    if (logLikelihoodSum === -Infinity) {
        return -Infinity;
    }
    
    return logLikelihoodSum - Math.log(this.n_samples);
};

// the basic customized gatekeepers
GaussianKDE.prototype.acceptance = function(current, proposal, temperature=1.0) {

    const density_current = this.density(current);
    const density_proposal = this.density(proposal);
    return Math.exp(density_proposal/temperature) / (Math.exp(density_current/temperature) + Math.exp(density_proposal/temperature))
}

GaussianKDE.prototype.processing = async function(current_state, proposal, table_name, proposal_cov) {
    var trial_number;
    var proposal = proposal;
    var acceptance_rate = this.acceptance(current_state, proposal);
    while (Math.random() > acceptance_rate) {
        trial_number = await pool.query(
            `UPDATE ${table_name} SET picked = true, gatekeeper = true WHERE id = (
            SELECT id FROM ${table_name} ORDER BY id DESC OFFSET 1 LIMIT 1) 
            RETURNING trial;`,
        );
        // console.log(trial_number.rows[0].trial);
        proposal = transformer.limit_array_in_range(sampling.gaussian_array(current_state, proposal_cov), min=lower_bound, max=upper_bound);
        acceptance_rate = this.acceptance(current_state, proposal);
        await pool.query(
            `INSERT INTO ${table_name} (trial, choices) 
            VALUES ($1, $2), ($3, $4)`,
            [trial_number.rows[0].trial+1, JSON.stringify(current_state), trial_number.rows[0].trial+1, JSON.stringify(proposal)]
        );
    }
    // console.log(proposal);
    return proposal;
}

GaussianKDE.prototype.sampling = function() {
    // 1. Pick a Kernel (one of the original data points)
    const randomIndex = Math.floor(Math.random() * this.n_samples);
    const selectedKernelCenter = this.tree_data[randomIndex];
    // 2. Sample from that Kernel (add Gaussian noise scaled by bandwidth)
    const newSample = [];
    for (let i = 0; i < this.dimensionality; i++) {
        const noise = sampling.gaussianRandom() * this.bandwidth;
        newSample.push(selectedKernelCenter[i] + noise);
    }
    return newSample;
}



module.exports = {GaussianKDE};