Has anyone figured out a way around plotting very large arrays?

[cwoolfo1]

I am writing a code in Jax and I need to plot 3D arrays every n timesteps. The only reason I haven't used the jax scan methods is because I am concerned about memory usage for large simulations. Has anyone found a way around this? I was thinking of maybe trying to convert the arrays into a sparse data format, but I'm not entirely sure if that is really going to help much...

[jakevdp]

Hi - could you provide more details about what the particular challenge is? Is your question about how to visualize arrays? About how to move data to host during training? Or is it a matter of the data from multiple steps being too large to store in memory or on disk?

[jakevdp]

Rather than storing the full extracted data at each timestep, could you generate the desired summary plots in real-time and save those as PNGs?

[cwoolfo1]

That is what I have been doing right now. The process of saving PNGs is slower than I would like and unfortunately prevents me from using the lax.scan method to speed up my code.

[jakevdp]

If you want to save images within a scan, you could use jax.experimental.io_callback. For an example of its use, see https://docs.jax.dev/en/latest/external-callbacks.html#exploring-io-callback. This should also be able to execute asynchronously, so depending on your computation it may lead to better performance of the overall workflow compared to executing eagerly on the host.

[cwoolfo1]

Interesting.. I attempted to use the io_callback feature and I was having issues getting it to write values to txt files. It kept trying to write traced values.

def write_to_file(filename, time, data):
    with open(filename, "a") as f:
     f.write(f"{time}, {data}\n")

io_callback(lambda _: write_to_file(filename, time, data), None, 0, ordered=True)

[jakevdp]

You're misusing io_callback here: the only value that is passed back to host here is 0 (via the _ argument of your callback function). If you want to send time and delta to the callback, you need to make them explicit arguments of io_callback and the function you pass to it.