def store_state_and_check_file(self, aliases): assert not os.path.isfile(self.ncfile) monitor = NetCDFMonitor(self.ncfile, aliases=aliases, write_on_store=True) monitor.store(state) assert os.path.isfile(self.ncfile)
def test_netcdf_monitor_multiple_times_batched_all_vars(): time_list = [ datetime(2013, 7, 20, 0), datetime(2013, 7, 20, 6), datetime(2013, 7, 20, 12), ] current_state = state.copy() try: assert not os.path.isfile('out.nc') monitor = NetCDFMonitor('out.nc') for time in time_list: current_state['time'] = time monitor.store(current_state) assert not os.path.isfile('out.nc') # not set to write on store monitor.write() assert os.path.isfile('out.nc') with xr.open_dataset('out.nc') as ds: assert len(ds.data_vars.keys()) == 2 assert 'air_temperature' in ds.data_vars.keys() assert ds.data_vars['air_temperature'].attrs['units'] == 'degK' assert tuple( ds.data_vars['air_temperature'].shape) == (len(time_list), nx, ny, nz) assert 'air_pressure' in ds.data_vars.keys() assert ds.data_vars['air_pressure'].attrs['units'] == 'Pa' assert tuple( ds.data_vars['air_pressure'].shape) == (len(time_list), nx, ny, nz) assert len(ds['time']) == len(time_list) assert np.all(ds['time'].values == [np.datetime64(time) for time in time_list]) finally: # make sure we remove the output file if os.path.isfile('out.nc'): os.remove('out.nc')
def test_netcdf_monitor_single_write_on_store(): try: assert not os.path.isfile('out.nc') monitor = NetCDFMonitor('out.nc', write_on_store=True) monitor.store(state) assert os.path.isfile('out.nc') with xr.open_dataset('out.nc') as ds: assert len(ds.data_vars.keys()) == 2 assert 'air_temperature' in ds.data_vars.keys() assert ds.data_vars['air_temperature'].attrs['units'] == 'degK' assert tuple(ds.data_vars['air_temperature'].shape) == (1, nx, ny, nz) assert 'air_pressure' in ds.data_vars.keys() assert ds.data_vars['air_pressure'].attrs['units'] == 'Pa' assert tuple(ds.data_vars['air_pressure'].shape) == (1, nx, ny, nz) assert len(ds['time']) == 1 assert ds['time'][0] == np.datetime64(state['time']) finally: # make sure we remove the output file if os.path.isfile('out.nc'): os.remove('out.nc')
def test_netcdf_monitor_raises_when_names_change_on_batch_write(): current_state = state.copy() try: assert not os.path.isfile('out.nc') monitor = NetCDFMonitor('out.nc') current_state['time'] = datetime(2013, 7, 20, 0) monitor.store(current_state) assert not os.path.isfile('out.nc') current_state['time'] = datetime(2013, 7, 20, 6) current_state['air_density'] = current_state['air_pressure'] monitor.store(current_state) try: monitor.write() except InvalidStateError: pass except Exception as err: raise err else: raise AssertionError( 'Expected InvalidStateError but was not raised.') finally: # make sure we remove the output file if os.path.isfile('out.nc'): os.remove('out.nc')
# Set initial/boundary conditions latitudes = my_state['latitude'].values longitudes = my_state['longitude'].values zenith_angle = np.radians(latitudes) surface_shape = [len(longitudes), len(latitudes)] my_state['zenith_angle'].values = zenith_angle my_state['eastward_wind'].values[:] = np.random.randn( *my_state['eastward_wind'].shape) my_state['ocean_mixed_layer_thickness'].values[:] = 50 surf_temp_profile = 290 - (40 * np.sin(zenith_angle)**2) my_state['surface_temperature'].values = surf_temp_profile for i in range(1500 * 24 * 6): diag, my_state = dycore(my_state, model_time_step) my_state.update(diag) my_state['time'] += model_time_step if i % (6 * 24) == 0: netcdf_monitor.store(my_state) monitor.store(my_state) print('max. zonal wind: ', np.amax(my_state['eastward_wind'].values)) print('max. humidity: ', np.amax(my_state['specific_humidity'].values)) print('max. surf temp: ', my_state['surface_temperature'].max(keep_attrs=True).values) print(my_state['time'])
def main(): # ============ Adjustable Variables ============ # Integration Options dt = timedelta(minutes=15) # timestep duration = '48_00:00' # run duration ('<days>_<hours>:<mins>')t linearized = True ncout_freq = 6 # netcdf write frequency (hours) plot_freq = 6 # plot Monitor call frequency (hours) ntrunc = 42 # triangular truncation for spharm (e.g., 21 --> T21) # Diffusion Options diff_on = True # Use diffusion? k = 2.338e16 # Diffusion coefficient for del^4 hyperdiffusion # Forcing Options forcing_on = True # Apply vort. tendency forcing? damp_ts = 14.7 # Damping timescale (in days) # I/O Options ncoutfile = os.path.join(os.path.dirname(__file__), 'sardeshmukh88.nc') append_nc = False # Append to an existing netCDF file? # ============================================== start = time() # Get the initial state state = super_rotation(linearized=linearized, ntrunc=ntrunc) # Set up the Timestepper with the desired Prognostics if linearized: dynamics_prog = LinearizedDynamics(ntrunc=ntrunc) diffusion_prog = LinearizedDiffusion(k=k, ntrunc=ntrunc) damping_prog = LinearizedDamping(tau=damp_ts) else: dynamics_prog = NonlinearDynamics(ntrunc=ntrunc) diffusion_prog = NonlinearDiffusion(k=k, ntrunc=ntrunc) damping_prog = NonlinearDamping(tau=damp_ts) prognostics = [TendencyInDiagnosticsWrapper(dynamics_prog, 'dynamics')] if diff_on: prognostics.append(TendencyInDiagnosticsWrapper(diffusion_prog, 'diffusion')) if forcing_on: # Get our suptropical RWS forcing (from equatorial divergence) rws, rlat, rlon = rws_from_tropical_divergence(state) prognostics.append(TendencyInDiagnosticsWrapper(Forcing.from_numpy_array(rws, rlat, rlon, ntrunc=ntrunc, linearized=linearized), 'forcing')) prognostics.append(TendencyInDiagnosticsWrapper(damping_prog, 'damping')) stepper = Leapfrog(prognostics) # Create Monitors for plotting & storing data plt_monitor = PlotFunctionMonitor(debug_plots.fourpanel) if os.path.isfile(ncoutfile) and not append_nc: os.remove(ncoutfile) aliases = get_component_aliases(*prognostics) nc_monitor = NetCDFMonitor(ncoutfile, write_on_store=True, aliases=aliases) # Figure out the end date of this run d, h, m = re.split('[_:]', duration) end_date = state['time'] + timedelta(days=int(d), hours=int(h), minutes=int(m)) # Begin the integration loop idate = state['time'] while state['time'] <= end_date: # Get the state at the next timestep using our Timestepper diagnostics, next_state = stepper(state, dt) # Add any calculated diagnostics to our current state state.update(diagnostics) # Write state to netCDF every <ncout_freq> hours fhour = (state['time'] - idate).days*24 + (state['time'] - idate).seconds/3600 if fhour % ncout_freq == 0: print(state['time']) nc_monitor.store(state) # Make plot(s) every <plot_freq> hours if fhour % plot_freq == 0: plt_monitor.store(state) # Advance the state to the next timestep next_state['time'] = state['time'] + dt state = next_state print('TOTAL INTEGRATION TIME: {:.02f} min\n'.format((time()-start)/60.))
# Figure out the end date of this run d, h, m = re.split('[_:]', duration) end_date = state['time'] + timedelta(days=int(d), hours=int(h), minutes=int(m)) # Begin the integration loop idate = state['time'] while state['time'] <= end_date: # Get the state at the next timestep using our Timestepper diagnostics, next_state = stepper(state, dt) # Add any calculated diagnostics to our current state state.update(diagnostics) # Write state to netCDF every <ncout_freq> hours fhour = (state['time'] - idate).days * 24 + (state['time'] - idate).seconds / 3600 if fhour % ncout_freq == 0: print(state['time']) nc_monitor.store(state) # Make plot(s) every <plot_freq> hours if fhour % plot_freq == 0: plt_monitor.store(state) # Advance the state to the next timestep next_state['time'] = state['time'] + dt state = next_state print('TOTAL INTEGRATION TIME: {:.02f} min\n'.format((time() - start) / 60.))
[simple_physics, convection, radiation_lw, radiation_sw, slab]) state['air_temperature'].values[:] = 270 state['surface_albedo_for_direct_shortwave'].values[:] = 0.5 state['surface_albedo_for_direct_near_infrared'].values[:] = 0.5 state['surface_albedo_for_diffuse_shortwave'].values[:] = 0.5 state['zenith_angle'].values[:] = np.pi / 2.5 state['surface_temperature'].values[:] = 300. state['ocean_mixed_layer_thickness'].values[:] = 5 state['area_type'].values[:] = 'sea' time_stepper = AdamsBashforth([convection, radiation_lw, radiation_sw, slab]) for i in range(20000): convection.current_time_step = timestep diagnostics, state = time_stepper(state, timestep) state.update(diagnostics) diagnostics, new_state = simple_physics(state, timestep) state.update(diagnostics) if (i + 1) % 20 == 0: monitor.store(state) netcdf_monitor.store(state) print(i, state['surface_temperature'].values) print(state['surface_upward_sensible_heat_flux']) print(state['surface_upward_latent_heat_flux']) state.update(new_state) state['time'] += timestep state['eastward_wind'].values[:] = 3.
class MyModel(): """Climate model class.""" def __init__(self, dt_seconds=1800, nx=64, ny=32, nz=10, state=None, input_fields_to_store=input_vars, output_fields_to_store=output_vars, input_save_fn=None, output_save_fn=None, save_interval=6, convection=None, extra_components=[]): """ Initialize model. Uses SSTs from Andersen and Kuang 2012. Creates initial state unless state is given. """ climt.set_constants_from_dict( {'stellar_irradiance': { 'value': 200, 'units': 'W m^-2' }}) self.model_time_step = timedelta(seconds=dt_seconds) self.step_counter = 0 self.save_interval = save_interval # Create components if convection is None: convection = climt.EmanuelConvection( tendencies_in_diagnostics=True) simple_physics = TimeDifferencingWrapper( climt.SimplePhysics(tendencies_in_diagnostics=True)) radiation = climt.GrayLongwaveRadiation(tendencies_in_diagnostics=True) components = [simple_physics, radiation, convection] + extra_components self.dycore = climt.GFSDynamicalCore(components, number_of_damped_levels=2) grid = climt.get_grid(nx=nx, ny=ny, nz=nz) if state is None: self.create_initial_state(grid) else: self.state = state if not input_save_fn is None: self.input_netcdf_monitor = NetCDFMonitor( input_save_fn, write_on_store=True, store_names=input_fields_to_store) if not output_save_fn is None: self.output_netcdf_monitor = NetCDFMonitor( output_save_fn, write_on_store=True, store_names=output_fields_to_store) def create_initial_state(self, grid): """Create initial state.""" # Create model state self.state = climt.get_default_state([self.dycore], grid_state=grid) # Set initial/boundary conditions latitudes = self.state['latitude'].values sst_k = and_kua_sst(latitudes) self.state['surface_temperature'] = DataArray(sst_k, dims=['lat', 'lon'], attrs={'units': 'degK'}) self.state['eastward_wind'].values[:] = np.random.randn( *self.state['eastward_wind'].shape) def step(self): """Take one time step forward.""" self.diag, self.state = self.dycore(self.state, self.model_time_step) self.state.update(self.diag) self.state['time'] += self.model_time_step if hasattr(self, 'input_netcdf_monitor'): if self.step_counter % self.save_interval == 0: self.input_netcdf_monitor.store(self.state) if hasattr(self, 'output_netcdf_monitor'): if (self.step_counter - 1) % self.save_interval == 0: self.output_netcdf_monitor.store(self.state) self.step_counter += 1 def iterate(self, steps, noprog=False): """Iterate over several time steps.""" for i in tqdm(range(steps), disable=noprog): self.step()