diff --git a/opendrift/models/basemodel/__init__.py b/opendrift/models/basemodel/__init__.py index e93725d19..523dc5428 100644 --- a/opendrift/models/basemodel/__init__.py +++ b/opendrift/models/basemodel/__init__.py @@ -34,6 +34,7 @@ import geojson import xarray as xr +import pandas as pd import numpy as np import scipy import pyproj @@ -1915,7 +1916,8 @@ def run(self, if export_buffer_length is None: self.export_buffer_length = self.expected_steps_output else: - self.export_buffer_length = export_buffer_length + # TODO: Not clear why +1 is needed below, as this is already taken into account above + self.export_buffer_length = np.minimum(export_buffer_length, self.expected_steps_output+1) if self.time_step.days < 0: # For backwards simulation, we start at last seeded element @@ -1962,6 +1964,22 @@ def run(self, self.history.mask = True self.steps_exported = 0 + # Create Xarray Dataset to hold result + #coords = { # Initialize for the part fitting in memory + # 'trajectory': ('trajectory', np.arange(len(self.elements_scheduled))), + # 'time': ('time', pd.date_range(self.start_time, periods=self.export_buffer_length, freq=self.time_step_output))} + #shape = (len(coords['trajectory'][1]), len(coords['time'][1])) + #dims = ('trajectory', 'time') # Presently, but shall also allow single dimension + + #element_vars = {varname: (dims, np.nan*np.ones(shape=shape, dtype=var['dtype']), + # {attr:var[attr] for attr in var if attr not in ['seed', 'dtype']}) + # for varname,var in self.ElementType.variables.items() + # if self.export_variables is None or varname in self.export_variables} + #environment_vars = {varname: (dims, np.nan*np.ones(shape=shape, dtype=np.float32)) + # for varname,var in self.required_variables.items() + # if self.export_variables is None or varname in self.export_variables} + #self.result = xr.Dataset(coords=coords, data_vars=element_vars | environment_vars) + if outfile is not None: self.io_init(outfile) else: @@ -2206,6 +2224,8 @@ def run(self, del self.environment_profiles self.io_import_file(outfile) + #return self.result + def increase_age_and_retire(self): """Increase age of elements, and retire if older than config setting.""" # Increase age of elements @@ -2234,9 +2254,36 @@ def increase_age_and_retire(self): self.deactivate_elements(self.elements.lat > N, reason='outside') + def state_to_xarray(self): + if pd.to_datetime(self.time) in self.result.time: # Output time step + ID_ind = self.elements.ID - 1 + element_ind = range(len(ID_ind)) + insert_time = pd.to_datetime(self.time) + else: # Deactivated elements must be written even if no output timestep + deactivated = np.where(self.elements.status != 0)[0] + if len(deactivated) == 0: + return # No deactivated elements this sub-timestep + ID_ind = self.elements.ID[deactivated] - 1 + element_ind = deactivated + insert_time = self.result['time'].sel(time=self.time, method='backfill').values + + for var in self.result.var(): + if var == 'ID': + continue + for source in (self.elements, self.environment): + d = getattr(source, var, None) + if d is not None: + self.result[var].loc[{'time': insert_time, + 'trajectory': ID_ind}] = d[element_ind] + + if pd.to_datetime(self.time) == self.result.time[-1]: + pass # write to file + def state_to_buffer(self): """Append present state (elements and environment) to recarray.""" + #self.state_to_xarray() # Not ready yet + steps_calculation_float = \ (self.steps_calculation * self.time_step.total_seconds() / self.time_step_output.total_seconds()) + 1 diff --git a/tests/models/test_readers.py b/tests/models/test_readers.py index 9ad5b2366..a79a5b30a 100644 --- a/tests/models/test_readers.py +++ b/tests/models/test_readers.py @@ -244,8 +244,8 @@ def test_lazy_reader_leeway_compare(self): o1 = Leeway(loglevel=0) #o1.set_config('environment:fallback:land_binary_mask', 0) - o1.required_variables = [r for r in o1.required_variables - if r != 'land_binary_mask'] + o1.required_variables = {n:r for n,r in o1.required_variables.items() + if n != 'land_binary_mask'} o1.add_readers_from_list(reader_list, lazy=False) time = o1.env.readers['roms native'].start_time o1.seed_elements(lat=67.85, lon=14, time=time) @@ -253,8 +253,8 @@ def test_lazy_reader_leeway_compare(self): o2 = Leeway(loglevel=20) #o2.set_config('environment:fallback:land_binary_mask', 0) - o2.required_variables = [r for r in o1.required_variables - if r != 'land_binary_mask'] + o2.required_variables = {n:r for n,r in o1.required_variables.items() + if n != 'land_binary_mask'} o2.add_readers_from_list(reader_list, lazy=True) o2.seed_elements(lat=67.85, lon=14, time=time) o2.run(steps=5)