create_psd.Rmd

---
title: "Make popcycle PSD"
output: html_notebook
---

```{r}
renv::activate("~/git/popcycle")
library(popcycle)
library(dplyr)
library(glue)
```

Notebook configuration. Inputs, outputs, population, quantile, etc.
Make sure you've downloaded a local copy of the Seaflow instrument log spreadsheet as a tsv file.
Here it's called `SeaFlow\ instrument\ log\ -\ metadata-2022-03-12.tsv`.
It should have a column `cruise` for cruise name and TRUE/FALSE column `Zenodo` to indicate if the data has been published to Zenodo. This file is only used to build a list of cruises to process. Modify the code in the block below to change this list of cruises.

```{r}
cruise_dir <- "./data"
inst_log_tsv <- "SeaFlow\ instrument\ log\ -\ metadata-2022-03-12.tsv"
inst_log <- readr::read_delim(inst_log_tsv, delim = "\t")
cruises <- inst_log %>% filter(Curation == TRUE) %>% pull(cruise)
pop <- "prochloro"
quantile_ <- "2.5"
bins <- 50
global_Qc_range_flag <- TRUE
cores <- 6
per_cruise_out_dir <- "psd-data"
combined_outfile_prefix <- "combined"
```

Declare functions used in this notebook.

```{r}
#' Add zero count rows to a sparse one dimensional, one population, PSD data frame
#' 
#' Removes rows where coord_col is NA.
unsparsify <- function(psd, bins) {
  coord_col <- names(psd)[endsWith(names(psd), "_coord")]
  if (length(coord_col) != 1) {
    stop("expected only one '*_coord' column in psd")
  }

  psd_with_zeros <- psd %>%
    dplyr::filter(!is.na(.data[[coord_col]])) %>%
    dplyr::group_by(date) %>%
    dplyr::group_modify(function(x, y) {
      date_counts <- tibble::tibble(
        !!coord_col := seq(bins),
        n = 0L
      )
      date_counts[x[[coord_col]], "n"] <- x[["n"]]
      return(date_counts)
    }) %>%
    dplyr::ungroup()
  return(psd_with_zeros)
}

get_delta_v_int_Qc_range <- function(Qc_range, bins) {
  # The growth model takes 1 / (log2 distance between bins) as an integer. Calculate the end of
  # the grid range as the closest such integer that creates a grid that contains the true grid range
  # from the previous step
  delta_log2 <- diff(seq(from=log2(Qc_range[1]), to=log2(Qc_range[2]), length=(bins+1)))[1]
  delta_log2_inv <- 1 / delta_log2
  delta_log2_inv_int <- as.integer(delta_log2_inv)
  # Now original range should be
  # c(Qc_range_orig[1], 2**(log2(Qc_range_orig[1]) + (bins * (1 / delta_log2_inv))))
  # Expressing delta_log2_inv as an int gives a little headroom at the top end
  result <- list(
    Qc_range = c(
      Qc_range[1],
      2**(log2(Qc_range[1]) + (bins * (1 / delta_log2_inv_int)))
    ),
    delta_log2 = delta_log2,
    delta_log2_inv = delta_log2_inv,
    delta_log2_inv_int = delta_log2_inv_int
  )
  return(result)
}

get_Qc_ranges <- function(cruise_paths, quantile_, pop, bins) {
  for (i in seq(nrow(cruise_paths))) {
    row <- cruise_paths[i, ]
    message("getting Qc_range for cruise = ", row$cruise)
    message("pop = ", pop)
    message(row)
    
    sfl_tbl <- get_sfl_table(row$db)
    
    flagged_dates <- sfl_tbl %>%
      dplyr::filter(flag != 0) %>%
      dplyr::pull(date)
    
    pop_refrac <- read_refraction_csv() %>%
      dplyr::filter(cruise == row$cruise) %>%
      pull(pop)

    data_cols <- paste0("Qc_", pop_refrac)
    
    vct_files <- list.files(row$vct_dir, "\\.parquet$", full.names=T)
    
    ptm <- proc.time()
    # Qc_range_orig <- popcycle::get_vct_quantile_range(
    #  vct_files, data_cols, quantile_, c(0.01, 0.99), pop = pop, ignore_dates = flagged_dates
    # )
    message(glue::glue("data_cols={data_cols} quantile={quantile_} pop={pop}"))
    message(glue::glue("length(flagged_dates) == {length(flagged_dates)} out of {length(sfl_tbl$date)} possible dates"))
    Qc_range_orig <- popcycle::get_vct_range(
     vct_files, data_cols, quantile_, pop = pop, ignore_dates = flagged_dates, cores = 1
    )
    if (any(!is.finite(Qc_range_orig))) {
      message("No cells of interest for this cruise")
      next
    }
    #Qc_range_orig <- get_vct_range(vct_files, data_cols, quantile = quantile_, pop = pop, cores = cores)
    # The growth model takes 1 / (log2 distance between bins) as an integer. Calculate the end of
    # the grid range as the closest such integer that creates a grid that contains the true grid range
    # from the previous step
    newrange <- get_delta_v_int_Qc_range(Qc_range_orig, bins)
    cruise_paths[i, "Qc_range_orig_start"] <- Qc_range_orig[1]
    cruise_paths[i, "Qc_range_orig_end"] <- Qc_range_orig[2] * 2  # to make sure to have empty bins at the end of the PSD
    cruise_paths[i, "delta_log2"] <- newrange$delta_log2
    cruise_paths[i, "delta_log2_inv"] <- newrange$delta_log2_inv
    cruise_paths[i, "delta_log2_inv_int"] <- newrange$delta_log2_inv_int
    cruise_paths[i, "Qc_range_start"] <- newrange$Qc_range[1]
    cruise_paths[i, "Qc_range_end"] <- newrange$Qc_range[2]
    invisible(gc())
    
    message("population = ", pop)
    message("refractive index = ", pop_refrac)
    message("quantile = ", quantile_)
    message("data columns = ", data_cols)
    message(glue("Qc range = {Qc_range_orig[1]}, {Qc_range_orig[2]}"))
    message("delta_log2 = log2 distance between bins = ", newrange$delta_log2)
    message("delta_log2_inv = 1 / (log2 distance between bins) = ", newrange$delta_log2_inv)
    message(glue("Qc range with integer value for delta_log2_inv ({newrange$delta_log2_inv_int}) = {newrange$Qc_range[1]}, {newrange$Qc_range[2]}"))
    
    deltat <- proc.time() - ptm
    message("vct range in ", lubridate::duration(deltat[["elapsed"]]))
    message(row$cruise, " finished")
    message("")
  }

  return(cruise_paths)
}

create_model_data <- function(cruise_paths, bins, quantile_, pop, sparse = FALSE) {
  for (i in seq(nrow(cruise_paths))) {
    row <- cruise_paths[i, ]
    Qc_range <- c(row$Qc_range_start, row$Qc_range_end)
    message("processing cruise = ", row$cruise)
    message("bins = ", bins)
    message("sparse = ", sparse)
    message("pop = ", pop)
    message("Qc_range =", Qc_range)
    out <- row$output
    
    sfl_tbl <- get_sfl_table(row$db)
    
    flagged_dates <- sfl_tbl %>%
      dplyr::filter(flag != 0) %>%
      dplyr::pull(date)
    
    par <- sfl_tbl %>%
      dplyr::filter(flag == 0) %>%
      dplyr::select(date, par, lat, lon)
    
    # Correct raw PAR values
    par_calib <- popcycle::read_par_csv() %>%
      dplyr::filter(!is.na(correction), !is.infinite(correction), cruise == row$cruise)
    if (nrow(par_calib) == 1) {
      message("Applying PAR correction value " ,par_calib$correction)
      par$par <- par$par * par_calib$correction[1]
    } else {
      message("No PAR correction value found for this cruise")
    }
  
    refracs <- read_refraction_csv() %>%
      dplyr::filter(cruise == row$cruise) %>%
      select(-c(cruise))
    pop_refrac <- refracs[[pop]]
    data_cols <- paste0("Qc_", pop_refrac)
    
    vct_files <- list.files(row$vct_dir, "\\.parquet$", full.names=T)
    #vct_files <- head(vct_files, 72)
    
    if (!any(is.infinite(Qc_range)) && !any(is.na(Qc_range))) {
      ptm <- proc.time()
      
      dir.create(dirname(out), recursive = TRUE, showWarnings = FALSE)
      
      # Make the grid
      Qc_range <- c(row$Qc_range_start, row$Qc_range_end)
      grid <- popcycle::create_grid(bins, log_base=2, log_answers=FALSE, Qc_range = Qc_range)
      #grid <- popcycle::create_grid(bins, log_base=2, log_answers=FALSE, Qc_range = c(0.0135, 0.11777484))
      grid <- grid["Qc"]
      grid_df <- tibble::tibble(cruise=row$cruise, Qc=grid$Qc)
      
      # Create the distribution
      psd <- popcycle::create_PSD(
        vct_files, quantile_, refracs, grid, ignore_dates = flagged_dates, pop = pop, cores = cores
      )
      if (all(is.na(psd$date))) {
        # No data for this cruise
        message("no data for ", row$cruise)
        next
      }
      
      # Remove counts out of grid range (coord is NA)
      # Remove Qc_sum column
      psd <- psd %>%
        dplyr::filter(!is.na(Qc_coord)) %>%
        dplyr::select(-c(Qc_sum))
      hourly_psd <- popcycle::group_psd_by_time(psd, time_expr = "1 hours")
      psd <- tibble::as_tibble(psd) # group_psd_by_time may convert psd to data.table
      
      # Remove population label since we only have only population
      psd$pop <- NULL
      hourly_psd$pop <- NULL
      # Add zero counts if necessary
      if (!sparse) {
        message("adding zero count rows")
        psd <- unsparsify(psd, bins)
        hourly_psd <- unsparsify(hourly_psd, bins)
      }
      # Add cruise column
      psd <- psd %>% dplyr::mutate(cruise = row$cruise, .before = 1)
      hourly_psd <- hourly_psd %>% dplyr::mutate(cruise = row$cruise, .before = 1)
      arrow::write_parquet(grid_df, paste0(out, ".psd-grid.parquet"))
      arrow::write_parquet(psd, paste0(out, ".psd-full.parquet"))
      arrow::write_parquet(hourly_psd, paste0(out, ".psd-hourly.parquet"))
      invisible(gc())
      
      deltat <- proc.time() - ptm
      message("Full PSD dim = ", stringr::str_flatten(dim(psd), " "), ", MB = ", object.size(psd) / 2**20)
      message("Hourly PSD dim = ", stringr::str_flatten(dim(hourly_psd), " "), ", MB = ", object.size(hourly_psd) / 2**20)
      message("psd in ", lubridate::duration(deltat[["elapsed"]]))
      
      # Only keep PAR dates that are in PSD
      # Average by hour
      par <- par %>%
        dplyr::filter(date %in% unique(psd$date))
      hourly_par <- par %>%
        dplyr::group_by(date = lubridate::floor_date(date, "hour")) %>%
        dplyr::summarise(par = mean(par, na.rm = T), lat = mean(lat, na.rm = T), lon = mean(lon, na.rm = T))
      # Add cruise column
      par <- par %>% dplyr::mutate(cruise = row$cruise, .before = 1)
      hourly_par <- hourly_par %>% dplyr::mutate(cruise = row$cruise, .before = 1)
      arrow::write_parquet(par, paste0(out, ".par-full.parquet"))
      arrow::write_parquet(hourly_par, paste0(out, ".par-hourly.parquet"))
    } else {
      message("psd range has infinite values")
    }
    message(row$cruise, " finished")
    message("")
  }
}
```

Prepare a dataframe of configuration for data processing.

```{r}
getwd()

cruise_paths <- tibble::tibble(
  cruise = cruises,
  vct_dir = file.path(cruise_dir, cruises, paste0(cruises, "_vct")),
  #vct_dir = "vct_test",
  db = file.path(cruise_dir, cruises, paste0(cruises, ".db")),
  output = file.path(per_cruise_out_dir, cruises)
)

t0 <- proc.time()
# Add Qc Range information to cruise_paths
cruise_paths <- get_Qc_ranges(cruise_paths, quantile_, pop, bins)
deltat <- proc.time() - t0
message("Calculated Qc_range for all  data in ", lubridate::duration(deltat[["elapsed"]]))

cruise_paths_orig <- cruise_paths  # keep a copy before altering

# Set a global Qc range for all cruises
if (global_Qc_range_flag) {
  global_Qc_range_orig <- c(
    min(cruise_paths$Qc_range_orig_start, na.rm=T),
    max(cruise_paths$Qc_range_orig_end, na.rm=T)
  )
  global_Qc_range_data <- get_delta_v_int_Qc_range(global_Qc_range_orig, bins)
  print(global_Qc_range_data)
  
  # To use the global range for all cruises, overriding their per-cruise
  # Qc ranges.
  cruise_paths$Qc_range_start <- global_Qc_range_data$Qc_range[1]
  cruise_paths$Qc_range_end <- global_Qc_range_data$Qc_range[2]
}
```

Generate model input data.

```{r}
t0 <- proc.time()
create_model_data(cruise_paths, bins, quantile_, pop, sparse = FALSE)
deltat <- proc.time() - t0
message("Calculated PSD for all  data in ", lubridate::duration(deltat[["elapsed"]]))
```

Create single combined Parquet files for each type of data: PSD, PSD grid, and PAR.
These files will be saved in the current working directory with a file name prefix of `combined`.

```{r}
par_full_files <- list.files(per_cruise_out_dir, pattern = "*\\.par-fully\\.parquet", full.names = TRUE)
par_hourly_files <- list.files(per_cruise_out_dir, pattern = "*\\.par-hourly\\.parquet", full.names = TRUE)
psd_full_files <- list.files(per_cruise_out_dir, pattern = "*\\.psd-full\\.parquet", full.names = TRUE)
psd_hourly_files <- list.files(per_cruise_out_dir, pattern = "*\\.psd-hourly\\.parquet", full.names = TRUE)
psd_grid_files <- list.files(per_cruise_out_dir, pattern = "*\\.psd-grid\\.parquet", full.names = TRUE)

combined_par_hourly <- dplyr::bind_rows(lapply(par_hourly_files, function(f) {
  arrow::read_parquet(f)
}))
arrow::write_parquet(combined_par_hourly, paste0(combined_outfile_prefix, ".par-hourly.parquet"))

combined_psd_full <- dplyr::bind_rows(lapply(psd_full_files, function(f) {
  arrow::read_parquet(f)
}))
arrow::write_parquet(combined_psd_full, paste0(combined_outfile_prefix, ".psd-full.parquet"))

combined_psd_hourly <- dplyr::bind_rows(lapply(psd_hourly_files, function(f) {
  arrow::read_parquet(f)
}))
arrow::write_parquet(combined_psd_hourly, paste0(combined_outfile_prefix, ".psd-hourly.parquet"))

combined_grid <- dplyr::bind_rows(lapply(psd_grid_files, function(f) {
  arrow::read_parquet(f)
}))
arrow::write_parquet(combined_grid, paste0(combined_outfile_prefix, ".psd-grid.parquet"))
```