Repository: WAM2layers/WAM2layers

Issue:
need for subsetting regions for backtracking
As it stands, the backtracking code will run over the entire domain of the era5 data. However, one should not have to re-download (or resave) the ERA5 data for each separate domain they may want to perform the backtracking on. In my view, the best way to use the WAM for future flexibility is to download the global ERA5 data and run the pre-processing (fluxes/states) on that data. Then, at the backtracking stage, have the user specify the domain over which to perform the backtracking. This could be a subset of the globe, or the entire globe. I have written code to do this, but am putting in an issue to start discussion.


Discussion:
Peter9192: This could still be polished a bit and might need more thorough testing, but the basis is ready, I think.
geek-yang: Now we can also update the ci to let it run all tests.
ruudvdent: Idea to use terminology "tagging_region" and "tracking_domain"
chrisweij: The 3d tagging can be read, however one issue is that it now continues after (i.e. earlier in time) with the last provided location. Either the tagging begin_date or the time stamp in the tagging region netcdf should be used. 

Patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index dd2b921..955bbb0 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -39,5 +39,5 @@ jobs:
           wam2layers --help
         shell: bash -el {0}
       - name: Run pytest
-        run: pytest tests/test_workflow.py
+        run: pytest
         shell: micromamba-shell {0}
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1f673d0..6bd4f6f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,10 +13,14 @@ All notable changes to this project will be documented in this file.
 - Add additional documentation and doctest for vertical advection terms ([#274](https://github.com/WAM2layers/WAM2layers/pull/274)).
 - Copy config yaml file to output path ([#249](https://github.com/WAM2layers/WAM2layers/pull/249)).
 - Add forward tracking ([#289](https://github.com/WAM2layers/WAM2layers/pull/289))
+- Support for time-dependent tagging region ([#297](https://github.com/WAM2layers/WAM2layers/pull/297))
+- tagging_region can be specified as a bounding box ([#297](https://github.com/WAM2layers/WAM2layers/pull/297))
+- tracking_domain can now be used to subselect preprocessed data ([#297](https://github.com/WAM2layers/WAM2layers/pull/297))
 
 ### Removed
 
 - Remove example test floodcase from CI ([#265](https://github.com/WAM2layers/WAM2layers/pull/265)).
+- Chunks argument in config ([#290](https://github.com/WAM2layers/WAM2layers/pull/290))
 
 ### Fixed
 
@@ -30,6 +34,7 @@ All notable changes to this project will be documented in this file.
 - There's now a dedicated function to compute the advection term. This makes it easier to apply different solvers. Padding is added to the edges of the domain, and boundary losses are added to e-track at the edges of the domain ([#266](https://github.com/WAM2layers/WAM2layers/pull/266)).
 - Seperated the vertical advection and dispersion terms: dispersion is modelled as `kvf * abs(Fv) * dS/dz` ([#274](https://github.com/WAM2layers/WAM2layers/pull/274)).
 - Splitted the backtracking module, moving functions to separate io, core, and preprocessing.shared modules ([#274](https://github.com/WAM2layers/WAM2layers/pull/274)).
+- `region` renamed to `tagging_region`; `track_start_date` to `tracking_start_date`; `track_end_date` to `tracking_end_date` ([#297](https://github.com/WAM2layers/WAM2layers/pull/297))
 
 ## Release v3.0.0-beta.5 (2023-07-21)
 
diff --git a/example_config.yaml b/example_config.yaml
index d450430..5e69e6a 100644
--- a/example_config.yaml
+++ b/example_config.yaml
@@ -1,22 +1,27 @@
 # General
 log_level: info
-preprocessed_data_folder: ~/floodcase_202107/preprocessed_data
-output_folder: ~/floodcase_202107/output_data
-chunks: null
+preprocessed_data_folder: tests/tmp/preprocessed_data
+output_folder: tests/tmp/output_data
 
 # Preprocessing
-filename_template: /ERA5data/{year}/{month:02}/ERA5_{year}-{month:02d}-{day:02d}{levtype}_{variable}.nc
+filename_template: tests/test_data/era5/ERA5_{year}-{month:02d}-{day:02d}{levtype}_{variable}.nc
 preprocess_start_date: "2021-07-01T00:00"
 preprocess_end_date: "2021-07-15T00:00"
 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
-region: /data/volume_2/era5_2021/source_region_global.nc
-track_start_date: "2021-07-01T00:00"
-track_end_date: "2021-07-15T00:00"
-event_start_date: "2021-07-13T00:00"
-event_end_date: "2021-07-15T00:00"
+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"
+tracking_end_date: "2021-07-15T00:00"
+
+# Tagging
+tagging_region: tests/test_data/test_region.nc
+# tagging_region: [0, 50, 10, 55]  # alternative to mask file; west south east north
+tagging_start_date: "2021-07-13T00:00"
+tagging_end_date: "2021-07-15T00:00"
+
 input_frequency: "1h"
 target_frequency: "10min"
 output_frequency: "1d"
diff --git a/src/wam2layers/analysis/visualization.py b/src/wam2layers/analysis/visualization.py
index 12c3d35..6091be9 100644
--- a/src/wam2layers/analysis/visualization.py
+++ b/src/wam2layers/analysis/visualization.py
@@ -8,7 +8,7 @@ from cmocean import cm
 
 from wam2layers.config import Config
 from wam2layers.preprocessing.shared import get_grid_info
-from wam2layers.tracking.io import input_path, load_region, output_path
+from wam2layers.tracking.io import input_path, load_tagging_region, output_path
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +38,7 @@ def polish(ax, region):
 def _plot_precip(config, ax):
     """Return subplot with precip."""
     # Load config and some usful stuf.
-    region = load_region(config)
+    region = load_tagging_region(config)
 
     # Load data
     dates = pd.date_range(
@@ -54,8 +54,8 @@ def _plot_precip(config, ax):
 
     ds = xr.open_mfdataset(input_files, combine="nested", concat_dim="time")
     # TODO: make region time-dependent
-    start = config.event_start_date
-    end = config.event_end_date
+    start = config.tagging_start_date
+    end = config.tagging_end_date
     subset = ds.precip.sel(time=slice(start, end))
     precip = (subset * region * 3600).sum("time").compute()
 
@@ -67,13 +67,13 @@ def _plot_precip(config, ax):
 
 def _plot_evap(config, ax):
     """Return subplot with tracked evaporation."""
-    region = load_region(config)
+    region = load_tagging_region(config)
     a_gridcell, lx, ly = get_grid_info(region)
 
     # Load data
     dates = pd.date_range(
-        start=config.track_start_date,
-        end=config.track_end_date,
+        start=config.tracking_start_date,
+        end=config.tracking_end_date,
         freq=config.output_frequency,
         inclusive="left",
     )[1:]
diff --git a/src/wam2layers/config.py b/src/wam2layers/config.py
index e4b2bdd..dd112f0 100644
--- a/src/wam2layers/config.py
+++ b/src/wam2layers/config.py
@@ -1,14 +1,45 @@
 import logging
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Literal, Optional, Union
+from typing import Dict, List, Literal, NamedTuple, Optional, Union
 
 import yaml
-from pydantic import BaseModel, ConfigDict, FilePath, field_validator, model_validator
+from pydantic import (
+    AfterValidator,
+    BaseModel,
+    ConfigDict,
+    Field,
+    FilePath,
+    field_validator,
+    model_validator,
+)
+from typing_extensions import Annotated
 
 logger = logging.getLogger(__name__)
 
 
+class BoundingBox(NamedTuple):
+    west: Annotated[float, Field(ge=-180, le=180)]
+    south: Annotated[float, Field(ge=-80, le=80)]
+    east: Annotated[float, Field(ge=-180, le=180)]
+    north: Annotated[float, Field(ge=-80, le=80)]
+
+
+def validate_region(region):
+    """Check if region is existing path or valid bbox."""
+    if isinstance(region, Path):
+        return region.absolute()
+
+    assert -180 <= region.west <= 180, "longitude should be between -180 and 180"
+    assert -180 <= region.east <= 180, "longitude should be between -180 and 180"
+    assert -80 <= region.north <= 80, "latitude should be between -80 and 80"
+    assert -80 <= region.south <= 80, "latitude should be between -80 and 80"
+    assert region.south < region.north, "south should be smaller than north"
+    if region.west > region.east:
+        logger.info("west > east, coordinates will be rolled around meridian")
+    return region
+
+
 class Config(BaseModel):
     # Pydantic configuration
     model_config = ConfigDict(validate_assignment=True)
@@ -49,16 +80,51 @@ class Config(BaseModel):
 
     """
 
-    region: FilePath
-    """Location where the region mask is stored.
+    tagging_region: Annotated[
+        Union[FilePath, BoundingBox], AfterValidator(validate_region)
+    ]
+    """Subdomain delimiting the source/sink regions for tagged moisture.
+
+    You can either specify a path that contains a netcdf file, or a bounding box
+    of the form [west, south, east, north].
+
+    The bounding box should be inside -180, -80, 180, 80; if west > south, the
+    coordinates will be rolled to retain a continous longitude.
+
+    The file should exist. If it has a time dimension, the nearest field will be
+    used as tagging region, and the time should still be between
+    tagging_start_date and tagging_end_date
+
+    For example:
+
+    .. code-block:: yaml
+
+        tagging_region: /data/volume_2/era5_2021/tagging_region_global.nc
+        tagging_region: [0, 50, 10, 55]
+    """
+
+    tracking_domain: Optional[
+        Annotated[BoundingBox, AfterValidator(validate_region)]
+    ] = None
+    """Subdomain delimiting the region considered during tracking.
+
+    This is useful when you have global pre-processed data but you don't need
+    global tracking.
+
+    You can specify a bounding box of the form [west, south, east, north].
+
+    The bounding box should be inside -180, -80, 180, 80; if west > south, the
+    coordinates will be rolled to retain a continous longitude.
 
-    The file should exist.
+    If it is set to `null`, then it will use full domain of preprocessed data.
+
+    Note that you should always set period to False if you use a subdomain.
 
     For example:
 
     .. code-block:: yaml
 
-        region: /data/volume_2/era5_2021/source_region_global.nc
+        tracking_domain: [0, 50, 10, 55]
 
     """
 
@@ -100,22 +166,22 @@ class Config(BaseModel):
         preprocess_end_date: "2021-07-15T23:00"
     """
 
-    track_start_date: datetime
+    tracking_start_date: datetime
     """Start date for tracking.
 
     Should be formatted as: `"YYYY-MM-DD[T]HH:MM"`. Start date < end date, even if
     backtracking.
-    When backward tracking the track_start_date is not given as output date.
+    When backward tracking the tracking_start_date is not given as output date.
 
     For example:
 
     .. code-block:: yaml
 
-        track_start_date: "2021-07-01T00:00"
+        tracking_start_date: "2021-07-01T00:00"
 
     """
 
-    track_end_date: datetime
+    tracking_end_date: datetime
     """Start date for tracking.
 
     Should be formatted as: `"YYYY-MM-DD[T]HH:MM"`. Start date < end date, even if
@@ -125,16 +191,16 @@ class Config(BaseModel):
 
     .. code-block:: yaml
 
-        track_end_date: "2021-07-15T23:00"
+        tracking_end_date: "2021-07-15T23:00"
     """
 
-    event_start_date: datetime
-    """Start date for event.
+    tagging_start_date: datetime
+    """Start date for tagging.
 
     For tracking individual (e.g. heavy precipitation) events, you can set the
-    start and end date of the event to something different than the total
-    tracking start and end date, you can also indicate the hours that you want to track.
-    The event_start_date is included in the event tracking.
+    start and end date to something different than the total tracking start and
+    end date, you can also indicate the hours that you want to track. The
+    start date is included.
 
     Should be formatted as: `"YYYY-MM-DD[T]HH:MM"`. Start date < end date, even if
     backtracking.
@@ -143,17 +209,17 @@ class Config(BaseModel):
 
     .. code-block:: yaml
 
-        event_start_date: "2021-07-13T00:00"
+        tagging_start_date: "2021-07-13T00:00"
 
     """
 
-    event_end_date: datetime
-    """Start date for event.
+    tagging_end_date: datetime
+    """End date for tagging.
 
     For tracking individual (e.g. heavy precipitation) events, you can set the
-    start and end date of the event to something different than the total
-    tracking start and end date, you can also indicate the hours that you want to track.
-    The event_end_date is included in the event tracking.
+    start and end date to something different than the total tracking start and
+    end date, you can also indicate the hours that you want to track. The
+    end date is included.
 
     Should be formatted as: `"YYYY-MM-DD[T]HH:MM"`. Start date < end date, even if
     backtracking.
@@ -162,7 +228,7 @@ class Config(BaseModel):
 
     .. code-block:: yaml
 
-        event_end_date: "2021-07-14T23:00"
+        tagging_end_date: "2021-07-14T23:00"
 
     """
 
@@ -289,11 +355,11 @@ class Config(BaseModel):
 
         return cls(**settings)
 
-    @field_validator("preprocessed_data_folder", "region", "output_folder")
+    @field_validator("preprocessed_data_folder", "output_folder")
     @classmethod
     def _expanduser(cls, path):
         """Resolve ~ in input paths."""
-        return path.expanduser()
+        return path.absolute()
 
     @field_validator("preprocessed_data_folder", "output_folder")
     @classmethod
@@ -306,13 +372,20 @@ class Config(BaseModel):
 
     @model_validator(mode="after")
     def check_date_order(self):
-        if self.track_start_date > self.track_end_date:
-            raise ValueError("track_end_date should be later than track_start_date")
-        if self.event_start_date > self.event_end_date:
-            raise ValueError("event_end_date should be later than event_start_date")
+        if self.tracking_start_date > self.tracking_end_date:
+            raise ValueError(
+                "tracking_end_date should be later than tracking_start_date"
+            )
+        if self.tagging_start_date > self.tagging_end_date:
+            raise ValueError("tagging_end_date should be later than tagging_start_date")
         if self.preprocess_start_date > self.preprocess_end_date:
             raise ValueError(
                 "preprocess_end_date should be later than preprocess_start_date"
             )
 
+        if self.tracking_domain is not None and self.periodic_boundary:
+            logger.warning(
+                "Periodic boundary set to true while using a subdomain. Are you sure?"
+            )
+
         return self
diff --git a/src/wam2layers/preprocessing/shared.py b/src/wam2layers/preprocessing/shared.py
index c04939f..8d3409a 100644
--- a/src/wam2layers/preprocessing/shared.py
+++ b/src/wam2layers/preprocessing/shared.py
@@ -342,6 +342,7 @@ def stabilize_fluxes(F, S, progress_tracker, config, t):
 
     CFL: Water cannot move further than one grid cell per timestep.
     """
+    coords = S.coords
     for level in ["upper", "lower"]:
         fx = F["fx_" + level]
         fy = F["fy_" + level]
@@ -358,7 +359,7 @@ def stabilize_fluxes(F, S, progress_tracker, config, t):
         fy_limit = 1 / 2 * fy_abs / ft_abs * s.values
         fy_stable = np.minimum(fy_abs, fy_limit)
 
-        progress_tracker.track_stability_correction(fy_limit, fy_abs, config, t)
+        progress_tracker.track_stability_correction(fy_limit, fy_abs, config, coords, t)
 
         # Get rid of any nan values
         fx_stable.fillna(0)
diff --git a/src/wam2layers/tracking/backtrack.py b/src/wam2layers/tracking/backtrack.py
index 601b0be..1cd6f0e 100644
--- a/src/wam2layers/tracking/backtrack.py
+++ b/src/wam2layers/tracking/backtrack.py
@@ -1,4 +1,5 @@
 import logging
+from pathlib import Path
 
 import numpy as np
 import pandas as pd
@@ -17,7 +18,13 @@ from wam2layers.tracking.core import (
     vertical_advection,
     vertical_dispersion,
 )
-from wam2layers.tracking.io import load_data, load_region, output_path, write_output
+from wam2layers.tracking.io import (
+    load_data,
+    load_tagging_region,
+    output_path,
+    write_output,
+)
+from wam2layers.tracking.shared import initialize_tagging_region, initialize_time
 from wam2layers.utils.profiling import ProgressTracker
 
 logger = logging.getLogger(__name__)
@@ -27,7 +34,7 @@ def backtrack(
     F,
     S1,
     S0,
-    region,
+    tagging_mask,
     output,
     config,
 ):
@@ -58,7 +65,7 @@ def backtrack(
         + vertical_dispersion(
             -f_vert, s_track_relative_lower, s_track_relative_upper, config.kvf
         )
-        + region * precip * s_lower / (s_upper + s_lower)
+        + tagging_mask * precip * s_lower / (s_upper + s_lower)
         - evap * s_track_relative_lower
     )
 
@@ -68,7 +75,7 @@ def backtrack(
         - vertical_dispersion(
             -f_vert, s_track_relative_lower, s_track_relative_upper, config.kvf
         )
-        + region * precip * s_upper / (s_upper + s_lower)
+        + tagging_mask * precip * s_upper / (s_upper + s_lower)
     )
 
     # down and top: redistribute unaccounted water that is otherwise lost from the sytem
@@ -102,7 +109,7 @@ def backtrack(
 
     output["s_track_lower_restart"].values = s_track_lower
     output["s_track_upper_restart"].values = s_track_upper
-    output["tagged_precip"] += region * precip
+    output["tagged_precip"] += tagging_mask * precip
 
 
 def initialize(config_file):
@@ -110,53 +117,29 @@ def initialize(config_file):
     logger.info(f"Initializing experiment with config file {config_file}")
 
     config = Config.from_yaml(config_file)
-    region = load_region(config)
-
-    output = initialize_outputs(region)
-    region = region.values
-
-    if config.restart:
-        # Reload last state from existing output
-        date = config.track_end_date
-        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
-
-    logger.info(f"Output will be written to {config.output_folder.absolute()}.")
-    return config, region, output
 
-
-def initialize_outputs(region):
-    """Allocate output arrays."""
-
-    proto = region
+    # Initialize outputs as empty fields based on the input coords
+    t = pd.Timestamp(config.tracking_end_date)
+    grid = load_data(t, config, "states").coords
     output = xr.Dataset(
         {
             # Keep last state for a restart
-            "s_track_upper_restart": xr.zeros_like(proto),
-            "s_track_lower_restart": xr.zeros_like(proto),
-            "e_track": xr.zeros_like(proto),
-            "tagged_precip": xr.zeros_like(proto),
+            "s_track_upper_restart": xr.DataArray(0, coords=grid).astype("float32"),
+            "s_track_lower_restart": xr.DataArray(0, coords=grid).astype("float32"),
+            "e_track": xr.DataArray(0, coords=grid).astype("float32"),
+            "tagged_precip": xr.DataArray(0, coords=grid).astype("float32"),
         }
     )
-    return output
-
 
-def initialize_time(config, direction="forward"):
-    dt = pd.Timedelta(config.target_frequency)
-
-    if direction == "forward":
-        t0 = pd.Timestamp(config.track_start_date)
-        th = t0 + dt / 2  # th is "half" time, i.e. between t0 and t1
-        t1 = t0 + dt
-    elif direction == "backward":
-        t1 = pd.Timestamp(config.track_end_date)
-        th = t1 - dt / 2  # th is "half" time, i.e. between t0 and t1
-        t0 = t1 - dt
-    else:
-        raise ValueError("Direction should be forward or backward")
+    if config.restart:
+        # Reload last state from existing output
+        date = config.tracking_end_date
+        ds = xr.open_dataset(output_path(date, config, mode="backtrack"))
+        output["s_track_upper_restart"].values = ds.s_track_upper_restart.values
+        output["s_track_lower_restart"].values = ds.s_track_lower_restart.values
 
-    return t0, th, t1, dt
+    logger.info(f"Output will be written to {config.output_folder.absolute()}.")
+    return config, output, grid
 
 
 #############################################################################
@@ -167,18 +150,34 @@ def initialize_time(config, direction="forward"):
 
 def run_experiment(config_file):
     """Run a backtracking experiment from start to finish."""
-    config, region, output = initialize(config_file)
+    config, output, grid = initialize(config_file)
 
-    event_start, event_end = config.event_start_date, config.event_end_date
     progress_tracker = ProgressTracker(output, mode="backtrack")
 
     t0, th, t1, dt = initialize_time(config, direction="backward")
 
-    while t0 >= config.track_start_date:
+    if isinstance(config.tagging_region, Path):
+        tagging_region = load_tagging_region(config)
+        tagging_region_stationary = tagging_region.ndim == 2
+    else:
+        lat, lon = grid["latitude"], grid["longitude"]
+        bbox = config.tagging_region
+        tagging_region = initialize_tagging_region(bbox, lat, lon)
+        tagging_region_stationary = True
+
+    while t0 >= config.tracking_start_date:
         S0 = load_data(t0, config, "states")
         F = load_data(th, config, "fluxes")
         S1 = load_data(t1, config, "states")
 
+        # Load/update tagging mask
+        tagging_mask = 0
+        if config.tagging_start_date <= t1 <= config.tagging_end_date:
+            if tagging_region_stationary:
+                tagging_mask = tagging_region
+            else:
+                tagging_mask = load_tagging_region(config, t=t1)
+
         # Convert data to volumes
         change_units(S0, config.target_frequency)
         change_units(F, config.target_frequency)
@@ -190,19 +189,15 @@ def run_experiment(config_file):
         # Determine the vertical moisture flux
         F["f_vert"] = calculate_fz(F, S0, S1, config.kvf)
 
-        # Only track the precipitation at certain timesteps
-        if not event_start <= t1 <= event_end:
-            F["precip"] = 0
-
         # Inside backtrack the "output" dictionary is updated
-        backtrack(F, S1, S0, region, output, config)
+        backtrack(F, S1, S0, tagging_mask, output, config)
         t1 -= dt
         th -= dt
         t0 -= dt
 
         # Daily output
         is_output_time = t1 == t1.floor(config.output_frequency)
-        is_final_step = t0 < config.track_start_date
+        is_final_step = t0 < config.tracking_start_date
         if is_output_time or is_final_step:
             progress_tracker.print_progress(t1, output)
             progress_tracker.store_intermediate_states(output)
diff --git a/src/wam2layers/tracking/forwardtrack.py b/src/wam2layers/tracking/forwardtrack.py
index 85eb3fa..f2d1902 100644
--- a/src/wam2layers/tracking/forwardtrack.py
+++ b/src/wam2layers/tracking/forwardtrack.py
@@ -1,4 +1,5 @@
 import logging
+from pathlib import Path
 
 import numpy as np
 import pandas as pd
@@ -17,7 +18,13 @@ from wam2layers.tracking.core import (
     vertical_advection,
     vertical_dispersion,
 )
-from wam2layers.tracking.io import load_data, load_region, output_path, write_output
+from wam2layers.tracking.io import (
+    load_data,
+    load_tagging_region,
+    output_path,
+    write_output,
+)
+from wam2layers.tracking.shared import initialize_tagging_region, initialize_time
 from wam2layers.utils.profiling import ProgressTracker
 
 logger = logging.getLogger(__name__)
@@ -132,54 +139,30 @@ def initialize(config_file):
     logger.info(f"Initializing experiment with config file {config_file}")
 
     config = Config.from_yaml(config_file)
-    region = load_region(config)
-
-    output = initialize_outputs(region)
-    region = region.values
-
-    if config.restart:
-        # Reload last state from existing output
-        date = config.track_end_date
-        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
-
-    logger.info(f"Output will be written to {config.output_folder.absolute()}.")
-    return config, region, output
 
-
-def initialize_outputs(region):
-    """Allocate output arrays."""
-
-    proto = region
+    # Initialize outputs as empty fields based on the input coords
+    t = pd.Timestamp(config.tracking_start_date)
+    grid = load_data(t, config, "states").coords
     output = xr.Dataset(
         {
             # Keep last state for a restart
-            "s_track_upper_restart": xr.zeros_like(proto),
-            "s_track_lower_restart": xr.zeros_like(proto),
-            "p_track_upper": xr.zeros_like(proto),
-            "p_track_lower": xr.zeros_like(proto),
-            "tagged_evap": xr.zeros_like(proto),
+            "s_track_upper_restart": xr.DataArray(0, coords=grid).astype("float32"),
+            "s_track_lower_restart": xr.DataArray(0, coords=grid).astype("float32"),
+            "p_track_upper": xr.DataArray(0, coords=grid).astype("float32"),
+            "p_track_lower": xr.DataArray(0, coords=grid).astype("float32"),
+            "tagged_evap": xr.DataArray(0, coords=grid).astype("float32"),
         }
     )
-    return output
-
 
-def initialize_time(config, direction="forward"):
-    dt = pd.Timedelta(config.target_frequency)
-
-    if direction == "forward":
-        t0 = pd.Timestamp(config.track_start_date)
-        th = t0 + dt / 2  # th is "half" time, i.e. between t0 and t1
-        t1 = t0 + dt
-    elif direction == "backward":
-        t1 = pd.Timestamp(config.track_end_date)
-        th = t1 - dt / 2  # th is "half" time, i.e. between t0 and t1
-        t0 = t1 - dt
-    else:
-        raise ValueError("Direction should be forward or backward")
+    if config.restart:
+        # Reload last state from existing output
+        date = config.tracking_start_date
+        ds = xr.open_dataset(output_path(date, config, mode="forwardtrack"))
+        output["s_track_upper_restart"].values = ds.s_track_upper_restart.values
+        output["s_track_lower_restart"].values = ds.s_track_lower_restart.values
 
-    return t0, th, t1, dt
+    logger.info(f"Output will be written to {config.output_folder.absolute()}.")
+    return config, output, grid
 
 
 #############################################################################
@@ -190,18 +173,34 @@ def initialize_time(config, direction="forward"):
 
 def run_experiment(config_file):
     """Run a forwardtracking experiment from start to finish."""
-    config, region, output = initialize(config_file)
+    config, output, grid = initialize(config_file)
 
-    event_start, event_end = config.event_start_date, config.event_end_date
     progress_tracker = ProgressTracker(output, mode="forwardtrack")
 
     t0, th, t1, dt = initialize_time(config, direction="forward")
 
-    while t1 <= config.track_end_date:
+    if isinstance(config.tagging_region, Path):
+        tagging_region = load_tagging_region(config)
+        tagging_region_stationary = tagging_region.ndim == 2
+    else:
+        lat, lon = grid["latitude"], grid["longitude"]
+        bbox = config.tagging_region
+        tagging_region = initialize_tagging_region(bbox, lat, lon)
+        tagging_region_stationary = True
+
+    while t1 <= config.tracking_end_date:
         S0 = load_data(t0, config, "states")
         F = load_data(th, config, "fluxes")
         S1 = load_data(t1, config, "states")
 
+        # Load/update tagging mask
+        tagging_mask = 0
+        if config.tagging_start_date <= t0 <= config.tagging_end_date:
+            if tagging_region_stationary:
+                tagging_mask = tagging_region
+            else:
+                tagging_mask = load_tagging_region(config, t=t0)
+
         # Convert data to volumes
         change_units(S0, config.target_frequency)
         change_units(F, config.target_frequency)
@@ -213,20 +212,16 @@ def run_experiment(config_file):
         # Determine the vertical moisture flux
         F["f_vert"] = calculate_fz(F, S0, S1, config.kvf)
 
-        # Only track the precipitation at certain timesteps
-        if not event_start <= t0 <= event_end:
-            F["evap"] = 0
-
         # TODO: consider changing signature to F, S0, S1
         # Inside forwardtrack the "output" dictionary is updated
-        forwardtrack(F, S1, S0, region, output, config)
+        forwardtrack(F, S1, S0, tagging_mask, output, config)
         t1 += dt
         th += dt
         t0 += dt
 
         # Daily output
         is_output_time = t0 == t0.floor(config.output_frequency)
-        is_final_step = t1 > config.track_end_date
+        is_final_step = t1 > config.tracking_end_date
         if is_output_time or is_final_step:
             progress_tracker.print_progress(t0, output)
             progress_tracker.store_intermediate_states(output)
diff --git a/src/wam2layers/tracking/io.py b/src/wam2layers/tracking/io.py
index 62bb8e8..d57bd52 100644
--- a/src/wam2layers/tracking/io.py
+++ b/src/wam2layers/tracking/io.py
@@ -1,4 +1,6 @@
 import logging
+from functools import lru_cache
+from pathlib import Path
 
 import xarray as xr
 
@@ -18,7 +20,10 @@ def output_path(date, config, mode):
 def read_data_at_date(d, config):
     """Load input data for given date."""
     file = input_path(d, config)
-    return xr.open_dataset(file, cache=False)
+    ds = xr.open_dataset(file, cache=False)
+    if config.tracking_domain is not None:
+        return select_subdomain(ds, config.tracking_domain)
+    return ds
 
 
 def read_data_at_time(t, config):
@@ -27,6 +32,25 @@ def read_data_at_time(t, config):
     return ds.sel(time=t, drop=True)
 
 
+def select_subdomain(ds, bbox):
+    """Limit the data to a subdomain."""
+    if bbox.west < bbox.east:
+        return ds.sel(
+            longitude=slice(bbox.west, bbox.east),  # -180 to 180
+            latitude=slice(bbox.north, bbox.south),  # 90 - -90 (decreasing)
+            drop=True,
+        )
+    else:
+        # Need to roll coordinates to maintain a continuous longitude
+        # Roll the first lon that's NOT in the bbox to the front of the array
+        # That way, all values that ARE in the bbox, must be contiguous.
+        in_bbox = (ds.longitude < bbox.east) | (ds.longitude > bbox.west)
+        n_roll = in_bbox.argmin().item()  # first false value
+        ds_rolled = ds.roll(longitude=-n_roll, roll_coords=True)
+        in_bbox = (ds_rolled.longitude < bbox.east) | (ds_rolled.longitude > bbox.west)
+        return ds_rolled.sel(longitude=in_bbox, latitude=slice(bbox.north, bbox.south))
+
+
 def load_data(t, config, subset="fluxes"):
     """Load variable at t, interpolate if needed."""
     variables = {
@@ -48,9 +72,15 @@ def load_data(t, config, subset="fluxes"):
     return da0 + (t - t0) / (t1 - t0) * (da1 - da0)
 
 
-def load_region(config):
-    # TODO: make variable name more generic
-    return xr.open_dataset(config.region).tagging_region
+def load_tagging_region(config, t=None):
+    tagging_region = xr.open_dataarray(config.tagging_region)
+
+    if config.tracking_domain is not None:
+        tagging_region = select_subdomain(tagging_region, config.tracking_domain)
+
+    if t is not None:
+        return tagging_region.sel(time=t, method="nearest")
+    return tagging_region
 
 
 def write_output(output, t, config, mode):
diff --git a/src/wam2layers/tracking/shared.py b/src/wam2layers/tracking/shared.py
new file mode 100644
index 0000000..04fc464
--- /dev/null
+++ b/src/wam2layers/tracking/shared.py
@@ -0,0 +1,30 @@
+import pandas as pd
+import xarray as xr
+
+
+def initialize_time(config, direction="forward"):
+    dt = pd.Timedelta(config.target_frequency)
+
+    if direction == "forward":
+        t0 = pd.Timestamp(config.tracking_start_date)
+        th = t0 + dt / 2  # th is "half" time, i.e. between t0 and t1
+        t1 = t0 + dt
+    elif direction == "backward":
+        t1 = pd.Timestamp(config.tracking_end_date)
+        th = t1 - dt / 2  # th is "half" time, i.e. between t0 and t1
+        t0 = t1 - dt
+    else:
+        raise ValueError("Direction should be forward or backward")
+
+    return t0, th, t1, dt
+
+
+def initialize_tagging_region(bbox, lat, lon):
+    """Build a new xr.DataArray with ones inside and zeros outside bbox."""
+    lat_in_bbox = (bbox.south <= lat) & (lat <= bbox.north)
+    unrolled_lon_in_bbox = (bbox.west <= lon) & (lon <= bbox.east)
+    rolled_lon_in_bbox = (bbox.west <= lon) | (lon <= bbox.east)
+    rolled_domain = bbox.west > bbox.east
+    lon_in_bbox = rolled_lon_in_bbox if rolled_domain else unrolled_lon_in_bbox
+    tagging_region = xr.where((lat_in_bbox & lon_in_bbox), 1, 0).values
+    return tagging_region
diff --git a/src/wam2layers/utils/profiling.py b/src/wam2layers/utils/profiling.py
index 7d69f05..6ac5e79 100644
--- a/src/wam2layers/utils/profiling.py
+++ b/src/wam2layers/utils/profiling.py
@@ -1,11 +1,12 @@
 import logging
 import time
 from datetime import datetime
+import xarray as xr
 
 import numpy as np
 import psutil
 
-from wam2layers.tracking.io import load_region
+from wam2layers.tracking.io import load_tagging_region
 
 logger = logging.getLogger(__name__)
 
@@ -82,7 +83,7 @@ class ProgressTracker:
             f"Time since start: {time}s, RAM: {memory:.2f} MB"
         )
 
-    def track_stability_correction(self, fy_corrected, fy_abs, config, t):
+    def track_stability_correction(self, fy_corrected, fy_abs, config, coords, t):
         """Issue warning if correction exceeds criterion.
 
         Warning advises to reduce the timestep.
@@ -114,8 +115,7 @@ class ProgressTracker:
         timestamp = t.strftime("%Y%m%d-%H%M%S")
         filename = debug_dir / f"stability_correction_{timestamp}.nc"
 
-        ncfile = load_region(config).rename("correction")
-        ncfile.values = correction
+        ncfile = xr.DataArray(correction, coords=coords, name="correction")
         ncfile.to_netcdf(filename)
 
         logger.warn(


Test patch:
diff --git a/tests/test_config.py b/tests/test_config.py
index 07bf70e..a2a29ee 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,18 +1,82 @@
+import logging
+from datetime import datetime
+
 import pytest
 
-from wam2layers.config import Config
+from wam2layers.config import BoundingBox, Config
 
 EXAMPLE_CONFIG = "example_config.yaml"
+TEST_CONFIG = "tests/test_data/test_config.yaml"
 
 
 def test_load_example_config():
+    """Verify that the example config file can be loaded"""
     Config.from_yaml(EXAMPLE_CONFIG)
 
 
-def test_config_fails_for_wrong_time():
-    config = Config.from_yaml(EXAMPLE_CONFIG)
+def test_load_test_config():
+    """Verify that the test config file can be loaded"""
+    Config.from_yaml(TEST_CONFIG)
+
+
+def test_check_date_order_invalid_tracking():
+    """Raise error when tracking_end_date is earlier than tracking_start_date."""
+    config = Config.from_yaml(TEST_CONFIG)
+    with pytest.raises(
+        ValueError, match="tracking_end_date should be later than tracking_start_date"
+    ):
+        config.tracking_end_date = datetime(2022, 7, 31)
+
+
+def test_check_date_order_invalid_tagging():
+    """Raise error when tagging_end_date is earlier than tagging_start_date."""
+    config = Config.from_yaml(TEST_CONFIG)
+    with pytest.raises(
+        ValueError, match="tagging_end_date should be later than tagging_start_date"
+    ):
+        config.tagging_end_date = datetime(2022, 7, 31)
+
+
+def test_check_date_order_invalid_preprocess():
+    """Raise error when preprocess_end_date is earlier than preprocess_start_date."""
+    config = Config.from_yaml(TEST_CONFIG)
+    with pytest.raises(
+        ValueError,
+        match="preprocess_end_date should be later than preprocess_start_date",
+    ):
+        config.preprocess_end_date = datetime(2022, 7, 31)
+
+
+class TestValidateRegion:
+    def test_path_expanded(self):
+        """Check path is expanded."""
+        config = Config.from_yaml(TEST_CONFIG)
+        assert config.tagging_region.is_absolute()
+        config.tagging_region = "tests/test_data/test_region.nc"
+        assert config.tagging_region.is_absolute()
+
+    def test_invalid_path(self):
+        """Raise error for invalid path."""
+        config = Config.from_yaml(TEST_CONFIG)
+        with pytest.raises(Exception, match="does not point to a file"):
+            config.tagging_region = "weird/path.exe"
+
+    def test_valid_bbox_ok(self):
+        """Check that we can set region to a bbox"""
+        config = Config.from_yaml(TEST_CONFIG)
+        config.tagging_region = BoundingBox(0, 50, 10, 55)
+        assert config.tagging_region == BoundingBox(0, 50, 10, 55)
 
-    with pytest.raises(ValueError) as excinfo:
-        config.preprocess_end_date = config.preprocess_start_date
+    def test_invalid_south(self):
+        """Raise error for invalid bbox."""
+        config = Config.from_yaml(TEST_CONFIG)
+        with pytest.raises(Exception, match="south should be smaller than north"):
+            config.tagging_region = BoundingBox(0, 60, -10, 55)
 
-    assert "End date should be later than start date" in str(excinfo.value)
+    def test_wrap_meridian(self, caplog):
+        """Warn for bbox crossing the antimeridian"""
+        config = Config.from_yaml(TEST_CONFIG)
+        with caplog.at_level(logging.INFO):
+            config.tagging_region = BoundingBox(10, 50, -10, 55)
+        assert "coordinates will be rolled" in caplog.text
+        assert config.tagging_region.west == 10
diff --git a/tests/test_data/test_config.yaml b/tests/test_data/test_config.yaml
index bac7bfb..579d990 100644
--- a/tests/test_data/test_config.yaml
+++ b/tests/test_data/test_config.yaml
@@ -1,29 +1,33 @@
 # Paths (relative to root of repo)
-filename_template: tests/test_data/era5/ERA5_{year}-{month:02d}-{day:02d}{levtype}_{variable}.nc
-region: tests/test_data/test_region.nc
 preprocessed_data_folder: tests/tmp/preprocessed_data
 output_folder: tests/tmp/output_data
 
-# Dates
+# Preprocessing
+filename_template: tests/test_data/era5/ERA5_{year}-{month:02d}-{day:02d}{levtype}_{variable}.nc
 preprocess_start_date: "2022-08-31T00:00"
 preprocess_end_date: "2022-08-31T23:00"
-track_start_date: "2022-08-31T00:00"
-track_end_date: "2022-08-31T23:00"
-event_start_date: "2022-08-31T12:00"
-event_end_date: "2022-08-31T23:00"
+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_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"
+tracking_end_date: "2022-08-31T23:00"
+
+# Tagging
+tagging_region: tests/test_data/test_region.nc
+# tagging_region: [0, 50, 10, 55]  # alternative to mask file
+tagging_start_date: "2022-08-31T12:00"
+tagging_end_date: "2022-08-31T23:00"
 
 # Time
 input_frequency: "1h"
 target_frequency: "10min"
 output_frequency: "6h"
 
-# Levels
-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]
-
 # Model settings
 log_level: info
 restart: False
 periodic_boundary: false
 kvf: 3
-chunks: null
diff --git a/tests/test_io.py b/tests/test_io.py
new file mode 100644
index 0000000..d7d3c6f
--- /dev/null
+++ b/tests/test_io.py
@@ -0,0 +1,38 @@
+import numpy as np
+import pytest
+import xarray as xr
+
+from wam2layers.config import BoundingBox
+from wam2layers.tracking.io import select_subdomain
+
+
+@pytest.fixture
+def ds():
+    lon = np.arange(-180, 180, 30)
+    lat = np.arange(90, -90, -30)
+    time = np.arange(3)
+
+    p = xr.DataArray(
+        np.random.randn(len(time), len(lat), len(lon)),
+        coords={"time": time, "latitude": lat, "longitude": lon},
+    )
+    e = xr.DataArray(
+        np.random.randn(len(time), len(lat), len(lon)),
+        coords={"time": time, "latitude": lat, "longitude": lon},
+    )
+    return xr.Dataset({"e": e, "p": p})
+
+
+def test_select_subdomain(ds):
+    """Check subdomain that fits in [-180, 180]."""
+    bbox = BoundingBox(0, -20, 90, 60)
+    sub = select_subdomain(ds, bbox)
+    assert np.array_equal(sub.latitude, [60, 30, 0])
+    assert np.array_equal(sub.longitude, [0, 30, 60, 90])
+
+
+def test_rolling_subdomain(ds):
+    """Check subdomain that does not fit in [-180, 180] or [0, 360]."""
+    bbox = BoundingBox(130, -20, 40, 60)
+    sub = select_subdomain(ds, bbox)
+    assert np.array_equal(sub.longitude, [150, -180, -150, -120, -90, -60, -30, 0, 30])
diff --git a/tests/test_workflow.py b/tests/test_workflow.py
index 970b60f..68c1432 100644
--- a/tests/test_workflow.py
+++ b/tests/test_workflow.py
@@ -37,7 +37,7 @@ def test_preprocess():
         i for i in Path("tests/tmp/preprocessed_data/").glob("wam2layers_*.log")
     ]
 
-    assert len(log_files) == 1
+    assert len(log_files) >= 1  # make sure the test passes when repeating
 
     # check config yaml
     config_path = Path("tests/tmp/preprocessed_data/test_config.yaml")
@@ -69,7 +69,7 @@ def test_backtrack():
     # check log file
     log_files = [i for i in Path("tests/tmp/output_data/").glob("wam2layers_*.log")]
 
-    assert len(log_files) == 1
+    assert len(log_files) >= 1  # make sure the test passes when repeating
 
     # check config yaml
     config_path = Path("tests/tmp/output_data/test_config.yaml")