replicated mecc

sims/consensus.py

@@ -0,0 +1,150 @@
+import jax
+import jax.numpy as jnp
+from dataclasses import dataclass
+from functools import partial
+import numpy as np
+
+
+class ConsensusConfig:
+    """
+    Config class for Consensus dynamics sims
+    """
+    num_sims: int = 500 # Number of consensus sims
+    num_agents: int = 5 # Number of agents in the consensus simulation
+    max_range: float = 1 # Max range of values each agent can take
+    step_size: float = 0.1 # Target range for length of simulation
+    directed: bool = False # Consensus graph directed?
+    weighted: bool = False
+    num_time_steps: int = 100
+
+
+def consensus_step(adj_matrix: jax.Array, agent_states: jax.Array, config: ConsensusConfig):
+    """
+    Takes a step given the adjacency matrix and the current agent state using consensus dynamics.
+
+
+    Parameters
+    -----------------------------
+    adj_matrix : jax.Array (num_agents, num_agents)
+        A jax array containing the adjacency matrix for the consensus step.
+        
+    
+    agent_states: jax.Array (num_agents)
+        A jax array containing the current agent state
+    
+    config: ConsensusConfig
+        Config class for Consensus Dynamics
+
+
+    Returns
+    ------------------------------
+    updated_agent_state: jax.Array (num_agents)
+        A jax array containing the updated agent state
+
+    """
+    L = adj_matrix
+    norms = jnp.sum(jnp.abs(L), axis=1, keepdims=True)
+    L = L / (norms)
+
+    return (agent_states + config.step_size * jnp.matmul(L , agent_states))/(1 + config.step_size)
+
+def generate_random_adjacency_matrix(key: jax.Array, config: ConsensusConfig):
+    """
+    Generates a random adjacency matrix in accordance with the config.
+    The diagonal of the matrix is ensured to be all ones.
+
+    Parameters
+    --------------------
+        key: jax.Array
+            A key for jax.random operations
+
+        config: ConsensusConfig
+            Config for Consensus dyanmics
+
+    Returns
+    ---------------------
+    adj_matrices: jax.Array (num_agents, num_agents)
+        Random matrix 
+    """    
+    rand_matrix =  jax.random.uniform(key, shape=(config.num_agents, config.num_agents))
+    # idxs = jnp.arange(config.num_agents)
+    # rand_matrix = rand_matrix.at[:, idxs, idxs].set(1)
+    rand_matrix = jnp.fill_diagonal(rand_matrix, 1, inplace=False) # Fill diagonal with ones
+    if not config.weighted:
+        rand_matrix = jnp.where(rand_matrix > 0.5, 1, 0)
+
+    if config.directed:
+        return rand_matrix
+
+    rand_matrix = jnp.tril(rand_matrix) + jnp.triu(rand_matrix.T, 1)
+    return rand_matrix
+
+
+def generate_random_agent_states(key: jax.Array, config: ConsensusConfig):
+    """
+    Generate a random initial state for the agents in accordance with the config.
+
+    Parameters
+    ---------------------
+        key: jax.Arrray
+            A key for jax.random operations
+
+        config: ConsensusConfig
+            Config for Consensus dynamics
+
+    Returns
+    ---------------------
+        rand_states: jax.Array (num_sims, num_agents)
+
+    """
+    rand_states = jax.random.uniform(key, shape=(config.num_sims, config.num_agents), minval=-config.max_range, maxval=config.max_range)
+
+    return rand_states
+
+@partial(jax.jit, static_argnames=["config"])
+def run_consensus_sim(adj_mat: jax.Array, initial_agent_state: jax.Array, config: ConsensusConfig):
+    """
+    Runs the consensus sim and returns a history of agent states.
+
+    Parameters
+    -------------------
+        adj_mat: jax.Array (num_agents, num_agents)
+            A jax array containing the adjacency matrix for the consensus step.
+
+        initial_agent_state: jax.Array (num_agents)
+            A jax array containing the initial agent state
+
+        config: ConsensusConfig
+            Config for Consensus dynamics
+    """ 
+    # batched consensus step (meant for many initial states)
+    batched_consensus_step = jax.vmap(consensus_step, in_axes=(None, 0, None), out_axes=0)
+    
+    def step(x_prev, _):
+        x_next = batched_consensus_step(adj_mat, x_prev, config)
+        return x_next, x_next
+
+    _, all_states = jax.lax.scan(step, initial_agent_state, None, config.num_time_steps)
+
+    return all_states.transpose(1, 0, 2)
+
+
+def plot_consensus(trajectory, config):
+    import matplotlib.pyplot as plt
+    import seaborn as sns
+    sns.set_theme()
+
+    states = np.array(trajectory)
+    timesteps, num_agents = states.shape
+    time = np.arange(timesteps)
+
+    plt.figure()
+    for agent_idx in range(num_agents):
+        plt.plot(time ,states[:, agent_idx], label=f"Agent {agent_idx}")
+
+    plt.xlabel("Timestep")
+    plt.ylabel("Agent statue")
+    plt.ylim(-config.max_range, config.max_range)
+    plt.title("Consensus simulation trajectories")
+    plt.legend()
+    plt.show()