- Add option to load netcdf in memory before to be open

- Use of safe_load for yaml
- replace some numpy operation by numba equivalent

Author: AntSimi

examples/02_eddy_identification/pet_eddy_detection.py

@@ -91,7 +91,7 @@ def update_axes(ax, mappable=None):
 update_axes(ax)
 
 # %%
-# Creteria for rejecting a contour
+# Criteria for rejecting a contour
 # 0. - Accepted (green)
 # 1. - Rejection for shape error (red)
 # 2. - Masked value within contour (blue)

notebooks/python_module/02_eddy_identification/pet_eddy_detection.ipynb

@@ -235,7 +235,7 @@
       },
       "outputs": [],
       "source": [
-        "ax = start_axes(\"Detected Eddies\")\na.display(ax, color=\"r\", linewidth=0.75, label=\"Anticyclonic ({nb_obs} eddies)\", ref=-10)\nc.display(ax, color=\"b\", linewidth=0.75, label=\"Cyclonic ({nb_obs} eddies)\", ref=-10)\nax.legend()\nupdate_axes(ax)"
+        "ax = start_axes(\"Detected Eddies\")\na.display(\n    ax, color=\"r\", linewidth=0.75, label=\"Anticyclonic ({nb_obs} eddies)\", ref=-10\n)\nc.display(ax, color=\"b\", linewidth=0.75, label=\"Cyclonic ({nb_obs} eddies)\", ref=-10)\nax.legend()\nupdate_axes(ax)"
       ]
     },
     {

setup.py

@@ -28,7 +28,6 @@
     scripts=[
         "src/scripts/EddySubSetter",
         "src/scripts/EddyTranslate",
-        "src/scripts/EddyTracking",
         "src/scripts/EddyFinalTracking",
         "src/scripts/EddyMergeCorrespondances",
     ],
@@ -43,6 +42,7 @@
             "EddyFrequency = py_eddy_tracker.appli.eddies:get_frequency_grid",
             "EddyInfos = py_eddy_tracker.appli.eddies:display_infos",
             "EddyCircle = py_eddy_tracker.appli.eddies:eddies_add_circle",
+            "EddyTracking = py_eddy_tracker.appli.eddies:eddies_tracking",
             # network
             "EddyNetworkGroup = py_eddy_tracker.appli.network:build_network",
             "EddyNetworkBuildPath = py_eddy_tracker.appli.network:divide_network",

src/py_eddy_tracker/__init__.py

@@ -85,6 +85,13 @@ def add_base_argument(self):
             help="Levels : DEBUG, INFO, WARNING," " ERROR, CRITICAL",
         )
 
+    def memory_arg(self):
+        self.add_argument(
+            "--memory",
+            action="store_true",
+            help="Load file in memory before to read with netCDF library",
+        )
+
     def parse_args(self, *args, **kwargs):
         logger = start_logger()
         # Parsing

src/py_eddy_tracker/appli/eddies.py

@@ -3,12 +3,24 @@
 Applications on detection and tracking files
 """
 import argparse
+import logging
+from datetime import datetime
+from glob import glob
+from os import mkdir
+from os.path import basename, dirname, exists
+from os.path import join as join_path
+from re import compile as re_compile
 
 from netCDF4 import Dataset
+from numpy import bytes_, empty, unique
+from yaml import safe_load
 
 from .. import EddyParser
 from ..observations.observation import EddiesObservations
 from ..observations.tracking import TrackEddiesObservations
+from ..tracking import Correspondances
+
+logger = logging.getLogger("pet")
 
 
 def eddies_add_circle():
@@ -41,24 +53,22 @@ def merge_eddies():
     parser.add_argument(
         "--include_var", nargs="+", type=str, help="use only listed variable"
     )
+    parser.memory_arg()
     args = parser.parse_args()
 
     if args.include_var is None:
         with Dataset(args.filename[0]) as h:
             args.include_var = h.variables.keys()
 
-    obs = TrackEddiesObservations.load_file(
-        args.filename[0], raw_data=True, include_vars=args.include_var
-    )
-    if args.add_rotation_variable:
-        obs = obs.add_rotation_type()
-    for filename in args.filename[1:]:
-        other = TrackEddiesObservations.load_file(
+    obs = list()
+    for filename in args.filename:
+        e = TrackEddiesObservations.load_file(
             filename, raw_data=True, include_vars=args.include_var
         )
         if args.add_rotation_variable:
-            other = other.add_rotation_type()
-        obs = obs.merge(other)
+            e = e.add_rotation_type()
+        obs.append(e)
+    obs = TrackEddiesObservations.concatenate(obs)
     obs.write_file(filename=args.out)
 
 
@@ -141,3 +151,229 @@ def display_infos():
             )
             e = e.extract_with_area(area)
         print(e)
+
+
+def eddies_tracking():
+    parser = EddyParser("Tool to use identification step to compute tracking")
+    parser.add_argument("yaml_file", help="Yaml file to configure py-eddy-tracker")
+    parser.add_argument("--correspondance_in", help="Filename of saved correspondance")
+    parser.add_argument("--correspondance_out", help="Filename to save correspondance")
+    parser.add_argument(
+        "--save_correspondance_and_stop",
+        action="store_true",
+        help="Stop tracking after correspondance computation,"
+        " merging can be done with EddyFinalTracking",
+    )
+    parser.add_argument(
+        "--zarr", action="store_true", help="Output will be wrote in zarr"
+    )
+    parser.add_argument("--unraw", action="store_true", help="Load unraw data")
+    parser.add_argument(
+        "--blank_period",
+        type=int,
+        default=0,
+        help="Nb of detection which will not use at the end of the period",
+    )
+    parser.memory_arg()
+    args = parser.parse_args()
+
+    # Read yaml configuration file
+    with open(args.yaml_file, "r") as stream:
+        config = safe_load(stream)
+
+    if "CLASS" in config:
+        classname = config["CLASS"]["CLASS"]
+        obs_class = dict(
+            class_method=getattr(
+                __import__(config["CLASS"]["MODULE"], globals(), locals(), classname),
+                classname,
+            ),
+            class_kw=config["CLASS"].get("OPTIONS", dict()),
+        )
+    else:
+        obs_class = dict()
+
+    c_in, c_out = args.correspondance_in, args.correspondance_out
+    if c_in is None:
+        c_in = config["PATHS"].get("CORRESPONDANCES_IN", None)
+    y_c_out = config["PATHS"].get(
+        "CORRESPONDANCES_OUT", "{path}/{sign_type}_correspondances.nc"
+    )
+    if c_out is None:
+        c_out = y_c_out
+
+    # Create ouput folder if necessary
+    save_dir = config["PATHS"].get("SAVE_DIR", None)
+    if save_dir is not None and not exists(save_dir):
+        mkdir(save_dir)
+
+    track(
+        pattern=config["PATHS"]["FILES_PATTERN"],
+        output_dir=save_dir,
+        c_out=c_out,
+        **obs_class,
+        virtual=int(config.get("VIRTUAL_LENGTH_MAX", 0)),
+        previous_correspondance=c_in,
+        memory=args.memory,
+        correspondances_only=args.save_correspondance_and_stop,
+        raw=not args.unraw,
+        zarr=args.zarr,
+        nb_obs_min=int(config.get("TRACK_DURATION_MIN", 10)),
+        blank_period=args.blank_period,
+    )
+
+
+def browse_dataset_in(
+    data_dir,
+    files_model,
+    date_regexp,
+    date_model,
+    start_date=None,
+    end_date=None,
+    sub_sampling_step=1,
+    files=None,
+):
+    pattern_regexp = re_compile(".*/" + date_regexp)
+    if files is not None:
+        filenames = bytes_(files)
+    else:
+        full_path = join_path(data_dir, files_model)
+        logger.info("Search files : %s", full_path)
+        filenames = bytes_(glob(full_path))
+
+    dataset_list = empty(
+        len(filenames),
+        dtype=[
+            ("filename", "S500"),
+            ("date", "datetime64[D]"),
+        ],
+    )
+    dataset_list["filename"] = filenames
+
+    logger.info("%s grids available", dataset_list.shape[0])
+    mode_attrs = False
+    if "(" not in date_regexp:
+        logger.debug("Attrs date : %s", date_regexp)
+        mode_attrs = date_regexp.strip().split(":")
+    else:
+        logger.debug("Pattern date : %s", date_regexp)
+
+    for item in dataset_list:
+        str_date = None
+        if mode_attrs:
+            with Dataset(item["filename"].decode("utf-8")) as h:
+                if len(mode_attrs) == 1:
+                    str_date = getattr(h, mode_attrs[0])
+                else:
+                    str_date = getattr(h.variables[mode_attrs[0]], mode_attrs[1])
+        else:
+            result = pattern_regexp.match(str(item["filename"]))
+            if result:
+                str_date = result.groups()[0]
+
+        if str_date is not None:
+            item["date"] = datetime.strptime(str_date, date_model).date()
+
+    dataset_list.sort(order=["date", "filename"])
+
+    steps = unique(dataset_list["date"][1:] - dataset_list["date"][:-1])
+    if len(steps) > 1:
+        raise Exception("Several days steps in grid dataset %s" % steps)
+
+    if sub_sampling_step != 1:
+        logger.info("Grid subsampling %d", sub_sampling_step)
+        dataset_list = dataset_list[::sub_sampling_step]
+
+    if start_date is not None or end_date is not None:
+        logger.info(
+            "Available grid from %s to %s",
+            dataset_list[0]["date"],
+            dataset_list[-1]["date"],
+        )
+        logger.info("Filtering grid by time %s, %s", start_date, end_date)
+        mask = (dataset_list["date"] >= start_date) * (dataset_list["date"] <= end_date)
+
+        dataset_list = dataset_list[mask]
+    return dataset_list
+
+
+def track(
+    pattern,
+    output_dir,
+    c_out,
+    nb_obs_min=10,
+    raw=True,
+    zarr=False,
+    blank_period=0,
+    correspondances_only=False,
+    **kw_c,
+):
+    kw = dict(date_regexp=".*_([0-9]*?).[nz].*", date_model="%Y%m%d")
+    if isinstance(pattern, list):
+        kw.update(dict(data_dir=None, files_model=None, files=pattern))
+    else:
+        kw.update(dict(data_dir=dirname(pattern), files_model=basename(pattern)))
+    datasets = browse_dataset_in(**kw)
+    if blank_period > 0:
+        datasets = datasets[:-blank_period]
+        logger.info("Last %d files will be pop", blank_period)
+
+    if nb_obs_min > len(datasets):
+        raise Exception(
+            "Input file number (%s) is shorter than TRACK_DURATION_MIN (%s)."
+            % (len(datasets), nb_obs_min)
+        )
+
+    c = Correspondances(datasets=datasets["filename"], **kw_c)
+    c.track()
+    logger.info("Track finish")
+    t0, t1 = c.period
+    kw_save = dict(
+        date_start=t0,
+        date_stop=t1,
+        date_prod=datetime.now(),
+        path=output_dir,
+        sign_type=c.current_obs.sign_legend,
+    )
+
+    c.save(c_out, kw_save)
+    if correspondances_only:
+        return
+
+    logger.info("Start merging")
+    c.prepare_merging()
+    logger.info("Longer track saved have %d obs", c.nb_obs_by_tracks.max())
+    logger.info(
+        "The mean length is %d observations for all tracks", c.nb_obs_by_tracks.mean()
+    )
+
+    kw_write = dict(path=output_dir, zarr_flag=zarr)
+
+    c.get_unused_data(raw_data=raw).write_file(
+        filename="%(path)s/%(sign_type)s_untracked.nc", **kw_write
+    )
+
+    short_c = c._copy()
+    short_c.shorter_than(size_max=nb_obs_min)
+    c.longer_than(size_min=nb_obs_min)
+
+    long_track = c.merge(raw_data=raw)
+    short_track = short_c.merge(raw_data=raw)
+
+    # We flag obs
+    if c.virtual:
+        long_track["virtual"][:] = long_track["time"] == 0
+        long_track.filled_by_interpolation(long_track["virtual"] == 1)
+        short_track["virtual"][:] = short_track["time"] == 0
+        short_track.filled_by_interpolation(short_track["virtual"] == 1)
+
+    logger.info("Longer track saved have %d obs", c.nb_obs_by_tracks.max())
+    logger.info(
+        "The mean length is %d observations for long track",
+        c.nb_obs_by_tracks.mean(),
+    )
+
+    long_track.write_file(**kw_write)
+    short_track.write_file(
+        filename="%(path)s/%(sign_type)s_track_too_short.nc", **kw_write
+    )

src/py_eddy_tracker/appli/network.py

@@ -32,9 +32,16 @@ def build_network():
         action="store_true",
         help="Use intern contour instead of outter contour",
     )
+
+    parser.memory_arg()
     args = parser.parse_args()
 
-    n = Network(args.identification_regex, window=args.window, intern=args.intern)
+    n = Network(
+        args.identification_regex,
+        window=args.window,
+        intern=args.intern,
+        memory=args.memory,
+    )
     group = n.group_observations(minimal_area=True)
     n.build_dataset(group).write_file(filename=args.out)
 

src/py_eddy_tracker/featured_tracking/area_tracker.py

@@ -1,6 +1,7 @@
 import logging
 
-from numpy import ma
+from numba import njit
+from numpy import empty, ma, ones
 
 from ..observations.observation import EddiesObservations as Model
 
@@ -29,11 +30,8 @@ def needed_variable(cls):
     def tracking(self, other):
         shape = (self.shape[0], other.shape[0])
         i, j, c = self.match(other, intern=False)
-        cost_mat = ma.empty(shape, dtype="f4")
-        cost_mat.mask = ma.ones(shape, dtype="bool")
-        m = c > self.cmin
-        i, j, c = i[m], j[m], c[m]
-        cost_mat[i, j] = 1 - c
+        cost_mat = ma.array(empty(shape, dtype="f4"), mask=ones(shape, dtype="bool"))
+        mask_cmin(i, j, c, self.cmin, cost_mat.data, cost_mat.mask)
 
         i_self, i_other = self.solve_function(cost_mat)
         i_self, i_other = self.post_process_link(other, i_self, i_other)
@@ -55,3 +53,13 @@ def propagate(
             if nb_virtual_extend > 0:
                 virtual[key][nb_dead:] = obs_to_extend[key]
         return virtual
+
+
+@njit(cache=True)
+def mask_cmin(i, j, c, cmin, cost_mat, mask):
+    for k in range(c.shape[0]):
+        c_ = c[k]
+        if c_ > cmin:
+            i_, j_ = i[k], j[k]
+            cost_mat[i_, j_] = 1 - c_
+            mask[i_, j_] = False

src/py_eddy_tracker/generic.py

@@ -83,7 +83,7 @@ def build_index(groups):
             first_index[group - i0 + 1 : next_group - i0 + 1] = i + 1
     last_index = zeros(amplitude, dtype=numba_types.int_)
     last_index[:-1] = first_index[1:]
-    last_index[-1] = i + 2
+    last_index[-1] = i + 1
     return first_index, last_index, i0