<reponame>oesst/pycog<filename>examples/models/multisensory.py
"""
A multisensory integration task, loosely based on

  A category-free neural population supports evolving demands during decision-making.
  <NAME>, <NAME>, & <NAME>, Nature Neurosci. 2014.

  http://www.nature.com/neuro/journal/v17/n12/full/nn.3865.html

"""
from __future__ import division

import numpy as np

from pycog import tasktools

#-----------------------------------------------------------------------------------------
# Network structure
#-----------------------------------------------------------------------------------------

Nin  = 5
N    = 150
Nout = 2

# E/I
ei, EXC, INH = tasktools.generate_ei(N)
Ne = len(EXC)
Ni = len(INH)

# Input labels
VISUAL_P   = 0 # Positively tuned visual input
AUDITORY_P = 1 # Positively tuned auditory input
VISUAL_N   = 2 # Negatively tuned visual input
AUDITORY_N = 3 # Negatively tuned auditory input
START      = 4 # Start cue

# Units receiving visual input
EXC_VISUAL = EXC[:Ne//3]
INH_VISUAL = INH[:Ni//3]

# Units receiving auditory input
EXC_AUDITORY = EXC[Ne//3:Ne*2//3]
INH_AUDITORY = INH[Ni//3:Ni*2//3]

#-----------------------------------------------------------------------------------------
# Input connectivity
#-----------------------------------------------------------------------------------------

Cin = np.zeros((N, Nin))
Cin[EXC_VISUAL   + INH_VISUAL,   VISUAL_P]   = 1
Cin[EXC_VISUAL   + INH_VISUAL,   VISUAL_N]   = 1
Cin[EXC_AUDITORY + INH_AUDITORY, AUDITORY_P] = 1
Cin[EXC_AUDITORY + INH_AUDITORY, AUDITORY_N] = 1

#-----------------------------------------------------------------------------------------
# Output connectivity: read out from excitatory units only
#-----------------------------------------------------------------------------------------

Cout = np.zeros((Nout, N))
Cout[:,EXC] = 1

#-----------------------------------------------------------------------------------------
# Baseline input
#-----------------------------------------------------------------------------------------

baseline_in = 0.2

#-----------------------------------------------------------------------------------------
# Task structure
#-----------------------------------------------------------------------------------------

modalities  = ['v', 'a', 'va']
freqs       = range(9, 16+1)
boundary    = 12.5
nconditions = len(modalities)*len(freqs)
pcatch      = 5/(nconditions + 1)

fmin = min(freqs)
fmax = max(freqs)

def scale_v_p(f):
    return 0.4 + 0.8*(f - fmin)/(fmax - fmin)

def scale_a_p(f):
    return 0.4 + 0.8*(f - fmin)/(fmax - fmin)

def scale_v_n(f):
    return 0.4 + 0.8*(fmax - f)/(fmax - fmin)

def scale_a_n(f):
    return 0.4 + 0.8*(fmax - f)/(fmax - fmin)

def generate_trial(rng, dt, params):
    #-------------------------------------------------------------------------------------
    # Select task condition
    #-------------------------------------------------------------------------------------

    catch_trial = False
    if params['name'] in ['gradient', 'test']:
        if params.get('catch', rng.rand() < pcatch):
            catch_trial = True
        else:
            modality = params.get('modality', rng.choice(modalities))
            freq     = params.get('freq',     rng.choice(freqs))
    elif params['name'] == 'validation':
        b = params['minibatch_index'] % (nconditions + 1)
        if b == 0:
            catch_trial = True
        else:
            k1, k2   = tasktools.unravel_index(b-1,
                                               (len(modalities), len(freqs)))
            modality = modalities[k1]
            freq     = freqs[k2]
    else:
        raise ValueError("Unknown trial type.")

    #-------------------------------------------------------------------------------------
    # Epochs
    #-------------------------------------------------------------------------------------

    if catch_trial:
        epochs = {'T': 2500}
    else:
        if params['name'] == 'test':
            fixation = 500
        else:
            fixation = 100
        stimulus = 1000
        decision = 300
        T        = fixation + stimulus + decision

        epochs = {
            'fixation': (0, fixation),
            'stimulus': (fixation, fixation + stimulus),
            'decision': (fixation + stimulus, T)
            }
        epochs['T'] = T

    #-------------------------------------------------------------------------------------
    # Trial info
    #-------------------------------------------------------------------------------------

    t, e  = tasktools.get_epochs_idx(dt, epochs) # Time, task epochs in discrete time
    trial = {'t': t, 'epochs': epochs}           # Trial

    if catch_trial:
        trial['info'] = {}
    else:
        # Correct choice
        if freq > boundary:
            choice = 0
        else:
            choice = 1

        # Trial info
        trial['info'] = {'modality':  modality, 'freq': freq, 'choice': choice}

    #-------------------------------------------------------------------------------------
    # Inputs
    #-------------------------------------------------------------------------------------

    X = np.zeros((len(t), Nin))
    if not catch_trial:
        if 'v' in modality:
            X[e['stimulus'],VISUAL_P] = scale_v_p(freq)
            X[e['stimulus'],VISUAL_N] = scale_v_n(freq)
        if 'a' in modality:
            X[e['stimulus'],AUDITORY_P] = scale_a_p(freq)
            X[e['stimulus'],AUDITORY_N] = scale_a_n(freq)
        X[e['stimulus'] + e['decision'],START] = 1
    trial['inputs'] = X

    #-------------------------------------------------------------------------------------
    # Target output
    #-------------------------------------------------------------------------------------

    if params.get('target_output', False):
        Y = np.zeros((len(t), Nout)) # Output matrix
        M = np.zeros_like(Y)         # Mask matrix

        # Hold values
        hi = 1
        lo = 0.2

        if catch_trial:
            Y[:] = lo
            M[:] = 1
        else:
            # Fixation
            Y[e['fixation'],:] = lo

            # Decision
            Y[e['decision'],choice]   = hi
            Y[e['decision'],1-choice] = lo

            # Only care about fixation and decision periods
            M[e['fixation']+e['decision'],:] = 1

        # Outputs and mask
        trial['outputs'] = Y
        trial['mask']    = M

    #-------------------------------------------------------------------------------------

    return trial

# Performance measure
performance = tasktools.performance_2afc

# Termination criterion
TARGET_PERFORMANCE = 90
def terminate(performance_history):
    return np.mean(performance_history[-5:]) > TARGET_PERFORMANCE

# Validation dataset
n_validation = 100*(nconditions + 1)

"""
A multisensory integration task, loosely based on

  A category-free neural population supports evolving demands during decision-making.
  <NAME>, <NAME>, & <NAME>, Nature Neurosci. 2014.

  http://www.nature.com/neuro/journal/v17/n12/full/nn.3865.html

"""
from __future__ import division

import numpy as np

from pycog import tasktools

#-----------------------------------------------------------------------------------------
# Network structure
#-----------------------------------------------------------------------------------------

Nin  = 5
N    = 150
Nout = 2

# E/I
ei, EXC, INH = tasktools.generate_ei(N)
Ne = len(EXC)
Ni = len(INH)

# Input labels
VISUAL_P   = 0 # Positively tuned visual input
AUDITORY_P = 1 # Positively tuned auditory input
VISUAL_N   = 2 # Negatively tuned visual input
AUDITORY_N = 3 # Negatively tuned auditory input
START      = 4 # Start cue

# Units receiving visual input
EXC_VISUAL = EXC[:Ne//3]
INH_VISUAL = INH[:Ni//3]

# Units receiving auditory input
EXC_AUDITORY = EXC[Ne//3:Ne*2//3]
INH_AUDITORY = INH[Ni//3:Ni*2//3]

#-----------------------------------------------------------------------------------------
# Input connectivity
#-----------------------------------------------------------------------------------------

Cin = np.zeros((N, Nin))
Cin[EXC_VISUAL   + INH_VISUAL,   VISUAL_P]   = 1
Cin[EXC_VISUAL   + INH_VISUAL,   VISUAL_N]   = 1
Cin[EXC_AUDITORY + INH_AUDITORY, AUDITORY_P] = 1
Cin[EXC_AUDITORY + INH_AUDITORY, AUDITORY_N] = 1

#-----------------------------------------------------------------------------------------
# Output connectivity: read out from excitatory units only
#-----------------------------------------------------------------------------------------

Cout = np.zeros((Nout, N))
Cout[:,EXC] = 1

#-----------------------------------------------------------------------------------------
# Baseline input
#-----------------------------------------------------------------------------------------

baseline_in = 0.2

#-----------------------------------------------------------------------------------------
# Task structure
#-----------------------------------------------------------------------------------------

modalities  = ['v', 'a', 'va']
freqs       = range(9, 16+1)
boundary    = 12.5
nconditions = len(modalities)*len(freqs)
pcatch      = 5/(nconditions + 1)

fmin = min(freqs)
fmax = max(freqs)

def scale_v_p(f):
    return 0.4 + 0.8*(f - fmin)/(fmax - fmin)

def scale_a_p(f):
    return 0.4 + 0.8*(f - fmin)/(fmax - fmin)

def scale_v_n(f):
    return 0.4 + 0.8*(fmax - f)/(fmax - fmin)

def scale_a_n(f):
    return 0.4 + 0.8*(fmax - f)/(fmax - fmin)

def generate_trial(rng, dt, params):
    #-------------------------------------------------------------------------------------
    # Select task condition
    #-------------------------------------------------------------------------------------

    catch_trial = False
    if params['name'] in ['gradient', 'test']:
        if params.get('catch', rng.rand() < pcatch):
            catch_trial = True
        else:
            modality = params.get('modality', rng.choice(modalities))
            freq     = params.get('freq',     rng.choice(freqs))
    elif params['name'] == 'validation':
        b = params['minibatch_index'] % (nconditions + 1)
        if b == 0:
            catch_trial = True
        else:
            k1, k2   = tasktools.unravel_index(b-1,
                                               (len(modalities), len(freqs)))
            modality = modalities[k1]
            freq     = freqs[k2]
    else:
        raise ValueError("Unknown trial type.")

    #-------------------------------------------------------------------------------------
    # Epochs
    #-------------------------------------------------------------------------------------

    if catch_trial:
        epochs = {'T': 2500}
    else:
        if params['name'] == 'test':
            fixation = 500
        else:
            fixation = 100
        stimulus = 1000
        decision = 300
        T        = fixation + stimulus + decision

        epochs = {
            'fixation': (0, fixation),
            'stimulus': (fixation, fixation + stimulus),
            'decision': (fixation + stimulus, T)
            }
        epochs['T'] = T

    #-------------------------------------------------------------------------------------
    # Trial info
    #-------------------------------------------------------------------------------------

    t, e  = tasktools.get_epochs_idx(dt, epochs) # Time, task epochs in discrete time
    trial = {'t': t, 'epochs': epochs}           # Trial

    if catch_trial:
        trial['info'] = {}
    else:
        # Correct choice
        if freq > boundary:
            choice = 0
        else:
            choice = 1

        # Trial info
        trial['info'] = {'modality':  modality, 'freq': freq, 'choice': choice}

    #-------------------------------------------------------------------------------------
    # Inputs
    #-------------------------------------------------------------------------------------

    X = np.zeros((len(t), Nin))
    if not catch_trial:
        if 'v' in modality:
            X[e['stimulus'],VISUAL_P] = scale_v_p(freq)
            X[e['stimulus'],VISUAL_N] = scale_v_n(freq)
        if 'a' in modality:
            X[e['stimulus'],AUDITORY_P] = scale_a_p(freq)
            X[e['stimulus'],AUDITORY_N] = scale_a_n(freq)
        X[e['stimulus'] + e['decision'],START] = 1
    trial['inputs'] = X

    #-------------------------------------------------------------------------------------
    # Target output
    #-------------------------------------------------------------------------------------

    if params.get('target_output', False):
        Y = np.zeros((len(t), Nout)) # Output matrix
        M = np.zeros_like(Y)         # Mask matrix

        # Hold values
        hi = 1
        lo = 0.2

        if catch_trial:
            Y[:] = lo
            M[:] = 1
        else:
            # Fixation
            Y[e['fixation'],:] = lo

            # Decision
            Y[e['decision'],choice]   = hi
            Y[e['decision'],1-choice] = lo

            # Only care about fixation and decision periods
            M[e['fixation']+e['decision'],:] = 1

        # Outputs and mask
        trial['outputs'] = Y
        trial['mask']    = M

    #-------------------------------------------------------------------------------------

    return trial

# Performance measure
performance = tasktools.performance_2afc

# Termination criterion
TARGET_PERFORMANCE = 90
def terminate(performance_history):
    return np.mean(performance_history[-5:]) > TARGET_PERFORMANCE

# Validation dataset
n_validation = 100*(nconditions + 1)
A multisensory integration task, loosely based on    A category-free neural population supports evolving demands during decision-making.   <NAME>, <NAME>, & <NAME>, Nature Neurosci. 2014.    http://www.nature.com/neuro/journal/v17/n12/full/nn.3865.html #----------------------------------------------------------------------------------------- # Network structure #----------------------------------------------------------------------------------------- # E/I # Input labels # Positively tuned visual input # Positively tuned auditory input # Negatively tuned visual input # Negatively tuned auditory input # Start cue # Units receiving visual input # Units receiving auditory input #----------------------------------------------------------------------------------------- # Input connectivity #----------------------------------------------------------------------------------------- #----------------------------------------------------------------------------------------- # Output connectivity: read out from excitatory units only #----------------------------------------------------------------------------------------- #----------------------------------------------------------------------------------------- # Baseline input #----------------------------------------------------------------------------------------- #----------------------------------------------------------------------------------------- # Task structure #----------------------------------------------------------------------------------------- #------------------------------------------------------------------------------------- # Select task condition #------------------------------------------------------------------------------------- #------------------------------------------------------------------------------------- # Epochs #------------------------------------------------------------------------------------- #------------------------------------------------------------------------------------- # Trial info #------------------------------------------------------------------------------------- # Time, task epochs in discrete time # Trial # Correct choice # Trial info #------------------------------------------------------------------------------------- # Inputs #------------------------------------------------------------------------------------- #------------------------------------------------------------------------------------- # Target output #------------------------------------------------------------------------------------- # Output matrix # Mask matrix # Hold values # Fixation # Decision # Only care about fixation and decision periods # Outputs and mask #------------------------------------------------------------------------------------- # Performance measure # Termination criterion # Validation dataset