# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:ohw20-proj-pyxpcm]
#     language: python
#     name: conda-env-ohw20-proj-pyxpcm-py
# ---

# ## Step 1 - Argo Data

# +
# load libraries
import numpy as np
import pyxpcm
from pyxpcm.models import pcm

import xarray as xr

from argopy import DataFetcher as ArgoDataFetcher

import cartopy.crs as ccrs 
import matplotlib.pyplot as plt
import cartopy.feature as cfeature

import gsw 


# -

def get_argo_region_data(llon,rlon,llat,ulat,depthmin,depthmax,time_in,time_f):
    
    """Function to get argo data for a given lat,lon box (using Argopy), 
       and return a 2D array collection of vertical profile for the given region
    
    Parameters
    ----------
    llon : int
        left longitude
    rlon : int
        right longtidue
    ulat : int
        upper latitude
    llat : int
        lower latitude
    time_in : str/datetime object
        the start time of desired range, formatted Y-m-d
    time_f : str/datetime object
        the end time of desired range, formatted Y-m-d
    
    
    Returns
    ---------
    xarray
        The result is a xarray of the vertical profile for the given range and region.
        
    
    """

    ds_points = ArgoDataFetcher(src='erddap').region([llon,rlon, llat,ulat, depthmin, depthmax,time_in,time_f]).to_xarray()
    ds_profiles = ds_points.argo.point2profile()
    return ds_profiles


def spliced_argo_region_data(llon,rlon,llat,ulat,depthmin,depthmax,time_start,time_end):
    """Function that gets the argo data for given latitude and longitude bounding box
       (using Argopy), and given start and end time range to return a 2D array collection of vertical
       profile for the given region and time frame
    
    Parameters
    ----------
    llon : int
        left longitude
    rlon : int
        right longtidue
    ulat : int
        upper latitude
    llat : int
        lower latitude
    time_in : str/datetime object
        the start time of desired range, formatted Y-m-d
    time_f : str/datetime object
        the end time of desired range, formatted Y-m-d
    
    
    Returns
    ---------
    xarray
        The result is a xarray of the vertical profile for the given range and region.
        
    
    """

    
    #step
    max_dt = timedelta(days = 10)
    
    if isinstance(time_start, str):
        time_start = datetime.strptime(time_start,"%Y-%m-%d")
    if isinstance(time_end, str):
        time_end = datetime.strptime(time_end,"%Y-%m-%d")
        
    if time_end - time_start <= max_dt:
        ds = get_argo_region_data(llon,rlon,llat,ulat,depthmin,depthmax,time_start,time_end)
        return ds
    else:
        early_end = time_start+max_dt
        ds = get_argo_region_data(llon,rlon,llat,ulat,depthmin,depthmax,time_start,early_end)
        print("Retrived data from " + str(time_start) + " to " + str(early_end) + ", retreived " + str(len(ds.N_PROF)) + " profiles")
        ds2 = spliced_argo_region_data(llon,rlon,llat,ulat,depthmin,depthmax, early_end,time_end)
        return xr.concat([ds,ds2],dim='N_PROF')


llon=-90;rlon=0
ulat=70;llat=0 
depthmin=0;depthmax=4000
# January 
time_in='2011-01'
time_f='2011-02'

# ## (i) Get data

ds = get_argo_region_data(llon,rlon,llat,ulat,depthmin,depthmax,time_in,time_f)
ds

ds.PRES

# ## Plot

# +
plt.figure(figsize=(15,8))
    
ax = plt.subplot(111, projection=ccrs.PlateCarree()
                 #Mercator()
                 )
ax.set_extent([llon,rlon,llat,ulat])
ax.add_feature(cfeature.LAND)
ax.add_feature(cfeature.COASTLINE)
sc = ax.scatter(ds['LONGITUDE'], ds['LATITUDE'], s=3)

plt.title('Floats selected')
plt.show()
# -

# ## (ii) Check QC
#
# According to Argo manual: 
# - QC flag values of 1, 2, 5, or 8 are GOOD data
# - QC flag values of 9 (missing) are NOT USED in the computation
# - All other QC flag values are BAD data 
#
# Page 57 in http://www.odip.org/documents/odip/downloads/20/argo-dm-user-manual.pdf

goodflags=[1,2,5,8]
badflags=[3,4,6,7,9]
QC=['POSITION_QC','PRES_QC','TEMP_QC','PSAL_QC','TIME_QC']
for q in QC:
    data=np.array(ds[q])
    for f in badflags:
        if np.any(data)==f: 
            print('Bad data in '+str(q)+' QC flag: '+str(f))

# ## (iii) Delayed X Real time Data

# DATA_MODE(PROF_N):
#
# R : real time data
#
# D : delayed mode data
#
# A : real time data with adjusted values

ds.DATA_MODE

ds2=ds.where(ds.DATA_MODE=='D',drop=True)
print('We selected '+str(len(ds2.N_PROF))+' from '+str(len(ds.N_PROF)))


ds2

#

# ## (iv) Vertical Interpolation
#
# ' We finally interpolated the data on a regular vertical grid with a 5 m resolution (the original resolution ranges from less than 10 m at the surface to 200 m at the bottom of the profile)' - maze et al. 2017

z=np.arange(0.,450,5.)
ds_interp = ds2.argo.interp_std_levels(z)
print(ds_interp)

# So, when we interpolate, the number of profiles decreases a lot, because only the profiles that have data in the interval we asked is kept!

print('The number of profiles changed from  '+str(len(ds2.N_PROF))+ ' to '+str(len(ds_interp.N_PROF)))


ds3=ds_interp


# ## (v) Pressure to Depth

p=np.array(ds3.PRES)
lat=np.array(ds3.LATITUDE)
z=np.ones_like(p)
nprof=np.array(ds3.N_PROF)
#z=gsw.z_from_p(p,lat)
for i in np.arange(0,len(nprof)):
    z[i,:]=gsw.z_from_p(p[i,:], lat[i])


p_interp=np.array(ds3.PRES_INTERPOLATED)
z_interp=gsw.z_from_p(p_interp, 35) # choosing a mid latitude of the basin

# ## (vi) Make a clean dataset 
#
# We could use ds.drop_vars() or ds.drop_dims() also

temp=np.array(ds3.TEMP)
sal=np.array(ds3.PSAL)
depth_var=z
depth=z_interp
lat=np.array(ds3.LATITUDE)
lon=np.array(ds3.LONGITUDE)
time=np.array(ds3.TIME)

da=xr.Dataset(data_vars={'TEMP':(('N_PROF','DEPTH'),temp),
                        'PSAL':(('N_PROF','DEPTH'),sal),
                        'TIME':(('N_PROF'),time),
                         'LATITUDE':(('N_PROF'),lat),
                         'LONGITUDE':(('N_PROF'),lon),
                         'DEPTH_PROF':(('N_PROF','DEPTH'),depth_var)
                        },
                         coords={'N_PROF':nprof,
                                'DEPTH':depth})

da

# ## Make model

z = np.arange(0.,-450,-10.) # depth array
pcm_features = {'temperature': z, 'salinity':z} #features that vary in function of depth
m = pcm(K=8, features=pcm_features) # create the 'basic' model
m

# ## Fit model

features_in_ds = {'temperature': 'TEMP', 'salinity': 'PSAL'}
features_zdim='DEPTH'
m.fit(da, features=features_in_ds, dim=features_zdim)
m

da['TEMP'].attrs['feature_name'] = 'temperature'
da['PSAL'].attrs['feature_name'] = 'salinity'
da['DEPTH'].attrs['axis'] = 'Z'
m.fit(da)

# ## Classify the data

da

m.predict(da, features=features_in_ds, dim=features_zdim,inplace=True)
da

m.predict_proba(da, features=features_in_ds, inplace=True)
da

for vname in ['TEMP', 'PSAL']:
    da = da.pyxpcm.quantile(m, q=[0.05, 0.5, 0.95], of=vname, outname=vname + '_Q', keep_attrs=True, inplace=True)
da

fig, ax = m.plot.quantile(da['TEMP_Q'], maxcols=4, figsize=(10, 8), sharey=True)


# +
ds=da
proj = ccrs.PlateCarree()
subplot_kw={'projection': proj, 'extent': np.array([-80,1,-1,66]) + np.array([-0.1,+0.1,-0.1,+0.1])}
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5,5), dpi=120, facecolor='w', edgecolor='k', subplot_kw=subplot_kw)

kmap = m.plot.cmap()
sc = ax.scatter(ds['LONGITUDE'], ds['LATITUDE'], s=3, c=ds['PCM_LABELS'], cmap=kmap, transform=proj, vmin=0, vmax=m.K)
cl = m.plot.colorbar(ax=ax)

gl = m.plot.latlongrid(ax, dx=10)
ax.add_feature(cfeature.LAND)
ax.add_feature(cfeature.COASTLINE)
ax.set_title('LABELS of the training set')
plt.show()
# -



StarcoderdataJupyterScriptsDedupFiltered