restart_copernicus_arc_simulation.py

""" Simulation script for NECCTON simulations using Parcels and PlasticParcels.
    This script runs a simulation for plastic particles in the 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
import xarray as xr

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, PlasticParticle
from kernels import ARC_kernel_Copernicus, PolyTEOS10_bsq, BiofoulingCopernicus, StokesDriftCopernicus, checkThroughSurfaceCopernicus, stopExecution, reflectAtSurfaceCopernicus, reflectAtBathymetry, checkErrorThroughSurface_2DAdvectionRK2, AdvectionRK2, unbeachingBySamplingAfterwards, belowLatitude, depth_delete, initialise_algae_amount

# 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)
# Create the particlefile and run the simulation
startdate_s = settings['simulation']['startdate'].isoformat()[:10]

try:
    pset = parcels.ParticleSet.from_particlefile(
        fieldset=fieldset,
        pclass=PlasticParticle,
        filename=output_dir+f'particles_{startdate_s}.zarr',
        restart=True  # Restart from the last time step
    )
except:
    print("Likely a memory error has occurred, construct particleset manually.")

    ds_particles = xr.open_zarr(output_dir+f'particles_{startdate_s}.zarr')

    global_start_times = (ds_particles.isel(obs=0).time.values - ds_particles.isel(obs=0, trajectory=0).time.values).astype('timedelta64[D]').astype(int)
    global_trajectory_id = ds_particles.trajectory.values

    x = xr.DataArray(global_trajectory_id, dims="traj")
    y = xr.DataArray(5450 - global_start_times, dims="traj") # 5450 was chosen in a similar way to what's shown in determine_fill_missing_trajectories_from_restart_files.ipynb

    print("__________________________________________")
    print("Time and number of particles at obs=5450:")
    sim_ds_day = ds_particles.sel(trajectory=x, obs=y)
    print(np.unique(sim_ds_day.time.values, return_counts=True))
    print("__________________________________________")
    # These should be a single time

    non_nan_id = ~np.isnat(sim_ds_day.time.values)

    last_lons = sim_ds_day.lon.values[non_nan_id]
    last_lats = sim_ds_day.lat.values[non_nan_id]
    last_times = sim_ds_day.time.values[non_nan_id]
    plastic_diameters = sim_ds_day.plastic_diameter.values[non_nan_id]
    plastic_densities = sim_ds_day.plastic_density.values[non_nan_id]
    plastic_amounts = sim_ds_day.plastic_amount.values[non_nan_id]
    release_classes = sim_ds_day.release_class.values[non_nan_id]
    release_ids = sim_ds_day.release_id.values[non_nan_id]

    # Manually create the particleset using the non-nan particles that remain!
    pset = parcels.ParticleSet(fieldset=fieldset,
                            pclass=PlasticParticle,
                            lon=last_lons,
                            lat=last_lats,
                            time=last_times,
                            plastic_diameter=plastic_diameters,
                            plastic_density=plastic_densities,
                            plastic_amount=plastic_amounts,
                            release_class=release_classes,
                            release_id=release_ids,
                            )

# Kernels for simulation
kernels = [initialise_algae_amount, # Restarted particles need to start with a non-zero algal amount otherwise they shoot to the surface.
           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']



pfilename = output_dir+f'restart_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}.")