Sparse Array Benchmarking¶

%load_ext autoreload
%autoreload 2

import sys
if "google.colab" in sys.modules:
    %pip install "inferactively-pymdp" -q

import jax.numpy as jnp
from jax import tree_util as jtu
from jax.experimental import sparse
from pymdp.agent import Agent
import matplotlib.pyplot as plt
import seaborn as sns
import time

from pymdp.inference import smoothing_ovf
import numpy as np

import tracemalloc

def sizeof(x):
    return np.prod(x.shape)


def sizeof_sparse(x):
    return np.prod(x.data.shape) + np.prod(x.indices.shape)


def get_matrices(n_batch, num_obs, n_states):

    A_1 = jnp.ones((num_obs[0], n_states[0]))
    A_1 = A_1.at[-1, :-1].set(0)
    A_2 = jnp.ones((num_obs[0], n_states[1]))
    A_2 = A_2.at[-1, 1:].set(0)

    A_tensor = A_1[..., None] * A_2[:, None]
    A_tensor /= A_tensor.sum(0)

    A = [jnp.broadcast_to(A_tensor, (n_batch, *A_tensor.shape))]

    # create two transition matrices, one for each state factor
    B_1 = jnp.eye(n_states[0])
    B_1 = B_1.at[:, 1:].set(B_1[:, :-1])
    B_1 = B_1.at[:, 0].set(0)
    B_1 = B_1.at[-1, 0].set(1)
    B_1 = jnp.broadcast_to(B_1, (n_batch, n_states[0], n_states[0]))

    B_2 = jnp.eye(n_states[1])
    B_2 = B_2.at[:, 1:].set(B_2[:, :-1])
    B_2 = B_2.at[:, 0].set(0)
    B_2 = B_2.at[-1, 0].set(1)
    B_2 = jnp.broadcast_to(B_2, (n_batch, n_states[1], n_states[1]))

    B = [B_1[..., None], B_2[..., None]]
    C = [jnp.zeros((n_batch, num_obs[0]))]  # flat preferences
    D = [jnp.ones((n_batch, n_states[0])) / n_states[0], jnp.ones((n_batch, n_states[1])) / n_states[1]]  # flat prior
    E = jnp.ones((n_batch, 1))

    return A, B, C, D, E

def profile(fun, *args): 
    tracemalloc.start()
    tracemalloc.reset_peak()
    bt = time.time()
    res = fun(*args)
    et = time.time()
    size, peak = tracemalloc.get_traced_memory()

    stats = {'time': et - bt}
    return res, stats

def experiment(n_states):
    results = {}

    n_batch = 1
    num_obs = [2]

    A, B, C, D, E = get_matrices(n_batch=n_batch, num_obs=num_obs, n_states=n_states)

    # for the single modality, a sequence over time of observations (one hot vectors)
    obs = [
        jnp.broadcast_to(
            jnp.array(
                [
                    [1.0, 0.0],  # observation 0 is ambiguous with respect state factors
                    [1.0, 0],  # observation 0 is ambiguous with respect state factors
                    [1.0, 0],  # observation 0 is ambiguous with respect state factors
                    [0.0, 1.0],
                ]
            )[:, None],
            (4, n_batch, num_obs[0]),
        )
    ]  # observation 1 provides information about exact state of both factors

    agents = Agent(
        A=A,
        B=B,
        C=C,
        D=D,
        E=E,
        pA=None,
        pB=None,
        policy_len=3,
        control_fac_idx=None,
        policies=None,
        gamma=16.0,
        alpha=16.0,
        use_utility=True,
        categorical_obs=True,
        action_selection="deterministic",
        sampling_mode="full",
        inference_algo="ovf",
        num_iter=16,
        learn_A=False,
        learn_B=False,
        batch_size=1,
    )

    jtu.tree_map(lambda b: sparse.BCOO.fromdense(b, n_batch=n_batch), agents.B)


    prior = agents.D
    qs_hist = None
    action_hist = []
    for t in range(len(obs[0])):
        first_obs = jtu.tree_map(lambda x: jnp.moveaxis(x[:t+1], 0, 1), obs)
        beliefs = agents.infer_states(first_obs, past_actions=None, empirical_prior=prior, qs_hist=qs_hist)
        actions = jnp.broadcast_to(agents.policies[0, 0], (n_batch, 2))
        prior = agents.update_empirical_prior(actions, beliefs)
        qs_hist = beliefs
        action_hist.append(actions)

    beliefs = jtu.tree_map(lambda x, y: jnp.concatenate([x[:, None], y], 1), agents.D, beliefs)

    take_first = lambda pytree: jtu.tree_map(lambda leaf: leaf[0], pytree)
    beliefs_single = take_first(beliefs)

    # ======
    # Dense implementation
    smoothed_beliefs_dense, run_stats = profile(
        smoothing_ovf, *(beliefs_single, take_first(agents.B), jnp.stack(action_hist, 1)[0])
    )
    results.update({k+'_dense': v for k, v in run_stats.items()})
    results["size_dense"] = sum([sizeof(sB) for sB in agents.B])
    # ======

    sparse_B_single = jtu.tree_map(lambda b: sparse.BCOO.fromdense(b[0]), agents.B)
    actions_single = jnp.stack(action_hist, 1)[0]

    # ======
    # Sparse implementation
    smoothed_beliefs_sparse, run_stats = profile(
        smoothing_ovf, *(beliefs_single, sparse_B_single, actions_single)
    )
    results.update({k+'_sparse': v for k, v in run_stats.items()})
    results["size_sparse"] = sum([sizeof_sparse(sB) for sB in sparse_B_single])
    # ======

    return results, [beliefs_single, smoothed_beliefs_dense, smoothed_beliefs_sparse]