Efficient way of combining xarray datasets (explode + combine_by_coords) with variables only sharing subset of coordinates

[adrien-perello]

(message edited for clarity)

Context

I'm running a simulation that generates multiple (networkx) graphs. (It's actually a mesa agent-based simulation inspired from the Virus on a Network example) and due to randomness, I run each set of parameters multiple times.

Minimal example

The following example hopefully gives an idea of what I'm trying to achieve

import numpy as np
import pandas as pd
import networkx as nx

def avg_degree(G):
    return 2 * G.number_of_edges() / G.number_of_nodes()

def degree_distribution(G):
    return np.array(nx.degree_histogram(G))

def network_metrics(G):
    return {
        "avg degree": avg_degree(G),
        "degree distribution": degree_distribution(G),
    }

def generate_data(step, run, n, p):
    G = nx.erdos_renyi_graph(n, p)
    dct = {
        'network_type': 'random',
        'run': run,
        'nb_agents': n,
        'probability': p,
        'infected': np.random.randint(0,100, step),
    }
    dct.update(network_metrics(G))
    return dct


def main(nb_steps):
    lst = []
    for run in range(4):
        for nb_nodes in [10, 20, 50]:
            for probability in [0.1, 0.5, 0.8]:
                dct = generate_data(nb_steps, run, nb_nodes, probability)
                lst.append(dct)

    result = pd.DataFrame(lst)
    indexes = ['network_type', 'run', 'nb_agents', 'probability']
    result.set_index(indexes, inplace=True, drop=True)
    return result

Which gives

Goal

1. Explode infected and degree distribution (!! note: infected has a fixed length (= nb_steps) but not degree distribution (whose length varies)).
2. Combine everything in a single dataset

Current solution

I use the following helper function to explode the different columns

def explode(df, mapping):
    new_df = df.reset_index()
    for col, idx in mapping.items():
        new_df.index.rename('_id', inplace=True)
        new_df = new_df.explode(col)
        new_df.insert(1, idx, new_df.groupby('_id').cumcount())
        new_df.reset_index(drop=True, inplace=True)

    idx = list(df.index.names)
    nested_idx = list(mapping.values())
    return new_df.set_index(idx + nested_idx)


def helper(df, mapping):
    sol = []
    for k,v in mapping.items():
        sol.append(explode(df[k], {k:v}).to_xarray())

    leftover_columns = list(df.columns.difference(mapping.keys()))
    sol.append(df[leftover_columns].to_xarray())
    return sol

which I use as follow:

mapping = {'infected': 'step', 'degree distribution': "nb_nodes_with_degree"}
lst = helper(result, mapping)
xr.combine_by_coords(lst)

Result

Limits

It works, but it is quite slow and really inelegant.
--> What would be a better way of handling it?

[max-sixty]

Coordinates can be any number of dimensions!

Can you give more detail if that doesn't answer your question, ref #5404

[adrien-perello]

(message deleted for clarity --> see original message)

[max-sixty]

Hi Adrien — I'd really encourage you to reduce your example much more to a targeted xarray question, use paste-able code — all the items in the linked issue.

[max-sixty]

Thanks for the edit! That's much clearer!

[max-sixty]

I'm going to respond quickly rather than completely, but we can go back & forth:

Working on data in lists and then using something like df.explode is very slow in python, relative to having the data in array form.

One of the advantages of xarray is being able to have an arbitrary number of dimensions — so infected shouldn't need to be stored as a python list IIUC, since it's always the same size — instead that can be stored as a variable with an additional dimension. degree_distribution may be more difficult, since it doesn't have a consistent size.

The profile you posted is excellent. The only surprising thing to me is how long to_xarray takes. If you can turn that specific operation into a minimal example (i.e. without nx, and ensuring none of the values are python lists) and you still see poor performance, that would be a bug.

[adrien-perello]

using cProfile and gprof2dot gives an idea of the bottlenecks. And it actually comes from the explode() together with to_xarray()