baroclinic initialization #31
Conversation
…tive to the fortran code
…e into feature/use_new_grid_in_dycore
…hrough the grid_data object instead. A future PR may orgnaize these to be computed at a different stage
… launching simulations that are not baroclinic
… on the MetricTerms class. To avoid a circular import with MirrorGrid, changed RIGHT_HAND_GRID to be an argument in the context of mirroring, set to to a constant in out calls of it
…k, but are still wrong on a few points
…ridData because the match to the variables used is better and it simplifies the code.
…pace into feature/baroclinic-initialization
| lat2, lat3, lat4, lat5 = compute_grid_edge_midpoint_latitude_components(lon, lat) | ||
| slice_3d = (slice(0, nx), slice(0, ny), slice(None)) | ||
| slice_2d = (slice(0, nx), slice(0, ny)) | ||
| # initialize temperature | ||
| t_mean = jablo_init.horizontally_averaged_temperature(eta) | ||
| pt[slice_3d] = cell_average_nine_components( | ||
| jablo_init.temperature, | ||
| [eta, eta_v, t_mean], | ||
| lat, | ||
| lat_agrid, | ||
| lat2, | ||
| lat3, | ||
| lat4, | ||
| lat5, | ||
| slice_2d, | ||
| ) | ||
|
|
||
| # initialize surface geopotential | ||
| phis[slice_2d] = cell_average_nine_components( | ||
| jablo_init.surface_geopotential_perturbation, | ||
| [], | ||
| lat, | ||
| lat_agrid, | ||
| lat2, | ||
| lat3, | ||
| lat4, | ||
| lat5, | ||
| slice_2d, | ||
| ) |
There was a problem hiding this comment.
cell_average_nine_components and compute_grid_edge_midpoint_latitude_components are doing something low level and pretty complex - I can't tell what they are doing. compute_grid_edge_midpoint_latitude_components is also only used for this second function.
Is there any change you can make to have it be more clear what's happening at this function's level, by pushing the low-level logic down one call level? For example, even if it's less efficient, by pushing the call to compute_grid_edge_midpoint_latitude_components down one level (removing many low-level call arguments from this function's level), and by pulling the slicing operation (which is a higher level operation you're already using here in the assignment slice) up to this level (e.g. by passing lat[slice_2d] and lat_agrid[slice_2d])?
There was a problem hiding this comment.
sure thing. this duplicates the calculation of these latitude factors, but does make it more understandable. I'll implement the suggestion. Pulling the slicing up is a little tricker, as the resulting shapes of the 9 points being averaged are not by default the same, as they are shifted in different directions. I can do this if you prefer, but want to double check before I iterate on it.
| assert not adjust_dry_mass | ||
| assert not hydrostatic |
There was a problem hiding this comment.
Does it make sense for us to just not implement these as keyword arguments? I don't think we have any plans to implement these, the code itself is easier to understand than the rest of the dycore in terms of if we were to implement a new option, and if we did implement a new option we'd probably do it using new functions rather than if flags?
There was a problem hiding this comment.
ok sure. The fortran code has a lot of if conditions for different test cases, so it could be potentially confusing to understand and implement the other options. But I am not worried about this as much as I was for the dycore, since this is just an idealized initial state.
There was a problem hiding this comment.
I'm also already not following this p_var subroutine to the letter, since it was recomputing pressure variables in the exact same way they were already computed.
There was a problem hiding this comment.
do you mean just for p_var, or for the whole module?
| # TODO: when the dycore state is updated to only include | ||
| # quantities and no storages, remove the "_quantity" from phis, u and v |
There was a problem hiding this comment.
You could remove this TODO here, when I go to do it this will be using vscode's refactor functionality so it will update everywhere at once (thanks to your use of a dataclass!). As it stands, this comment would probably remain after the refactor unless a reviewer catches it or I've thought long enough about this particular comment location while writing this message.
| # TODO: when the dycore state is updated to only include | |
| # quantities and no storages, remove the "_quantity" from phis, u and v |
| def init_from_serialized_data(cls, serializer, grid, quantity_factory): | ||
| savepoint_in = serializer.get_savepoint("FVDynamics-In")[0] | ||
| translate_object = fv3core.testing.TranslateFVDynamics([grid]) | ||
| input_data = translate_object.collect_input_data(serializer, savepoint_in) | ||
| # making just storages for the moment, revisit when making them all | ||
| # quantities (maybe use state_from_inputs) | ||
| translate_object._base.make_storage_data_input_vars(input_data) | ||
| # used for the translate test as inputs, but are generated by the | ||
| # MetricsTerms class and are not part of this data class | ||
| for delvar in ["ak", "bk", "ptop", "ks"]: | ||
| del input_data[delvar] | ||
| return cls(**input_data, quantity_factory=quantity_factory) |
There was a problem hiding this comment.
This method should only be getting used by the tests, and currently fv3core doesn't depend on serialbox other than for testing. Can this get moved to the testing code?
There was a problem hiding this comment.
It can. I was thinking it might be useful for performance testing if we end up wanting to test non-baroclinic namelists before we implement reading initial state from disk. But someone wanting to do that can probably figure it out.
fv3core/fv3core/testing/translate.py
Outdated
| @@ -293,6 +296,12 @@ def make_grid_storage(self, pygrid): | |||
| if k in self.data: | |||
| self.make_composite_var_storage(k, self.data[k], shape) | |||
| del self.data[k] | |||
| for k in TranslateGrid.ee_vars: | |||
There was a problem hiding this comment.
nit: This was pre-existing, but can you replace k here and below with key? When I read these my mind doesn't stop telling me it's indexing on the vertical direction.
| @@ -282,15 +284,27 @@ def compute_parallel(self, inputs, communicator): | |||
| grid_data.ptop = inputs["ptop"] | |||
| grid_data.ks = inputs["ks"] | |||
|
|
|||
| state = self.state_from_inputs(inputs) | |||
| input_storages = self.state_from_inputs(inputs) | |||
| # making sure we init DyCOreState with the exact set of variables | |||
There was a problem hiding this comment.
| # making sure we init DyCOreState with the exact set of variables | |
| # making sure we init DyCoreState with the exact set of variables |
…king one. we will revisit in a future PR to pull this up to the driver
| lat2, lat3, lat4, lat5 = compute_grid_edge_midpoint_latitude_components(lon, lat) | ||
| slice_3d = (slice(0, nx), slice(0, ny), slice(None)) | ||
| slice_2d = (slice(0, nx), slice(0, ny)) | ||
| # initialize temperature | ||
| t_mean = jablo_init.horizontally_averaged_temperature(eta) | ||
| pt[slice_3d] = cell_average_nine_components( | ||
| jablo_init.temperature, | ||
| [eta, eta_v, t_mean], | ||
| lat, | ||
| lat_agrid, | ||
| lat2, | ||
| lat3, | ||
| lat4, | ||
| lat5, | ||
| slice_2d, | ||
| ) | ||
|
|
||
| # initialize surface geopotential | ||
| phis[slice_2d] = cell_average_nine_components( | ||
| jablo_init.surface_geopotential_perturbation, | ||
| [], | ||
| lat, | ||
| lat_agrid, | ||
| lat2, | ||
| lat3, | ||
| lat4, | ||
| lat5, | ||
| slice_2d, | ||
| ) |
There was a problem hiding this comment.
sure thing. this duplicates the calculation of these latitude factors, but does make it more understandable. I'll implement the suggestion. Pulling the slicing up is a little tricker, as the resulting shapes of the 9 points being averaged are not by default the same, as they are shifted in different directions. I can do this if you prefer, but want to double check before I iterate on it.
| assert not adjust_dry_mass | ||
| assert not hydrostatic |
There was a problem hiding this comment.
ok sure. The fortran code has a lot of if conditions for different test cases, so it could be potentially confusing to understand and implement the other options. But I am not worried about this as much as I was for the dycore, since this is just an idealized initial state.
| assert not adjust_dry_mass | ||
| assert not hydrostatic |
There was a problem hiding this comment.
I'm also already not following this p_var subroutine to the letter, since it was recomputing pressure variables in the exact same way they were already computed.
| assert not adjust_dry_mass | ||
| assert not hydrostatic |
There was a problem hiding this comment.
do you mean just for p_var, or for the whole module?
| islice, jslice, slice_3d, slice_2d = compute_slices(nx, ny) | ||
| # Slices with extra buffer points in the horizontal dimension | ||
| # to accomodate averaging over shifted calculations on the grid | ||
| isliceb, jsliceb, slice_3db, slice_2db = compute_slices(nx + 1, ny + 1) |
There was a problem hiding this comment.
sure thing, good point.
| eta_s = 1.0 # surface level | ||
| eta_t = 0.2 # tropopause |
| # maximum windspeed amplitude - close to windspeed of zonal-mean time-mean | ||
| # jet stream in troposphere | ||
| u0 = 35.0 # From Table VI of DCMIP2016 | ||
| # [lon, lat] of zonal wind perturbation centerpoint at 20E, 40N | ||
| pcen = [math.pi / 9.0, 2.0 * math.pi / 9.0] # From Table VI of DCMIP2016 | ||
| u1 = 1.0 | ||
| pt0 = 0.0 | ||
| eta_0 = 0.252 | ||
| eta_s = 1.0 # surface level | ||
| eta_t = 0.2 # tropopause | ||
| t_0 = 288.0 | ||
| delta_t = 480000.0 | ||
| lapse_rate = 0.005 # From Table VI of DCMIP2016 | ||
| surface_pressure = 1.0e5 # units of (Pa), from Table VI of DCMIP2016 | ||
| # NOTE RADIUS = 6.3712e6 in FV3 vs Jabowski paper 6.371229e6 | ||
| R = constants.RADIUS / 10.0 # Perturbation radiusfor test case 13 |
There was a problem hiding this comment.
thanks! paid close attention to the paper on this one, which thankfully matches pretty well to what the fortran code does!
| def init_from_serialized_data(cls, serializer, grid, quantity_factory): | ||
| savepoint_in = serializer.get_savepoint("FVDynamics-In")[0] | ||
| translate_object = fv3core.testing.TranslateFVDynamics([grid]) | ||
| input_data = translate_object.collect_input_data(serializer, savepoint_in) | ||
| # making just storages for the moment, revisit when making them all | ||
| # quantities (maybe use state_from_inputs) | ||
| translate_object._base.make_storage_data_input_vars(input_data) | ||
| # used for the translate test as inputs, but are generated by the | ||
| # MetricsTerms class and are not part of this data class | ||
| for delvar in ["ak", "bk", "ptop", "ks"]: | ||
| del input_data[delvar] | ||
| return cls(**input_data, quantity_factory=quantity_factory) |
There was a problem hiding this comment.
It can. I was thinking it might be useful for performance testing if we end up wanting to test non-baroclinic namelists before we implement reading initial state from disk. But someone wanting to do that can probably figure it out.
fv3core/fv3core/testing/translate.py
Outdated
| @@ -293,6 +296,12 @@ def make_grid_storage(self, pygrid): | |||
| if k in self.data: | |||
| self.make_composite_var_storage(k, self.data[k], shape) | |||
| del self.data[k] | |||
| for k in TranslateGrid.ee_vars: | |||
| lon, | ||
| lat, | ||
| lat_agrid, | ||
| grid_slice, |
There was a problem hiding this comment.
The type of this argument is non-intuitive and would be particularly helpful to type hint.
| grid_slice, | |
| grid_slice: Tuple[slice, slice], |
There was a problem hiding this comment.
If you do still want to pull grid_slice up one level, the way you would do it is by passing in lon and lat sliced on their own compute domains - that is, up to nx+1 and ny+1 since they're defined on cell corners. Then below, every access to lat or lon should be clipping off 1 point. pt6 becomes lat[:-1, :-1] while pt9 becomes lat[:-1, 1:].
Co-authored-by: Jeremy McGibbon <jeremym@allenai.org>
…slice as needed, slicing at a level up
Purpose
To add the computation of the baroclinic perturbation test case (Jablonowski & Williamson JRMS 2006) to enable running performance scripts without reading serialized data.
This PR: