ising_efficient.py

import numpy as np
import jax.numpy as jnp
import jax
import typing
import sparse_adj_mat

State = jnp.ndarray
"""For each node, a number in {-1, 0, 1} indicating it's belief"""
Field = typing.Callable[[int, tuple[int]], float]
"""A function that takes a timestep 't' and (an optionally empty) tuple of what node we are looking at and generates some float that is the field strength at that node at that time"""


class BeliefNetwork:
    def __init__(
        self,
        sparse_adj: sparse_adj_mat.Sparse_Adjacency_Matrix,
        external_field: Field,
        init_state: State,
        µ: float = 0.5,
        beta: float = 0.5,
        µ_is_weighted_according_to_neighborhood_size: bool = False,
    ):
        """
        :param sparse_adj: Can be generated by helper functions. Represents network topology and connection strengths.
        :type sparse_adj: sparse_adj_mat.Sparse_Adjacency_Matrix
        :param external_field: A function that specifies, for each integer timestep and node idx, what the strength of the external field is. See "Field" type docstring for more.
        :type external_field: Field
        :param init_state: The initial belief state of the nodes.
        :type init_state: State | None
        :param μ: The memory coefficient, in [0, 1]
        :type μ: float
        :param beta: The inverse temperature parameter. Should be nonnegative and below 100.
        :type beta: float
        :param µ_is_weighted_according_to_neighborhood_size: If True, the memory coefficient µ is scaled by the
                                                             size of the neighborhood for each node. This allows the
                                                             influence of past states to be weighted by the connectivity
                                                             of the node.
        :type µ_is_weighted_according_to_neighborhood_size: bool
        """

        self.external_field: Field = external_field
        self.init_state = init_state
        self.state = self.init_state
        self.prev_state = self.state
        self.graph_adjacency_mat = sparse_adj
        self.all_states = [
            self.init_state
        ]  # added this as it's used later for diameter
        self.µ_is_weighted_according_to_neighborhood_size = (
            µ_is_weighted_according_to_neighborhood_size
        )
        assert 0 <= µ <= 1, (
            "The memory-parameter µ should be in [0, 1] to remain interpretable"
        )
        self.µ = µ
        assert 100 >= beta >= 0, (
            "Beta, the inverse temperature parameter, should be nonnegative and not too big (leads to underflow)"
        )
        self.beta = beta

        self.neighbors, self.weights = (
            self.graph_adjacency_mat.precompute_neighbors_and_weights()
        )

    def run_for_steps(self, steps: int, seed: int):
        """Let s = steps and let n = number of nodes. This function runs the simulation for s steps and return a (s + 1, n) integer array which holds, for all s+1 timesteps (including time zero) the state for each node."""
        assert steps > 0, "Run for a positive amount of steps"
        return BeliefNetwork._run_for_steps(
            steps,
            self.neighbors,
            self.weights,
            self.init_state,
            self.external_field,
            self.µ_is_weighted_according_to_neighborhood_size,
            seed,
            # *params
            self.µ,
            self.beta,
        )

    @staticmethod
    def _run_for_steps(
        steps: int,
        neighbors,
        weights,
        init_state: State,
        field: Field,
        µ_is_weighted_according_to_neighborhood_size: bool,
        seed: int,
        *params,
    ):
        """
        Internal JIT compiled function to be called from the initialized class. Let 'n' be the size of the network.

        :param steps: The amount of timesteps to run the simulation for.
        :type steps: int
        :param neighbors: A (n, b) int array where b is an upper bound on the max network degree. Holds the value '-1' to represent None
        :param weights: A (n, b) int array where b is an upper bound on the max network degree. Holds the value 0.0
        :param init_state: An (n,) int array holding the initial state of each node.
        :type init_state: State
        :param field: The (potentially) time-varying field that will nudge the nodes to change belief.
        :type field: Field
        :param μ_is_weighted_according_to_neighborhood_size: Default False. If True, will multiply the effect of the past state by the neighborhood size (per node), which allows the nodes to "hear themselves", even in networks with high average degrees.
        :type μ_is_weighted_according_to_neighborhood_size: bool
        :param seed: For reproducibility, determines how the next state is sampled.
        :type seed: int
        :param params: Holds eventual extra parameters.
        """

        def __step(carry: tuple[State, State], xs) -> tuple[tuple[State, State], State]:
            state, prev_state = carry
            t, key = xs

            new_state = BeliefNetwork._step(
                state,
                prev_state,
                neighbors,
                weights,
                field,
                t,
                key,
                µ_is_weighted_according_to_neighborhood_size,
                *params,
            )
            return (new_state, state), new_state

        (last_state, prev_state), all_states = jax.lax.scan(
            f=__step,
            init=(init_state, init_state),
            xs=(jnp.arange(steps), jax.random.split(jax.random.PRNGKey(seed), steps)),
        )

        # Here we are adding the first state manually because it is dropped otherwise.
        return jnp.concat((init_state[None, :], all_states))

    @staticmethod
    def _step(
        state: State,
        prev_state: State,
        neighbors,
        weights,
        field: Field,
        t: int,
        key,
        µ_is_weighted_according_to_neighborhood_size: bool,
        *params,
    ) -> State:
        """
        See docstring for `_run_for_steps` for remaining parameter meanings.

        :param state: The current state.
        :type state: State
        :param prev_state: The previous state.
        :type prev_state: State
        :param t: The current time step. This is only used as information for the time-varying external field.
        :type t: int
        :param key: The JAX PRNGKey used to sample the next state.
        :return: Returns the next sampled state
        :rtype: State
        """
        (µ, beta) = params
        n = neighbors.shape[0]
        new_state = jnp.zeros_like(state)

        ss = jnp.array([-1, 0, 1])

        def update_i(i: int, key) -> int:
            nbs = neighbors[i]
            wght = weights[i]
            nbs_states = state[nbs]

            if µ_is_weighted_according_to_neighborhood_size:
                n_neighbors = jnp.sum(jnp.clip(jnp.sign(nbs + 1), a_min=0, a_max=1))
                memory_multiplier = n_neighbors
            else:
                memory_multiplier = 1.0

            ext = field(t, (i,))
            eff_i = wght @ nbs_states + ext + memory_multiplier * µ * prev_state[i]

            sample_p = jnp.exp(beta * ss * eff_i)
            sample_p /= jnp.sum(sample_p)
            sample = jax.random.choice(a=ss, p=sample_p, key=key)

            return sample

        new_state = jax.vmap(update_i, in_axes=(0, 0))(
            np.arange(n), jax.random.split(key, n)
        )

        return new_state