SIPNET restarts

[ashiklom]

Summarizing today's discussion of SIPNET restarts (i.e., what I'll be working on next).

We need a way for SIPNET to handle exogenous changes to its internal state, such as crop rotation or land use changes.

The PEcAn standard way to describe these "events" is via events.json files. PEcAn then converts these to model-specific input files.

In ensemble simulations, we are expecting one events.json file per ensemble member.

There is some machinery for stopping models, adjusting the state, and then restarting the model in the PEcAn state data assimilation code. We should see if we can adapt this first, before doing new implementation.

One consideration is whether the logic for restarting should be inside a job.sh file or outside of it (i.e., in the parent process / calling code). The current SDA implementation seems to do the latter, modifying job.sh at each restart step. Pseudocode for a job.sh file implementing the former approach might look something like this:

# job.sh

./sipnet --inputs inputs_20180101_20180730... --from 2018-01-01 --to 2018-07-30
./sipnet --inputs inputs_20180730_20190801... --from 2018-07-30 --to 2019-08-01
# etc...

(A cleverer implementation may be possible with loops / arrays, etc.).

Separately from this workflow, @Alomir noted that SIPNET does not currently have the ability to export a snapshot of its entire current state, which makes it tricky to "perfectly" restart SIPNET. We collectively agreed that the cleanest approach was to separate SIPNET time-series outputs for analysis (current .out files) from "restart" files that contain a complete snapshot of SIPNET's state but only at the indicated timestep.

Let me know if I missed anything!

[mdietze]

You are correct that the current SDA code does the restart outside the job.sh, which we have to do because in the Analysis step we're reading in the current state, modifying it, and writing it back out on the fly. While the CARB workflow restarts don't need this functionality now, I'd argue that retaining this approach, while possibly a tad slower, may make it easier to eventually bring the SDA and events codes together in the future (either to do SDA for CARB and/or to include events in the North American SDA)

[Alomir]

I believe that summary is correct for the SIPNET side of this. For those interested, see PecanProject/sipnet#271 for my tracking of where SIPNET is currently falling short in outputting state. (Though, as @ashiklom points out, our current proposed approach is to output a "restart" file that contains a complete state snapshot.)

[dlebauer]

That lines up with #180.

To write down what I recall from our discussions of how to do this:

• The motivation for doing it inside job.sh was to contain what is otherwise a hack around SIPNET’s inability to change PFT parameterizations mid-run. But this was not informed by the existing SDA architecture.
• We considered and rejected the alternative approach of having SIPNET swap parameters mid-run.
• Following this discussion/thread, and after looking at the SDA workflow, it makes sense to me to follow the SDA approach.

So it sounds like the path forward is:

1. Keep restart orchestration outside job.sh, following SDA implementation.
2. Treat events.json as workflow input (one per ensemble member) that defines restart boundaries + management changes.
3. SIPNET is responsible for writing/reading restart state.
4. PEcAn handles timing, PFT/parameter/event/clim updates between segments, and output stitching.

In the current implementation, it seems assim.sequential builds restart args and reruns model segments in a loop. Met splitting is handled in the SDA/restart workflow and then passed into write.configs.SIPNET for each run; the analogous approach would be to split the events.json files when an event changes the PFT (ie: planting a different PFT). This makes sense from the perspective of having model-agnostic workflow components.

[dlebauer]

How to Restart?

To support parameter and state updates, we have at least two options:

1. Implement a portable restart format like NetCDF
2. C binary restarts and apply explicit overrides after loading the checkpoint with existing config files (ic and parameters)

I think 1) is more standard and 2) is less work / less change from current workflow.

[mdietze]

Option 3 would be a tabular ASCII restart format, which might be easier to implement than netCDF (option 1). I'd strongly recommend against a custom binary restart (option 2) as reading/writing/modifying such a file in R is very difficult and slow (this is what we have to do for lpj-guess and it's a major pain). I'm also not convinced that a dedicated restart file is actually required to meet the MVP

[ashiklom]

My vote would be for Mike's option 3 (ASCII tabular). Easy to write, easy to read, easy to implement.

[infotroph]

I vote ASCII as well -- one key-value pair per line, like the param file?

[dlebauer]

That makes sense, one key-value pair per line

[dlebauer]

Implemented in PecanProject/sipnet#276

[ashiklom]

Progress update: To get restarts working outside of a job.sh, my sense is that the logic looks something like this. This code will successfully execute two SIPNET simulations back-to-back (though the second one will immediately fail because the first segment writes nan values of Soil C into the restart; that's probably a parameter/IC issue with my quick and dirty test cases).

Note that this depends on:

• Function to parse restart dates out of event files PecanProject/pecan#3828
• Update events schema: require crop id when planting, remove pft from site PecanProject/pecan#3836
• SIP279 SIPNET Restart MVP PecanProject/sipnet#276
• Some (currently) local changes to write.configs.SIPNET that allow it to use the restart_out and restart_in options; this interface will evolve based on outcomes discussed here: Allow SIPNET to take command line arguments PecanProject/pecan#3861

I'll be returning to this on Friday, but I wanted to put this here for discussion and input to (1) confirm that I'm basically on the right track; (2) I'm not duplicating anyone else's effort; (3) to confirm that, at a high level, this is the approach we want; and (4) to give folks an early chance to weigh in on design preferences/tweaks before I resume work on this in a few days.

#!/usr/bin/env Rscript

# devtools::install("~/projects/pecan/sipnet-events/modules/data.remote", upgrade = FALSE)
# devtools::install("~/projects/pecan/sipnet-events/base/workflow", upgrade = FALSE)
devtools::install("~/projects/pecan/sipnet-events/models/sipnet", upgrade = FALSE)

devtools::load_all("~/projects/pecan/sipnet-events/modules/data.land")
devtools::load_all("~/projects/pecan/sipnet-events/models/sipnet")

site1 <- "~/projects/pecan/sipnet-events/modules/data.land/inst/events_fixtures/events_site1.json"
site1_multi <- "~/projects/pecan/sipnet-events/modules/data.land/inst/events_fixtures/events_site1_multi.json"
site_12 <- "~/projects/pecan/sipnet-events/modules/data.land/inst/events_fixtures/events_site1_site2.json"

binary <- path.expand("~/projects/pecan/_helpers/sipnet/sipnet.restart")
met <- path.expand("~/data/ERA5-site-processed/8773c58306e8d9a0/ERA5.1.2016-01-01.2023-12-31.clim")

start_date <- as.Date("2022-01-02")
end_date <- as.Date("2023-07-01")

site_id <- "8773c58306e8d9a0"

################################################################################
events_json <- site1

outdir <- normalizePath("_test/segments")
dir.create(outdir, showWarnings = FALSE)

settings <- PEcAn.settings::as.Settings(list(
  outdir = file.path(outdir, "out"),
  rundir = file.path(outdir, "run"),
  modeloutdir = file.path(outdir, "out"),
  pfts = list(list(
    name = "grassland",
    constants = list(num = 1)
  )),
  model = list(
    type = "SIPNET",
    binary = binary,
    revision = "v2"
  ),
  run = list(
    site = list(
      id = site_id,
      name = site_id,
      lat = 32.71585,
      lon = -115.47163
    ),
    start.date = start_date,
    end.date = end_date,
    inputs = list(met = list(path = met))
  ),
  host = list(
    name = "localhost"
  )
))

events <- jsonlite::fromJSON(events_json, simplifyVector = FALSE)
# TODO: Iterate over events
site_events_obj <- events[[1]]

site_id <- site_events_obj[["site_id"]]
site_events_list <- site_events_obj[["events"]]
site_events_common <- site_events_obj
site_events_common[["events"]] <- NULL

crop_cycles <- events_to_crop_cycle_starts(site1_multi)

# Empty example
# crop_cycles <- tibble::tibble(
#   site_id = character(0),
#   date = as.Date(NULL),
#   crop_code = character(0)
# )

# Get segments
segments <- tibble::tibble(
  start_date = c(start_date, crop_cycles[["date"]]),
  end_date = c(crop_cycles[["date"]] - 1, end_date)
) |>
  dplyr::mutate(segment_id = dplyr::row_number())

################################################################################

for (isegment in seq_len(nrow(segments))) {
  # isegment <- 1
  segment <- segments[isegment, ]
  segment_id <- sprintf("%03d", isegment)
  dstart <- segment[["start_date"]]
  dend <- segment[["end_date"]]

  segment_dir <- file.path(outdir, paste0("segment_", segment_id))
  if (dir.exists(segment_dir)) {
    unlink(segment_dir, recursive = TRUE)
  }
  dir.create(segment_dir, showWarnings = FALSE, recursive = TRUE)

  # Filter events to relevant dates
  events_sub <- site_events_list |>
    purrr::keep(~as.Date(.x[["date"]]) >= dstart) |>
    purrr::keep(~as.Date(.x[["date"]]) <= dend)

  # Segment-separated events file
  eventfile <- file.path(segment_dir, "events.json")
  segment_event_obj <- list(c(site_events_common, events = list(events_sub)) )
  jsonlite::write_json(segment_event_obj, eventfile, auto_unbox = TRUE)

  segment_eventfile <- PEcAn.SIPNET::write.events.SIPNET(eventfile, segment_dir)

  # Subset the met to only the dates in this segment. SIPNET does not respect
  # start/end date, only the dates in the .clim file.
  met_orig <- read.table(settings[[c("run", "inputs", "met", "path")]])
  met_segment <- met_orig |>
    # Create a date from the year + DOY
    dplyr::mutate(date = as.Date(paste0(V2, "-01-01")) + V3) |>
    dplyr::filter(date >= dstart, date <= dend) |>
    dplyr::select(-c("date"))
  met_segment_file <- file.path(segment_dir, "met.clim")
  write.table(
    met_segment,
    met_segment_file,
    quote = FALSE,
    sep = "\t",
    row.names = FALSE,
    col.names = FALSE
  )

  # Segment-specific settings
  segment_outdir <- file.path(segment_dir, "out")
  dir.create(segment_outdir, showWarnings = FALSE, recursive = TRUE)
  segment_rundir <- file.path(segment_dir, "run")
  dir.create(segment_rundir, showWarnings = FALSE, recursive = TRUE)
  segment_rundir_withid <- file.path(segment_rundir, segment_id)
  dir.create(segment_rundir_withid, showWarnings = FALSE, recursive = TRUE)

  segment_settings <- settings
  segment_settings[["outdir"]] <- segment_outdir
  segment_settings[["rundir"]] <- segment_rundir
  segment_settings[["modeloutdir"]] <- segment_rundir
  segment_settings[[c("run", "start.date")]] <- dstart
  segment_settings[[c("run", "end.date")]] <- dend
  segment_settings[[c("run", "inputs", "met", "path")]] <- met_segment_file
  segment_settings[[c("run", "inputs", "events")]] <- list(path = segment_eventfile)

  if (isegment > 1) {
    # For isegment > 1, we restart from the *previous* segment's restart.out
    segment_settings[[c("model", "restart_in")]] <- restart_out
  }
  # ...and now, define a new restart.out for *this* segment
  restart_out <- file.path(segment_dir, "restart.out")
  segment_settings[[c("model", "restart_out")]] <- restart_out

  # Write runs file
  writeLines(segment_id, file.path(segment_rundir, "runs.txt"))

  # TODO: Logic to get the trait values corresponding to the segment's PFT.
  # 1. Cross-reference crop_code against PFT
  # 2. Get traits from PFT posterior file.
  segment_traits <- list(list())

  config <- PEcAn.SIPNET::write.config.SIPNET(
    defaults = settings[["pfts"]],
    trait.values = segment_traits,
    settings = segment_settings,
    run.id = segment_id
  )

  runs <- PEcAn.workflow::start_model_runs(segment_settings, write = FALSE)
}

# TODO: Post processing. Combine all the segments together and return output in
# PEcAn standard.

[ashiklom]

Chris pointed out that a lot of this functionality already exists in write_restart.SIPNET. So I need to refactor this to use that instead.

[infotroph]

Here's an attempted hop test using the existing restart tools from PEcAn.SIPNET (and not yet calling Sipnet's binary restarts). This runs top to bottom for me, but is buggy -- for reasons I haven't yet tracked down, the ICs never get updated even though I can see their new values being passed in to write_restart. See s3://carb/tmp/hop_demo.zip for the full working directory I ran this in.

#!/usr/bin/env Rscript

library(tidyverse)
library(PEcAn.SIPNET)
PEcAn.logger::logger.setLevel("DEBUG")

settings <- PEcAn.settings::read.settings("output/settings.xml")

# Debug mode
settings <- settings[1:3]

if (!dir.exists(settings$outdir)) dir.create(settings$outdir, recursive = TRUE)

design <- PEcAn.uncertainty::generate_joint_ensemble_design(
  settings = settings[[1]],
  ensemble_size = settings$ensemble$size
)$X
write.csv(design, file = file.path(settings$outdir, "input_design.csv"))

settings <- PEcAn.workflow::runModule.run.write.configs(settings, input_design = design)

runs_mfst <- read.csv(file.path(settings$outdir, "runs_manifest.csv")) |>
  # TODO seems like the manifest ought to have a column for this...
  dplyr::mutate(ens_idx = str_extract(run_id, "\\d+") |> as.numeric())

# full-length run
PEcAn.workflow::runModule_start_model_runs(settings)

# Break each site's simulation into fixed 6-month chunks for a hop test.
# With an events file for production restart this would be something like
# PEcAn.data.land::events_to_crop_cycle_starts("events.json") |>
#   filter(
#       site_id == settings$run$site$id,
#       date >= settings$run$start.date,
#       date <= settings$run$end.date
#   )
#
# NB each end date is same as next start date, as expected by split_inputs
# (which compares with start >= date < stop)
make_breaks <- function(site_id, start, end, crop) {
  data.frame(
    site_id = site_id,
    crop = crop,
    start_date = seq(as.Date(start), as.Date(end), by = "6 months")
  ) |>
    dplyr::mutate(
      stop_date = dplyr::lead(start_date, default = as.Date(end)),
    )
}
site_breaks <- PEcAn.settings::papply(
  settings,
  \(s) make_breaks(s$run$site$id, s$run$start.date, s$run$end.date, s$run$site$site.pft$veg)
) |>
  purrr::list_rbind()

runs_mfst <- dplyr::left_join(runs_mfst, site_breaks) |>
  # TODO 1: pft_name is blank in current ensemble manifests. Change?
  # TODO 2: Need better crop -> pft mapping here. This one completely bogus.
  dplyr::mutate(
    pft_name = dplyr::case_when(
      crop %in% c("temperate.deciduous", "almonds", "oranges") ~ "temperate.deciduous",
      crop %in% c("grass", "switchgrass", "hay") ~ "grass",
      TRUE ~ "PFT_NOT_DEFINED"))
write.csv(runs_mfst, file.path(settings$outdir, "runs_manifest_for_restart.csv"))

# settings for whole run with many paths per input ->
# settings for one ensemble member with one path per input
# @param settings single-site settings object (not Multisettings)
# @param inputs named list of input indices (one row of the sample design)
subset_paths <- function(settings, path_nums) {
  for (input in names(path_nums)) {
    if (!(input %in% names(settings$run$inputs))) {
      next
    }
    path_idx <- path_nums[[input]]
    all_paths <- settings$run$inputs[[input]]$path
    if (path_idx > length(all_paths)) {
      stop("No path at input ", sQuote(input), " index ", path_idx)
    }
    settings$run$inputs[[input]]$path <- all_paths[[path_idx]]
  }
  settings
}

split_events.SIPNET <- function(inpath, outpath, start, stop) {
  evts <- readLines(inpath)
  dates <- evts |>
    grep(pattern = r"(^\d+\s+\d+)", value = TRUE) |>
    strptime("%Y%j")
  filename <- file.path(
    outpath,
    sub(".in$", paste0(start, "_", stop, ".in"), basename(inpath))
  )
  if(!dir.exists(outpath)) dir.create(outpath, recursive =TRUE)
  writeLines(
    evts[dates >= start & dates < stop],
    filename
  )

  filename
}

# TODO Maybe write.ensemble.configs with a restart argument can handle some of this?
restart_one_run <- function(site_id, run_id, ens_idx,
                            start_date, stop_date,
                            pft_name,
                            ...,
                            run_settings = settings, run_design = design) {
  input_row_nums <- run_design[ens_idx,]
  all_sites <- PEcAn.settings::papply(run_settings, \(s)s$run$site$id)
  settings <- run_settings[all_sites == site_id]
  settings <- subset_paths(settings, input_row_nums)
  # TODO met and event paths could (should?) be passed into inputs arg of
  #  write_restart instead of altering settings
  settings$run$inputs$met$path <- split_inputs.SIPNET(
    start.time = start_date,
    stop.time = stop_date,
    inputs = settings$run$inputs$met$path,
    outpath = file.path(settings$outdir, "met")
  )
  settings$run$inputs$events$path <- split_events.SIPNET(
    inpath = settings$run$inputs$events$path,
    outpath = file.path(settings$outdir, "events"),
    start = start_date,
    stop = stop_date
  )
  ens_samples <- PEcAn.utils::load_local(
    file.path(settings$outdir, paste0("ensemble.samples.", settings$ensemble$ensemble.id, ".RData"))
  )
  params <- ens_samples$ens.samples[pft_name] |>
    lapply(\(x) x[input_row_nums$param,, drop = FALSE])
  
  state <- read_restart.SIPNET(
    outdir = settings$modeloutdir,
    runid = run_id,
    stop.time = stop_date,
    settings = settings,
    var.names = c("NEE",
                  "Qle",
                  "TotSoilCarb",
                  "AbvGrndWood",
                  "abvGrndWoodFrac",
                  "coarseRootFrac",
                  "fineRootFrac"),
    params = params
  )
  PEcAn.logger::logger.debug("########################################################")
  str(state)
  PEcAn.logger::logger.debug("########################################################")
  dput(state$params$restart)
  #
  # TODO:
  #  - IC values not being adjusted between steps
  #  - write.config.SIPNET accepts but so far ignores a `restart` arg.
  #    Do we want to make write_restart use it? 
  #
  write_restart.SIPNET(
    outdir = settings$modeloutdir, # TODO probably only works if this == settings$host$outdir -- make write_restart use that internally?
    runid = run_id,
    start.time = start_date,
    stop.time = stop_date,
    settings = settings,
    new.state = state$X,
    new.params = state$params,
    inputs = NULL
  )
  PEcAn.settings::write.settings(
    settings,
    outputfile = paste0(run_id, "_", start_date, "_", stop_date, ".xml"))

  # TODO return settings too?
}

# Testing just one step
# head(runs_mfst)
# runs_mfst |>
#  slice_min(start_date) |> # TODO once this works, figure out how to apply to later slices
#  purrr::pwalk(restart_one_run)
# PEcAn.workflow::runModule_start_model_runs(settings)


# TODO this loop strategy is only efficient if steps are coordinated across sites --
# otherwise this will operate one site-step at a time without parallelizing across directories
for (step in unique(runs_mfst$start_date)) {
  PEcAn.logger::logger.info("Starting date:", step)
  runs_mfst |>
    filter(start_date == step) |>
    purrr::pwalk(restart_one_run)
  PEcAn.workflow::runModule_start_model_runs(settings)
}



# Plotting from raw sipnet output
# Would use netcdf, but don't see a convenient way to tell which datapoints
# are from which restart -- consider tracking this?
plt <- list.files(settings$outdir, pattern = "sipnet.*.out", full.names=TRUE, recursive = TRUE) |>
  setNames(nm = _) |>
  map(\(x)read.table(x, header = TRUE)) |>
  bind_rows(.id = "path") |>
  separate_wider_regex(
    path,
    c(".*ENS-", ens_num = "\\d+", "-", site_id = "[a-f0-9]+", "/sipnet\\.?", hop_date = "[\\d-]*", ".out")
  ) |>
  mutate(posix_time = PEcAn.SIPNET:::sipnet2datetime(year, day, time)) |>
  ggplot() +
  aes(posix_time, soil, color = hop_date) +
  facet_grid(ens_num~site_id) +
  geom_point() +
  theme_bw()
ggsave("hop.png", plot = plt, height = 8, width = 10)

[infotroph]

Note also that this needs the patches in PecanProject/pecan#3905