From 6b64b1bafe39e804a3a8f676712ad54bfb45e428 Mon Sep 17 00:00:00 2001
From: Gaute Hope <gauteh@met.no>
Date: Wed, 20 Nov 2024 09:44:50 +0100
Subject: [PATCH] traj: allow specifying the trajectory, observation and time
 dimensions/variables

Not 100% sure if this is necessary. I was trying to avoid the problem
where there's another dimension to time (say 'campaign'), but that's
another issue. This is mostly to allow users to specify the dimensions
manually when there is ambiguity. Maybe it would be better to allow
specifying the longitude and latitude variables (e.g. in the case when
there are interpolated values for the IMU, sometimes you want one,
sometimes the other).
---
 tests/test_dims.py |  8 ++++++++
 trajan/traj.py     | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/tests/test_dims.py b/tests/test_dims.py
index 0943f56..0fdd8d4 100644
--- a/tests/test_dims.py
+++ b/tests/test_dims.py
@@ -12,3 +12,11 @@ def test_no_time_dim(barents):
     # can we use trajan without
 
     print(b.traj)
+
+def test_custom_dims(barents: xr.DataArray):
+    # b = barents.expand_dims('campaign')
+    b = barents.rename(trajectory='campaign')
+    print(b)
+
+    g = b.traj(trajectory_dim='campaign').gridtime('1H')
+    print(g)
diff --git a/trajan/traj.py b/trajan/traj.py
index e0408d4..a899048 100644
--- a/trajan/traj.py
+++ b/trajan/traj.py
@@ -50,6 +50,39 @@ def __init__(self, ds, trajectory_dim, obs_dim, time_varname):
         self.obs_dim = obs_dim  # dimension along which time increases
         self.time_varname = time_varname
 
+    def __call__(self, trajectory_dim=None, obs_dim=None, time_varname=None):
+        """
+        Specify the trajectory, observation and time dimensions or variable names.
+
+
+        Parameters
+        ----------
+        trajectory_dim : string
+            Name of trajectory dimension (usually `trajectory` or `drifter`).
+
+        obs_dim : string
+            Name of observation dimension (usually `obs` or `time`).
+
+        time_varname : string
+            Name of time variable or dimension coordinate (usually `time`).
+
+
+        Returns
+        -------
+        traj : Traj
+            Traj accessor with manually specified coordinates and dimensions.
+        """
+        if trajectory_dim is not None:
+            self.trajectory_dim = trajectory_dim
+
+        if obs_dim is not None:
+            self.obs_dim = obs_dim
+
+        if time_varname is not None:
+            self.time_varname = time_varname
+
+        return self
+
     def __repr__(self):
         output = '=======================\n'
         output += 'TrajAn info:\n'