This is nice, but I believe a simpler design could work better. Simply make a model which transforms a 3d section of an image to an embedding vector. Make an

This is nice, but I believe a simpler design could work better. Simply make a model which transforms a 3d section of an image to an embedding vector. Make another model which can reverse the process (ie. encoder-decoder). Do that for every tile of a starting state. Make an 'upscale' and 'downscale' model which can take a grid of embedding vectors and return a new vector representing the whole. Then make an 'advance time' model, which takes an embedding vector and advances time by a given number of seconds/microseconds/days. Now train all the models end to end to ensure that all combinations of upscaling/downscaling/advancing/encoding/decoding produce similar outputs to traditional physics models. Use an ensemble of models or a sampling scheme to find places where outputs do not closely match, and insert more training data from the physical simulation at those points.