graph_cluster.py

import numpy as np
import networkx as nx
from system import DrMoagiArchitecture


class GraphMoagiCluster:
    def __init__(self, nodes=6, dim=16):
        self.graph = nx.erdos_renyi_graph(nodes, 0.5)
        self.nodes = {i: DrMoagiArchitecture(dim=dim) for i in self.graph.nodes}

    def graph_step(self):
        new_states = {}

        for i in self.graph.nodes:
            neighbors = list(self.graph.neighbors(i))
            Psi = np.random.randn(self.nodes[i].n)
            self.nodes[i].step(Psi)

            if neighbors:
                neighbor_states = [self.nodes[j].Xi for j in neighbors]
                consensus = np.mean(neighbor_states, axis=0)
                new_states[i] = 0.7 * self.nodes[i].Xi + 0.3 * consensus
            else:
                new_states[i] = self.nodes[i].Xi

        for i in new_states:
            self.nodes[i].Xi = new_states[i]