Skip to content

Commit

Permalink
REF: generalize psd code
Browse files Browse the repository at this point in the history
  • Loading branch information
Konrad Mayer committed Sep 12, 2023
1 parent fa2a4ed commit c36136b
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 69 deletions.
31 changes: 24 additions & 7 deletions dev/postprocessing/plot_results.R
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,11 @@
#!/usr/bin/env Rscript

##################################################################
#Description : plot fields (era5, cerra, downscaled and
# difference) and PSD for individual timesteps
#Author : Konrad Mayer
##################################################################

library(here)
library(lubridate)
library(stars)
Expand All @@ -24,6 +32,14 @@ replace_na <- function(m, replacement = 0) {
m[is.na(m)] <- replacement
m
}

do_psd <- function(x) {
tmp <- replace_na(truncate_to_square(as.matrix(x[[1]])))
psd <- radial.psd(tmp, plot = FALSE, scaled = FALSE, normalized = FALSE)
max_dist <- sqrt(sum((ceiling(dim(tmp)/2)-1)^2))
mutate(psd, wavenumber = wavenumber / max_dist)
}

# main

per_lead_time <- function(lead_time) {
Expand Down Expand Up @@ -59,20 +75,21 @@ per_lead_time <- function(lead_time) {
theme_minimal()
),
list(map(
list(era5_regridded, cerra, downscaled) %>% map(~filter(.x %>%.[1, 3:238,6:160], time == eval_datetime)),
~{radial.psd(replace_na(truncate_to_square(as.matrix((st_apply(.x, c("x", "y"), mean)[[1]])))), plot = FALSE, normalized = FALSE, scaled = FALSE)}
list(era5_regridded, cerra, downscaled) %>% map(~filter(.x %>%.[1, 3:238,6:160], time == eval_datetime)[drop = TRUE]),
~{do_psd(st_apply(.x, c("x", "y"), mean))}
) %>%
setNames(c("era5", "cerra", "samos")) %>%
bind_rows(.id = "model") %>%
ggplot(aes(wavenumber, r_spectrum, lty = model)) +
geom_line() +
theme_minimal() +
scale_y_log10() + annotation_logticks() + scale_x_log10()
geom_line() +
theme_minimal(20) +
scale_y_log10() + annotation_logticks() + scale_x_log10(sec.axis = sec_axis(trans = ~ (.^-1), name = "wavelength")) +
labs(y = "Power")
))


wrap_plots(plts[1:4], ncol = 2) / plts[[5]] + plot_annotation(title = eval_datetime)
ggsave(here(glue("plt/SAMOS/{format(eval_datetime, '%Y%m%dT%H%M%S')}_testset.pdf")), height = 14, width = 10)
wrap_plots(plts[1:4], ncol = 2) | plts[[5]] + plot_annotation(title = eval_datetime)
ggsave(here(glue("plt/SAMOS/{format(eval_datetime, '%Y%m%dT%H%M%S')}_testset.pdf")), height = 6, width = 20)

}

Expand Down
120 changes: 58 additions & 62 deletions dev/verification/spatial_structures/compare_pds.R
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,10 @@
#!/usr/bin/env Rscript

##################################################################
#Description : power spectral density analysis
#Author : Konrad Mayer
##################################################################

library(stars)
library(purrr)
library(here)
Expand Down Expand Up @@ -33,102 +40,91 @@ replace_na <- function(m, replacement = 0) {
m
}

do_psd <- function(x) {
tmp <- replace_na(truncate_to_square(as.matrix(x[[1]])))
psd <- radial.psd(tmp, plot = FALSE, scaled = FALSE, normalized = FALSE)
max_dist <- sqrt(sum((ceiling(dim(tmp)/2)-1)^2))
mutate(psd, wavenumber = wavenumber / max_dist)
}

drop_units <- function(x) {
tryCatch(
units::drop_units(x),
error = function(e) {
return(x)
}
)
}
# main -------------------------------------------------------------------------

dothis <- function(lead_time) {
lead_time <- str_pad(lead_time, 2, pad = "0")

# load datasets (crop a bit as SAMOS baseline has some values NA around)
cerra <- read_stars(here(glue("dat/TESTING/PREPROCESSED/CERRA/t2m_cerra_{lead_time}.nc"))) %>%
setNames("cerra") %>%
.[1, 3:238,6:160]
downscaled <- read_stars(here(glue("dat/TESTING/SAMOS/postprocessed/samos-postprocessed_{lead_time}.nc"))) %>%
setNames("samos") %>%
.[1, 3:238,6:160]
era5 <- read_stars(here(glue("dat/TESTING/PREPROCESSED/ERA5_regridded/t2m_era5_{lead_time}.nc"))) %>%
setNames("era5") %>%
.[1, 3:238,6:160]

timesteps <- seq_len(dim(cerra)["time"])
datasets <- c(
cerra = "dat/TESTING/PREPROCESSED/CERRA/t2m_cerra_{lead_time}.nc",
downscaled = "dat/TESTING/SAMOS/postprocessed/samos-postprocessed_{lead_time}.nc",
era5 = "dat/TESTING/PREPROCESSED/ERA5_regridded/t2m_era5_{lead_time}.nc"
)
dat <- map2(datasets, names(datasets),
~setNames(drop_units(read_stars(here(glue(.x)))), .y)[1, 3:238,6:160]) # data gets truncated to get rid of NA values on two of the edges caused by CDO bilinear interpolation

timesteps <- seq_len(dim(dat[[1]])["time"])

#overall

psd_cerra_overall <- radial.psd(replace_na(truncate_to_square(as.matrix((st_apply(cerra, c("x", "y"), mean)[[1]])))), plot = FALSE, scaled = FALSE, normalized = FALSE)
psd_downscaled_overall <- radial.psd(replace_na(truncate_to_square(as.matrix((st_apply(downscaled, c("x", "y"), mean)[[1]])))), plot = FALSE, scaled = FALSE, normalized = FALSE)
psd_era5_overall <- radial.psd(replace_na(truncate_to_square(as.matrix((st_apply(era5, c("x", "y"), mean)[[1]])))), plot = FALSE, scaled = FALSE, normalized = FALSE)
psd_overall <- bind_rows(list(cerra = psd_cerra_overall, downscaled = psd_downscaled_overall, era5 = psd_era5_overall), .id = "model")

psd_overall <- map(dat,
~do_psd(st_apply(.x, c("x", "y"), mean))) %>%
bind_rows(.id = "model")

psd_overall %>%
ggplot(aes(wavenumber, r_spectrum, lty = model)) +
geom_line() +
theme_minimal() +
scale_y_log10() + scale_x_log10() + annotation_logticks()
ggsave(here(glue("plt/PSD/radialPSD_samos_leadtime{lead_time}_overall.pdf")))
theme_minimal(20) +
scale_y_log10() + scale_x_log10(sec.axis = sec_axis(trans = ~ (.^-1), name = "wavelength")) + annotation_logticks() +
labs(y = "Power")

ggsave(here(glue("plt/PSD/radialPSD_samos_leadtime{lead_time}_overall.pdf")), width = 10, height = 6)

# loop over timesteps and calculate mean psd
psd_cerra <- list()
psd_ds <- list()
psd_era5 <- list()

for (i in seq_along(timesteps)){
psd_cerra[[i]] <- radial.psd(replace_na(truncate_to_square(as.matrix(units::drop_units(cerra[[1]][,,timesteps[i]])))), plot = FALSE)#, scaled = FALSE, normalized = FALSE)
psd_ds[[i]] <- radial.psd(replace_na(truncate_to_square(as.matrix(downscaled[[1]][,,timesteps[i]]))), plot = FALSE)#, scaled = FALSE, normalized = FALSE)
psd_era5[[i]] <- radial.psd(replace_na(truncate_to_square(as.matrix(era5[[1]][,,timesteps[i]]))), plot = FALSE)#, scaled = FALSE, normalized = FALSE)
}

# loop over timesteps and calculate mean psd
psd_timesteps <- map(dat, function(ds) {
map(timesteps, function(ts) do_psd(ds[,,,ts, drop = TRUE]))
})
# common dataframe
plotdat <- bind_rows(list(
cerra = psd_cerra |>
setNames(timesteps) |>
bind_rows(.id = "timesteps"),
samos = psd_ds |>
setNames(timesteps) |>
bind_rows(.id = "timesteps"),
era5 = psd_era5 |>
setNames(timesteps) |>
bind_rows(.id = "timesteps")
), .id = "dataset")

plotdat <- psd_timesteps %>%
map(~bind_rows(setNames(.x, timesteps), .id = "timesteps")) %>%
bind_rows(.id = "dataset")

# plot
plotdat %>%
ggplot(aes(x = wavenumber, y = r_spectrum, color = dataset, group = dataset)) +
# stat_smooth(
# geom="ribbon",
# aes(ymax = after_stat(y) + 1,
# ymin = after_stat(y) - 1,
# fill = dataset),
# alpha = 0.1, linetype = 0
# ) +
# geom_smooth(se = FALSE)+
geom_interval() +
theme_minimal() +
scale_y_log10() + scale_x_log10() + annotation_logticks()
theme_minimal(20) +
scale_y_log10() + scale_x_log10(sec.axis = sec_axis(trans = ~ (.^-1), name = "wavelength")) + annotation_logticks() +
labs(y = "Power")

ggsave(here(glue("plt/PSD/radialPSD_samos_leadtime{lead_time}_distribution.pdf")))
ggsave(here(glue("plt/PSD/radialPSD_samos_leadtime{lead_time}_distribution.pdf")), width = 10, height = 6)

# aggregate to seasons and calculate PSD per season
seasons <- list(DJF = c(12, 1, 2), MAM = 3:5, JJA = 6:8, SON = 9:11)

per_season <- function(months) {
cerra_season <- st_apply(cerra %>% filter(lubridate::month(time) %in% months), c("x", "y"), mean)
downscaled_season <- st_apply(downscaled %>% filter(lubridate::month(time) %in% months), c("x", "y"), mean)
era5_season <- st_apply(era5 %>% filter(lubridate::month(time) %in% months), c("x", "y"), mean)
psd_cerra_season <- radial.psd(replace_na(truncate_to_square(as.matrix((cerra_season[[1]])))), plot = FALSE)#, scaled = FALSE, normalized = FALSE)
psd_downscaled_season <- radial.psd(replace_na(truncate_to_square(as.matrix((downscaled_season[[1]])))), plot = FALSE)#, scaled = FALSE, normalized = FALSE)
psd_era5_season <- radial.psd(replace_na(truncate_to_square(as.matrix((era5_season[[1]])))), plot = FALSE)#, scaled = FALSE, normalized = FALSE)
bind_rows(list(cerra = psd_cerra_season, downscaled = psd_downscaled_season, era5 = psd_era5_season), .id = "model")
map(dat, ~do_psd(st_apply(.x %>% filter(lubridate::month(time) %in% months), c("x", "y"), mean))) %>%
bind_rows(.id = "model")
}

psd_season <- map(seasons, per_season)

psd_season %>%
bind_rows(.id = "season") %>%
ggplot(aes(wavenumber, r_spectrum, color = season, lty = model)) +
geom_line() +
theme_minimal() +
scale_y_log10() + scale_x_log10() + annotation_logticks()
ggsave(here(glue("plt/PSD/radialPSD_samos_leadtime{lead_time}_season.pdf")))
theme_minimal(20) +
scale_y_log10() + scale_x_log10(sec.axis = sec_axis(trans = ~ (.^-1), name = "wavelength")) + annotation_logticks() +
labs(y = "Power")

ggsave(here(glue("plt/PSD/radialPSD_samos_leadtime{lead_time}_season.pdf")), width = 10, height = 6)
}

lead_times <- seq(0, 21, by = 3)
Expand Down

0 comments on commit c36136b

Please sign in to comment.