Merge pull request #1223 from knutfrode/dev

Dev

opendrift/models/oceandrift.py

@@ -199,11 +199,38 @@ def update(self):
         # Optional machine learning correction
         self.machine_learning_correction()
 
+    def simulate_trajectories(self, outfile, trajectories,
+                              wind_drift_factors=None, current_drift_factors=None,
+                              time_step=None, time_step_output=None,
+                              simulation_duration=None, simulation_interval=None):
+
+        import pandas as pd
+        time_step_output = pd.Timedelta(time_step_output)
+        simulation_interval = pd.Timedelta(simulation_interval)
+        # Interpolate trajectories to output time step
+        trajectories = trajectories.traj.gridtime(time_step_output)
+        # Find all seed combinations: position, time, wdf, cdf
+        #  Loop også over trajektorier -> origin_marker
+        step = int(simulation_interval / time_step_output)
+        tind = np.arange(0, trajectories.sizes['time'], step)
+        start_lons = trajectories.isel(time=tind).isel(trajectory=0).lon.values
+        start_lats = trajectories.isel(time=tind).isel(trajectory=0).lat.values
+        start_times = trajectories.time[tind]
+        self.set_config('drift:max_age_seconds', simulation_duration.total_seconds())
+        for (lo,la,ti) in zip(start_lons, start_lats, start_times):
+            if np.isnan(lo):
+                continue
+            ti = pd.Timestamp(ti.values).to_pydatetime()
+            self.seed_elements(lon=lo, lat=la, time=ti)
+        print(self)
+        self.run(outfile=outfile, end_time=pd.Timestamp(start_times[-1].values).to_pydatetime()+simulation_duration)
+        # Simulate and save to file
+
     def wind_drift_factor_from_trajectory_lw(self, drifters, wind_drift_factors,
                                              simulation_length, simulation_interval):
         """Perform simulations and use skillscore to optimize wind_drift_factor
 
-        drifters: list of disctionaries with numpy arrays of 'lon' and 'lat'
+        drifters: list of dictionaries with numpy arrays of 'lon' and 'lat'
                     and list of datetimes
         wind_drift_factors: the wind_drift_factors to use for simulations/optimalizations
         """

opendrift/readers/basereader/variables.py

@@ -58,6 +58,15 @@ def center(self):
 
     def rotate_vectors(self, reader_x, reader_y, u_component, v_component,
                        proj_from, proj_to):
+        if isinstance(u_component, list):  # Looping recursively over ensemble members
+            uout = []
+            vout = []
+            for ucomp, vcomp in zip(u_component, v_component):
+                ucomprot, vcomprot = self.rotate_vectors(
+                    reader_x, reader_y, ucomp, vcomp, proj_from, proj_to)
+                uout.append(ucomprot)
+                vout.append(vcomprot)
+            return uout, vout
         """Rotate vectors from one crs to another."""
 
         if type(proj_from) is str:

opendrift/readers/reader_netCDF_CF_generic.py

@@ -119,6 +119,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
         lat_var_name = None
         self.unitfactor = 1
         self.realizations = None
+        self.ensemble_dimension = None
+        self.dimensions = {}
         for var_name in self.Dataset.variables:
             var = self.Dataset.variables[var_name]
 
@@ -156,6 +158,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
             if (axis == 'X' or standard_name == 'projection_x_coordinate' or standard_name == 'grid_longitude') \
                     and var.ndim == 1:
                 self.xname = var_name
+                if len(var.dims)==1:
+                    self.dimensions['x'] = var.dims[0]
                 # Fix for units; should ideally use udunits package
                 if units == 'km':
                     self.unitfactor = 1000
@@ -166,6 +170,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
             if (axis == 'Y' or standard_name == 'projection_y_coordinate' or standard_name == 'grid_latitude') \
                     and var.ndim == 1:
                 self.yname = var_name
+                if len(var.dims)==1:
+                    self.dimensions['y'] = var.dims[0]
                 # Fix for units; should ideally use udunits package
                 if units == 'km':
                     self.unitfactor = 1000
@@ -175,6 +181,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
                 y = var_data*self.unitfactor
             if (standard_name == 'depth' or axis == 'Z') and var.ndim==1:
                 var_data = var.values
+                if len(var.dims)==1:
+                    self.dimensions['z'] = var.dims[0]
                 if var_data.ndim == 1:  # Earlier this was not a requirement above
                     if 'positive' not in var.attrs or \
                             var.attrs['positive'] == 'up':
@@ -186,6 +194,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
                 var_data = var.values
                 time = var_data
                 time_units = units
+                if len(var.dims)==1:
+                    self.dimensions['time'] = var.dims[0]
 
                 if isinstance(time[0], np.bytes_):
                     # This hack is probably only necessary for CERSAT/GELOBCURRENT
@@ -215,6 +225,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
                     self.realizations = var_data
                     logger.debug('%i ensemble members available'
                                 % len(self.realizations))
+                if len(var.dims)==1:
+                    self.ensemble_dimension = var.dims[0]
 
         # Temporary workaround for Barents EPS model
         if self.realizations is None and 'ensemble_member' in self.Dataset.dims:
@@ -241,6 +253,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
                 y = lat_var.data
                 self.xname = lon_var_name
                 self.yname = lat_var_name
+                self.dimensions['x'] = lon_var.dims[0]
+                self.dimensions['y'] = lat_var.dims[0]
                 if self.proj4 is None:
                     self.proj4 = '+proj=latlong'
             elif lon_var.ndim == 2:
@@ -288,6 +302,8 @@ def __init__(self, filename=None, zarr_storage_options=None, name=None, proj4=No
             self.xmax -= 360
             self.x -= 360
 
+        logger.info(f'Detected dimensions: {self.dimensions}')
+
         ##########################################
         # Find all variables having standard_name
         ##########################################
@@ -401,18 +417,31 @@ def get_variables(self, requested_variables, time=None,
 
             ensemble_dim = None
             if continuous is True:
-                if var.ndim == 2:
-                    variables[par] = var[indy, indx]
-                elif var.ndim == 3:
-                    variables[par] = var[indxTime, indy, indx]
-                elif var.ndim == 4:
-                    variables[par] = var[indxTime, indz, indy, indx]
-                elif var.ndim == 5:  # Ensemble data
-                    variables[par] = var[indxTime, indz, indrealization, indy, indx]
-                    ensemble_dim = 0  # Hardcoded ensemble dimension for now
-                else:
-                    raise Exception('Wrong dimension of variable: ' +
-                                    self.variable_mapping[par])
+                if True:  # new dynamic way
+                    dimindices = {'x': indx, 'y': indy, 'time': indxTime, 'z': indz}
+                    subset = {vdim:dimindices[dim] for dim,vdim in self.dimensions.items() if vdim in var.dims}
+                    variables[par] = var.isel(subset)
+                    # Remove any unknown dimensions
+                    for dim in variables[par].dims:
+                        if dim not in self.dimensions.values() and dim != self.ensemble_dimension:
+                            logger.debug(f'Removing unknown dimension: {dim}')
+                            variables[par] = variables[par].squeeze(dim=dim)
+                    if self.ensemble_dimension is not None and self.ensemble_dimension in variables[par].dims:
+                        ensemble_dim = 0  # hardcoded, may not work for MEPS
+                else:  # old hardcoded way
+                    if var.ndim == 2:
+                        variables[par] = var[indy, indx]
+                    elif var.ndim == 3:
+                        variables[par] = var[indxTime, indy, indx]
+                    elif var.ndim == 4:
+                        variables[par] = var[indxTime, indz, indy, indx]
+                    elif var.ndim == 5:  # Ensemble data
+                        variables[par] = var[indxTime, indz, indrealization, indy, indx]
+                        ensemble_dim = 0  # Hardcoded ensemble dimension for now
+                    else:
+                        raise Exception('Wrong dimension of variable: ' +
+                                        self.variable_mapping[par])
+            # The below should also be updated to dynamic subsetting
             else:  # We need to read left and right parts separately
                 if var.ndim == 2:
                     left = var[indy, indx_left]
@@ -446,7 +475,7 @@ def get_variables(self, requested_variables, time=None,
             # Ensemble blocks are split into lists
             if ensemble_dim is not None:
                 num_ensembles = variables[par].shape[ensemble_dim]
-                logger.debug('Num ensembles: %i ' % num_ensembles)
+                logger.debug(f'Num ensembles for {par}: {num_ensembles}')
                 newvar = [0]*num_ensembles
                 for ensemble_num in range(num_ensembles):
                     newvar[ensemble_num] = \
@@ -482,7 +511,7 @@ def get_variables(self, requested_variables, time=None,
                 from opendrift.readers.basereader import vector_pairs_xy
                 for vectorpair in vector_pairs_xy:
                     if vectorpair[0] in self.rotate_mapping and vectorpair[0] in variables.keys():
-                        logger.debug(f'Rotating vector from east/north to xy orientation: {vectorpair}')
+                        logger.debug(f'Rotating vector from east/north to xy orientation: {vectorpair[0:2]}')
                         variables[vectorpair[0]], variables[vectorpair[1]] = self.rotate_vectors(
                             lon, lat, variables[vectorpair[0]], variables[vectorpair[1]],
                             pyproj.Proj('+proj=latlong'), self.proj)