brain/test-blocks.ts at main · RLabs-Inc/brain · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import { core as mx } from '@frost-beta/mlx'

console.log('=== Testing Building Blocks ===\n')

// Import our building blocks
import {
  allocatePopulation,
  voltage,
  current,
  populationSize,
  integrate,
  THRESHOLD,
} from './dist/src/core/neuron.svelte.js'

import {
  allocateSynapseGroup,
  createAllToAllConnectivity,
  weights,
  transmit,
} from './dist/src/core/synapse.svelte.js'

import {
  allocateNetwork,
  addPopulationToNetwork,
  addSynapseGroupToNetwork,
  step,
} from './dist/src/core/network.svelte.js'

async function runTest() {
  try {
    console.log('1. Creating populations...')
    const pop1 = allocatePopulation('input', 5, 'RS')
    const pop2 = allocatePopulation('output', 5, 'RS')
    console.log('   Created pop1=' + pop1 + ', pop2=' + pop2)
    console.log('   Sizes: ' + populationSize[pop1] + ', ' + populationSize[pop2])

    console.log('\n2. Checking initial voltage...')
    await mx.asyncEval(voltage[pop1])
    const avgV1 = mx.mean(voltage[pop1])
    await mx.asyncEval(avgV1)
    console.log('   Pop1 avg voltage: ' + avgV1.item())

    console.log('\n3. Creating network...')
    const net = allocateNetwork('test')
    addPopulationToNetwork(net, pop1)
    addPopulationToNetwork(net, pop2)
    console.log('   Network ' + net + ' created')

    console.log('\n4. Creating synapses...')
    const conn = createAllToAllConnectivity(5, 5)
    await mx.asyncEval(conn.preIndices, conn.postIndices)
    console.log('   Connectivity shape: pre=' + conn.preIndices.shape + ', post=' + conn.postIndices.shape)

    const syn = allocateSynapseGroup('syn1', pop1, pop2, conn.preIndices, conn.postIndices)
    addSynapseGroupToNetwork(net, syn)
    await mx.asyncEval(weights[syn])
    const avgW = mx.mean(weights[syn])
    await mx.asyncEval(avgW)
    console.log('   Synapse group ' + syn + ', avg weight: ' + avgW.item())

    console.log('\n5. Injecting current...')
    const inputCurrent = mx.multiply(mx.random.uniform(0, 1, [5]), mx.array(20))
    current[pop1] = mx.add(current[pop1], inputCurrent)
    await mx.asyncEval(current[pop1])
    const avgCurrent = mx.mean(current[pop1])
    await mx.asyncEval(avgCurrent)
    console.log('   Injected current, avg: ' + avgCurrent.item())

    console.log('\n6. Running simulation step...')
    step(net, 1.0)
    await mx.asyncEval(voltage[pop1], voltage[pop2])
    const newAvgV1 = mx.mean(voltage[pop1])
    const newAvgV2 = mx.mean(voltage[pop2])
    await mx.asyncEval(newAvgV1, newAvgV2)
    console.log('   After step - Pop1 avg V: ' + newAvgV1.item() + ', Pop2 avg V: ' + newAvgV2.item())

    console.log('\n7. Checking for spikes...')
    const firing1 = mx.greaterEqual(voltage[pop1], mx.array(THRESHOLD))
    const firing2 = mx.greaterEqual(voltage[pop2], mx.array(THRESHOLD))
    const spikes1 = mx.sum(firing1)
    const spikes2 = mx.sum(firing2)
    await mx.asyncEval(spikes1, spikes2)
    console.log('   Pop1 spikes: ' + spikes1.item() + ', Pop2 spikes: ' + spikes2.item())

    console.log('\n8. Running more steps to see activity...')
    for (let i = 0; i < 10; i++) {
      // Inject more current
      const noise = mx.multiply(mx.random.uniform(0, 1, [5]), mx.array(15))
      current[pop1] = mx.add(current[pop1], noise)

      // Step
      step(net, 1.0)

      // Check spikes
      await mx.asyncEval(voltage[pop1], voltage[pop2])
      const f1 = mx.greaterEqual(voltage[pop1], mx.array(THRESHOLD))
      const f2 = mx.greaterEqual(voltage[pop2], mx.array(THRESHOLD))
      const s1 = mx.sum(f1)
      const s2 = mx.sum(f2)
      await mx.asyncEval(s1, s2)
      console.log('   Step ' + (i+2) + ': Pop1=' + s1.item() + ' spikes, Pop2=' + s2.item() + ' spikes')
    }

    console.log('\n=== Test Complete ===')
  } catch (e) {
    console.error('ERROR:', e)
  }
}

runTest()