Include a has_aux-like parameter in the dataset generation pipeline

[DSilva27]

Sometimes there are quantities that are computed during the image generation process that could be useful to have, but that normally don't go into a starfile. For example, when generating data from ensembles it would be useful to have the index of the configuration used for each image. There are more cases, right now I can only think of images with different SNR, but I am sure there are more.

What I mean by hax-aux-like is to have a functionality similar to the one in jax.value_and_grad function, where one can specify that the objective function outputs an image, and some extra variable. In this case we would have to carry this output, and then give it as an output at the end of the dataset generation process.

[mjo22]

I think I see what you mean—you mean to say that we allow for some auxiliary data to be returned in the compute_image function for the write_simulated_image_stack_from_starfile function, and return this auxiliary data at the end of dataset generation?