pool.go

// Code generated by "goal build"; DO NOT EDIT.
//line pool.goal:1
// Copyright (c) 2019, The Emergent Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package axon

import (
	"cogentcore.org/core/base/atomicx"
	"cogentcore.org/core/math32"
	"github.com/emer/axon/v2/fsfffb"
	"log"
	"sync/atomic"
)

//gosl:start
//gosl:import "github.com/emer/axon/v2/fsfffb"

// PoolIntVars are int32 pool variables, for computing fsfffb inhibition etc.
// Note that we use int32 instead of uint32 so that overflow errors can be detected.
// See [PoolVars] for float32 variables.
type PoolIntVars int32 //enums:enum

const (
	// PoolNeurSt is the starting layer-wise index within the list
	// of neurons in this pool.
	// Add layer starting neuron index (NeurSt) to get index into global
	// network neurons list.
	PoolNeurSt PoolIntVars = iota

	// PoolNeurEd is the ending (exclusive) layer-wise index within the list
	// of neurons in this pool.
	// Add layer starting neuron index (NeurSt) to get index into global
	// network neurons list.
	PoolNeurEd

	// PoolLayerIdx is the layer index for this pool.
	PoolLayerIdx

	// PoolIsLayer is true (> 0) if this pool represents the entire layer,
	// which is always the first pool in the list of pools for a layer.
	PoolIsLayer

	// Clamped if true (!=0), this layer is hard-clamped and should
	// use GeExts exclusively for PV.
	Clamped

	// PoolGated is true (> 0) if this pool gated (for [MatrixLayer], [BGThalLayer])
	PoolGated

	// FFsRawInt is the int32 atomic add compatible integration of [fsfffb.FFsRaw].
	FFsRawInt

	// FBsRawInt is the int32 atomic add compatible integration of [fsfffb.FBsRaw].
	FBsRawInt

	// GeExtRawInt is the int32 atomic add compatible integration of [fsfffb.GeExtRaw].
	GeExtRawInt

	// PoolIntAvgMaxStart is the starting point for int32 AvgMax variables.
	// Use AvgMaxIntVarIndex to get the relevant variable index.
	// There are only values for Cycle phase, for the different variables.
	PoolIntAvgMaxStart
)

// AvgMax are Avg and Max
type AvgMax int32 //enums:enum

const (
	Avg AvgMax = iota
	Max
)

// AvgMaxPhases are the different Phases over which AvgMax values are tracked.
type AvgMaxPhases int32 //enums:enum -trim-prefix AM

const (
	// Cycle is the current cycle, which is the source for the rest.
	AMCycle AvgMaxPhases = iota

	// Minus is at the end of the minus phase.
	AMMinus

	// Plus is at the end of the plus phase.
	AMPlus

	// Prev is at the end of the previous plus phase.
	AMPrev
)

// AvgMaxVars are the different Neuron variables for which [AvgMaxPhases]
// is computed.
type AvgMaxVars int32 //enums:enum -trim-prefix AM

const (
	// CaSpkP is the primary variable for tracking overall pool activity
	// over a recent timescale, integrated at roughly 40 msec time constant.
	AMCaSpkP AvgMaxVars = iota

	// CaSpkD is a slower moving activation signal, capable of reflecting
	// activity over the entire trial.
	AMCaSpkD

	// SpkMax is the maximum CaSpkP over the trial of processing.
	AMSpkMax

	// Act is the computed rate-code equivalent of current spike rate.
	AMAct

	// GeInt is the integrated running-average value of excitatory conductance.
	AMGeInt

	// GiInt is the integrated running-average value of inhibitory conductance.
	AMGiInt

	// AvgDif is the integrated AvgDif between ActPct - TrgAvg.
	// Only the Plus phase is used.
	AMAvgDif
)

const (
	// poolFloatAvgMaxStart is the starting index for AvgMax float32 variables.
	poolFloatAvgMaxStart = fsfffb.InhibVarsN

	PoolVarsN = poolFloatAvgMaxStart + fsfffb.InhibVars(int32(AvgMaxVarsN)*int32(AvgMaxN)*int32(AvgMaxPhasesN))

	PoolIntVarsTot = PoolIntAvgMaxStart + PoolIntVars(int32(AvgMaxVarsN)*int32(AvgMaxN))
)

// avgMaxToNeuron is the mapping from AvgMaxVars to neuron vars.
var avgMaxToNeuron = [AMAvgDif]NeuronVars{CaSpkP, CaSpkD, SpkMax, Act, GeInt, GiInt}

// AvgMaxVarIndex returns the variable index for accessing
// [Pools] AvgMax float32 variables.
func AvgMaxVarIndex(vr AvgMaxVars, phase AvgMaxPhases, am AvgMax) uint32 {
	return uint32(poolFloatAvgMaxStart) + uint32(vr)*uint32(AvgMaxN)*uint32(AvgMaxPhasesN) + uint32(phase)*uint32(AvgMaxN) + uint32(am)
}

// AvgMaxIntVarIndex returns the variable index for accessing
// [Pools] AvgMax int32 variables.  Avg = Sum actually.
// There are only values for the Cycle phase level.
func AvgMaxIntVarIndex(vr AvgMaxVars, am AvgMax) uint32 {
	return uint32(PoolIntAvgMaxStart) + uint32(vr)*uint32(AvgMaxN) + uint32(am)
}

// PoolAvgMax returns an AvgMax value for given variable, phase,
// and Avg or Max, for given pool index and data index.
func PoolAvgMax(vr AvgMaxVars, phase AvgMaxPhases, am AvgMax, pi, di uint32) float32 {
	return Pools.Value(int(AvgMaxVarIndex(vr, phase, am)), int(pi), int(di))
}

// PoolNNeurons returns the number of neurons in the given pool.
// pi = global pool index.
func PoolNNeurons(pi uint32) int32 {
	return PoolsInt.Value(int(PoolNeurEd), int(pi), int(0)) - PoolsInt.Value(int(PoolNeurSt), int(pi), int(0))
}

// PoolAvgMaxInit initializes the AvgMax Int accumulators for Cycle vals
// for update start. always left init'd so generally unnecessary.
// pi = global pool index.
func PoolAvgMaxInit(pi, di uint32) {
	for vr := range AMAvgDif {
		PoolsInt.Set(0, int(AvgMaxIntVarIndex(vr, Avg)), int(pi), int(di))
		PoolsInt.Set(0, int(AvgMaxIntVarIndex(vr, Max)), int(pi), int(di))
	}
}

// PoolAvgMaxZero initializes all the AvgMax values to zero.
// pi = global pool index.
func PoolAvgMaxZero(pi, di uint32) {
	PoolAvgMaxInit(pi, di)
	for vr := range AMAvgDif {
		for ph := range AvgMaxPhasesN {
			Pools.Set(0, int(AvgMaxVarIndex(vr, ph, Avg)), int(pi), int(di))
			Pools.Set(0, int(AvgMaxVarIndex(vr, ph, Max)), int(pi), int(di))
		}
	}
}

// PoolAvgMaxUpdateVar updates the AvgMax value based on given value.
// pi = global pool index.
func PoolAvgMaxUpdateVar(vr AvgMaxVars, pi, di uint32, val float32) {
	n := float32(PoolNNeurons(pi))
	floatToInt := float32(uint32(1) << 20)
	floatToSum := floatToInt / n
	vis := AvgMaxIntVarIndex(vr, Avg)
	vim := AvgMaxIntVarIndex(vr, Max)
	atomic.AddInt32(PoolsInt.ValuePtr(int(vis), int(pi), int(di)), int32(val*floatToSum))
	atomicx.MaxInt32(PoolsInt.ValuePtr(int(vim), int(pi), int(di)), int32(val*floatToInt))
}

// PoolAvgMaxUpdate updates the AvgMax values based on current neuron values.
// pi = global pool index.
func PoolAvgMaxUpdate(pi, di, ni uint32) {
	PoolAvgMaxUpdateVar(AMCaSpkP, pi, di, math32.Abs(Neurons.Value(int(avgMaxToNeuron[AMCaSpkP]), int(ni), int(di))))
	PoolAvgMaxUpdateVar(AMCaSpkD, pi, di, math32.Abs(Neurons.Value(int(avgMaxToNeuron[AMCaSpkD]), int(ni), int(di))))
	PoolAvgMaxUpdateVar(AMSpkMax, pi, di, math32.Abs(Neurons.Value(int(avgMaxToNeuron[AMSpkMax]), int(ni), int(di))))
	PoolAvgMaxUpdateVar(AMAct, pi, di, math32.Abs(Neurons.Value(int(avgMaxToNeuron[AMAct]), int(ni), int(di))))
	PoolAvgMaxUpdateVar(AMGeInt, pi, di, math32.Abs(Neurons.Value(int(avgMaxToNeuron[AMGeInt]), int(ni), int(di))))
	PoolAvgMaxUpdateVar(AMGiInt, pi, di, math32.Abs(Neurons.Value(int(avgMaxToNeuron[AMGiInt]), int(ni), int(di))))
}

// PoolAvgMaxCalcVar does Calc on Cycle level, and re-inits, for given Var
func PoolAvgMaxCalcVar(vr AvgMaxVars, pi, di uint32) {
	floatFromInt := float32(1.0) / float32(uint32(1)<<20)
	vim := AvgMaxIntVarIndex(vr, Max)
	sum := PoolsInt.Value(int(vim), int(pi), int(di))
	if sum < 0 {
		//gosl:end
		log.Println("PoolAvgMaxCalc overflow in Sum", "pi:", pi, "di:", di, "sum:", sum)
		//gosl:start
		sum = int32(uint32(1) << 20)
	}
	Pools.Set(float32(sum)*floatFromInt, int(AvgMaxVarIndex(vr, AMCycle, Avg)), int(pi), int(di))
	mx := PoolsInt.Value(int(AvgMaxIntVarIndex(vr, Max)), int(pi), int(di))
	if mx < 0 {
		//gosl:end
		log.Println("PoolAvgMaxCalc overflow in Max", "pi:", pi, "di:", di, "max:", mx)
		//gosl:start
		mx = int32(uint32(1) << 20)
	}
	Pools.Set(float32(mx)*floatFromInt, int(AvgMaxVarIndex(vr, AMCycle, Max)), int(pi), int(di))
}

// PoolAvgMaxCalc does Calc on Cycle level, and re-inits
func PoolAvgMaxCalc(pi, di uint32) {
	for vr := range AMAvgDif { // don't do AvgDif
		PoolAvgMaxCalcVar(vr, pi, di)
	}
}

// PoolAvgDifInit initializes the AvgMax AvgDif Int accumulators for Cycle vals
// for update start. always left init'd so generally unnecessary.
// pi = global pool index.
func PoolAvgDifInit(pi, di uint32) {
	PoolsInt.Set(0, int(AvgMaxIntVarIndex(AMAvgDif, Avg)), int(pi), int(di))
	PoolsInt.Set(0, int(AvgMaxIntVarIndex(AMAvgDif, Max)), int(pi), int(di))
}

// PoolAvgDifUpdate updates the AvgMax values for AvgDif Var.
// pi = global pool index.
func PoolAvgDifUpdate(pi, di uint32, avdif float32) {
	PoolAvgMaxUpdateVar(AMAvgDif, pi, di, avdif)
}

// PoolAvgDifCalc does Calc on Cycle level, and re-inits
func PoolAvgDifCalc(pi, di uint32) {
	PoolAvgMaxCalcVar(AMAvgDif, pi, di)
}

// PoolCycleToMinus grabs current Cycle values into the Minus phase values,
// and Plus values into Prev.
func PoolCycleToMinus(pi, di uint32) {
	for vr := range AMAvgDif { // don't do AvgDif
		Pools.Set(Pools.Value(int(AvgMaxVarIndex(vr, AMCycle, Avg)), int(pi), int(di)), int(AvgMaxVarIndex(vr, AMMinus, Avg)), int(pi), int(di))
		Pools.Set(Pools.Value(int(AvgMaxVarIndex(vr, AMCycle, Max)), int(pi), int(di)), int(AvgMaxVarIndex(vr, AMMinus, Max)), int(pi), int(di))

		Pools.Set(Pools.Value(int(AvgMaxVarIndex(vr, AMPlus, Avg)), int(pi), int(di)), int(AvgMaxVarIndex(vr, AMPrev, Avg)), int(pi), int(di))
		Pools.Set(Pools.Value(int(AvgMaxVarIndex(vr, AMPlus, Max)), int(pi), int(di)), int(AvgMaxVarIndex(vr, AMPrev, Max)), int(pi), int(di))
	}
}

// PoolCycleToPlus grabs current Cycle values into the Plus phase values.
func PoolCycleToPlus(pi, di uint32) {
	for vr := range AMAvgDif { // don't do AvgDif
		Pools.Set(Pools.Value(int(AvgMaxVarIndex(vr, AMCycle, Avg)), int(pi), int(di)), int(AvgMaxVarIndex(vr, AMPlus, Avg)), int(pi), int(di))
		Pools.Set(Pools.Value(int(AvgMaxVarIndex(vr, AMCycle, Max)), int(pi), int(di)), int(AvgMaxVarIndex(vr, AMPlus, Max)), int(pi), int(di))
	}
}

// PoolInit is callled during InitActs
func PoolInit(pi, di uint32) {
	PoolInhibInit(pi, di)
	PoolsInt.Set(0, int(PoolGated), int(pi), int(di))
	PoolAvgMaxZero(pi, di)
}

// PoolPoolGi computes the total inhibitory conductance for the pool.
func PoolPoolGi(ctx *Context, pi, di uint32) {
	if PoolsInt.Value(int(PoolIsLayer), int(pi), int(di)) > 0 {
		return
	}
	li := PoolsInt.Value(int(PoolLayerIdx), int(pi), int(di))
	PoolAvgMaxCalc(pi, di)
	PoolInhibIntToRaw(pi, di)
	ly := GetLayers(uint32(li))
	giMult := LayerStates.Value(int(LayerGiMult), int(li), int(di))
	lyIsOn := (ly.Inhib.Layer.On == 1)
	lpi := ly.PoolIndex(uint32(0))
	ly.SubPoolGiFromSpikes(ctx, lpi, pi, di, lyIsOn, giMult)
}

//gosl:end

// IndexToAvgMaxVar returns the AvgMaxVar indexes from overall Pool variable index.
func IndexToAvgMaxVar(vi uint32) (vr AvgMaxVars, phase AvgMaxPhases, am AvgMax) {
	vi -= uint32(poolFloatAvgMaxStart)
	vr = AvgMaxVars(vi / (uint32(AvgMaxN) * uint32(AvgMaxPhasesN)))
	rmdr := vi % (uint32(AvgMaxN) * uint32(AvgMaxPhasesN))
	phase = AvgMaxPhases(rmdr / uint32(AvgMaxN))
	am = AvgMax(rmdr % uint32(AvgMaxN))
	return
}

func PoolVarName(vi uint32) string {
	if vi < uint32(fsfffb.InhibVarsN) {
		return fsfffb.InhibVars(vi).String()
	}
	vr, phase, am := IndexToAvgMaxVar(vi)
	return vr.String() + "_" + phase.String() + "_" + am.String()
}

// IndexToAvgMaxIntVar returns the AvgMaxVar indexes from overall PoolInt variable index.
func IndexToAvgMaxIntVar(vi uint32) (vr AvgMaxVars, am AvgMax) {
	vi -= uint32(PoolIntAvgMaxStart)
	vr = AvgMaxVars(vi / uint32(AvgMaxN))
	am = AvgMax(vi % uint32(AvgMaxN))
	return
}

func PoolIntVarName(vi uint32) string {
	if vi < uint32(PoolIntAvgMaxStart) {
		return PoolIntVars(vi).String()
	}
	vr, am := IndexToAvgMaxIntVar(vi)
	return vr.String() + "_" + am.String()
}

// TestValues returns a map of CaSpkD.Avg, which provides an
// integrated summary of pool activity for testing
func PoolTestValues(pi, di uint32, layKey string, vals map[string]float32) {
	vals[layKey+" CaSpkD Avg"] = PoolAvgMax(AMCaSpkD, AMCycle, Avg, pi, di)
}