To integrate a reservoir models into fv3gfs this PR introduces a new steppers that functions across different TimeLoop checkpoints (i.e., the beginning and end of a timestep) with the use of a shared model object and a state machine for consistency.
Added public API:
get_reservoir_steppers: return a tuple of steppers, one for the increment of the underlying reservoir (ReservoirIncrementOnlyStepper) and one for prediction (ReservoirPredictStepper)
ReservoirConfig: Simple configuration to load a reservoir model with a specified number of steps to synchronize the reservoir and the reservoir timestep duration.
Significant internal changes:
loop.py now includes the reservoir steppers in the loop. The increment takes place before any state updates such that it can grab the input state to increment the reservoir to t+1. The prediction step takes place after all other updates such that we see the t+1 of the final model outputs (including any ML corrections) for the readout layer. The reservoir increment will take place on the subsequent timestep.
To integrate a reservoir models into fv3gfs this PR introduces a new steppers that functions across different
TimeLoop
checkpoints (i.e., the beginning and end of a timestep) with the use of a shared model object and a state machine for consistency.Added public API:
get_reservoir_steppers
: return a tuple of steppers, one for the increment of the underlying reservoir (ReservoirIncrementOnlyStepper
) and one for prediction (ReservoirPredictStepper
)ReservoirConfig
: Simple configuration to load a reservoir model with a specified number of steps to synchronize the reservoir and the reservoir timestep duration.Significant internal changes:
loop.py
now includes the reservoir steppers in the loop. The increment takes place before any state updates such that it can grab the input state to increment the reservoir to t+1. The prediction step takes place after all other updates such that we see the t+1 of the final model outputs (including any ML corrections) for the readout layer. The reservoir increment will take place on the subsequent timestep.[x] Tests added
Coverage reports (updated automatically):