SGS probabilistic modeling implementation log

[jiaxi98]

We use this thread to record the bug and issue we encountered during the implementation of the probabilistic SGS modeling in the KS equation for future reference.

[jiaxi98]

There is huge caveat in calculation using array of shape (DIM, ) and (DIM, 1), basically they will broadcast to (DIM, DIM). See details in 565c3ae

[jiaxi98]

Batch normalization:
the is_training should be specified when the UNet model is constructed. Within the framework of train_state, one needs to init a new UNet with training = False for inference and validation use.
see bc58497

[jiaxi98]

float64 training:
It turns out that training with params of dtype float64 does not very much improve the performance, i.e. decrease the a-priori training error. Fairly speaking, we are not comparing to float32 training since previously both the input and output are of dtype float64, only the params of dtype 32.

[jiaxi98]

Writing a unified pipeline for both local and global modeling:

• No much issue in the case of PBC.
• With Dirichlet and Neumann BC, padding is used to make sure the shape of the input is a power of 2. This has to be done both loading the data and inference time.
• It is recommended to not include the padded data as the training data for the local model, as if the mesh size is small, e.g. 32, including one padded input for each mesh will results in large bias.
• To provide a unified interface for the a-posteriori simulation, both the local and global model should receive a padded input. In our KS case, the input should be a (1, 256, 1) tensor and for local model the augment_inputs function should operate on this function.
• In implementation, we currently adopt the following convention:
  • Input passed by simulator will be padded, i.e. shape (B, 256, 1).
  • In the wrapped function forward_fn, the input is reduced to size (B, 255, 1), augmented by augment_inputs with shape (B, 255, F) where F is the number of features, and padded back to (B, 256, F). Notice that the reduction of padding and repadding are handled explicitly in forward_fn.
  • The (B, 256, F) tensor is finally fed into the network apply method to obtain the SGS stresses.