restart_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 xarray as xr

from argparse import ArgumentParser
import os

# Lagrangian analysis
import parcels
import plasticparcels as pp

from datetime import date, timedelta

# Import custom functions
from functions import create_full_fieldset, add_additional_fields, create_particleset_from_release_files, PlasticParticle
from kernels import PolyTEOS10_bsq, Biofouling, StokesDrift, checkThroughSurface, stopExecution, reflectAtSurface, reflectAtBathymetry, checkErrorThroughSurface_2DAdvectionRK2, AdvectionRK2_3D, unbeachingBySamplingAfterwards, belowLatitude, initialise_algae_amount, 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/initial_simulations/'
release_folder = '/nethome/denes001/Projects/NECCTONsimulations/release_files_initial/'

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

# Set model indices for faster loading
settings['ocean']['indices'] = {'lat': range(1800,3059),
                                'depth': range(0,20)}
settings['bgc']['indices'] = {'lat': range(590,1022),
                                'depth': range(0,20)}
settings['stokes']['indices'] = {'lat': range(230, 361)}

# Create the simulation settings - used to create the fieldset
settings['simulation'] = {
    'startdate': starting_time,      # Start date of simulation
    'endtime': endtime,              # End time of simulation
    'runtime': runtime,              # Runtime of simulation
    'outputdt': timedelta(days=int(args['outputdt'])),       # Timestep of output
    'dt': timedelta(minutes=int(args['dt'])),                # Timestep of advection
    }

# Create the fieldset
fieldset = create_full_fieldset(settings)

# Add constants to the fieldset
fieldset.add_constant('use_mixing', settings['use_mixing'])
fieldset.add_constant('use_biofouling', settings['use_biofouling'])
fieldset.add_constant('use_stokes', settings['use_stokes'])
fieldset.add_constant('use_wind', settings['use_wind'])
fieldset.add_constant('G', 9.81)  # Gravitational constant [m s-1]
fieldset.add_constant('use_3D', settings['use_3D'])

# Add Stokes and biogeochemistry and unbeaching
fieldset = add_additional_fields(fieldset, settings)

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

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 memory error thrown! Manually creating particleset.")
    # For some reason, the 2009-01-01 file will throw a memory error...
    # Let's manually construct the pset.
    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(3779 - global_start_times, dims="traj") # See determine_fill_missing_trajectories_from_restart_files.ipynb for explanation of 3779

    print("__________________________________________")
    print("Time and number of particles at obs=3779:")
    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.
           PolyTEOS10_bsq,
           AdvectionRK2_3D,
           checkErrorThroughSurface_2DAdvectionRK2,
           StokesDrift,
           Biofouling,
           pp.VerticalMixing,
           reflectAtSurface,
           reflectAtBathymetry,
           checkThroughSurface,
           unbeachingBySamplingAfterwards,
           stopExecution,
           belowLatitude,
           depth_delete]

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

# Create the particlefile and run the simulation


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("Simulation complete.")