|
3 | 3 | from credit.postblock.reconstruct import Reconstruct |
4 | 4 | from credit.postblock.wet_mask_samudra import WetMaskBlock |
5 | 5 | from credit.postblock.scaler import BridgeScalerTransformer |
6 | | - |
| 6 | +from credit.postblock.gen1 import TracerFixer, GlobalMassFixer, GlobalWaterFixer, GlobalEnergyFixer |
| 7 | +from credit.postblock.geopotential import GeopotentialDiagnostic |
7 | 8 |
|
8 | 9 | POSTBLOCK_REGISTRY = { |
9 | 10 | "reconstruct": Reconstruct, |
10 | 11 | "bridgescaler_transform": BridgeScalerTransformer, |
11 | 12 | "wet_mask_samudra": WetMaskBlock, |
| 13 | + "tracer_fixer": TracerFixer, |
| 14 | + "global_mass_fixer": GlobalMassFixer, |
| 15 | + "global_water_fixer": GlobalWaterFixer, |
| 16 | + "global_energy_fixer": GlobalEnergyFixer, |
| 17 | + "geopotential_diagnostic": GeopotentialDiagnostic, |
12 | 18 | } |
13 | 19 |
|
| 20 | +_VALID_SECTIONS = {"per_step", "post_rollout"} |
14 | 21 |
|
15 | | -def build_postblocks(postblock_cfg: dict | None = None) -> nn.ModuleDict: |
16 | | - """Instantiates all postblocks from the config's ``postblocks`` section. |
17 | | -
|
18 | | - ``per_step`` defaults to ``False`` for every block and can be set per-block |
19 | | - in the config as a signal to the trainer about call timing. |
20 | 22 |
|
21 | | - Args: |
22 | | - postblock_cfg: the full postblocks dict from the config, e.g.:: |
23 | | -
|
24 | | - postblocks: |
25 | | - reconstruct: |
26 | | - type: reconstruct |
27 | | - inverse_scale: |
28 | | - type: bridgescaler_transform |
29 | | - args: |
30 | | - method: inverse_transform |
31 | | - scaler_path: /path/to/scaler.json |
32 | | - per_step: false |
33 | | - mass_fixer: |
34 | | - type: global_mass_fixer |
35 | | - args: ... |
36 | | -
|
37 | | - Returns: |
38 | | - ``nn.ModuleDict`` of instantiated postblocks, ordered as in config. |
39 | | - """ |
| 23 | +def _build_postblock_section(section_cfg: dict) -> nn.ModuleDict: |
40 | 24 | modules = {} |
41 | | - for name, block_cfg in (postblock_cfg or {}).items(): |
42 | | - block_type = block_cfg["type"] |
43 | | - args = block_cfg.get("args") or {} |
44 | | - instance = POSTBLOCK_REGISTRY[block_type](**args) |
45 | | - instance.per_step = block_cfg.get("per_step", False) |
46 | | - modules[name] = instance |
| 25 | + for name, block_cfg in section_cfg.items(): |
| 26 | + modules[name] = POSTBLOCK_REGISTRY[block_cfg["type"]](**(block_cfg.get("args") or {})) |
47 | 27 | return nn.ModuleDict(modules) |
48 | 28 |
|
49 | 29 |
|
50 | | -def apply_postblocks(postblocks: nn.ModuleDict, batch_dict: dict) -> dict: |
51 | | - """Applies all postblocks sequentially on a shared batch dict. |
| 30 | +def build_postblocks(postblock_cfg: dict | None = None, phase: str = "per_step") -> nn.ModuleDict: |
| 31 | + """Instantiate postblocks for a single phase from a two-section config. |
| 32 | +
|
| 33 | + Config format:: |
| 34 | +
|
| 35 | + postblocks: |
| 36 | + per_step: # run after every forward pass in the rollout loop |
| 37 | + reconstruct: |
| 38 | + type: reconstruct |
| 39 | + inverse_scale: |
| 40 | + type: bridgescaler_transform |
| 41 | + args: |
| 42 | + method: inverse_transform |
| 43 | + scaler_path: /path/to/scaler.json |
| 44 | + post_rollout: # run once after all rollout steps complete |
| 45 | + mass_fixer: |
| 46 | + type: global_mass_fixer |
| 47 | + args: ... |
| 48 | +
|
| 49 | + Typical usage — build once per phase, store separately:: |
52 | 50 |
|
53 | | - The caller is responsible for adding ``"prediction"`` (flat model output |
54 | | - tensor) and ``"meta"`` (metadata from ``apply_preblocks``) to ``batch_dict`` |
55 | | - before calling. Any additional data needed by postblocks (e.g. ``"input"``, |
56 | | - ``"target"``, ``"_raw"``) should also be added by the caller beforehand. |
| 51 | + step_postblocks = build_postblocks(cfg, phase="per_step") |
| 52 | + rollout_postblocks = build_postblocks(cfg, phase="post_rollout") |
57 | 53 |
|
58 | | - ``Reconstruct`` must be the first registered postblock — it converts |
59 | | - ``batch_dict["prediction"]`` from a flat tensor into a nested variable dict. |
60 | | - Subsequent postblocks operate on that nested dict via their ``key=`` arg. |
| 54 | + # inside rollout loop, after each forward pass: |
| 55 | + full_data_dict = apply_postblocks(step_postblocks, full_data_dict) |
| 56 | +
|
| 57 | + # once after rollout loop completes: |
| 58 | + apply_postblocks(rollout_postblocks, full_data_dict) |
61 | 59 |
|
62 | 60 | Args: |
63 | | - postblocks: ``nn.ModuleDict`` built by ``build_postblocks``. |
64 | | - batch_dict: Dict containing at minimum ``"prediction"`` (flat tensor) |
65 | | - and ``"meta"`` (with ``_channel_map["output"]``). |
| 61 | + postblock_cfg: the full ``postblocks`` config dict (both sections). |
| 62 | + phase: which section to build — ``"per_step"`` or ``"post_rollout"``. |
66 | 63 |
|
67 | 64 | Returns: |
68 | | - The same ``batch_dict`` after all postblocks have run. ``"prediction"`` |
69 | | - will be a nested variable dict after ``Reconstruct`` runs. |
| 65 | + ``nn.ModuleDict`` of instantiated blocks for the requested phase. |
70 | 66 |
|
71 | 67 | Raises: |
72 | | - RuntimeError: if ``postblocks`` is empty or ``Reconstruct`` is not first. |
| 68 | + ValueError: if the config contains keys other than ``"per_step"`` / ``"post_rollout"``, |
| 69 | + or if ``phase`` is not one of those values. |
73 | 70 | """ |
74 | | - blocks = list(postblocks.values()) |
| 71 | + cfg = postblock_cfg or {} |
| 72 | + unknown = set(cfg) - _VALID_SECTIONS |
| 73 | + if unknown: |
| 74 | + raise ValueError( |
| 75 | + f"build_postblocks: unexpected top-level keys {sorted(unknown)}. " |
| 76 | + "Expected only 'per_step' and/or 'post_rollout'. " |
| 77 | + "If you are using the old flat postblock format, migrate to the two-section layout." |
| 78 | + ) |
| 79 | + if phase not in _VALID_SECTIONS: |
| 80 | + raise ValueError(f"build_postblocks: phase must be one of {sorted(_VALID_SECTIONS)}, got {phase!r}.") |
| 81 | + return _build_postblock_section(cfg.get(phase) or {}) |
75 | 82 |
|
76 | | - for block in blocks: |
77 | | - batch_dict = block(batch_dict) |
78 | 83 |
|
| 84 | +def apply_postblocks(postblocks: nn.ModuleDict, batch_dict: dict) -> dict: |
| 85 | + """Apply a postblock group built by ``build_postblocks``. |
| 86 | +
|
| 87 | + Args: |
| 88 | + postblocks: ``nn.ModuleDict`` built by ``build_postblocks`` for a single phase. |
| 89 | + batch_dict: dict containing at minimum ``"y_pred"`` and ``"metadata"``. |
| 90 | +
|
| 91 | + Returns: |
| 92 | + The same ``batch_dict`` after all blocks in the group have run. |
| 93 | + """ |
| 94 | + for block in postblocks.values(): |
| 95 | + batch_dict = block(batch_dict) |
79 | 96 | return batch_dict |
0 commit comments