copernicus_arc_simulation.py

""" Simulation script for NECCTON simulations using Parcels and PlasticParcels.
    This script runs a simulation for plastic particles in the North West European Continental Shelf and Arctic Ocean.
    This script should be run on an HPC cluster, using the simulation.sh submission script.
    The submission script takes 3 arguments, which are passed to this script:
        --startrelease: The first date of particle release (format: YYYY-MM-DD)
        --endrelease: The last date of particle release (format: YYYY-MM-DD)
        --endsimulation: The end date of the simulation (format: YYYY-MM-DD)
    Example usage: sbatch simulation.sh --startrelease 2020-01-01 --endrelease 2020-02-01 --endsimulation 2021-01-01
"""

## Library imports
# Data handling
import numpy as np
import pandas as pd

from argparse import ArgumentParser

# Lagrangian analysis
import os
os.chdir('/storage/home/denes001/Projects/PlasticParcels/')
import plasticparcels as pp
import parcels

from datetime import date, timedelta

# Import custom functions
folder_with_functions = '/nethome/denes001/Projects/NECCTONsimulations/'
os.chdir(folder_with_functions)

from functions import create_full_fieldset_copernicus, add_additional_fields_copernicus, create_particleset_from_release_files, select_files
from kernels import ARC_kernel_Copernicus, PolyTEOS10_bsq, BiofoulingCopernicus, StokesDriftCopernicus, checkThroughSurfaceCopernicus, stopExecution, reflectAtSurfaceCopernicus, reflectAtBathymetry, checkErrorThroughSurface_2DAdvectionRK2, AdvectionRK2, unbeachingBySamplingAfterwards, belowLatitude, depth_delete

# Parse the arguments
p = ArgumentParser(description="""NECCTON Simulation using Parcels""")
p.add_argument('-startrelease', '--startrelease', default='2020-01-01', help='First release date of the particles')
p.add_argument('-endrelease', '--endrelease', default='2020-02-01', help='Last release date of the particles')
p.add_argument('-endsimulation', '--endsimulation', default='2021-01-01', help='Time the simulation ends')
parsed_args = p.parse_args()

# Simulation date settings
args = {'startrelease': parsed_args.startrelease+'T00:00:00.000000000',
        'endrelease': parsed_args.endrelease+'T00:00:00.000000000',
        'endsimulation': parsed_args.endsimulation+'T00:00:00.000000000',
        'outputdt': 1,              # Daily output [d]
        'dt': 20}                   # Timestep choice in minutes

starttime = np.datetime64(args['startrelease'])
endtime = np.datetime64(args['endsimulation'])
runtime = timedelta(days=int((endtime-starttime).astype('timedelta64[D]').astype('int')))

# Create start date of simulation
start_year, start_month, start_day = [int(dateportion) for dateportion in args['startrelease'][:10].split('-')]
starting_time = date(start_year, start_month, start_day)

# Location to store output data, and the location containing release information
output_dir = '/storage/shared/oceanparcels/output_data/data_Michael/NECCTONsimulations/data/copernicus_simulations/'
release_folder = '/nethome/denes001/Projects/NECCTONsimulations/release_files_copernicus_ARC/' #TODO

# Create fieldset
settings = pp.utils.load_settings('NECCTON_copernicus_downloaded.json')

# Create the simulation settings - used to create the fieldset
settings['simulation'] = {
    'startdate': pd.Timestamp(starttime),      # Start date of simulation
    'endtime': pd.Timestamp(endtime) + pd.Timedelta(days=5),              # End time of simulation # buffer at end for copernicusmarine
    'runtime': runtime,              # Runtime of simulation
    'outputdt': timedelta(days=int(args['outputdt'])),       # Timestep of output
    'dt': timedelta(minutes=int(args['dt'])),                # Timestep of advection
}

fieldset = create_full_fieldset_copernicus(settings)
fieldset = add_additional_fields_copernicus(fieldset, settings)

z_start = 0 # Top layer for the simulation, based on whichever field has the deepest top layer
for field in fieldset.get_fields():
    try:
        if field.grid.depth[0] > z_start:
            z_start = field.grid.depth[0]
    except: # Not a field, but likely a vectorfield
        pass

fieldset.add_constant('z_start', z_start)

# Create the particle set
pset = create_particleset_from_release_files(args['startrelease'], args['endrelease'], release_folder, settings, fieldset)

# Kernels for simulation
kernels = [ARC_kernel_Copernicus,
           PolyTEOS10_bsq,
           AdvectionRK2,
           checkErrorThroughSurface_2DAdvectionRK2, # likely will never run
           StokesDriftCopernicus,
           BiofoulingCopernicus,
           reflectAtSurfaceCopernicus,
           reflectAtBathymetry,
           checkThroughSurfaceCopernicus,
           unbeachingBySamplingAfterwards,
           stopExecution,
           depth_delete]

# Simulation parameters
dt = settings['simulation']['dt']
outputdt = settings['simulation']['outputdt']
endtime = settings['simulation']['endtime']

# Create the particlefile and run the simulation
startdate_s = settings['simulation']['startdate'].isoformat()[:10]

pfilename = output_dir+f'particles_{startdate_s}.zarr'

pfile = pp.ParticleFile(pfilename, pset, settings=settings, outputdt=outputdt, chunks=(len(pset), 7)) #7 obs per chunk = 1 week if daily output
pset.execute(kernels, endtime=endtime, dt=dt, output_file=pfile)

# Write the latest locations at the end of the simulation
pfile.write_latest_locations(pset, time=np.max(pset.time_nextloop))

print(f"Simulation complete for start release date: {startdate_s}.")