graph_net_layer_nodes: gn_input_dict_renamer_layer_nodes
| 332 | |
| 333 | |
| 334 | class LevelEdgesAndLayerNodesGraph(EdgesAndNodesTransforms, OnlyLayersAreNodesTransforms): |
| 335 | """ |
| 336 | graph_net_layer_nodes: gn_input_dict_renamer_layer_nodes |
| 337 | """ |
| 338 | |
| 339 | def __init__(self, exp_type: str): |
| 340 | super().__init__(exp_type) |
| 341 | self._n_edges = self.spatial_input_dim[LEVELS] * 2 # bi-directional |
| 342 | self._out_dim = {NODES: self.input_dim[LAYERS], EDGES: self.input_dim[LEVELS], GLOBALS: self.input_dim[GLOBALS]} |
| 343 | |
| 344 | def transform(self, x: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: |
| 345 | x[NODES] = x.pop(LAYERS) |
| 346 | x[EDGES] = x.pop(LEVELS)[1:-1, :] # remove the surface and toa level |
| 347 | x[EDGES] = einops.repeat(x[EDGES], "e d -> (repeat e) d", repeat=2) # bidirectional edges |
| 348 | return x |
| 349 | |
| 350 | def batched_transform(self, x: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: |
| 351 | x[NODES] = x.pop(LAYERS) |
| 352 | x[EDGES] = x.pop(LEVELS)[:, 1:-1, :] # remove the surface and toa level |
| 353 | x[EDGES] = einops.repeat(x[EDGES], "b e d -> b (repeat e) d", repeat=2) # bidirectional edges |
| 354 | return x |
| 355 | |
| 356 | |
| 357 | class ColumnPreprocesser: |
nothing calls this directly
no outgoing calls
no test coverage detected