vic1309 / opendrift Goto Github PK

import xarray as xr
import matplotlib.pyplot as plt 
import trajan as ta
import numpy as np
import pandas as pd
import cf_xarray as _
import cftime
import pyproj
from datetime import datetime, timedelta

# Work on output aggregation
def open_mfdataset_overlap(url_base, time_series, timedim='time', variables=None):
    urls = [t.strftime(url_base) for t in time_series]
    time_step = time_series[1] - time_series[0]
    print('Opening individual URLs...')
    #datasets = [xr.open_dataset(u, chunks='auto').sel({timedim: slice(t, t+time_step-timedelta(seconds=1))})
    #            for u,t in zip(urls, time_series)]
    datasets = []
    
    i = 0
    for u, t in zip(urls, time_series):
        try:
            print(f'Opening {u}')
            
            #dF = xr.open_dataset(u, engine='netcdf4', decode_times=True, chunks={'time':1})

            #print(f'Converting dates')
            #testdate = cftime.num2pydate(dF.time[0], units='seconds since 1970-01-01 00:00:00')
            #makedate = pd.date_range(testdate, periods=67, freq='1H')

            #dF = dF.assign_coords({'time':('time', makedate.values, dF.time.attrs)})

            print(f'Slicing dataset..')
            T0 = time_series[-1]
            TF = T0 + timedelta(hours=24)

            d = xr.open_dataset(u, engine='netcdf4', decode_times=True, chunks={'time':1}).sel({timedim: slice(T0, TF-timedelta(seconds=1))})

            #d = dF.sel({timedim: slice(T0, TF-timedelta(seconds=1))})

            if variables is not None:
                try:
                    gm = d.cf['grid_mapping']
                    beps_crs = pyproj.CRS.from_cf(gm.attrs)
                    
                    d = d[variables]
                except Exception as e:
                    raise ValueError(f'Not containing relevant variables: {e}')

            # Adjust ensemble_member coordinate values
            size_ts = len(time_series)
            ens_members = np.arange(0, 48,1)
            members = ens_members.reshape((8, 6))
            
            d.coords['ensemble_member'] = np.array(members[i,:], dtype='int32')
            datasets.append(d)
            
            i += 1
        except:
            print(f'Could not open {u}')
                  
    print('Concatenating...')    

    '''
    ds = xr.concat(datasets, dim=timedim,
                   compat='override', combine_attrs='override', join='override', coords='all')
    '''
    ds = xr.concat(datasets, dim='ensemble_member', data_vars='minimal',
           compat='override', combine_attrs='override', join='override', coords='all')
    return ds, beps_crs

# Define original path
path_orig = '/lustre/storeB/project/fou/hi/barents_eps/eps/barents_eps_%Y%m%dT%HZ.nc'

datasets = []

poi = pd.date_range('2022-08-06', '2022-08-10')

for i in poi:

    year = i.year; month = i.month
    day = i.day; hour = i.hour
    
    t0 = datetime(year,month,day,hour)
    tf = t0 - timedelta(days=2) + timedelta(hours=6)
    
    # Define files to be opened
    times = pd.date_range(tf, t0, freq='6H')
  
    dM, beps_crs = open_mfdataset_overlap(path_orig, times, variables=['u', 'v', 'Uwind', 'Vwind'])

    # Append
    datasets.append(dM)


# Ensure that all steps have 48 ensemble members (8 runs x 6 members)
j = 0
for i in datasets:
    if len(i.ensemble_member.values) != 48:        
        # Work arround for missing ensemble members
        number_ens = len(i.ensemble_member.values)
        members = np.arange(0,48,1)
        
        # Choose a random number of EPS members to suplement
        sup = 48 - number_ens # 48 = 6 members times 8 lagged series
        X1 = np.random.randint(low=0, high=number_ens, size=(sup,))
        
        # Assign the 'new' members to the dataset
        teste = i.isel(ensemble_member= X1)
        teste.coords['ensemble_member'] = np.array(members[-sup:], dtype='int32')

        print('Merging data..')
        #datasets[j] = xr.merge([i,teste], compat='override', combine_attrs='override', join='override')
        datasets[j] = xr.concat([i,teste], dim='ensemble_member', data_vars='minimal',
           compat='override', combine_attrs='override', join='override', coords='minimal')

    j += 1
        
# Concatenate
print('Concatenating...')
dS = xr.concat(datasets, dim='time', data_vars='minimal',
           compat='override', combine_attrs='override', join='override', coords='minimal')

dS.ensemble_member.attrs['long_name'] = 'ensemble run member'
dS.ensemble_member.attrs['standard_name'] = 'realization'

ds = xr.open_dataset('/lustre/storeB/project/fou/hi/barents_eps/eps/barents_eps_20220908T06Z.nc')

gm = ds.cf['grid_mapping']
beps_crs = pyproj.CRS.from_cf(gm.attrs)

tes = gm.attrs['proj4']

##################################
### Start a loop for opendrift ###
##################################
import datetime
from opendrift.readers import reader_netCDF_CF_generic
from opendrift.models.oceandrift import OceanDrift


t0 = pd.to_datetime(dS.time[0].values)

# Start Opendrift
o = OceanDrift(loglevel=0) 

reader_eps = reader_netCDF_CF_generic.Reader(dS, proj4=tes)

# WAM Model
reader_wave = reader_netCDF_CF_generic.Reader('https://thredds.met.no/thredds/dodsC/sea/mywavewam4/mywavewam4_be')
o.add_reader(reader_wave, variables=['sea_surface_wave_stokes_drift_x_velocity', 'sea_surface_wave_stokes_drift_y_velocity'])

# Workaround to handle EPS products
o.add_reader(reader_eps, variables=['x_sea_water_velocity', 'y_sea_water_velocity', \
    'x_wind', 'y_wind'])

# Seeding location
seed_lon = 21.990812976496922
seed_lat = 73.14869327644097

seeding_time = 2022-08-06 00:00:00

o.seed_elements(lon=seed_lon, lat=seed_lat, radius=100, number=480,
    time=seeding_time, wind_drift_factor=0.02)

# ----- Define processes ----- # 
o.set_config('drift:stokes_drift',  True)
o.set_config('drift:vertical_mixing', False)
o.set_config('drift:vertical_advection', False)

# ----- Perform simulation ----- #
save_path = 'your_path/'
#o.run(duration=datetime.timedelta(days=2), time_step=3600, time_step_output=3600, outfile=save_path+'drifter_2_sl%i.nc'%(i))
o.run(duration=datetime.timedelta(days=1), time_step=3600, time_step_output=3600, outfile=save_path+'drifter_test.nc')

import trajan as ta

file_path = 'your_path/drifter_test.nc'

tf = xr.open_dataset(file_path)

number_ens = 48
ensemble = np.remainder(np.arange(tf.dims['trajectory']), number_ens) + 1 # Ensemble members

for i in range(1,49,1):
    index = np.where(ensemble==i)
    speed = tf.isel(trajectory=index[0]).traj.speed()

[0.24844839, 0.28512105, 0.32085344, 0.34196123, 0.36232346,
        0.35882047, 0.32987076, 0.29624984, 0.26153788, 0.24456808,
        0.24667085, 0.26559073, 0.28721333, 0.31005155, 0.31487129,
        0.3104002 , 0.30435823, 0.30027105, 0.28364826, 0.26265018,
        0.23458891, 0.201583  , 0.16002346, *0.11709443*, *0.11709443*]