Skip to content

Commit

Permalink
Merge pull request #27 from axiom-data-science/mostly_kerchunk_improv…
Browse files Browse the repository at this point in the history
…ements

Mostly kerchunk improvements
  • Loading branch information
kthyng authored Jul 26, 2024
2 parents 05ab77b + e5b1536 commit 66d6047
Show file tree
Hide file tree
Showing 7 changed files with 465 additions and 16 deletions.
8 changes: 8 additions & 0 deletions docs/whats_new.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,13 @@
# What's New

## v0.9.0 (July 26, 2024)

* Added utilities to generate kerchunk files on the fly for the time period of the simulation length for CIOFS and NWGOA. This has majorly sped up CIOFS simulations and modestly sped up NWGOA simulations.
* depth z should be negative! Fixed this in tests.
* added `start_time_end`, which adds OpenDrift capability for starting drifters over linear time frame
* fixed so unique log file is output for each simulation even if run in a script, and has the same name as `output_file`.
* small fix to histogram plot

## v0.8.4 (April 24, 2024)

* updated the `ptm_level` of a bunch of config parameters
Expand Down
44 changes: 35 additions & 9 deletions particle_tracking_manager/models/opendrift/opendrift.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@

from ...cli import is_None
from ...the_manager import _KNOWN_MODELS, ParticleTrackingManager
from .utils import make_ciofs_kerchunk, make_nwgoa_kerchunk


# from .cli import is_None
Expand Down Expand Up @@ -661,8 +662,11 @@ def run_add_reader(
"hraw",
"snow_thick",
]
start = f"{self.start_time.year}-{str(self.start_time.month).zfill(2)}-{str(self.start_time.day).zfill(2)}"
end = f"{self.end_time.year}-{str(self.end_time.month).zfill(2)}-{str(self.end_time.day).zfill(2)}"
loc_local = make_nwgoa_kerchunk(start=start, end=end)

loc_local = "/mnt/depot/data/packrat/prod/aoos/nwgoa/processed/nwgoa_kerchunk.parq"
# loc_local = "/mnt/depot/data/packrat/prod/aoos/nwgoa/processed/nwgoa_kerchunk.parq"
loc_remote = (
"http://xpublish-nwgoa.srv.axds.co/datasets/nwgoa_all/zarr/"
)
Expand All @@ -674,8 +678,14 @@ def run_add_reader(
"wetdry_mask_psi",
]
if self.ocean_model == "CIOFS":

loc_local = "/mnt/vault/ciofs/HINDCAST/ciofs_kerchunk.parq"
start = f"{self.start_time.year}_{str(self.start_time.dayofyear - 1).zfill(4)}"
end = (
f"{self.end_time.year}_{str(self.end_time.dayofyear).zfill(4)}"
)
loc_local = make_ciofs_kerchunk(start=start, end=end)
# loc_local = make_ciofs_kerchunk(start="2005_0052", end="2005_0068")
# loc_local = "/mnt/vault/ciofs/HINDCAST/ciofs_kerchunk_2005.parq"
# loc_local = "/mnt/vault/ciofs/HINDCAST/ciofs_kerchunk.parq"
loc_remote = "http://xpublish-ciofs.srv.axds.co/datasets/ciofs_hindcast/zarr/"

elif self.ocean_model == "CIOFSOP":
Expand Down Expand Up @@ -762,11 +772,19 @@ def run_add_reader(
dt_model = float(
ds.ocean_time[1] - ds.ocean_time[0]
) # time step of the model output in seconds
start_time_num = (self.start_time - units_date).total_seconds()
# want to include the next ocean model output before the first drifter simulation time
# in case it starts before model times
start_time_num = (
self.start_time - units_date
).total_seconds() - dt_model
# want to include the next ocean model output after the last drifter simulation time
end_time_num = (self.end_time - units_date).total_seconds() + dt_model
ds = ds.sel(ocean_time=slice(start_time_num, end_time_num))
self.logger.info("Narrowed model output to simulation time")
if len(ds.ocean_time) == 0:
raise ValueError(
"No model output left for simulation time. Check start_time and end_time."
)
else:
raise ValueError(
"start_time and end_time must be set to narrow model output to simulation time"
Expand Down Expand Up @@ -823,10 +841,19 @@ def seed_kws(self):
"drift:truncate_ocean_model_below_m",
]

if self.start_time_end is not None:
# time can be a list to start drifters linearly in time
time = [
self.start_time.to_pydatetime(),
self.start_time_end.to_pydatetime(),
]
elif self.start_time is not None:
time = self.start_time.to_pydatetime()
else:
time = None

_seed_kws = {
"time": self.start_time.to_pydatetime()
if self.start_time is not None
else None,
"time": time,
"z": self.z,
}

Expand Down Expand Up @@ -861,7 +888,6 @@ 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":
Expand Down Expand Up @@ -900,7 +926,7 @@ def run_drifters(self):

output_file_initial = (
f"{self.output_file}_initial"
or f"output-results_{datetime.datetime.utcnow():%Y-%m-%dT%H%M:%SZ}.nc"
or f"output-results_{datetime.datetime.now():%Y-%m-%dT%H%M:%SZ}.nc"
)

self.o.run(
Expand Down
Loading

0 comments on commit 66d6047

Please sign in to comment.