r2bScore.cpp

#include <vector>
#include <random>
#include <numeric>
#include <pybind11/stl.h>
#include <pybind11/stl_bind.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <Python.h>
#include <iostream>
using namespace std;
namespace py = pybind11;
#include "tcHeader.h"
PYBIND11_MAKE_OPAQUE(std::vector<double>);


//////////////////////////////////////////////////////////////////////////
// stuff for r2b calculations.
//////////////////////////////////////////////////////////////////////////

/**
 * trialAve: Mean over ALL trials for each frame and cell.
 * Returns the peak frame index for each cell.
 * ret[cell][frame]
 */
vector< unsigned int > trialAve( const double* data,
		vector< vector< double > >& ret,
		unsigned int numTrials, unsigned int numCells,
		const AnalysisParams& ap )
{
	unsigned int nf = ap.circShuffleFrames;
	ret.resize( numCells );
	for ( unsigned int cell = 0; cell < numCells; cell++ )
		ret[cell].assign( nf, 0.0 );

	vector< unsigned int > pk;
	for ( unsigned int cell = 0; cell < numCells; cell++ ) {
		vector< double >& mean = ret[cell];
		for ( unsigned int ff = 0; ff < nf; ++ff ) {
			const double* d = data + cell + (ff + ap.csOnsetFrame - ap.circPad) * numCells * numTrials;
			for ( unsigned int tt = 0; tt < numTrials; tt++ )
				mean[ff] += d[ tt * numCells ];
		}
		pk.push_back( max_element( mean.begin(), mean.end() ) - mean.begin() );
		for ( unsigned int ff = 0; ff < nf; ++ff )
			mean[ff] /= numTrials;
	}
	return pk;
}

/**
 * shuffleTrialAve: Circularly shuffle ALL trials, return the mean and the
 * peak frame of that shuffled mean. ret[cell][frame].
 */
vector< unsigned int > shuffleTrialAve( const double* data,
		vector< vector< double > >& ret,
		unsigned int numTrials, unsigned int numCells,
		std::uniform_int_distribution<std::mt19937::result_type>& shuffler,
		std::mt19937& rng, const AnalysisParams& ap )
{
	unsigned int nf = ap.circShuffleFrames;
	ret.resize( numCells );
	for ( auto& r : ret ) r.assign( nf, 0.0 );
	unsigned int ncxnt = numCells * numTrials;

	vector< unsigned int > shuff( numTrials );
	for ( unsigned int tt = 0; tt < numTrials; ++tt )
		shuff[tt] = shuffler( rng );

	vector< unsigned int > pk;
	for ( unsigned int cell = 0; cell < numCells; cell++ ) {
		vector< double >& mean = ret[cell];
		for ( unsigned int tt = 0; tt < numTrials; tt++ ) {
			for ( unsigned int ff = 0; ff < nf; ++ff ) {
				unsigned int circff = ap.csOnsetFrame - ap.circPad + (ff + shuff[tt]) % nf;
				mean[ff] += data[ cell + circff * ncxnt + tt * numCells ];
			}
		}
		pk.push_back( max_element( mean.begin(), mean.end() ) - mean.begin() );
		for ( unsigned int ff = 0; ff < nf; ++ff )
			mean[ff] /= numTrials;
	}
	return pk;
}

double r2b( const vector< double >& ave, unsigned int pkfr) {
	unsigned int f0 = pkfr > 0 ? pkfr-1 : 0;
	unsigned int f2 = pkfr < ave.size()-1 ? pkfr+1 : pkfr;
	double ridge = ave[f0] + ave[pkfr] + ave[f2];
	double background = std::accumulate( ave.begin(), ave.end(), 0.0 ) - ridge;
	if ( ridge < 0.0 || background <= 0.0 ) {
		return 0.0;
	}
	return ridge/background;
}

vector< CellScore > r2bScore( py::array_t<double> xs, const AnalysisParams& ap, double r2bThresh, double r2bPercentile)
{
	unsigned int SEED = 1234;
	py::buffer_info info = xs.request();
	auto data = static_cast< double* >( info.ptr);
	assert( info.shape[0] == ap.numFrames * ap.numCells * ap.numTrials );

	std::mt19937 rng( SEED );
	std::uniform_int_distribution<std::mt19937::result_type> shuffler(0, ap.circShuffleFrames);

	vector< CellScore > ret( ap.numCells );
	vector< vector< double > > tave;
	trialAve( data, tave, ap.numTrials, ap.numCells, ap );

	unsigned int minSepFrames = (unsigned int)round( ap.minPeakSep / ap.frameDt );
	if ( minSepFrames < 1 ) minSepFrames = 1;
	unsigned int nf = ap.circShuffleFrames;

	// Peel candidate peaks from the full mean trace.
	// Real R2B is also computed on the full mean.
	vector< vector< unsigned int > > candidates( ap.numCells );
	vector< vector< double > >       realR2B( ap.numCells );

	for ( unsigned int cc = 0; cc < ap.numCells; cc++ ) {
		ret[cc].meanTrace = tave[cc];
		vector< bool > available( nf, true );
		while ( true ) {
			unsigned int pkFrame = nf; // sentinel
			double pkVal = -1e30;
			for ( unsigned int ff = 0; ff < nf; ff++ )
				if ( available[ff] && tave[cc][ff] > pkVal ) {
					pkVal = tave[cc][ff]; pkFrame = ff;
				}
			if ( pkFrame == nf || pkVal <= 0.0 ) break;
			candidates[cc].push_back( pkFrame );
			realR2B[cc].push_back( r2b( tave[cc], pkFrame ) );
			unsigned int lo = (pkFrame >= minSepFrames) ? pkFrame - minSepFrames : 0;
			unsigned int hi = min( pkFrame + minSepFrames, nf - 1 );
			for ( unsigned int ff = lo; ff <= hi; ff++ )
				available[ff] = false;
		}
	}

	// Bootstrap.
	// Primary peak uses the max-statistic: shuffled data finds its own peak
	// and we compare its R2B against the real primary-peak R2B.
	// Secondary peaks use fixed-location tests on the shuffled mean.
	vector< vector< double > > bootCounts( ap.numCells );
	vector< double >           sumShuff( ap.numCells, 0.0 );
	for ( unsigned int cc = 0; cc < ap.numCells; cc++ )
		bootCounts[cc].assign( candidates[cc].size(), 0.0 );

	vector< vector< double > > shuffTave;
	for ( unsigned int ii = 0; ii < ap.numShuffle; ++ii ) {
		auto shuffPk = shuffleTrialAve( data, shuffTave, ap.numTrials, ap.numCells, shuffler, rng, ap );
		for ( unsigned int cc = 0; cc < ap.numCells; cc++ ) {
			if ( candidates[cc].empty() ) continue;
			// Primary: max-statistic — shuffled R2B at shuffled peak.
			double sr0 = r2b( shuffTave[cc], shuffPk[cc] );
			if ( !isnan( sr0 ) ) {
				sumShuff[cc] += sr0;
				bootCounts[cc][0] += ( realR2B[cc][0] > sr0 );
			}
			// Secondary peaks: fixed location on shuffled mean.
			for ( unsigned int pk = 1; pk < candidates[cc].size(); pk++ ) {
				double sr = r2b( shuffTave[cc], candidates[cc][pk] );
				if ( !isnan( sr ) )
					bootCounts[cc][pk] += ( realR2B[cc][pk] > sr );
			}
		}
	}

	// Fill CellScore for each cell.
	for ( unsigned int cc = 0; cc < ap.numCells; cc++ ) {
		if ( candidates[cc].empty() ) continue;

		ret[cc].meanPkIdx       = candidates[cc][0];
		ret[cc].baseScore       = realR2B[cc][0];
		ret[cc].meanScore       = sumShuff[cc] / ap.numShuffle;
		ret[cc].percentileScore = bootCounts[cc][0] / ap.numShuffle;
		ret[cc].sigMean         = ( ret[cc].baseScore > r2bThresh * ret[cc].meanScore );
		ret[cc].sigBootstrap    = ( ret[cc].percentileScore > r2bPercentile / 100.0 );

		// Dip threshold: mean + dipSdev * sdev of this cell's mean trace.
		double dipThresh;
		{
			double sum = 0.0, sumSq = 0.0;
			for ( auto v : tave[cc] ) { sum += v; sumSq += v * v; }
			double mn = sum / nf;
			double sd = sqrt( sumSq / nf - mn * mn );
			dipThresh = mn + ap.dipSdev * sd;
		}

		for ( unsigned int pk = 0; pk < candidates[cc].size(); pk++ ) {
			double pctile = bootCounts[cc][pk] / ap.numShuffle;
			if ( pctile <= r2bPercentile / 100.0 ) continue;

			unsigned int pkFrame = candidates[cc][pk];

			// Dip check: require dip between this peak and the nearest already-accepted peak.
			if ( !ret[cc].allPkIndices.empty() ) {
				unsigned int nearest = ret[cc].allPkIndices[0];
				unsigned int minDist = (pkFrame > nearest) ? pkFrame - nearest : nearest - pkFrame;
				for ( auto prevIdx : ret[cc].allPkIndices ) {
					unsigned int d = (pkFrame > prevIdx) ? pkFrame - prevIdx : prevIdx - pkFrame;
					if ( d < minDist ) { minDist = d; nearest = prevIdx; }
				}
				if ( !hasDipBetween( tave[cc], nearest, pkFrame, dipThresh, ap.dipFrames ) )
					continue;
			}

			ret[cc].allPkIndices.push_back( pkFrame );
			ret[cc].allPkScores.push_back( realR2B[cc][pk] );
			ret[cc].allPkPvalues.push_back( 1.0 - pctile );
		}
	}
	return ret;
}