tiScore.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>

using namespace std;

namespace py = pybind11;
#include "tcHeader.h"
PYBIND11_MAKE_OPAQUE(std::vector<double>);

/// Returns the mean dfbf[ frame# ] for the specified shuffle[ trial# ]
/// using data[ trial# ][ frame# ]. Operates on data for a single cell.
vector< double > aveOfTrials( 
				const vector< vector< double > >& data, 
				const vector< unsigned int >& shuff, 
				unsigned int numFrames )
{
	unsigned int numTrials = data.size();
	vector< double > aveBin( numFrames, 0.0 );
	assert( numTrials == data.size() && numTrials > 0  );
	for ( unsigned int ii = 0; ii < numTrials; ++ii ) {
		vector< double >::iterator aptr = aveBin.begin();
		auto dd = data[ii].begin();
		for ( unsigned int bb = shuff[ii]; bb < shuff[ii] + numFrames; bb++, aptr++) {
			*aptr += *( dd + bb%numFrames );
		}
	}
	for( vector< double >::iterator aptr = aveBin.begin(); aptr != aveBin.end(); aptr++ ) {
		*aptr /= numFrames;
	}
	return aveBin;
}

/// returns binned data for a single trial.
vector< double > binner( const vector< double >& orig, unsigned int binFrames )
{
	if (binFrames <= 1 || binFrames >= orig.size()/2 ) return orig;
	unsigned int numRet = orig.size() / binFrames;
	vector< double > ret( numRet, 0.0 );
	for ( unsigned int idx = 0; idx < numRet; idx++ ) {
		for ( unsigned int bf = 0; bf < binFrames; bf++ )
			ret[idx] += orig[ idx * binFrames + bf ];
	}
	return ret;
}

/// returns binned data for a single cell, all trials.
vector< vector< double > > binnedCellData( const vector< vector< double > >& data, unsigned int numFrames, unsigned int binFrames )
{
	unsigned int numTrials = data.size();
	vector< vector< double > > ret( numTrials );
	for( unsigned int trialIdx = 0; trialIdx < numTrials; trialIdx++ ) {
		ret[ trialIdx ] = binner( data[ trialIdx ], binFrames );
	}
	return ret;
}


/// Return the average tuning curve of the cell. Uses peeling to find all
/// significant peaks separated by at least ap.minPeakSep seconds.
void tuningCurve( const vector< vector< double >>& data, const vector< vector< unsigned int > >& shuff, unsigned int numFrames, const AnalysisParams& ap, CellScore& cs )
{
	unsigned int numTrials = data.size();
	vector< vector< double >> binnedCellData_ = binnedCellData( data, numFrames, ap.binFrames );
	unsigned int numBins = binnedCellData_[0].size();
	vector< unsigned int > zeros( numTrials, 0 );
	cs.meanTrace = aveOfTrials( binnedCellData_, zeros, numBins );

	// Pre-compute all shuffled mean traces once; reused for every peak.
	vector< vector< double > > shuffMeans( ap.numShuffle );
	for ( unsigned int ii = 0; ii < ap.numShuffle; ii++ )
		shuffMeans[ii] = aveOfTrials( binnedCellData_, shuff[ii], numBins );

	unsigned int minSepBins = (unsigned int)round( ap.minPeakSep / (ap.frameDt * ap.binFrames) );
	if ( minSepBins < 1 ) minSepBins = 1;

	// Dip threshold: mean + dipSdev * sdev of the mean trace.
	double dipThresh;
	{
		double sum = 0.0, sumSq = 0.0;
		for ( auto v : cs.meanTrace ) { sum += v; sumSq += v * v; }
		double mn = sum / numBins;
		double sd = sqrt( sumSq / numBins - mn * mn );
		dipThresh = mn + ap.dipSdev * sd;
	}

	// Availability mask: bins set to false are inside a previously accepted
	// peak window and are excluded from future searches.
	vector< bool > available( numBins, true );
	bool firstPeak = true;

	while ( true ) {
		// Find the highest available bin in the original mean trace.
		unsigned int pkBin = numBins; // sentinel
		double pkVal = -1e30;
		for ( unsigned int bb = 0; bb < numBins; bb++ ) {
			if ( available[bb] && cs.meanTrace[bb] > pkVal ) {
				pkVal = cs.meanTrace[bb];
				pkBin = bb;
			}
		}
		if ( pkBin == numBins || pkVal <= 0.0 ) break;

		// Adjacent bin for two-bin significance test.
		unsigned int adjBin;
		if ( pkBin == 0 ) adjBin = 1;
		else if ( pkBin == numBins - 1 ) adjBin = numBins - 2;
		else adjBin = (cs.meanTrace[pkBin-1] > cs.meanTrace[pkBin+1]) ? pkBin-1 : pkBin+1;

		unsigned int numOK = 0;
		for ( unsigned int ii = 0; ii < ap.numShuffle; ii++ )
			numOK += ( cs.meanTrace[pkBin]  > shuffMeans[ii][pkBin]  ) &&
			         ( cs.meanTrace[adjBin] > shuffMeans[ii][adjBin] );

		if ( (numOK * 100) <= (99 * ap.numShuffle) ) break;

		// Dip check: for 2nd+ peaks the trace must drop below dipThresh for
		// at least dipFrames consecutive bins between this and the nearest accepted peak.
		unsigned int lo = (pkBin >= minSepBins) ? pkBin - minSepBins : 0;
		unsigned int hi = min( pkBin + minSepBins, numBins - 1 );
		if ( !firstPeak ) {
			unsigned int nearest = cs.allPkIndices[0];
			unsigned int minDist = (pkBin > nearest) ? pkBin - nearest : nearest - pkBin;
			for ( auto prevIdx : cs.allPkIndices ) {
				unsigned int d = (pkBin > prevIdx) ? pkBin - prevIdx : prevIdx - pkBin;
				if ( d < minDist ) { minDist = d; nearest = prevIdx; }
			}
			if ( !hasDipBetween( cs.meanTrace, nearest, pkBin, dipThresh, ap.dipFrames ) ) {
				for ( unsigned int bb = lo; bb <= hi; bb++ ) available[bb] = false;
				continue;
			}
		}

		double pval = 1.0 - (double)numOK / ap.numShuffle;
		cs.allPkIndices.push_back( pkBin );
		cs.allPkScores.push_back( cs.meanTrace[pkBin] );
		cs.allPkPvalues.push_back( pval );

		if ( firstPeak ) {
			cs.meanPkIdx = pkBin;
			cs.meanScore = cs.meanTrace[pkBin];
			cs.sigMean   = true;
			firstPeak    = false;
		}

		// Mark the separation window as unavailable.
		for ( unsigned int bb = lo; bb <= hi; bb++ )
			available[bb] = false;
	}

	if ( firstPeak ) {
		cs.meanPkIdx = (unsigned int)(
			max_element( cs.meanTrace.begin(), cs.meanTrace.end() ) - cs.meanTrace.begin() );
		cs.meanScore = cs.meanTrace[cs.meanPkIdx];
		cs.sigMean   = false;
	}
}

/// args: data[ frame# ][ trial# ][ cell# ]
/// Pass-back arg: ret[ cell# ][ trial# ][ frame# ]
/// Fills in only the data in the window around the stimuli, from 
/// CS_ONSET_FRAME - CIRC_PAD to US_ONSET_FRAME + CIRC_PAD.
void reorderData( const double* data, unsigned int numCells, unsigned int numTrials, unsigned int numFrames, vector< vector< vector< double >>>& ret, const AnalysisParams& ap ) 
{
	ret.clear();
	ret.resize( numCells );
	unsigned int ncxnt = numCells * numTrials;
	unsigned int nf = ap.circShuffleFrames;
	for ( unsigned int cell = 0; cell < numCells; cell++ ) {
		vector< vector< double > >& rc = ret[cell];
		rc.resize( numTrials );
		for( unsigned int tt = 0; tt < numTrials; tt++ ) {
			vector< double >& rct = rc[ tt ];
			rct.resize( nf,  0.0 );
			for( unsigned int ff = 0; ff < nf; ++ff ) {
				rct[ff] = data[ cell + tt * numCells + (ff + ap.csOnsetFrame - ap.circPad) * ncxnt];
			}
		}
	}
}

// Returns threshold for sig peak for data[trial#][frame#]. The
// data has already been sub-selected for a given cell.
double findCellThresh(
		const vector< vector< double > >& data, unsigned int numFrames, double transientThresh )
{
	double sum = 0.0;
	double sq = 0.0;
	for( auto trialD = data.begin(); trialD != data.end(); ++trialD ) {
		for( auto frameD = trialD->begin(); frameD != trialD->end(); ++frameD ) {
			sum += *frameD;
			sq += *frameD * *frameD;
		}
	}
	double numSamples = numFrames * data.size(); //numFrames * numTrials
	double mean = sum/numSamples;
	return mean + transientThresh * sqrt( sq / numSamples - mean*mean );
}

/// Returns time stamps of transients trial# ][ transient# ]. 
// given the reordered dfbf data[ trial# ][ frame# ]
// A given trial may have zero or more transients.
vector< vector< unsigned int > > findTransients(
		const vector< vector< double > >& data, double cellThresh, unsigned int numFrames )
{
	unsigned int numTrials = data.size();
	vector< vector< unsigned int > > ret ( numTrials );
	vector< vector< unsigned int > >::iterator trialT = ret.begin();
	for( auto trialD = data.begin(); trialD != data.end(); ++trialD, ++trialT ) {
		trialT->clear();
		double lastFrame = 0.0;
		bool refractory = 0;
		for( unsigned int ii = 0; ii < numFrames; ++ii ) {
			double frameD = (*trialD)[ii];
			// Don't permit another transient till signal goes < cellThresh
			if (frameD > cellThresh && frameD > lastFrame && !refractory) {
				trialT->push_back( ii );
				refractory = 1;
			} else {
				refractory = (frameD > cellThresh);
			}
			lastFrame = frameD;
		}
	}
	return ret;
}

// Returns transient rates for the current
double findAveTransientRate( const vector< vector< unsigned int > >& transients, double trialDuration, unsigned int numTrials ) 
{
	double ret = 0.0;
	for( auto trial: transients ) {
		for ( auto bin: trial ) {
			ret += bin;
		}
	}
	return ret/ ( trialDuration * numTrials );
}

/// Returns temporal information for the current cell.
double findTemporalInformation( const vector< vector< unsigned int > >& transients, double aveTransientRate, const vector< unsigned int >& shuff, unsigned int numFrames )
{
	double ret = 0.0;
	unsigned int numTrials = transients.size();

	vector< double > binAve( numFrames, 0.0 );
	for( unsigned int trialIdx = 0; trialIdx < numTrials; ++trialIdx ) {
		for ( unsigned int bin: transients[trialIdx] ) {
			unsigned int b = ( bin + shuff[trialIdx] ) % numFrames;
			binAve[b] += 1.0;
		}
	}
	for ( auto bb: binAve ) {
		if ( bb > 0.0 ) {
			ret += bb * log2( bb / aveTransientRate );
		}
	}
	if ( aveTransientRate > 0.0 )
		return ret / ( aveTransientRate * numFrames );
	return 0.0;
}

// Returns the specified percentile value of shuffled TIs, for current cell
double percentileTI( const vector< vector< unsigned int > >& transients, double aveTransientRate, const AnalysisParams& ap, double percentile )
{
	std::random_device dev;
	// std::mt19937 rng(dev());
	std::mt19937 rng( 1234 );
	std::uniform_int_distribution<std::mt19937::result_type> shuffler(0, ap.circShuffleFrames );
	unsigned int numTrials = transients.size();
	assert( numTrials > 0 );

	// allocate tiRank[iter#], for sorting.
	vector< double > tiRank( ap.numShuffle ); 

	for ( unsigned int iter = 0; iter < ap.numShuffle; ++iter ) {
		vector< unsigned int > shuff( numTrials );
		for ( unsigned int tt = 0; tt < numTrials; ++tt ) {
			shuff[tt] = shuffler( rng );
		}
		tiRank[iter] = findTemporalInformation( transients, aveTransientRate, shuff, ap.circShuffleFrames  );
	}

	unsigned int percIdx = (percentile * ap.numShuffle ) / 100;
	sort( tiRank.begin(), tiRank.end() );
	return tiRank[percIdx];
}

// Fraction of trials in which there was >= 1 transient. frac[ cell# ]
double fracTrialsFired( const vector< vector< unsigned int > >& transients )
{
	assert( transients.size() > 0 );
	double ret = 0.0;
	for( auto tt : transients ) {
		ret += (tt.size() > 0 );
	}
	return ret / transients.size();
}

/// Computes TI score for a given neuron. zero means it fails a criterion.
// Should really return a lot of stuff: the mean bin vector, the
// baseTI, the PTI, the frac fired, the percentile bin height.
// Suggest return a 2-D vector or a small struct.
CellScore cellTIscore( const vector< vector< double > >& data, const AnalysisParams& ap, const TiAnalysisParams& tip )
{
	CellScore cs;
	unsigned int numTrials = data.size();
	assert( numTrials > 0 );
	unsigned int numFrames = data[0].size();
	vector< unsigned int > nonShuff( numTrials, 0 );
	vector< double > aveOfTrials_ = aveOfTrials( data, nonShuff, ap.circShuffleFrames );

	double cellThresh = findCellThresh( data, ap.circShuffleFrames, tip.transientThresh );
	vector< vector< unsigned int > > transients =
			findTransients( data, cellThresh, ap.circShuffleFrames  );

	double trialDuration = tip.frameDt * ap.circShuffleFrames;
	double aveTransientRate = findAveTransientRate( transients, trialDuration, numTrials );
	cs.baseScore = findTemporalInformation( transients, aveTransientRate, nonShuff, ap.circShuffleFrames  );
	cs.percentileScore = percentileTI( transients, aveTransientRate, ap, tip.tiPercentile );
	cs.fracTrialsFired = fracTrialsFired( transients );

	// Fill up the RNGs
	std::random_device dev;
	std::mt19937 rng(dev());
	std::uniform_int_distribution<std::mt19937::result_type> shuffler(0, numFrames / ap.binFrames );
	vector< vector< unsigned int > > binShuff( ap.numShuffle );
	for ( unsigned int ii = 0; ii < ap.numShuffle; ii++ ) {
		binShuff[ii].resize( numTrials );
		for ( unsigned int jj = 0; jj < numTrials; jj++ ) {
			binShuff[ii][jj] = shuffler( rng );
		}
	}
	tuningCurve( data, binShuff, numFrames, ap, cs );

	cs.sigBootstrap = ( cs.baseScore > cs.percentileScore ) && (cs.fracTrialsFired > tip.fracTrialsFiredThresh );

	return cs;
}


// Returns the reliability index.
vector< CellScore > tiScore( py::array_t<double> xs, const AnalysisParams& ap, const TiAnalysisParams& tip )
{
	py::buffer_info info = xs.request();
	auto data = static_cast< double* >( info.ptr);
	// unsigned int numDat = info.itemsize;

	vector< vector< vector< double >>> reorderedData;
	reorderData( data, ap.numCells, ap.numTrials, ap.numFrames, reorderedData, ap );

	vector< CellScore > ret( ap.numCells );
	for( unsigned int cellIdx = 0; cellIdx < ap.numCells; cellIdx++ ) {
		ret[cellIdx] = cellTIscore( reorderedData[cellIdx], ap, tip );
	}

	return ret;
}