ai2cm / fv3net

explore the FV3 data for parameterization
MIT License
16 stars 3 forks source link

Make RankDivider agnostic to time dimension #2240

Closed AnnaKwa closed 1 year ago

AnnaKwa commented 1 year ago

The RankDivider class is used in reservoir computing to handle the division and reconstruction of cubedsphere ranks into subdomains. Previously, its initialization was assumed to include information about the time dim and its extent. This PR changes RankDivider to only allow ordering and extent of spatial dimensions to be provided. Since the RankDivider only cares about spatial dimensions when slicing data, it didn't make sense for it to have the time extent in its attributes since this dimension could change or be absent in various use cases.

Refactored public API:

Added TimeSeriesRankDivider, which inherits from `RankDivider and has all of its functionality, the following methods differ from the parent class in that they operate on data with dims [time, x, y, z] as opposed to [x, y, z].

e.g. for training, the data being split and reshaped is a time series, so the training function uses a TimeSeriesRankDivider. Predictions are a single time step and don't have a time dimension, so they are reconstituted into the original [x, y, z] dims using a RankDivider.

Resolves # [JIRA-TAG]

Coverage reports (updated automatically):

AnnaKwa commented 1 year ago

Thanks, removing the time dimension handling entirely from RankDivider was a good idea. I added a TimeSeriesRankDivider child class which redefines the relevant methods to assume the first dim is a time dimension.

frodre commented 1 year ago

Awesome. I agree, this was a nice simplifying update!