#!/usr/bin/env python
# coding: utf-8

# # NOAA UFS Replay on Google Cloud Storage (GCS)

# We have put the atmosphere and ocean model components from the UFS Replay on GCS, which includes a model
# run natively at 1/4 degree for ~30 years and at 1 degree for ~5.5 years. 
# Additionally, for both components we have datasets that are spatially subsampled from 1/4 degree to
# approximately 1 degree by selecting every 4th point in latitude and longitude.
# 
# The atmosphere dataset is at 3 hour frequency, and ocean dataset is at 6 hour frequency.
# The time stamps selected correspond to the "analysis" portion of the Replay, coinciding with the IAU.
# While the atmosphere fields are snapshots, the ocean fields are a 6 hour temporal average over the IAU period.
# 
# The ocean fields were conservatively regridded from the native Tripolar grid to the same Gaussian grid as the atmosphere model output.
# The atmosphere fields were regridded from the cubed sphere FV3 grid to Gaussian grid.
# 
# Note that there are many land surface and some ice fields in the "atmosphere" (fv3) dataset.

# ## Locations on GCS
# 
# One can browse the datasets here: https://console.cloud.google.com/storage/browser/noaa-ufs-gefsv13replay
# 
# ### Atmosphere
# 
# - 1/4 degree: `gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree/03h-freq/zarr/fv3.zarr`
# - 1/4 degree subsampled: `gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree-subsampled/03h-freq/zarr/fv3.zarr`
# - 1 degree: `gs://noaa-ufs-gefsv13replay/ufs-hr1/1.00-degree/03h-freq/zarr/fv3.zarr`
# 
# ### Ocean
# 
# - 1/4 degree: `gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree/06h-freq/zarr/mom6.zarr`
# - 1/4 degree subsampled: `gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree-subsampled/06h-freq/zarr/mom6.zarr`
# - 1 degree: `gs://noaa-ufs-gefsv13replay/ufs-hr1/1.00-degree/06h-freq/zarr/mom6.zarr`

# <div class="alert alert-block alert-info">
# 
#  <header>
#     <h3>⚠️ Warning about the Atmosphere Vertical Coordinate ⚠️</h1>
#   </header>
# 
# The atmosphere data is on the native model vertical coordinate.
# Despite the fact that the coordinate is named `pfull` and has units hPa, the vertical coordinate is **not** 
# pressure levels.
# The coordinate refers to a "reference" pressure level, if the surface pressure at this particular
# point in time and horizontal space were 1000 hPa. 
# 
# The python package <a href="https://ufs2arco.readthedocs.io/en/latest/">ufs2arco</a> can be used to interpolate from the native FV3 vertical coordinate to pressure levels.
# Please see <a href="https://ufs2arco.readthedocs.io/en/latest/example_pressure_interpolation.html">this documentation page</a> for example usage.
# </div>

# ## 1/4 Degree Data
# 
# 30 years (1994 - 2023) of 1/4 degree at 3 hour frequency for atmosphere fields and 6 hour frequency for ocean fields

# In[1]:


import xarray as xr


# In[2]:


atm = xr.open_zarr(
    "gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree/03h-freq/zarr/fv3.zarr",
    storage_options={"token": "anon"},
)


# In[3]:


atm


# In[4]:


ocn = xr.open_zarr(
    "gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree/06h-freq/zarr/mom6.zarr",
    storage_options={"token": "anon"},
)


# In[5]:


ocn


# Data and chunksizes are as follows

# 3D Atmosphere fields

# In[6]:


print("chunks = ", atm.tmp.encoding["preferred_chunks"])
atm.tmp.data


# 3D Ocean fields

# In[7]:


print("chunks = ", ocn.temp.encoding["preferred_chunks"])
ocn.temp.data


# 2D Atmosphere fields

# In[8]:


print("chunks = ", atm.tmp2m.encoding["preferred_chunks"])
atm.tmp2m.data


# 2D Ocean fields

# In[9]:


print("chunks = ", ocn.SSH.encoding["preferred_chunks"])
ocn.SSH.data


# Some example plots...

# In[10]:


get_ipython().run_cell_magic('time', '', 'atm.tmp2m.isel(time=slice(8)).plot(col="time", col_wrap=4)\n')


# In[11]:


get_ipython().run_cell_magic('time', '', 'ocn.SSH.isel(time=slice(8)).plot(col="time", col_wrap=4)\n')


# ## Subsampled 1/4 Degree Dataset
# 
# 30 years (1994 - 2023) of 1/4 degree data, subsampled from the 1/4 degree data to ~1 degree resolution
# by selecting every 4th point in latitude/longitude (dimensions grid_yt/grid_xt in the atmosphere 
# dataset)

# In[12]:


atm = xr.open_zarr(
    "gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree-subsampled/03h-freq/zarr/fv3.zarr",
    storage_options={"token": "anon"},
)


# In[13]:


atm


# In[14]:


ocn = xr.open_zarr(
    "gs://noaa-ufs-gefsv13replay/ufs-hr1/0.25-degree-subsampled/06h-freq/zarr/mom6.zarr",
    storage_options={"token": "anon"},
)


# In[15]:


ocn


# Data and chunksizes are as follows

# 3D Atmosphere fields

# In[16]:


print("chunks = ", atm.tmp.encoding["preferred_chunks"])
atm.tmp.data


# 3D Ocean fields

# In[17]:


print("chunks = ", ocn.temp.encoding["preferred_chunks"])
ocn.temp.data


# 2D Atmosphere fields

# In[18]:


print("chunks = ", atm.tmp2m.encoding["preferred_chunks"])
atm.tmp2m.data


# 2D Ocean fields

# In[19]:


print("chunks = ", ocn.SSH.encoding["preferred_chunks"])
ocn.SSH.data


# Some example plots...

# In[20]:


get_ipython().run_cell_magic('time', '', 'atm.tmp2m.isel(time=slice(8)).plot(col="time", col_wrap=4)\n')


# In[21]:


get_ipython().run_cell_magic('time', '', 'ocn.SSH.isel(time=slice(8)).plot(col="time", col_wrap=4)\n')


# ## 1 Degree Data
# 
# 5.5 years (1994 - 1999) of 1 degree at 3 hour frequency for atmosphere fields and 6 hour frequency for ocean fields

# In[22]:


atm = xr.open_zarr(
    "gs://noaa-ufs-gefsv13replay/ufs-hr1/1.00-degree/03h-freq/zarr/fv3.zarr",
    storage_options={"token": "anon"},
)


# In[23]:


atm


# In[24]:


ocn = xr.open_zarr(
    "gs://noaa-ufs-gefsv13replay/ufs-hr1/1.00-degree/06h-freq/zarr/mom6.zarr",
    storage_options={"token": "anon"},
)


# In[25]:


ocn


# Data and chunksizes are as follows

# 3D Atmosphere fields

# In[26]:


print("chunks = ", atm.tmp.encoding["preferred_chunks"])
atm.tmp.data


# 3D Ocean fields

# In[27]:


print("chunks = ", ocn.temp.encoding["preferred_chunks"])
ocn.temp.data


# 2D Atmosphere fields

# In[28]:


print("chunks = ", atm.tmp2m.encoding["preferred_chunks"])
atm.tmp2m.data


# 2D Ocean fields

# In[29]:


print("chunks = ", ocn.SSH.encoding["preferred_chunks"])
ocn.SSH.data


# Some example plots...

# In[30]:


get_ipython().run_cell_magic('time', '', 'atm.tmp2m.isel(time=slice(8)).plot(col="time", col_wrap=4)\n')


# In[31]:


get_ipython().run_cell_magic('time', '', 'ocn.SSH.isel(time=slice(8)).plot(col="time", col_wrap=4)\n')


# In[ ]: