Merge pull request #24 from kthyng/updates

improved drift_model_config(), include PTM metadata with output file

docs/conf.py

@@ -78,7 +78,13 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ["_build", "**.ipynb_checkpoints", "Thumbs.db", ".DS_Store"]
+exclude_patterns = [
+    "*.ipynb",
+    "_build",
+    "**.ipynb_checkpoints",
+    "Thumbs.db",
+    ".DS_Store",
+]
 
 
 # -- Options for HTML output -------------------------------------------------

docs/configuration.md

@@ -173,8 +173,8 @@ import xroms
 
 m = ptm.OpenDriftModel(lon=-90, lat=28.7, number=1, steps=2)
 url = xroms.datasets.CLOVER.fetch("ROMS_example_full_grid.nc")
-reader_kwargs = dict(loc=url, kwargs_xarray={})
-m.add_reader(**reader_kwargs)
+ds = xr.open_dataset(url, decode_times=False)
+m.add_reader(ds=ds)
 m.run_all()
 ```
 
@@ -184,7 +184,7 @@ To run an idealized scenario, no reader should be added (`ocean_model` should be
 ```
 from datetime import datetime
 m = ptm.OpenDriftModel(lon=4.0, lat=60.0, start_time=datetime(2015, 9, 22, 6),
-                       use_auto_landmask=True,)
+                       use_auto_landmask=True, steps=5)
 
 # idealized simulation, provide a fake current
 m.o.set_config('environment:fallback:y_sea_water_velocity', 1)
@@ -201,7 +201,7 @@ For testing purposes, all steps can be run (including added a "reader") with the
 ```
 from datetime import datetime
 m = ptm.OpenDriftModel(lon=4.0, lat=60.0, start_time=datetime(2015, 9, 22, 6),
-                       use_auto_landmask=True, ocean_model="test")
+                       use_auto_landmask=True, ocean_model="test", steps=5)
 
 m.run_all()
 ```
@@ -246,7 +246,7 @@ The default list of `export_variables` is set in `config_model` but is modified
 
 #### How to modify details for Stokes Drift
 
-Turn on (on by default):
+Turn on (on by default, drift model-dependent):
 
 ```
 m = ptm.OpenDriftModel(stokes_drift=True)

docs/quick_start.md

@@ -30,7 +30,8 @@ Run directly from the Lagrangian model you want to use, which will inherit from
 ```
 import particle_tracking_manager as ptm
 
-m = ptm.OpenDriftModel(ocean_model="NWGOA", lon=-151, lat=59)
+m = ptm.OpenDriftModel(ocean_model="NWGOA", lon=-151, lat=59, steps=1)
+# Can modify `m` between these steps, or look at `OpenDrift` config with `m.drift_model_config()`
 m.run_all()
 ```
 
@@ -43,7 +44,13 @@ Then find results in file `m.outfile_name`.
 The equivalent for the set up above for using the command line is:
 
 ```
-ptm lon=-151 lat=59 ocean_model=NWGOA
+ptm lon=-151 lat=59 ocean_model=NWGOA steps=1
+```
+
+To just initialize the simulation and print the `OpenDrift` configuration to screen without running the simulation, add the `--dry-run` flag:
+
+```
+ptm lon=-151 lat=59 ocean_model=NWGOA steps=1 --dry-run
 ```
 
 `m.outfile_name` is printed to the screen after the command has been run. `ptm` is installed as an entry point with `particle-tracking-manager`.
@@ -122,11 +129,17 @@ m.reader
 Get reader/ocean model properties (gathered metadata about model):
 
 ```
-m.reader_metadata(key)
+m.reader_metadata(<key>)
 ```
 
 Show configuration details — many more details on this in {doc}`configuration`:
 
 ```
 m.show_config()
 ```
+
+Show `OpenDrift` configuration for selected `drift_model`:
+
+```
+m.drift_model_config()
+```

docs/whats_new.md

@@ -1,5 +1,12 @@
 # What's New
 
+## v0.8.2 (April 10, 2024)
+
+* updated docs
+* improved `drift_model_config()`
+* updated tests
+* now include PTM metadata with output file
+
 ## v0.8.1 (April 5, 2024)
 
 * updated docs

particle_tracking_manager/cli.py

@@ -106,6 +106,13 @@ def main():
         help="Input keyword arguments for running PTM. Available options are specific to the `catalog_type`. Dictionary-style input, e.g. `case='Oil'`. Format for list items is e.g. standard_names='[sea_water_practical_salinity,sea_water_temperature]'.",
     )
 
+    parser.add_argument(
+        "--dry-run",
+        help="Return configuration parameters without running the model.",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+
     args = parser.parse_args()
 
     to_bool = {
@@ -115,14 +122,20 @@ def main():
     }
     args.kwargs.update(to_bool)
 
-    # # set default
-    # if "model" not in args:
-    #     args.kwargs["model"] = "opendrift"
+    m = ptm.OpenDriftModel(**args.kwargs)
 
-    # if args.kwargs["ocean_model"] is None and args.kwargs["start_time"] is None:
-    #     raise KeyError("Need to either use a reader or input a start_time to avoid error.")
+    if args.dry_run:
 
-    m = ptm.OpenDriftModel(**args.kwargs)
-    m.run_all()
+        # run this to make sure everything is updated fully
+        m.add_reader()
+        print(m.drift_model_config())
+
+    else:
+
+        m.add_reader()
+        print(m.drift_model_config())
+
+        m.seed()
+        m.run()
 
-    print(m.outfile_name)
+        print(m.outfile_name)

particle_tracking_manager/models/opendrift/opendrift.py

@@ -1,8 +1,12 @@
 """Using OpenDrift for particle tracking."""
 import copy
 import datetime
+import gc
 import json
 import logging
+import os
+import platform
+import tempfile
 
 from pathlib import Path
 from typing import Optional, Union
@@ -735,11 +739,17 @@ def run_add_reader(
                 else:
                     if ".nc" in loc_remote:
                         ds = xr.open_dataset(
-                            loc_remote, chunks={}, drop_variables=drop_vars
+                            loc_remote,
+                            chunks={},
+                            drop_variables=drop_vars,
+                            decode_times=False,
                         )
                     else:
                         ds = xr.open_zarr(
-                            loc_remote, chunks={}, drop_variables=drop_vars
+                            loc_remote,
+                            chunks={},
+                            drop_variables=drop_vars,
+                            decode_times=False,
                         )
 
             # For NWGOA, need to calculate wetdry mask from a variable
@@ -811,8 +821,12 @@ def run_add_reader(
         else:
             raise ValueError("reader did not set an ocean_model")
 
-    def run_seed(self):
-        """Seed drifters for model."""
+    @property
+    def seed_kws(self):
+        """Gather seed input kwargs.
+
+        This could be run more than once.
+        """
 
         already_there = [
             "seed:number",
@@ -824,22 +838,28 @@ def run_seed(self):
             "seed:droplet_diameter_sigma",
             "seed:droplet_diameter_max_subsea",
             "seed:object_type",
+            "seed_flag",
+            "drift:use_tabularised_stokes_drift",
+            "drift:vertical_advection",
+            "drift:truncate_ocean_model_below_m",
         ]
 
-        seed_kws = {
-            "time": self.start_time.to_pydatetime(),
+        _seed_kws = {
+            "time": self.start_time.to_pydatetime()
+            if self.start_time is not None
+            else None,
             "z": self.z,
         }
 
         # update seed_kws with drift_model-specific seed parameters
         seedlist = self.drift_model_config(prefix="seed")
         seedlist = [(one, two) for one, two in seedlist if one not in already_there]
         seedlist = [(one.replace("seed:", ""), two) for one, two in seedlist]
-        seed_kws.update(seedlist)
+        _seed_kws.update(seedlist)
 
         if self.seed_flag == "elements":
             # add additional seed parameters
-            seed_kws.update(
+            _seed_kws.update(
                 {
                     "lon": self.lon,
                     "lat": self.lat,
@@ -848,20 +868,34 @@ def run_seed(self):
                 }
             )
 
-            self.o.seed_elements(**seed_kws)
+        elif self.seed_flag == "geojson":
+
+            # geojson needs string representation of time
+            _seed_kws["time"] = (
+                self.start_time.isoformat() if self.start_time is not None else None
+            )
+
+        self._seed_kws = _seed_kws
+        return self._seed_kws
+
+    def run_seed(self):
+        """Actually seed drifters for model."""
+
+        if self.seed_flag == "elements":
+
+            self.o.seed_elements(**self.seed_kws)
 
         elif self.seed_flag == "geojson":
 
             # geojson needs string representation of time
-            seed_kws["time"] = self.start_time.isoformat()
-            self.geojson["properties"] = seed_kws
+            self.seed_kws["time"] = self.start_time.isoformat()
+            self.geojson["properties"] = self.seed_kws
             json_string_dumps = json.dumps(self.geojson)
             self.o.seed_from_geojson(json_string_dumps)
 
         else:
             raise ValueError(f"seed_flag {self.seed_flag} not recognized.")
 
-        self.seed_kws = seed_kws
         self.initial_drifters = self.o.elements_scheduled
 
     def run_drifters(self):
@@ -885,8 +919,8 @@ def run_drifters(self):
 
         self.o._config = config_input_to_opendrift  # only OpenDrift config
 
-        output_file = (
-            self.output_file
+        output_file_initial = (
+            f"{self.output_file}_initial"
             or f"output-results_{datetime.datetime.utcnow():%Y-%m-%dT%H%M:%SZ}.nc"
         )
 
@@ -895,11 +929,37 @@ def run_drifters(self):
             time_step_output=self.time_step_output,
             steps=self.steps,
             export_variables=self.export_variables,
-            outfile=output_file,
+            outfile=output_file_initial,
         )
 
         self.o._config = full_config  # reinstate config
 
+        # open outfile file and add config to it
+        # config can't be present earlier because it breaks opendrift
+        ds = xr.open_dataset(output_file_initial)
+        for k, v in self.drift_model_config():
+            if isinstance(v, (bool, type(None), pd.Timestamp, pd.Timedelta)):
+                v = str(v)
+            ds.attrs[f"ptm_config_{k}"] = v
+
+        # Make new output file
+        output_file = (
+            self.output_file
+            or f"output-results_{datetime.datetime.utcnow():%Y-%m-%dT%H%M:%SZ}.nc"
+        )
+
+        ds.to_netcdf(output_file)
+
+        # update with new path name
+        self.o.outfile_name = output_file
+
+        try:
+            # remove initial file to save space
+            os.remove(output_file_initial)
+        except PermissionError:
+            # windows issue
+            pass
+
     @property
     def _config(self):
         """Surface the model configuration."""
@@ -1011,7 +1071,8 @@ def drift_model_config(self, ptm_level=[1, 2, 3], prefix=""):
 
         This shows all PTM-controlled parameters for the OpenDrift
         drift model selected and their current values, at the selected ptm_level
-        of importance.
+        of importance. It includes some additional configuration parameters
+        that are indirectly controlled by PTM parameters.
 
         Parameters
         ----------
@@ -1022,14 +1083,37 @@ def drift_model_config(self, ptm_level=[1, 2, 3], prefix=""):
             prefix to search config for, only for OpenDrift parameters (not PTM).
         """
 
-        return [
+        outlist = [
             (key, value_dict["value"])
             for key, value_dict in self.show_config(
                 substring=":", ptm_level=ptm_level, level=[1, 2, 3], prefix=prefix
             ).items()
             if "value" in value_dict
         ]
 
+        # also PTM config parameters that are separate from OpenDrift parameters
+        outlist2 = [
+            (key, value_dict["value"])
+            for key, value_dict in self.show_config(
+                ptm_level=ptm_level, prefix=prefix
+            ).items()
+            if "od_mapping" not in value_dict and "value" in value_dict
+        ]
+
+        # extra parameters that are not in the config_model but are set by PTM indirectly
+        extra_keys = [
+            "drift:vertical_advection",
+            "drift:truncate_ocean_model_below_m",
+            "drift:use_tabularised_stokes_drift",
+        ]
+        outlist += [
+            (key, self.show_config(key=key)["value"])
+            for key in extra_keys
+            if "value" in self.show_config(key=key)
+        ]
+
+        return outlist + outlist2
+
     def get_configspec(self, prefix, substring, excludestring, level, ptm_level):
         """Copied from OpenDrift, then modified."""