Training and Inference on sub-domains

[chelseajohn]

This PR contains code to perform sampling along with training and evaluation using patches(sub-domain) of samples.

Training

Training on sub-domains is enabled by the arguments in config.data

data:
  use_domainSampling: true
  use_fixedPatchSize: false        # fixed domain with overlapping
  use_orderedSampling: false       # fixed domain w/o overlapping 
  pad_to_fullGrid: false
  use_fixedPatch_startIdx: true
  nPatch_per_sample: 100
  use_minLimit: true
  padding : [1,1,1,1]
  # slices: []
  # patch_startIdx: []

• With a combination of the above parameters the following sampling can be done:

  1. nPatch_per_sample random sized patches (sX,sY) randomly sampled from (nX,nY) grid (can contain overlapping patches)

     use_domainSampling: true 
     nPatch_per_sample: $value
     

  2. nPatch_per_sample fixed sized patch of (sX,sY) randomly sampled from (nX,nY) grid (can contain overlapping patches)

     use_domainSampling: true 
     nPatch_per_sample: $value
     use_fixedPatchSize: true
     

  3. To sample fixed sized patch of (sX,sY) without overlapping from (nX,nY) grid

     use_domainSampling: true 
     nPatch_per_sample: $value
     use_fixedPatchSize: true
     use_orderedSampling: true
     

     nPatch_per_sample will be overwritten to (nX//sX)*(nY//sY).

     slices: [[sX,sY]]
     

• Use padding argument the inflow/outflow of patch can be set.

• Use use_minLimit: true to restrict (sX,sY) to be > (2 * (kX-1), 2 * (kY-1)), it must be set to True when pad_to_fullGrid is False

• Use use_fixedPatch_startIdx to sample slices for every epoch from the same index

• Use pad_to_fullGrid to embed patches (sX,sY) into a (nX,nY) grid of zeroes

• Additionally, to make training runs reproducible slices and patch_startIdx can be passed

Inference

To perform evaluation such that FNO predicts patches of the input grid , the model filters out the output patch before passing it into the Q projection layer. It can be enabled using the following arguments in config.model

model:
  get_subdomain_output: true
  iXBeg: 5
  iYBeg: 12
  iXEnd: 30
  iYEnd: 60

Keep in mind that the slice size i{X,Y}End - i{X,Y}Beg has to be > 2*(fourier_modes -1).

Contour Plotting

11_modelOutput.py has been modified to allow evaluation and plotting contours for :

1. Full grid input (nx,ny) ---> full grid output (nx,ny) (--use_full_input)
2. Full grid input (nx,ny) ---> sub-domain grid output (sx,sy) (--use_full_input --get_subdomain_output + above mentioned config.model args)
3. patch grid input (sx,sy)---> patch grid output (sx,sy) (--iXBeg, --iXEnd, --iYBeg, --iYEnd)

[tlunet]

Hi @chelseajohn, thanks for the PR. There is many things that are not very clear, sometime confusing, that need to be properly defined. In particular, this parameters mixing between input patches and subdomain output. I believe this can be made much more simpler, which would avoid a lot of confusion later on.