sup3r.preprocessing.conditional_moment_batch_handling.BatchMom2SepSF

class BatchMom2SepSF(low_res, high_res, output, mask)[source]

Bases: BatchMom1SF

Batch of low_res, high_res and output data when learning second moment of subfilter vel separate from first moment

Stores low, high res, output and mask data

Parameters:
  • low_res (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

  • high_res (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

  • output (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

  • mask (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

Methods

get_coarse_batch(high_res, s_enhance[, ...])

Coarsen high res data and return Batch with high res and low res data

make_mask(high_res[, s_padding, t_padding, ...])

Make mask for output.

make_output(low_res, high_res[, s_enhance, ...])

Make custom batch output

Attributes

high_res

Get the high-resolution data for the batch.

low_res

Get the low-resolution data for the batch.

mask

Get the mask for the batch.

output

Get the output for the batch.

shape

Get the (low_res_shape, high_res_shape) shapes.

static make_output(low_res, high_res, s_enhance=None, t_enhance=None, model_mom1=None, hr_features_ind=None, t_enhance_mode='constant')[source]

Make custom batch output

Parameters:
  • low_res (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

  • high_res (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

  • s_enhance (int | None) – Spatial enhancement factor

  • t_enhance (int | None) – Temporal enhancement factor

  • model_mom1 (Sup3rCondMom | None) – Model used to modify the make the batch output

  • hr_features_ind (list | np.ndarray | None) – List/array of feature channel indices that are used for generative output, without any feature indices used only for training.

  • t_enhance_mode (str) – Enhancing mode for temporal subfilter. Can be either constant or linear

Returns:

SF**2 (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features) SF is subfilter, HR is high-res and LR is low-res SF = HR - LR

classmethod get_coarse_batch(high_res, s_enhance, t_enhance=1, temporal_coarsening_method='subsample', temporal_enhancing_method='constant', hr_features_ind=None, features=None, smoothing=None, smoothing_ignore=None, model_mom1=None, s_padding=None, t_padding=None, end_t_padding=False)

Coarsen high res data and return Batch with high res and low res data

Parameters:
  • high_res (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

  • s_enhance (int) – Factor by which to coarsen spatial dimensions of the high resolution data

  • t_enhance (int) – Factor by which to coarsen temporal dimension of the high resolution data

  • temporal_coarsening_method (str) – Method to use for temporal coarsening. Can be subsample, average, or total

  • temporal_enhancing_method (str) – [constant, linear] Method to enhance temporally when constructing subfilter. At every temporal location, a low-res temporal data is substracted from the high-res temporal data predicted. constant will assume that the low-res temporal data is constant between landmarks. linear will linearly interpolate between landmarks to generate the low-res data to remove from the high-res.

  • hr_features_ind (list | np.ndarray | None) – List/array of feature channel indices that are used for generative output, without any feature indices used only for training.

  • features (list | None) – Ordered list of low-resolution training features input to the generative model

  • smoothing (float | None) – Standard deviation to use for gaussian filtering of the coarse data. This can be tuned by matching the kinetic energy of a low resolution simulation with the kinetic energy of a coarsened and smoothed high resolution simulation. If None no smoothing is performed.

  • smoothing_ignore (list | None) – List of features to ignore for the smoothing filter. None will smooth all features if smoothing kwarg is not None

  • model_mom1 (Sup3rCondMom | None) – Model used to modify the make the batch output

  • s_padding (int | None) – Width of spatial padding to predict only middle part. If None, no padding is used

  • t_padding (int | None) – Width of temporal padding to predict only middle part. If None, no padding is used

  • end_t_padding (bool | False) – Zero pad the end of temporal space. Ensures that loss is calculated only if snapshot is surrounded by temporal landmarks. False by default

Returns:

Batch – Batch instance with low and high res data

property high_res

Get the high-resolution data for the batch.

property low_res

Get the low-resolution data for the batch.

static make_mask(high_res, s_padding=None, t_padding=None, end_t_padding=False, t_enhance=None)

Make mask for output. The mask is used to ensure consistency when training conditional moments. Consider the case of learning E(HR|LR) where HR is the high_res and LR is the low_res. In theory, the conditional moment estimation works if the full LR is passed as input and predicts the full HR. In practice, only the LR data that overlaps and surrounds the HR data is useful, ie E(HR|LR) = E(HR|LR_nei) where LR_nei is the LR data that surrounds the HR data. Physically, this is equivalent to saying that data far away from a region of interest does not matter. This allows learning the conditional moments on spatial and temporal chunks only if one restricts the high_res output as being overlapped and surrounded by the input low_res. The role of the mask is to ensure that the input low_res always surrounds the output high_res.

Parameters:
  • high_res (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

  • s_padding (int | None) – Spatial padding size. If None or 0, no padding is applied. None by default

  • t_padding (int | None) – Temporal padding size. If None or 0, no padding is applied. None by default

  • end_t_padding (bool | False) – Zero pad the end of temporal space. Ensures that loss is calculated only if snapshot is surrounded by temporal landmarks. False by default

  • t_enhance (int | None) – Temporal enhancement factor to define end padding. None by default

  • model_mom1 (Sup3rCondMom | None) – Model used to modify the make the batch output

  • hr_features_ind (list | np.ndarray | None) – List/array of feature channel indices that are used for generative output, without any feature indices used only for training.

Returns:

mask (np.ndarray) – 4D | 5D array (batch_size, spatial_1, spatial_2, features) (batch_size, spatial_1, spatial_2, temporal, features)

property mask

Get the mask for the batch.

property output

Get the output for the batch. Output predicted by the neural net can be different than the high_res when doing moment estimation. For ex: output may be (high_res)**2 We distinguish output from high_res since it may not be possible to recover high_res from output.

property shape

Get the (low_res_shape, high_res_shape) shapes.