Skip to content
Prev Previous commit
Next Next commit
changes of extract_light_with_mask
- possibility to select extra variables to extract
  • Loading branch information
ludwigVonKoopa committed Dec 8, 2021
commit f4a6fee5a6ba93550c487b9a44759451c6a6df27
16 changes: 13 additions & 3 deletions src/py_eddy_tracker/observations/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -1301,7 +1301,7 @@ def extract_with_period(self, period):

return self.extract_with_mask(self.get_mask_with_period(period))

def extract_light_with_mask(self, mask):
def extract_light_with_mask(self, mask, track_extra_variables=[]):
"""extract data with mask, but only with variables used for coherence, aka self.array_variables

:param mask: mask used to extract
Expand All @@ -1319,7 +1319,7 @@ def extract_light_with_mask(self, mask):
variables = ["time"] + self.array_variables
new = self.__class__(
size=nb_obs,
track_extra_variables=[],
track_extra_variables=track_extra_variables,
track_array_variables=self.track_array_variables,
array_variables=self.array_variables,
only_variables=variables,
Expand All @@ -1333,9 +1333,19 @@ def extract_light_with_mask(self, mask):
f"{nb_obs} observations will be extracted ({nb_obs / self.shape[0]:.3%})"
)

for field in variables:
for field in variables + track_extra_variables:
logger.debug("Copy of field %s ...", field)
new.obs[field] = self.obs[field][mask]

if "previous_obs" in track_extra_variables and "next_obs" in track_extra_variables:
# n & p must be re-index
n, p = self.next_obs[mask], self.previous_obs[mask]
# we add 2 for -1 index return index -1
translate = -ones(len(self) + 1, dtype="i4")
translate[:-1][mask] = arange(nb_obs)
new.next_obs[:] = translate[n]
new.previous_obs[:] = translate[p]

return new

def extract_with_mask(self, mask):
Expand Down