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'