Class Normalizer

climart/data_transform/normalization.py:212–354

class Normalizer:
    def __init__(
            self,
            datamodule_config: DictConfig,
            input_normalization: Optional[str] = None,
            output_normalization: Optional[str] = None,
            spatial_normalization_in: bool = False,
            spatial_normalization_out: bool = False,
            log_scaling: Union[bool, List[str]] = False,
            data_dir: Optional[str] = None,
            verbose: bool = True
    ):
        """
        input_normalization (str): "z" for z-scaling (zero mean and unit standard deviation)
        """
        if not verbose:
            log.setLevel(logging.WARNING)

        if data_dir is None:
            data_dir = datamodule_config.get("data_dir") or constants.DATA_DIR
        exp_type = datamodule_config.get("exp_type")
        target_type = datamodule_config.get("target_type")
        target_variable = datamodule_config.get("target_variable")

        self._layer_mask = 45 if exp_type == constants.CLEAR_SKY else 14
        self._recover_meta_info(data_dir)
        self._input_normalizer: Dict[str, NormalizationMethod] = dict()
        self._output_normalizer: Optional[Dict[str, NormalizationMethod]] = None

        self._target_variables = get_target_variable_names(target_type, target_variable)
        if input_normalization is not None:
            norma_type = '_spatial' if spatial_normalization_in else ''
            info_msg = f" Applying {norma_type.lstrip('_')} {input_normalization} normalization to input data," \
                       f" based on pre-computed stats."
            log.info(info_msg)

            precomputed_stats = get_statistics(data_dir)
            precomputed_stats = {k: precomputed_stats[k] for k in precomputed_stats.keys() if
                                 (('spatial' in k and spatial_normalization_in) or ('spatial' not in k))}
            if isinstance(log_scaling, list) or log_scaling:
                log.info(' Log scaling pressure and height variables! (no other normalization is applied to them)')
                post_log_vals = dict(pressg=(11.473797, 0.10938317),
                                     layer_pressure=(9.29207, 2.097411),
                                     dz=(6.5363674, 1.044927),
                                     layer_thickness=(6.953938313568889, 1.3751644503732554),
                                     level_pressure=(9.252319, 2.1721559))
                vars_to_log_scale = ['pressg', 'layer_pressure', 'dz', 'layer_thickness', 'level_pressure']
                self._layer_log_mask = torch.tensor([2, 5, 12])
                for var in vars_to_log_scale:
                    dtype = self._variables[var]['data_type']
                    s, e = self.feature_by_var[dtype][var]['start'], self.feature_by_var[dtype][var]['end']
                    # precomputed_stats[f'{dtype}{prefix}_mean'][..., s:e] = 0
                    # precomputed_stats[f'{dtype}{prefix}_std'][..., s:e] = 1
                    precomputed_stats[f'{dtype}{norma_type}_mean'][..., s:e] = post_log_vals[var][0]
                    precomputed_stats[f'{dtype}{norma_type}_std'][..., s:e] = post_log_vals[var][1]

                def log_scaler(X: Dict[str, Tensor]) -> Dict[str, Tensor]:
                    # layer_log_mask = torch.tensor([2, 5, 12])