Class ColumnPreprocesser

climart/data_transform/transforms.py:357–433

class ColumnPreprocesser:
    ONLY_LAYER_NODES = ['duplication', 'graph_net_layer_nodes', 'identity']
    ONLY_LEVEL_NODES = ['graph_net_level_nodes']

    def __init__(self,
                 preprocessing: str,
                 exp_type: str,
                 projector_hidden_dim: int = 128,  # only if preprocessing == 'mlp'
                 projector_n_layers: int = 1,  # only if preprocessing == 'mlp'
                 projector_net_normalization: str = 'layer_norm',  # only if preprocessing == 'mlp'
                 use_level_features: bool = True,
                 ):
        self.preprocessing_type = preprocessing.lower()
        self.projector_hidden_dim = projector_hidden_dim
        self.projector_n_layers = projector_n_layers
        self.projector_net_normalization = projector_net_normalization
        self.use_level_features = use_level_features

    def get_preprocesser(self, batched: bool = False, verbose: bool = True):
        if self.preprocessing_type in ['mlp', 'mlp_projection']:
            in_dims = self.input_dims.copy()
            if not self.use_level_features:
                in_dims.pop(LEVELS)
            preprocesser = FeatureProjector(
                input_name_to_feature_dim=in_dims,
                projector_n_layers=self.projector_n_layers,
                projection_dim=self.projector_hidden_dim,
                projector_activation_func='Gelu',
                projector_net_normalization=self.projector_net_normalization,
                output_normalization=True,
                output_activation_function=False,
                projections_aggregation=self.intersperse if self.use_level_features else self.intersperse_no_levels)
            self.out_dim = self.projector_hidden_dim
            self.as_string = f'All var types are MLP-projected to a {self.projector_hidden_dim} hidden dimension'

        return preprocesser

    def intersperse(self,
                    name_to_array: Optional[Dict[str, Tensor]] = None,
                    global_node: Optional[Tensor] = None,  # shape (b, #feats)
                    levels: Optional[Tensor] = None,  # shape (b, #levels, #feats)
                    layers: Optional[Tensor] = None  # shape (b, #layers, #feats)
                    ) -> Tensor:
        global_node, levels, layers = self.get_data_types(name_to_array, global_node, levels, layers)

        if global_node.shape[-1] != layers.shape[-1] or levels.shape[-1] != layers.shape[-1]:
            raise ValueError("Expected all node types to have same dimensions. Project them first or pad them instead!")

        batch_size, _, n_feats = levels.shape

        interspersed_data = torch.empty((batch_size, self.n_nodes, n_feats))
        interspersed_data[:, self.GLOBAL_NODE, :] = global_node
        interspersed_data[:, self.LEVEL_NODES, :] = levels
        interspersed_data[:, self.LAYER_NODES, :] = layers

        interspersed_data = interspersed_data.to(global_node.device)
        return interspersed_data