Make it possible to save velocity components along a particle track

doc/source/examples/configs/stokes_drift_and_leeway_template.cfg

@@ -450,3 +450,24 @@ windage_calculator = none
 # The wind factor, which should be a fraction of the windspeed
 wind_factor = 0.01
 
+[OUTPUT]
+
+# Save ocean current velocity components
+save_ocean_current_variables = False
+
+# Save surface wind velocity components (only valid if using atmospheric forcing)
+save_surface_wind_variables = False
+
+# Save stokes drift velocity components
+save_stokes_drift_variables = False
+
+# List of extra grid variables. Support for the following extra grid variables
+# is included.
+#
+# zeta - Sea surface elevation at the particle's position.0
+# h - The depth of the sea floor at the particle's location.
+#
+# Variables should be given as a comma separated list. If neither of these variables are
+# needed, leave this blank of comment the line.
+extra_grid_variables = h
+

pylag/model.pyx

@@ -66,8 +66,8 @@ cdef class OPTModel:
     cdef NumMethod num_method
     cdef ParticleStateNumMethod particle_state_num_method
     cdef SettlingVelocityCalculator settling_velocity_calculator
-    cdef object environmental_variables
     cdef object extra_grid_variables
+    cdef object environmental_variables
     cdef object particle_seed_smart_ptrs
     cdef object particle_smart_ptrs
     cdef vector[Particle*] particle_ptrs
@@ -86,6 +86,11 @@ cdef class OPTModel:
     cdef bint use_bio_model
     cdef BioModel bio_model
 
+    # Output flags
+    cdef bint save_ocean_current_vars
+    cdef bint save_surface_wind_vars
+    cdef bint save_stokes_drift_vars
+
     # Data precision, used when constructing diagnostic variable arrays
     cdef object precision
 
@@ -143,15 +148,7 @@ cdef class OPTModel:
             raise PyLagValueError(f"Unsupported model coordinate "
                                   f"system `{coordinate_system}`")
 
-        # Save a list of environmental variables to be returned as diagnostics
-        try:
-            var_names = self.config.get("OUTPUT",
-                    "environmental_variables").strip().split(',')
-            self.environmental_variables = [var_name.strip() for var_name in var_names]
-        except (configparser.NoSectionError, configparser.NoOptionError) as e:
-            self.environmental_variables = []
-
-        # Save a list of extra grid variables to be returned as diagnostics
+        # Save extra grid variables; to be returned as diagnostics
         try:
             var_names = self.config.get("OUTPUT",
                     "extra_grid_variables").strip().split(',')
@@ -160,6 +157,38 @@ cdef class OPTModel:
         except (configparser.NoSectionError, configparser.NoOptionError) as e:
             self.extra_grid_variables = []
 
+        # Save ocean current variables; to be returned as diagnostics
+        try:
+            save_ocean_current_vars = self.config.getboolean(
+                    "OUTPUT", "save_ocean_current_variables")
+        except (configparser.NoSectionError, configparser.NoOptionError) as e:
+            save_ocean_current_vars = False
+        self.save_ocean_current_vars = save_ocean_current_vars
+
+        # Save surface wind variables; to be returned as diagnostics
+        try:
+            save_surface_wind_vars = self.config.getboolean(
+                    "OUTPUT", "save_surface_wind_variables")
+        except (configparser.NoSectionError, configparser.NoOptionError) as e:
+            save_surface_wind_vars = False
+        self.save_surface_wind_vars = save_surface_wind_vars
+
+        # Save Stokes drift variables; to be returned as diagnostics
+        try:
+            save_stokes_drift_vars = self.config.getboolean(
+                    "OUTPUT", "save_stokes_drift_variables")
+        except (configparser.NoSectionError, configparser.NoOptionError) as e:
+            save_stokes_drift_vars = False
+        self.save_stokes_drift_vars = save_stokes_drift_vars
+
+        # Save environmental variables; to be returned as diagnostics
+        try:
+            var_names = self.config.get("OUTPUT",
+                    "environmental_variables").strip().split(',')
+            self.environmental_variables = [var_name.strip() for var_name in var_names]
+        except (configparser.NoSectionError, configparser.NoOptionError) as e:
+            self.environmental_variables = []
+
         self._global_time_step = get_global_time_step(self.config)
 
         # Are we running with biology?
@@ -609,6 +638,11 @@ cdef class OPTModel:
         """
         cdef ParticleSmartPtr particle_smart_ptr
 
+        # Velocity components
+        cdef DTYPE_FLOAT_t ocean_velocity[3]
+        cdef DTYPE_FLOAT_t wind_velocity[2]
+        cdef DTYPE_FLOAT_t stokes_drift_velocity[2]
+
         # Grid offsets
         if self.coordinate_system == "cartesian":
             xmin = self.data_reader.get_xmin()
@@ -641,15 +675,33 @@ cdef class OPTModel:
             dtype = variable_library.get_data_type(var_name, self.precision)
             diags[var_name] = np.empty(self._n_particles, dtype)
 
+        # Number of boundary encounters
+        dtype = variable_library.get_data_type('land_boundary_encounters',
+                                               self.precision)
+        diags['land_boundary_encounters'] = np.empty(self._n_particles, dtype)
+
         # Extra grid variables
         for var_name in self.extra_grid_variables:
             dtype = variable_library.get_data_type(var_name, self.precision)
             diags[var_name] = np.empty(self._n_particles, dtype)
 
-        # Number of boundary encounters
-        dtype = variable_library.get_data_type('land_boundary_encounters',
-                                               self.precision)
-        diags['land_boundary_encounters'] = np.empty(self._n_particles, dtype)
+        # Ocean current variables
+        if self.save_ocean_current_vars:
+            for var_name in ['uo', 'vo', 'wo']:
+                dtype = variable_library.get_data_type(var_name, self.precision)
+                diags[var_name] = np.empty(self._n_particles, dtype)
+
+        # Surface wind variables
+        if self.save_surface_wind_vars:
+            for var_name in ['u10', 'v10']:
+                dtype = variable_library.get_data_type(var_name, self.precision)
+                diags[var_name] = np.empty(self._n_particles, dtype)
+
+        # Stokes drift variables
+        if self.save_stokes_drift_vars:
+            for var_name in ['usd', 'vsd']:
+                dtype = variable_library.get_data_type(var_name, self.precision)
+                diags[var_name] = np.empty(self._n_particles, dtype)
 
         # Environmental variables
         for var_name in self.environmental_variables:
@@ -708,6 +760,47 @@ cdef class OPTModel:
 
                 diags['zeta'][i] = zeta
 
+            # Ocean current variables
+            if self.save_ocean_current_vars:
+                if particle_smart_ptr.in_domain:
+                    self.data_reader.get_velocity(time, particle_smart_ptr.get_ptr(), ocean_velocity)
+                    diags['uo'][i] = ocean_velocity[0]
+                    diags['vo'][i] = ocean_velocity[1]
+                    diags['wo'][i] = ocean_velocity[2]
+                else:
+                    diags['uo'][i] = variable_library.get_invalid_value(
+                        diags['uo'].dtype)
+                    diags['vo'][i] = variable_library.get_invalid_value(
+                        diags['vo'].dtype)
+                    diags['wo'][i] = variable_library.get_invalid_value(
+                        diags['wo'].dtype)
+
+            # Surface wind variables
+            if self.save_surface_wind_vars:
+                if particle_smart_ptr.in_domain:
+                    self.data_reader.get_ten_meter_wind_velocity(time,
+                            particle_smart_ptr.get_ptr(), wind_velocity)
+                    diags['u10'][i] = wind_velocity[0]
+                    diags['v10'][i] = wind_velocity[1]
+                else:
+                    diags['u10'][i] = variable_library.get_invalid_value(
+                        diags['u10'].dtype)
+                    diags['v10'][i] = variable_library.get_invalid_value(
+                        diags['v10'].dtype)
+
+            # Stokes drift variables
+            if self.save_stokes_drift_vars:
+                if particle_smart_ptr.in_domain:
+                    self.data_reader.get_surface_stokes_drift_velocity(time, 
+                            particle_smart_ptr.get_ptr(), stokes_drift_velocity)
+                    diags['usd'][i] = stokes_drift_velocity[0]
+                    diags['vsd'][i] = stokes_drift_velocity[1]
+                else:
+                    diags['usd'][i] = variable_library.get_invalid_value(
+                        diags['usd'].dtype)
+                    diags['vsd'][i] = variable_library.get_invalid_value(
+                        diags['vsd'].dtype)
+
             # Number of boundary encounters
             diags['land_boundary_encounters'][i] = \
                 particle_smart_ptr.land_boundary_encounters

pylag/netcdf_logger.py

@@ -119,6 +119,30 @@ def __init__(self, config, file_name, start_datetime, n_particles,
         except (configparser.NoSectionError, configparser.NoOptionError) as e:
             self.extra_grid_variables = []
 
+        # Save ocean current components
+        self._ocean_current_vars = {}
+        try:
+            self.save_ocean_current_vars = self.config.getboolean("OUTPUT",
+                                                       "save_ocean_current_variables")
+        except (configparser.NoSectionError, configparser.NoOptionError) as e:
+            self.save_ocean_current_vars = False
+
+        # Save 10 meter wind velocity components
+        self._wind_velocity_vars = {}
+        try:
+            self.save_wind_velocity_vars = self.config.getboolean("OUTPUT",
+                                                       "save_surface_wind_variables")
+        except (configparser.NoSectionError, configparser.NoOptionError) as e:
+            self.save_wind_velocity_vars = False
+
+        # Save Stokes drift components
+        self._stokes_drift_vars = {}
+        try:
+            self.save_stokes_drift_vars = self.config.getboolean("OUTPUT",
+                                                       "save_stokes_drift_variables")
+        except (configparser.NoSectionError, configparser.NoOptionError) as e:
+            self.save_stokes_drift_vars = False
+
         # Grid names
         self.grid_names = grid_names
 
@@ -226,6 +250,48 @@ def _create_file_structure(self, n_particles):
             self._zeta.invalid = \
                     f'{variable_library.get_invalid_value(data_type)}'
 
+        # Ocean current variables
+        if self.save_ocean_current_vars:
+            for comp in ['uo', 'vo', 'wo']:
+                data_type = variable_library.get_data_type(comp,
+                                                           self._precision)
+                self._ocean_current_vars[comp] = self._ncfile.createVariable(
+                        comp, data_type, ('time', 'particles',), **self._ncopts)
+                self._ocean_current_vars[comp].units = \
+                    variable_library.get_units(comp)
+                self._ocean_current_vars[comp].long_name = \
+                    variable_library.get_long_name(comp)
+                self._ocean_current_vars[comp].invalid = \
+                    f'{variable_library.get_invalid_value(data_type)}'
+
+        # Wind velocity variables
+        if self.save_wind_velocity_vars:
+            for comp in ['u10', 'v10']:
+                data_type = variable_library.get_data_type(comp,
+                                                           self._precision)
+                self._wind_velocity_vars[comp] = self._ncfile.createVariable(
+                        comp, data_type, ('time', 'particles',), **self._ncopts)
+                self._wind_velocity_vars[comp].units = \
+                    variable_library.get_units(comp)
+                self._wind_velocity_vars[comp].long_name = \
+                    variable_library.get_long_name(comp)
+                self._wind_velocity_vars[comp].invalid = \
+                    f'{variable_library.get_invalid_value(data_type)}'
+
+        # Stokes drift variables
+        if self.save_stokes_drift_vars:
+            for comp in ['usd', 'vsd']:
+                data_type = variable_library.get_data_type(comp,
+                                                           self._precision)
+                self._stokes_drift_vars[comp] = self._ncfile.createVariable(
+                        comp, data_type, ('time', 'particles',), **self._ncopts)
+                self._stokes_drift_vars[comp].units = \
+                    variable_library.get_units(comp)
+                self._stokes_drift_vars[comp].long_name = \
+                    variable_library.get_long_name(comp)
+                self._stokes_drift_vars[comp].invalid = \
+                    f'{variable_library.get_invalid_value(data_type)}'
+
         # Add number of land boundary encounters
         data_type = variable_library.get_data_type('land_boundary_encounters',
                                                    self._precision)
@@ -326,6 +392,22 @@ def write(self, time, particle_data):
         if 'zeta' in self.extra_grid_variables:
             self._zeta[tidx, :] = particle_data['zeta']
 
+        # Add ocean current variables
+        if self.save_ocean_current_vars:
+            self._ocean_current_vars['uo'][tidx, :] = particle_data['uo']
+            self._ocean_current_vars['vo'][tidx, :] = particle_data['vo']
+            self._ocean_current_vars['wo'][tidx, :] = particle_data['wo']
+
+        # Add wind velocity variables
+        if self.save_wind_velocity_vars:
+            self._wind_velocity_vars['u10'][tidx, :] = particle_data['u10']
+            self._wind_velocity_vars['v10'][tidx, :] = particle_data['v10']
+
+        # Add Stokes drift variables
+        if self.save_stokes_drift_vars:
+            self._stokes_drift_vars['usd'][tidx, :] = particle_data['usd']
+            self._stokes_drift_vars['vsd'][tidx, :] = particle_data['vsd']
+
         # Add environmental variables
         for var_name in self.environmental_variables:
             self._env_vars[var_name][tidx, :] = particle_data[var_name]

pylag/variable_library.py

@@ -230,6 +230,44 @@ def get_coordinate_variable_name(coordinate_system, variable_name):
 _variable_units['depth'] = 'm'
 _variable_long_names['depth'] = 'depth'
 
+# Velocity components
+# -------------------
+
+# u current component
+_variable_data_types['uo'] = 'REAL'
+_variable_units['uo'] = 'm s-1'
+_variable_long_names['uo'] = 'Eastward Current Velocity in the Water Column'
+
+# v current component
+_variable_data_types['vo'] = 'REAL'
+_variable_units['vo'] = 'm s-1'
+_variable_long_names['vo'] = 'Northward Current Velocity in the Water Column'
+
+# w current component
+_variable_data_types['wo'] = 'REAL'
+_variable_units['wo'] = 'm s-1'
+_variable_long_names['wo'] = 'Vertical Current Velocity in the Water Column'
+
+# u10 wind component
+_variable_data_types['u10'] = 'REAL'
+_variable_units['u10'] = 'm s-1'
+_variable_long_names['u10'] = '10 meter U wind component'
+
+# v10 wind component
+_variable_data_types['v10'] = 'REAL'
+_variable_units['v10'] = 'm s-1'
+_variable_long_names['v10'] = '10 meter V wind component'
+
+# usd Stokes drift u-component
+_variable_data_types['usd'] = 'REAL'
+_variable_units['usd'] = 'm s-1'
+_variable_long_names['usd'] = 'Sea surface wave stokes drift x velocity'
+
+# vsd Stokes drift v-component
+_variable_data_types['vsd'] = 'REAL'
+_variable_units['vsd'] = 'm s-1'
+_variable_long_names['vsd'] = 'Sea surface wave stokes drift y velocity'
+
 
 # Extra grid variables
 # --------------------

resources/pylag_fvcom.cfg

@@ -261,12 +261,24 @@ vert_bound_cond = reflecting
 
 [OUTPUT]
 
+# Save ocean current velocity components
+save_ocean_current_variables = False
+
+# Save surface wind velocity components (only valid if using atmospheric forcing)
+save_surface_wind_variables = False
+
+# Save stokes drift velocity components
+save_stokes_drift_variables = False
+
 # List of extra grid variables. Support for the following extra grid variables
 # is included.
 #
 # zeta - Sea surface elevation at the particle's position.0
 # h - The depth of the sea floor at the particle's location.
-
+#
+# Variables should be given as a comma separated list. If neither of these variables are
+# needed, leave this blank of comment the line.
+extra_grid_variables = h
 
 # List of environmental variables to be saved to file along with particle data.
 # This facility allows one to analyse changes in a given environmental variable