More methods now return Xarray Dataarrays instead of numpy arrays

trajan/traj.py

@@ -42,7 +42,21 @@ def ensure_time_dim(ds, time_dim):
         return ds
 
 def grid_area(lons, lats):
-    """Calculate the area of each grid cell"""
+    """
+    Calculate the area of each grid cell.
+
+    Parameters
+    ----------
+    lons : array-like
+        Longitudes of the grid.
+    lats : array-like
+        Latitudes of the grid.
+
+    Returns
+    -------
+    xarray.DataArray
+        Grid cell areas with dimensions ('lat', 'lon').
+    """
     from shapely.geometry import Polygon
 
     if lons.ndim == 1:
@@ -59,7 +73,7 @@ def grid_area(lons, lats):
             polygon = Polygon([(lon[0], lat[0]), (lon[1], lat[1]), (lon[2], lat[2]), (lon[3], lat[3])])
             grid_areas[i, j] = abs(geod.geometry_area_perimeter(polygon)[0])
 
-    return grid_areas
+    return xr.DataArray(grid_areas, name="grid_area", attrs={'units': 'm^2', 'description': 'Area of each grid cell'})
 
 
 class Traj:
@@ -595,15 +609,17 @@ def assign_cf_attrs(self,
         return ds
 
     def index_of_last(self):
-        """Find index of last valid position along each trajectory.
+        """
+        Find the index of the last valid position along each trajectory.
 
         Returns
         -------
-        array-like
-            Array of the index of the last valid position along each trajectory.
+        xarray.DataArray
+            Index of the last valid position for each trajectory.
+            Dimensions: ('trajectory',).
         """
-        return np.ma.notmasked_edges(np.ma.masked_invalid(self.ds.lon.values),
-                                     axis=1)[1][1]
+        last_indices = np.ma.notmasked_edges(np.ma.masked_invalid(self.ds.lon.values), axis=1)[1][1]
+        return xr.DataArray(last_indices, dims=[self.trajectory_dim], name="index_of_last")
 
     @abstractmethod
     def speed(self) -> xr.DataArray:
@@ -855,35 +871,33 @@ def convex_hull_contains_point(self, lon, lat):
         return p.contains_points(point)[0]
 
     def get_area_convex_hull(self):
-        """Return the area [m2] of the convex hull spanned by all positions.
+        """
+        Calculate the area [m2] of the convex hull spanned by all positions.
 
         Returns
         -------
-        scalar
-            Area [m2] of convex hull around all positions.
+        xarray.DataArray
+            Area of the convex hull in square meters.
         """
-
         from scipy.spatial import ConvexHull
 
         lon = self.ds.lon.where(self.ds.status == 0)
         lat = self.ds.lat.where(self.ds.status == 0)
         fin = np.isfinite(lat + lon)
         if np.sum(fin) <= 3:
-            return 0
+            return xr.DataArray(0, name="convex_hull_area", attrs={"units": "m2"})
         if len(np.unique(lat)) == 1 and len(np.unique(lon)) == 1:
-            return 0
+            return xr.DataArray(0, name="convex_hull_area", attrs={"units": "m2"})
         lat = lat[fin]
         lon = lon[fin]
-        # An equal area projection centered around the particles
         aea = pyproj.Proj(
             f'+proj=aea +lat_0={lat.mean().values} +lat_1={lat.min().values} +lat_2={lat.max().values} +lon_0={lon.mean().values} +x_0=0 +y_0=0 +datum=NAD83 +units=m +no_defs'
         )
-
         x, y = aea(lat, lon, inverse=False)
         fin = np.isfinite(x + y)
         points = np.vstack((y.T, x.T)).T
         hull = ConvexHull(points)
-        return np.array(hull.volume)  # volume=area for 2D as here
+        return xr.DataArray(hull.volume, name="convex_hull_area", attrs={"units": "m2"})
 
     @abstractmethod
     def gridtime(self, times, time_varname=None) -> xr.Dataset:
@@ -1277,7 +1291,7 @@ def make_grid(self, dx, dy=None, z=None,
 
         x = np.arange(xmin, xmax + dx*2, dx)  # One extra row/column
         y = np.arange(ymin, ymax + dy*2, dy)
-        area = grid_area(x, y)
+        area = grid_area(x, y).data
 
         # Create Xarray Dataset
         data_vars = {}

trajan/traj1d.py

@@ -19,9 +19,18 @@ def __init__(self, ds, trajectory_dim, obs_dim, time_varname):
         super().__init__(ds, trajectory_dim, obs_dim, time_varname)
 
     def timestep(self):
-        """Time step between observations in seconds."""
-        return ((self.ds.time[1] - self.ds.time[0]) /
-                np.timedelta64(1, 's')).values
+        """
+        Calculate the time step between observations in seconds.
+
+        Returns
+        -------
+        xarray.DataArray
+            Time step between observations with a single value.
+            Attributes:
+            - units: seconds
+        """
+        timestep = ((self.ds.time[1] - self.ds.time[0]) / np.timedelta64(1, 's')).values
+        return xr.DataArray(timestep, name="timestep", attrs={"units": "seconds"})
 
     def is_1d(self):
         return True
@@ -30,54 +39,78 @@ def is_2d(self):
         return False
 
     def to_2d(self, obs_dim='obs'):
+        """
+        Convert the dataset to a 2D representation.
+
+        Parameters
+        ----------
+        obs_dim : str, optional
+            Name of the observation dimension in the 2D representation, by default 'obs'.
+
+        Returns
+        -------
+        xarray.Dataset
+            Dataset with a 2D representation of trajectories.
+        """
         ds = self.ds.copy()
-        time = ds[self.time_varname].rename({
-            self.time_varname: obs_dim
-        }).expand_dims(
-            dim={
-                self.trajectory_dim: ds.sizes[self.trajectory_dim]
-            }).assign_coords({self.trajectory_dim: ds[self.trajectory_dim]})
-        # TODO should also add cf_role here
+        time = ds[self.time_varname].rename({self.time_varname: obs_dim}).expand_dims(
+            dim={self.trajectory_dim: ds.sizes[self.trajectory_dim]}
+        ).assign_coords({self.trajectory_dim: ds[self.trajectory_dim]})
         ds = ds.rename({self.time_varname: obs_dim})
         ds[self.time_varname] = time
-        ds[obs_dim] = np.arange(0, ds.sizes[obs_dim])
-
+        ds[obs_dim] = xr.DataArray(np.arange(0, ds.sizes[obs_dim]), dims=[obs_dim])
         return ds
 
     def to_1d(self):
         return self.ds.copy()
 
     def time_to_next(self):
+        """
+        Calculate the time difference to the next observation.
+
+        Returns
+        -------
+        xarray.DataArray
+            Time difference to the next observation with the same dimensions as the dataset.
+            Attributes:
+            - units: seconds
+        """
         time_step = self.ds.time[1] - self.ds.time[0]
-        return time_step
+        return xr.DataArray(time_step, name="time_to_next", attrs={"units": "seconds"})
 
     def velocity_spectrum(self):
-
+        """
+        Calculate the velocity spectrum for a single trajectory.
+
+        Returns
+        -------
+        xarray.DataArray
+            Velocity spectrum with dimensions ('period').
+            Attributes:
+            - units: power
+        """
         if self.ds.sizes[self.trajectory_dim] > 1:
-            raise ValueError(
-                'Spectrum can only be calculated for a single trajectory')
+            raise ValueError('Spectrum can only be calculated for a single trajectory')
 
         u, v = self.velocity_components()
         u = u.squeeze()
         v = v.squeeze()
         u = u[np.isfinite(u)]
         v = v[np.isfinite(v)]
 
-        timestep_h = (self.ds.time[1] - self.ds.time[0]) / np.timedelta64(
-            1, 'h')  # hours since start
+        timestep_h = (self.ds.time[1] - self.ds.time[0]) / np.timedelta64(1, 'h')  # hours since start
 
         ps = np.abs(np.fft.rfft(np.abs(u + 1j * v)))
         freq = np.fft.rfftfreq(n=u.size, d=timestep_h.values)
         freq[0] = np.nan
 
         da = xr.DataArray(
             data=ps,
-            name='velocity spectrum',
+            name='velocity_spectrum',
             dims=['period'],
-            coords={'period': (['period'], 1 / freq, {
-                'units': 'hours'
-            })},
-            attrs={'units': 'power'})
+            coords={'period': (['period'], 1 / freq, {'units': 'hours'})},
+            attrs={'units': 'power'}
+        )
 
         return da
 
@@ -146,9 +179,27 @@ def skill_matching(traj, expected):
 
         return s
 
-    def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.Dataset:
-
-        expected = expected.traj  # Normalise
+    def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.DataArray:
+        """
+        Calculate the skill score for trajectories.
+
+        Parameters
+        ----------
+        expected : Traj1d
+            Expected trajectory dataset.
+        method : str, optional
+            Skill score method, by default 'liu-weissberg'.
+        **kwargs : dict
+            Additional arguments for the skill score calculation.
+
+        Returns
+        -------
+        xarray.DataArray
+            Skill score with dimensions matching the dataset.
+            Attributes:
+            - method: Skill score calculation method.
+        """
+        expected = expected.traj  # Normalize
         expected_trajdim = expected.trajectory_dim
         self_trajdim = self.trajectory_dim
 
@@ -157,7 +208,8 @@ def skill(self, expected, method='liu-weissberg', **kwargs) -> xr.Dataset:
         if numtraj_self > 1 and numtraj_expected > 1 and numtraj_self != numtraj_expected:
             raise ValueError(
                 'Datasets must have the same number of trajectories, or a single trajectory. '
-                f'This dataset: {numtraj_self}, expected: {numtraj_expected}.')
+                f'This dataset: {numtraj_self}, expected: {numtraj_expected}.'
+            )
 
         numobs_self = self.ds.sizes[self.obs_dim]
         numobs_expected = expected.ds.sizes[expected.obs_dim]