Initial commit

consensus.py

@@ -0,0 +1,46 @@
+import jax
+import jax.numpy as jnp
+
+
+class ConsensusConfig:
+    """
+    Config class for Consensus dynamics sims
+    """
+
+    num_agents: int # Number of agents in the consensus simulation
+    max_range: float = 100 # Max range of values each agent can take
+    step_size: float = 1 # Target range for length of simulation
+
+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 = jnp.diag(adj_matrix.sum(axis=0)) - adj_matrix
+    return agent_states + config.step_size * L * agent_states
+
+def generate_random_adjacency_matrix(key: jax.Array, config: ConsensusConfig):
+    """
+    Generates a random adjacency matrix when given 
+
+    """    
+