eliocamp / metR

Tools for Easier Analysis of Meteorological Fields
https://eliocamp.github.io/metR/
139 stars 22 forks source link

Stream Function and Velocity Potential #178

Open pascaloettli opened 1 year ago

pascaloettli commented 1 year ago

Hello,

Not an issue but a request.

I am wondering if there is a way to add the calculation of both the stream function and the velocity potential, either at the global scale or the regional scale. Or maybe it already exists and I missed it.

eliocamp commented 1 year ago

For this paper I created a function that computes the stream function from the vorticity by solving the Poisson equation $\nabla^2 \psi = - \omega$. The problem is that it only works with global fields and it uses some crusty-old Fortran I found online and don't really understand.

If you can point me to any existing code that does this I'd be glad to port it, though.

pascaloettli commented 1 year ago

Thank you,

I have checked the linked paper and saw you are using the hwsssp function part of the fishpack library. If I am not mistaken, this is used by the OpenGrADS extension fish.gex.

As you mentioned, this only works on global fields, similarly to cdo's dv2ps, which requires, in addition, a spectral transformation, as well as the spherepack 3.0 library.

The problem is that calculating SF and VP on regional/limited area is non-trivial, as said in this NCL thread. In the same thread, there is a link to a 2006 paper, but I couldn't find any associated script.

There is also a link to a Matlab script. I know this script, but I am not sure if it is reliable.

Finally, I tried to use your script on another dataset, but it gave inconsistent results. I'll try to provide a reproducible example later.

pascaloettli commented 1 year ago

Here is the reproducible example, using the following data: gfs.t00z.atmanl.360x180.nc psi_opengrads.nc psi_cdo.nc

libs <- c("metR", "data.table", "ggplot2", "magrittr", "shceof", "ggthemes")
lapply(libs, library, character.only = TRUE)
#> 
#> Attaching package: 'shceof'
#> The following object is masked from 'package:metR':
#> 
#>     WaveFlux
#> [[1]]
#> [1] "metR"      "stats"     "graphics"  "grDevices" "utils"     "datasets" 
#> [7] "methods"   "base"     
#> 
#> [[2]]
#> [1] "data.table" "metR"       "stats"      "graphics"   "grDevices" 
#> [6] "utils"      "datasets"   "methods"    "base"      
#> 
#> [[3]]
#>  [1] "ggplot2"    "data.table" "metR"       "stats"      "graphics"  
#>  [6] "grDevices"  "utils"      "datasets"   "methods"    "base"      
#> 
#> [[4]]
#>  [1] "magrittr"   "ggplot2"    "data.table" "metR"       "stats"     
#>  [6] "graphics"   "grDevices"  "utils"      "datasets"   "methods"   
#> [11] "base"      
#> 
#> [[5]]
#>  [1] "shceof"     "magrittr"   "ggplot2"    "data.table" "metR"      
#>  [6] "stats"      "graphics"   "grDevices"  "utils"      "datasets"  
#> [11] "methods"    "base"      
#> 
#> [[6]]
#>  [1] "ggthemes"   "shceof"     "magrittr"   "ggplot2"    "data.table"
#>  [6] "metR"       "stats"      "graphics"   "grDevices"  "utils"     
#> [11] "datasets"   "methods"    "base"
Sys.setenv(TZ = "GMT")

psi <- 
  ReadNetCDF(file = "~/Downloads/gfs.t00z.atmanl.360x180.nc") %>%
  normalise_coords() %>% 
  setorder(-lat) %>% 
  .[, lev := 200] %>% 
  rbind(., .[ lon == 0 ][ , lon := 360][]) %>% 
  .[, f := coriolis(lat) ] %>% 
  .[, c("du.dx","du.dy","dv.dx","dv.dy") := Derivate(ugrd + vgrd ~ lon + lat, cyclical = c(TRUE, FALSE), fill = TRUE, sphere = TRUE) ] %>% 
  .[, vort := (-du.dy + dv.dx) ] %>% 
  .[, divg := du.dx + dv.dy ] %>% 
  .[, psi := solve_poisson(vort, lon, lat)$value, by = .(lev, time)] %>%
  .[, psi := psi - mean(psi), by = .(time, lev)]

ggplot(psi, aes(lon,lat)) +
  labs(title = bquote("Vorticity from "*bold(metR*"{"*Derivate*"}")~"function")) +
  geom_contour_fill(aes(z=vort*1e5), binwidth = 5e-05*1e5) +
  # geom_vector(data = psi[lat > -88 & lat < 88 ], aes(dx = dlon(ugrd, lat), dy = dlat(vgrd)), skip.x = 2, skip.y = 2) +
  geom_vector(data = psi[lat > -88 & lat < 88 ], aes(dx = ugrd, dy = vgrd), skip.x = 3, skip.y = 2) +
  scale_mag(max_size = 0.8, guide = "none") +
  scale_fill_viridis_c(direction = 1, option = "H", name = bquote("["*s^-1~x*1^5*"]")) +
  scale_x_longitude() +
  scale_y_latitude() +
  coord_equal(xlim = c(0,360), ylim = c(-90,90), expand = FALSE) +
  theme_light(base_size = 20, base_family = "Albert Sans") +
  theme(legend.key.height = unit(0.5, 'cm'), legend.key.width = unit(5, 'cm'), legend.position="bottom")


ggplot(psi, aes(lon,lat)) +
  labs(title = bquote("Stream function from "*bold(shceof*"{"*solve_poisson*"}")~"function")) +
  geom_contour_fill(aes(z=psi*1e06), binwidth = 0.2) +
  scale_fill_viridis_c(direction = 1, option = "H", name = bquote("["*x*1^6*"]")) +
  geom_vector(data = psi[lat > -88 & lat < 88 ], aes(dx = ugrd, dy = vgrd), skip.x = 3, skip.y = 2) +
  scale_mag(max_size = 0.8, guide = "none") +
  scale_x_longitude() +
  scale_y_latitude() +
  coord_equal(xlim = c(0,360), ylim = c(-90,90), expand = FALSE) +
  theme_light(base_size = 20, base_family = "Albert Sans") +
  theme(legend.key.height = unit(0.5, 'cm'), legend.key.width = unit(5, 'cm'), legend.position="bottom")


data <- data.table::data.table(forcing = psi$vort, lon = psi$lon, lat = psi$lat)
forcing <- metR:::.tidy2matrix(data, lat ~ lon, value.var = "forcing")
f <- forcing$matrix
colat <- (90 - forcing$rowdims$lat) * pi/180
lon <- forcing$coldims$lon * pi/180
M <- length(colat) - 1L
N <- length(lon) - 1L
TS <- as.single(min(colat))
TF <- as.single(max(colat))
TF_is_pi <- isTRUE(all.equal(as.single(pi), TF, check.attributes = FALSE))
TS_is_zero <- isTRUE(all.equal(0, TS, check.attributes = FALSE))
MBDCND <- 3L
BDTS <- as.single(rep(0, N + 1))
BDTF <- as.single(rep(0, N + 1))
PS <- as.single(min(lon))
PF <- as.single(max(lon))
NBDCND <- 0L
BDPS <- 1L
BDPF <- 1L
ELMBDA <- as.single(0)
f <- as.single(f)
IDIMF <- M + 1L
W <- rep(0, 4 * (N + 1) + (16 + (ceiling(log2(N + 1)))) * (M + 1))
PERTRB <- 1L
IERROR <- 0L
result <- .Fortran("hwsssp_", TS, TF, M, MBDCND, BDTS, BDTF, 
                   PS, PF, N, NBDCND, BDPS, BDPF, ELMBDA, F = f, IDIMF, 
                   PERTRB, IERROR = IERROR, W)
if (result$IERROR != 0) {
  warning("error ", result$IERROR)
  return(NULL)
}else{
  forcing$matrix[, ] <- result[["F"]]
  OUT <- data.table::CJ(lon = forcing$coldims$lon, lat = forcing$rowdims$lat, sorted = FALSE)[, `:=`(value, c(forcing$matrix))][] %>% 
    setorder(-lat)

  ggplot(OUT, aes(lon,lat)) +
    labs(title = bquote("Stream function from the code inside "*bold(shceof*"{"*solve_poisson*"}")~"function")) +
    geom_contour_fill(aes(z=value*1e06), binwidth = 0.2) +
    scale_fill_viridis_c(direction = 1, option = "H", name = bquote("["*x*1^6*"]")) +
    geom_vector(data = psi[lat > -88 & lat < 88 ], aes(dx = ugrd, dy = vgrd), skip.x = 3, skip.y = 2) +
    scale_mag(max_size = 0.8, guide = "none") +
    scale_x_longitude() +
    scale_y_latitude() +
    coord_equal(xlim = c(0,360), ylim = c(-90,90), expand = FALSE) +
    theme_light(base_size = 20, base_family = "Albert Sans") +
    theme(legend.key.height = unit(0.5, 'cm'), legend.key.width = unit(5, 'cm'), legend.position="bottom")
}


# Stream function calculated by OpenGrADS
psi_opengrads <- ReadNetCDF("/home/oettli/Downloads/psi_opengrads.nc")

ggplot(psi_opengrads, aes(lon,lat)) +
  labs(title = bquote("Stream Function from OpenGrADS "*bold(fish_vor))) +
  geom_contour_fill(aes(z=psi*1e-06), binwidth = 10) +
  scale_fill_viridis_c(direction = 1, option = "H", name = bquote("["*x*1^-6*"]")) +
  geom_vector(data = psi[lat > -88 & lat < 88 ], aes(dx = ugrd, dy = vgrd), skip.x = 3, skip.y = 2) +
  scale_mag(max_size = 0.8, guide = "none") +
  scale_x_longitude() +
  scale_y_latitude() +
  coord_equal(xlim = c(0,360), ylim = c(-90,90), expand = FALSE) +
  theme(legend.key.height = unit(2, 'cm'), legend.key.width = unit(1, 'cm'), legend.position="bottom") +
  theme_light(base_size = 20, base_family = "Albert Sans") +
  theme(legend.key.height = unit(0.5, 'cm'), legend.key.width = unit(5, 'cm'), legend.position="bottom")


# Stream function calculated by cdo
psi_cdo <- ReadNetCDF("/home/oettli/Downloads/psi_cdo.nc", "stream") %>% rbind(., .[ lon == 0 ][ , lon := 360][])

ggplot(psi_cdo, aes(lon,lat)) +
  labs(title = bquote("Stream Function from cdo "*bold(dv2ps))) +
  geom_contour_fill(aes(z=stream*1e-06), binwidth = 10) +
  scale_fill_viridis_c(direction = 1, option = "H", name = bquote("["*x*1^-6*"]")) +
  geom_vector(data = psi[lat > -88 & lat < 88 ], aes(dx = ugrd, dy = vgrd), skip.x = 3, skip.y = 2) +
  scale_mag(max_size = 0.8, guide = "none") +
  scale_x_longitude() +
  scale_y_latitude() +
  coord_equal(xlim = c(0,360), ylim = c(-90,90), expand = FALSE) +
  theme(legend.key.height = unit(2, 'cm'), legend.key.width = unit(1, 'cm'), legend.position="bottom") +
  theme_light(base_size = 20, base_family = "Albert Sans") +
  theme(legend.key.height = unit(0.5, 'cm'), legend.key.width = unit(5, 'cm'), legend.position="bottom")


ggplot(merge(psi_opengrads, psi_cdo), aes(lon,lat)) +
  labs(title = bquote("Stream Function (OpenGrADS - cdo)")) +
  geom_contour_fill(aes(z=(psi-stream)*1e-6), binwidth = 0.1) +
  scale_fill_viridis_c(direction = 1, option = "H", name = bquote("["*x*1^-6*"]")) +
  scale_x_longitude() +
  scale_y_latitude() +
  coord_equal(xlim = c(0,360), ylim = c(-90,90), expand = FALSE) +
  theme(legend.key.height = unit(2, 'cm'), legend.key.width = unit(1, 'cm'), legend.position="bottom") +
  theme_light(base_size = 20, base_family = "Albert Sans") +
  theme(legend.key.height = unit(0.5, 'cm'), legend.key.width = unit(5, 'cm'), legend.position="bottom")


ggplot(merge(OUT, psi_opengrads), aes(lon,lat)) +
  labs(title = bquote("Stream Function (shceof - OpenGrADS)")) +
  geom_contour_fill(aes(z=(value-psi)*1e-6), binwidth = 5) +
  scale_fill_viridis_c(direction = 1, option = "H", name = bquote("["*x*1^-6*"]")) +
  scale_x_longitude() +
  scale_y_latitude() +
  coord_equal(xlim = c(0,360), ylim = c(-90,90), expand = FALSE) +
  theme(legend.key.height = unit(2, 'cm'), legend.key.width = unit(1, 'cm'), legend.position="bottom") +
  theme_light(base_size = 20, base_family = "Albert Sans") +
  theme(legend.key.height = unit(0.5, 'cm'), legend.key.width = unit(5, 'cm'), legend.position="bottom")

Created on 2023-07-25 with reprex v2.0.2

eliocamp commented 9 months ago

The solve_poisson() function does not necessarily return the values in the same order, so you can't do psi := solve_poisson(..)$value. You need to return the whole data.table (it's not super great, but it is what worked for the paper).

I don't know about those comparisons, but for the paper I tested the function against "ground truth" by comparing psi from NCEP and psi derived from vorticity in NCEP. There is a constant factor difference that I computed and then added back, but otherwise the result fits very well:

library(metR)
library(ggplot2)
psi_file <- tempfile()
download.file("ftp://ftp2.psl.noaa.gov/Datasets/ncep.reanalysis.derived/sigma/psi.mon.mean.nc", 
              psi_file)
vor_file <- tempfile()
download.file("ftp://ftp2.psl.noaa.gov/Datasets/ncep.reanalysis.derived/sigma/vor.mon.mean.nc", 
              vor_file)

psi <- ReadNetCDF(psi_file, subset = list(level =  0.2101, 
                             time = c("1979-01-01", "1979-01-01"))) 

vor <- ReadNetCDF(vor_file, subset = list(level =  0.2101, 
                             time = c("1979-01-01", "1979-01-01")))

psi_derived <- vor[, shceof::solve_poisson(vor, lon, lat), by = .(level, time)]

ggplot(psi, aes(lon, lat)) +
    geom_contour_filled(aes(z = psi)) +
    labs(title = "Original psi")

ggplot(psi_derived, aes(lon, lat)) +
    geom_contour_filled(aes(z = value)) +
    labs(title = "psi derived from vor")

psi[psi_derived, on = c("lon", "lat")][, cor(value, psi)]
#> [1] 0.9999979

ggplot(psi[psi_derived, on = c("lon", "lat")], aes(value, psi)) +
    geom_point() +
    labs(x = "psi from solve_poisson", 
         y = "psi from NCEP")

Created on 2023-11-09 with reprex v2.0.2