Conversation
…l but needs further testing with batching.
… are too many right now
Greptile OverviewGreptile SummaryThis PR adds an in-situ cyclone tracking pipeline that integrates TempestExtremes with Earth2Studio's ensemble forecasting capabilities. The implementation supports asynchronous execution where CPU-based cyclone detection runs in parallel with GPU-based predictions. Key Changes
Critical IssueImport Error: Additional Issues
|
| Filename | Overview |
|---|---|
| earth2studio/models/dx/init.py | added TempestExtremes exports to __all__ without importing them, causing ImportError in tests |
| recipes/tc_tracking/src/tempest_extremes.py | implemented TempestExtremes and AsyncTempestExtremes classes for cyclone tracking with subprocess execution and thread pooling |
| recipes/tc_tracking/src/modes/generate_ensembles.py | added ensemble generation mode with cyclone tracking integration and stability checks |
| earth2studio/models/px/fcn3.py | added version check for torch_harmonics with duplicated logic |
| test/models/dx/test_tempest_extremes.py | added comprehensive tests for TempestExtremes, but will fail due to missing imports in dx/init.py |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| __all__ = [ | ||
| "AsyncTempestExtremes", | ||
| "ClimateNet", | ||
| "CorrDiff", | ||
| "CorrDiffCMIP6", | ||
| "CorrDiffTaiwan", | ||
| "PrecipitationAFNO", | ||
| "PrecipitationAFNOv2", | ||
| "SolarRadiationAFNO1H", | ||
| "SolarRadiationAFNO6H", | ||
| "TempestExtremes", | ||
| "WindgustAFNO", |
There was a problem hiding this comment.
TempestExtremes and AsyncTempestExtremes are exported in __all__ but never imported
| from earth2studio.models.dx.tempest_extremes import ( | |
| AsyncTempestExtremes, | |
| TempestExtremes, | |
| ) | |
| from earth2studio.models.dx.tc_tracking import ( | |
| TCTrackerVitart, | |
| TCTrackerWuDuan, | |
| ) | |
| from earth2studio.models.dx.wind_gust import WindgustAFNO # noqa | |
| __all__ = [ | |
| "AsyncTempestExtremes", | |
| "ClimateNet", | |
| "CorrDiff", | |
| "CorrDiffCMIP6", | |
| "CorrDiffTaiwan", | |
| "PrecipitationAFNO", | |
| "PrecipitationAFNOv2", | |
| "SolarRadiationAFNO1H", | |
| "SolarRadiationAFNO6H", | |
| "TempestExtremes", | |
| "WindgustAFNO", | |
| ] |
| """ | ||
| for cmd in (self.detect_cmd[0], self.stitch_cmd[0]): | ||
| try: | ||
| run(cmd, capture_output=True) # noqa: S603 |
There was a problem hiding this comment.
running arbitrary commands without validating executable path creates a security risk - user-provided detect_cmd and stitch_cmd are executed via subprocess without any sanitization
| ram_path = "/dev/shm" # noqa: S108 | ||
| if not os.path.exists(ram_path): | ||
| raise ValueError( | ||
| "/dev/shm not available. Set use_ram=False and use store_dir to specify a different location." # noqa: S108 | ||
| ) |
There was a problem hiding this comment.
hardcoded path /dev/shm may not exist on all systems (e.g., some containers or non-Linux systems)
| if Version(torch_harmonics.__version__) >= Version("0.8.1"): | ||
| from torch_harmonics.disco import cuda_kernels_is_available | ||
|
|
||
| _cuda_extension_available = cuda_kernels_is_available() | ||
| else: | ||
| from importlib.util import find_spec | ||
|
|
||
| _cuda_extension_available = find_spec("disco_cuda_extension") is not None |
There was a problem hiding this comment.
duplicated version check logic - same code appears in lines 38-50 inside the try block
| # exit() | ||
|
|
||
| # set random state or apply perturbation | ||
| if (not "model" in cfg) or (cfg.model == "fcn3"): |
There was a problem hiding this comment.
missing whitespace around in operator
| if (not "model" in cfg) or (cfg.model == "fcn3"): | |
| if (not "model" in cfg) or (cfg.model == "fcn3"): |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| If TempestExtremes command fails (detected by "EXCEPTION" in output) | ||
| """ | ||
| # detect nodes | ||
| out = run(command, capture_output=True) # noqa: S603 |
There was a problem hiding this comment.
verify that TempestExtremes executables come from trusted sources, as they are executed via subprocess
|
|
||
|
|
||
| def tile_xx_to_yy( | ||
| xx: torch.Tensor, xx_coords: CoordSystem, yy: torch.Tensor, yy_coords: CoordSystem |
There was a problem hiding this comment.
You shouldn't need both yy and yy_coords if I am understanding the function correctly? Only yy_coords should be sufficient I think.
Can the parameters match the same as map coords:
x: torch.Tensor,
input_coords: CoordSystem,
output_coords: CoordSystem,| return CoordSystem(adjusted_coords), mapping | ||
|
|
||
|
|
||
| def tile_xx_to_yy( |
There was a problem hiding this comment.
Not sure if tile is the right word we want here as I think this would imply repeating functionality like pyt. Perhaps something like broadcast_coords would be good, and perhaps this can expand in functionality as needed
| xx: torch.Tensor, xx_coords: CoordSystem, yy: torch.Tensor, yy_coords: CoordSystem | ||
| ) -> tuple[torch.Tensor, CoordSystem]: | ||
| """Tile tensor xx to match the leading dimensions of tensor yy and update coords accordingly. | ||
| eg if xx has shape (4, 721, 1440) and yy has shape (8, 2, 35, 721, 1440), |
There was a problem hiding this comment.
This example is good, and definitely needed. Can you make this and example section of the API doc?
https://github.com/NVIDIA/earth2studio/blob/main/earth2studio/models/dx/corrdiff_cmip6.py#L126
So we get the render in the API docs if needed.
| if dim not in coy: | ||
| raise ValueError(f"dim {dim} is not in coords: {list(coy)}.") | ||
|
|
||
| # fix difference in latitude |
There was a problem hiding this comment.
What is the intent of this mapping here? I think it would be maybe better to just having the cat only doing a cat, not intermix some mapping, but unsure about the use case
There was a problem hiding this comment.
So this would mean just force the user to first map_coords then cat_coords if needed.
| handshake_size, | ||
| ) | ||
| from earth2studio.utils.coords import map_coords, split_coords | ||
| from earth2studio.utils.coords import ( |
There was a problem hiding this comment.
This PR is too large to happen in one go... lets split it up a little. Can you start with opening a smaller PR with just the addition of the coord utils.
There was a problem hiding this comment.
And the respective updates to the hens recipe with the imports. Thanks!!!
3 participants
Earth2Studio Pull Request
Description
This recipe demonstrates ensemble forecasting and tropical cyclone tracking using FCN3 or AIFS-ENS with TempestExtremes integration. Additionally, it serves as an example of how to integrate custom downstream analysis tools into an inference pipeline.
This workflow implements an asynchronous execution mode where the CPU runs cyclone detection on completed predictions while the GPU continues producing predictions, enabling tracking with virtually no computational overhead. The pipeline supports three modes: generating ensemble forecasts with inline TC tracking, reproducing individual ensemble members (FCN3 only) to extract atmospheric fields for interesting tracks, and extracting reference tracks from ERA5 reanalysis for validation. Furthermore, it showcases how to integrate custom downstream analysis tools and deal with constraints - such as in-memory file handling and subprocess execution—that can be adapted for other downstream tasks.
Checklist
Dependencies