Repository: WAM2layers/WAM2layers

Issue:
Finalize cli visualize commands, support forwardtracking
Before v3 the interface for visualization should be finalized.

A test for the visualization of forward tracking data should also be added.

Discussion:
Peter9192: I'm all for this. Just wonder if 'tracking direction' is the best name for the new config setting. I was thinking of 'tracking mode' as an alternative. 

@ruudvdent are time and distance tracking alternative modes or additional options for a forward or backward run?
ruudvdent: I think tracking direction is clear, mode would be fine as well, so no particular preference from my side. Time and distance tracking would rely on the basic tracking in either direction for the calculations, so they would be additional options.
Peter9192: ok let's stick with direction then. Can this be merged?
Peter9192: > I tried running the example floodcase again, and this works nicely. Note that we need to update the example config [here](https://rebrand.ly/wam2layers-example-data) as well.

I have updated this file

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b04f283..045bf0e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,7 @@ All notable changes to this project will be documented in this file.
 - The stability correction warning will now occur only when 10% more of the grid is corrected compared to last warning, or the correction factor is 10% stronger  ([#332](https://github.com/WAM2layers/WAM2layers/pull/332)).
 - The output files now contain a time dimension (incl. the coordinate) ([#334](https://github.com/WAM2layers/WAM2layers/pull/334)).
 - The package version is now defined in `src/wam2layers/__init__.py` ([#334](https://github.com/WAM2layers/WAM2layers/pull/334)).
+- The command line interface for tracking has changed. Tracking experiments are now started by doing `wam2layers track config.yml`. The tracking direction is retrieved from the new configuration entry "tracking_direction" ([#338](https://github.com/WAM2layers/WAM2layers/pull/338)).
 
 ## Release v3.0.0-beta.6 (2023-12-22)
 
diff --git a/README.md b/README.md
index 01f5217..26f2b9e 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,7 @@ Ruud van der Ent (Delft University of Technology)
 Imme Benedict (Wageningen University)
 Chris Weijenborg (Wageningen University)
 Peter Kalverla (Netherlands eScienceCenter)
+Bart Schilperoort (Netherlands eScience Center)
 
 # How to use
 See the [documentation](https://wam2layers.readthedocs.io/en/latest/) for a more detailed description. 
diff --git a/docs/userguide.md b/docs/userguide.md
index 83c3efc..6e59820 100644
--- a/docs/userguide.md
+++ b/docs/userguide.md
@@ -30,15 +30,13 @@ interace (CLI), much like pip, conda, git, cdo, et cetera. It looks like this:
 ```
 # Get help
 wam2layers --help
-wam2layers backtrack --help
-wam2layers forwardtrack --help
+wam2layers track --help
 
 # Preprocess data (you need to have downloaded the data)
 wam2layers preprocess era5 floodcase_202107.yaml
 
 # Perform tracking
-wam2layers backtrack floodcase_202107.yaml
-wam2layers forwardtrack floodcase_202107.yaml
+wam2layers track floodcase_202107.yaml
 
 # Make some basic plots
 wam2layers visualize output floodcase_202107.yaml
@@ -59,8 +57,9 @@ here:
    `wam2layers` should point to the file `src/wam2layers/cli.py`.
 1. We use [`click`](https://click.palletsprojects.com/en/8.1.x/) to further
    process everything that comes after the `wam2layers` command. For example,
-   `wam2layers backtrack` is redirected to
-   `src/wam2layers/tracking/backtrack.py:cli`. This function takes the
+   `wam2layers track` is redirected to the appropriate module
+   (`src/wam2layers/tracking/backtrack.py:run_experiment` or
+   `src/wam2layers/tracking/forwardtrack.py:run_experiment`). This function takes the
    configuration file as input and passes it on to a function called
    `run_experiment`.
 ```
@@ -231,14 +230,14 @@ Assuming you have a preprocessed dataset and prepared a configuration file for
 your experiment, running WAM2layers is as simple as:
 
 ```
-wam2layers backtrack config-file.yaml
+wam2layers track config-file.yaml
 ```
 
 where `config-file.yaml` is the path to your configuration file. Among others,
-this file should have settings on the date range for which you want to run
-track, and also about the location where the preprocessed data are stored and
-where the output will be stored. For more information, see [](./config) or have
-a look at the example config file
+this file should have settings on the tracking direction (forward/backward), 
+the date range for which you want to run track, and also about the location where the 
+preprocessed data are stored and where the output will be stored.
+For more information, see [](./config) or have a look at the example config file
 [here](https://github.com/WAM2layers/WAM2layers/blob/main/example_config.yaml).
 
 ```{tip}
diff --git a/example_config.yaml b/example_config.yaml
index dd15a95..dc98563 100644
--- a/example_config.yaml
+++ b/example_config.yaml
@@ -11,6 +11,7 @@ level_type: model_levels
 levels: [20,40,60,80,90,95,100,105,110,115,120,123,125,128,130,131,132,133,134,135,136,137]
 
 # Tracking
+tracking_direction: backward
 tracking_domain: null  # use full domain of preprocessed data
 # tracking_domain: [0, 50, 10, 55]  # use subdomain of preprocessed data; west south east north
 tracking_start_date: "2021-07-01T00:00"
diff --git a/src/wam2layers/analysis/visualization.py b/src/wam2layers/analysis/visualization.py
index 232dfb3..a8c1499 100644
--- a/src/wam2layers/analysis/visualization.py
+++ b/src/wam2layers/analysis/visualization.py
@@ -36,16 +36,6 @@ def polish(ax, region):
     ax.set_ylim(region.latitude.min(), region.latitude.max())
 
 
-def infer_mode(output_folder: Path) -> Literal["forwardtrack", "backtrack"]:
-    output_file = next(output_folder.glob("*.nc"))
-    if "forwardtrack" in str(output_file):
-        return "forwardtrack"
-    if "backtrack" in str(output_file):
-        return "backtrack"
-    else:
-        raise FileNotFoundError(f"No output files found in {output_folder}")
-
-
 def _plot_input(config: Config, ax):
     """
     Return subplot with input.
@@ -65,12 +55,11 @@ def _plot_input(config: Config, ax):
         inclusive="left",
     )
 
-    mode = infer_mode(config.output_folder)
     input_files = []
     for date in dates[:-1]:
         input_files.append(input_path(date, config))
     ds = xr.open_mfdataset(input_files, combine="nested", concat_dim="time")
-    if mode == "backtrack":
+    if config.tracking_direction == "backward":
         subset = ds.precip.sel(time=slice(start, end))
     else:
         subset = ds.evap.sel(time=slice(start, end))
@@ -101,11 +90,10 @@ def _plot_output(config: Config, ax):
         inclusive="left",
     )
 
-    mode = infer_mode(config.output_folder)
-    if mode == "backtrack":
-        output_files = [output_path(date, config, mode=mode) for date in dates[:-1]]
+    if config.tracking_direction == "backward":
+        output_files = [output_path(date, config) for date in dates[:-1]]
     else:
-        output_files = [output_path(date, config, mode=mode) for date in dates[1:]]
+        output_files = [output_path(date, config) for date in dates[1:]]
 
     if not all([Path(file).exists() for file in output_files]):
         raise FileNotFoundError(f"Could not find all files: {output_files}.")
@@ -113,7 +101,7 @@ def _plot_output(config: Config, ax):
         raise ValueError("Too few output files to visualize.")
 
     ds = xr.open_mfdataset(output_files, combine="nested", concat_dim="time")
-    if mode == "backtrack":
+    if config.tracking_direction == "backward":
         out_track = ds.e_track.sum("time") * 1000 / a_gridcell[:, None]
     else:
         out_track = (
diff --git a/src/wam2layers/cli.py b/src/wam2layers/cli.py
index 254a341..3ce082b 100644
--- a/src/wam2layers/cli.py
+++ b/src/wam2layers/cli.py
@@ -3,7 +3,7 @@
 This module contains all the functionality that makes it possible to run wam2layers on the command line, like so:
 
     wam2layers preprocess era5 floodcase.yaml
-    wam2layers backtrack floodcase.yaml
+    wam2layers track floodcase.yaml
 
 et cetera. It is built with [click](https://click.palletsprojects.com/en/8.1.x/).
 """
@@ -57,51 +57,57 @@ def cli():
     pass
 
 
-# Command line setup for backtrack
+# Command line setup for tracking
 
 
 @cli.command()
 @click.argument("config_file", type=click.Path(exists=True))
-def backtrack(config_file):
-    """Run WAM2layers backtrack experiment.
+def track(config_file):
+    """Run WAM2layers tracking experiment.
 
     CONFIG_FILE: Path to WAM2layers experiment configuration file.
 
     Usage examples:
 
         \b
-        - wam2layers backtrack path/to/cases/era5_2021.yaml
+        - wam2layers track path/to/cases/era5_2021.yaml
     """
-    log_path = Config.from_yaml(config_file).output_folder
+    config = Config.from_yaml(config_file)
+    log_path = config.output_folder
     setup_logging(log_path)
     _copy_config_yaml(config_file, log_path)
     logger.info("Welcome to WAM2layers.")
-    logger.info("Starting backtrack experiment.")
-    run_backtrack_experiment(config_file)
-
-
-# Command line setup for forwardtrack
-
+    logger.info(f"Starting {config.tracking_direction} tracking experiment.")
+    if config.tracking_direction == "backward":
+        run_backtrack_experiment(config_file)
+    else:
+        run_forwardtrack_experiment(config_file)
 
 @cli.command()
-@click.argument("config_file", type=click.Path(exists=True))
-def forwardtrack(config_file):
-    """Run WAM2layers forwardtrack experiment.
-
-    CONFIG_FILE: Path to WAM2layers experiment configuration file.
-
-    Usage examples:
-
-        \b
-        - wam2layers forwardtrack path/to/cases/era5_2021.yaml
-    """
-    log_path = Config.from_yaml(config_file).output_folder
-    setup_logging(log_path)
-    _copy_config_yaml(config_file, log_path)
-    logger.info("Welcome to WAM2layers.")
-    logger.info("Starting forwardtrack experiment.")
-    run_forwardtrack_experiment(config_file)
+def backtrack():
+    msg = (
+        "The `backtrack` and `forwardtrack` commands have been removed in favor of "
+        "`track`.\nPlease specify the tracking direction in your config file and use "
+        "the `track` command instead."
+        "\n"
+        "\nFor example, in your config.yaml file add:"
+        "\n    # Tracking"
+        "\n    tracking_direction: backward"
+    )
+    raise ValueError(msg)
 
+@cli.command()
+def forwardtrack():
+    msg = (
+        "The `backtrack` and `forwardtrack` commands have been removed in favor of "
+        "`track`.\nPlease specify the tracking direction in your config file and use "
+        "the `track` command instead."
+        "\n"
+        "\nFor example, in your config.yaml file add:"
+        "\n    # Tracking"
+        "\n    tracking_direction: forward"
+    )
+    raise ValueError(msg)
 
 # Command line setup for preprocess
 @click.group()
@@ -140,7 +146,15 @@ def visualize_cli():
 @visualize_cli.command()
 @click.argument("config_file", type=click.Path(exists=True))
 def input(config_file):
-    """Visualize input data for experiment."""
+    """Visualize input data for experiment.
+
+    CONFIG_FILE: Path to WAM2layers experiment configuration file.
+
+    Usage examples:
+
+        \b
+        - wam2layers visualize input path/to/cases/era5_2021.yaml
+    """
     log_path = Config.from_yaml(config_file).output_folder
     setup_logging(log_path)
     _copy_config_yaml(config_file, log_path)
@@ -152,7 +166,15 @@ def input(config_file):
 @visualize_cli.command()
 @click.argument("config_file", type=click.Path(exists=True))
 def output(config_file):
-    """Visualize output data for experiment."""
+    """Visualize output data for experiment.
+
+    CONFIG_FILE: Path to WAM2layers experiment configuration file.
+
+    Usage examples:
+
+        \b
+        - wam2layers visualize output path/to/cases/era5_2021.yaml
+    """
     log_path = Config.from_yaml(config_file).output_folder
     setup_logging(log_path)
     _copy_config_yaml(config_file, log_path)
diff --git a/src/wam2layers/config.py b/src/wam2layers/config.py
index 2c32638..647290e 100644
--- a/src/wam2layers/config.py
+++ b/src/wam2layers/config.py
@@ -80,6 +80,19 @@ class Config(BaseModel):
 
     """
 
+    tracking_direction: Literal["forward", "backward"]
+    """The tracking direction, either forward or backward.
+
+    You have to specify if either forward or backward tracking should be performed.
+
+    For example:
+
+    .. code-block:: yaml
+    
+        tracking_direction: backward
+
+    """
+
     tagging_region: Annotated[
         Union[FilePath, BoundingBox], AfterValidator(validate_region)
     ]
diff --git a/src/wam2layers/preprocessing/era5.py b/src/wam2layers/preprocessing/era5.py
index fd49db9..fdb2c8a 100644
--- a/src/wam2layers/preprocessing/era5.py
+++ b/src/wam2layers/preprocessing/era5.py
@@ -277,10 +277,6 @@ def prep_experiment(config_file):
                 logger, f"Total column water not available; using water vapour only"
             )
 
-        # Determine the fluxes
-        fx = u * cw  # eastward atmospheric moisture flux (kg m-1 s-1)
-        fy = v * cw  # northward atmospheric moisture flux (kg m-1 s-1)
-
         # Integrate fluxes and states to upper and lower layer
         if config.level_type == "model_levels":
             # TODO: Check if this is a reasonable choice for boundary
diff --git a/src/wam2layers/tracking/backtrack.py b/src/wam2layers/tracking/backtrack.py
index def7554..6058a94 100644
--- a/src/wam2layers/tracking/backtrack.py
+++ b/src/wam2layers/tracking/backtrack.py
@@ -159,7 +159,7 @@ def initialize(config_file):
     if config.restart:
         # Reload last state from existing output
         date = config.tracking_end_date
-        ds = xr.open_dataset(output_path(date, config, mode="backtrack"))
+        ds = xr.open_dataset(output_path(date, config))
         output["s_track_upper_restart"].values = ds.s_track_upper_restart.values
         output["s_track_lower_restart"].values = ds.s_track_lower_restart.values
 
diff --git a/src/wam2layers/tracking/forwardtrack.py b/src/wam2layers/tracking/forwardtrack.py
index 23bd17c..7fe0ac7 100644
--- a/src/wam2layers/tracking/forwardtrack.py
+++ b/src/wam2layers/tracking/forwardtrack.py
@@ -164,7 +164,7 @@ def initialize(config_file):
     if config.restart:
         # Reload last state from existing output
         date = config.tracking_start_date
-        ds = xr.open_dataset(output_path(date, config, mode="forwardtrack"))
+        ds = xr.open_dataset(output_path(date, config))
         output["s_track_upper_restart"].values = ds.s_track_upper_restart.values
         output["s_track_lower_restart"].values = ds.s_track_lower_restart.values
 
diff --git a/src/wam2layers/tracking/io.py b/src/wam2layers/tracking/io.py
index 19ad594..66ecbce 100644
--- a/src/wam2layers/tracking/io.py
+++ b/src/wam2layers/tracking/io.py
@@ -18,8 +18,9 @@ def input_path(date, config):
     return f"{input_dir}/{date.strftime('%Y-%m-%d')}_fluxes_storages.nc"
 
 
-def output_path(date, config, mode):
+def output_path(date, config):
     output_dir = config.output_folder
+    mode = "backtrack" if config.tracking_direction == "backward" else "forwardtrack"
     return f"{output_dir}/{mode}_{date.strftime('%Y-%m-%dT%H-%M')}.nc"
 
 
@@ -96,7 +97,7 @@ def write_output(
     mode: Literal["forwardtrack", "backtrack"],
 ) -> None:
     # TODO: add back (and cleanup) coordinates and units
-    path = output_path(t, config, mode)
+    path = output_path(t, config)
     logger.info(f"{t} - Writing output to file {path}")
 
     output = output.assign_coords({"time": t})


Test patch:
diff --git a/tests/test_data/config_rhine.yaml b/tests/test_data/config_rhine.yaml
index daa2654..fbd3d7a 100644
--- a/tests/test_data/config_rhine.yaml
+++ b/tests/test_data/config_rhine.yaml
@@ -10,6 +10,7 @@ level_type: model_levels
 levels: [20,40,60,80,90,95,100,105,110,115,120,123,125,128,130,131,132,133,134,135,136,137]
 
 # Tracking
+tracking_direction: backward
 tracking_domain: null  # use full domain of preprocessed data
 # tracking_domain: [0, 50, 10, 55]  # use subdomain of preprocessed data
 tracking_start_date: "2022-08-31T00:00"
diff --git a/tests/test_data/config_west_africa.yaml b/tests/test_data/config_west_africa.yaml
index 4bd443c..b096474 100644
--- a/tests/test_data/config_west_africa.yaml
+++ b/tests/test_data/config_west_africa.yaml
@@ -10,6 +10,7 @@ level_type: model_levels
 levels: [20,40,60,80,90,95,100,105,110,115,120,123,125,128,130,131,132,133,134,135,136,137]
 
 # Tracking
+tracking_direction: forward
 tagging_region: tests/test_data/region_lake_volta.nc
 tracking_start_date: "1979-07-01T00:00"
 tracking_end_date: "1979-07-01T23:00"
diff --git a/tests/test_workflow.py b/tests/test_workflow.py
index df571e1..66cba41 100644
--- a/tests/test_workflow.py
+++ b/tests/test_workflow.py
@@ -102,7 +102,7 @@ def preprocess(test_config_backward: str) -> None:
 @pytest.fixture(scope="session")
 def backtrack(test_config_backward: str, preprocess) -> None:
     runner = CliRunner()
-    runner.invoke(cli, ["backtrack", test_config_backward], catch_exceptions=False)
+    runner.invoke(cli, ["track", test_config_backward], catch_exceptions=False)
 
 
 @pytest.fixture(scope="session")
@@ -124,7 +124,7 @@ def visualize_forward(test_config_forward: str, forwardtrack) -> None:
 @pytest.fixture(scope="session")
 def forwardtrack(test_config_forward: str) -> None:
     runner = CliRunner()
-    runner.invoke(cli, ["forwardtrack", test_config_forward], catch_exceptions=False)
+    runner.invoke(cli, ["track", test_config_forward], catch_exceptions=False)
 
 
 def plot_difference(