post_process_arc.py

# Script to post-process arctic simulation data and compute histograms for ARC
import numpy as np
import os
import xarray as xr
import pandas as pd
import os
#import gc

from argparse import ArgumentParser

# Parse the arguments
p = ArgumentParser(description="""Arctic post-process of parcels simulations""")
p.add_argument('-processtype', '--processtype', default='normal_0', help='Type of processing: normal_0, normal_1, normal_2, normal_3')
p.add_argument('-startdate', '--startdate', default='2000-01-01', help='Start date for processing (YYYY-MM-DD)')
p.add_argument('-order', '--order', default='forward', help='Order of processing: forward or backward')

parsed_args = p.parse_args()

processtype = parsed_args.processtype
startdate = parsed_args.startdate
order = parsed_args.order

# Locations of trajectory data and output folder
data_path = "/storage/shared/oceanparcels/output_data/data_Michael/NECCTONsimulations/data/copernicus_simulations/"
output_dir = '/storage/shared/oceanparcels/output_data/data_Michael/NECCTONsimulations/data/histograms_arc/'

# Construct the ARC grid for the histograms by shifting the bathymetry grid
bathy = xr.open_dataset("/storage/shared/oceanparcels/output_data/data_Michael/NECCTONsimulations/data/copernicus_data/cmems_mod_glo_phy_my_static_fulldomain.nc")
bathy_ARC = bathy.sel(latitude=slice(59.9, 90))


# Handle the north pole...
bathy_ARC_latitude_values = np.empty(bathy_ARC.latitude.size+1)
bathy_ARC_latitude_values[:-1] = bathy_ARC.latitude.values
bathy_ARC_latitude_values[-1] = bathy_ARC_latitude_values[-2] + 1/12

# Handle the periodic BC
bathy_ARC_longitude_values = np.empty(bathy_ARC.longitude.size+3)
bathy_ARC_longitude_values[1:-2] = bathy_ARC.longitude.values
bathy_ARC_longitude_values[0] = bathy_ARC_longitude_values[1] - 1/12
bathy_ARC_longitude_values[-2] = bathy_ARC_longitude_values[-3] + 1/12
#We will create an extra column on the far right, and later add these values to the first column, and remove
bathy_ARC_longitude_values[-1] = bathy_ARC_longitude_values[-2] + 1/12

# Construct bin edges and cell centers
lat_bins = (bathy_ARC_latitude_values[:-1] + bathy_ARC_latitude_values[1:]) / 2
lat_centers = bathy_ARC_latitude_values[1:-1]

lon_bins = (bathy_ARC_longitude_values[:-1] + bathy_ARC_longitude_values[1:]) / 2
lon_centers = bathy_ARC_longitude_values[1:-2]

# Create a grid to compute the densitites over
arc_x_edges = lon_bins
arc_y_edges = lat_bins

surface_threshhold = 5  # meters

# Save to file
if not os.path.isfile(output_dir + "arc_grid.npz"):
    np.savez(output_dir + "arc_grid.npz", arc_x_edges=arc_x_edges, arc_y_edges=arc_y_edges, lon_centers=lon_centers, lat_centers=lat_centers)


# # List of start times #'2008-01-01',
# start_dates = ['2009-01-01',
#                '2010-01-01',
#                '2011-01-01',
#                '2012-01-01',
#                '2013-01-01',
#                '2014-01-01',
#                '2015-01-01',
#                '2016-01-01',
#                '2017-01-01',
#                '2018-01-01',
#                '2019-01-01',
#                '2020-01-01',
#                '2021-01-01',
#                '2022-01-01']


# if processtype == 'normal_0':
#     start_dates_to_process = start_dates[0::4]
# elif processtype == 'normal_1':
#     start_dates_to_process = start_dates[1::4]
# elif processtype == 'normal_2':
#     start_dates_to_process = start_dates[2::4]
# elif processtype == 'normal_3':
#     start_dates_to_process = start_dates[3::4]
# elif processtype == 'restart':
#     start_dates_to_process = ['2008-01-01']
# else:
#     raise ValueError("Invalid processtype argument. Choose from: normal_0, normal_1, normal_2, normal_3, restart")
start_dates_to_process = [startdate]

print("Processing ARC for processtype:", processtype)

# Loop over release dates
if processtype == 'restart':
    # Loop over release dates without restarts
    for starting_day_str in start_dates_to_process:
        print(f"Processing ARC starting day: {starting_day_str}")

        # Load the dataset
        sim_ds = xr.open_zarr(data_path + f"particles_{starting_day_str}.zarr")
        sim_restart_ds = xr.open_zarr(data_path+ f"restart_particles_{starting_day_str}.zarr")

        # Compute how many days we have to process the initial file
        normal_startday = pd.to_datetime(starting_day_str)
        restart_startday = pd.to_datetime(np.unique(sim_restart_ds.isel(obs=0).time.values)[0])
        ndays_to_process = (restart_startday - normal_startday).days


        # pre-processing to forward fill missing data (when particles get deleted). Assumption -> particles deleted remain at the last known position!
        sim_ds['lon'] = sim_ds.lon.ffill(dim='obs') # Forward fill longitude
        sim_ds['lat'] = sim_ds.lat.ffill(dim='obs') # Forward fill latitude
        sim_ds['z'] = sim_ds.z.ffill(dim='obs') # Forward fill depth

        # Construct a list of "starting times" for each trajectory, for when the day index is >=, can use the traj
        global_start_times = (sim_ds.isel(obs=0).time.values - sim_ds.isel(obs=0, trajectory=0).time.values).astype('timedelta64[D]').astype(int)

        # Construct a list of trajectory IDS
        global_trajectory_id = sim_ds.trajectory.values

        # List of plastic sizes we model
        plastic_sizes = np.unique(sim_ds.plastic_diameter.values)

        # List of release types (0=river, 1=coastal, 2=fisheries)
        release_classes = np.unique(sim_ds.release_class.values).astype(int)


        # List of trajectories for each plastic size
        plastic_class_traj = {}
        for p_size in plastic_sizes:
            mask = ((sim_ds.plastic_diameter == p_size)).rename("plastic_mask")
            plastic_class_traj[p_size] = sim_ds.sel(trajectory=mask).trajectory.values

        # List of trajectories for each release type
        release_class_traj = {}
        for r_class in release_classes:
            mask = ((sim_ds.release_class == r_class)).rename("release_mask")
            release_class_traj[r_class] = sim_ds.sel(trajectory=mask).trajectory.values


        # Big loop to construct histograms for each day
        loop_ndays_to_process = range(ndays_to_process)
        if order == 'backward':
            loop_ndays_to_process = reversed(loop_ndays_to_process)

        for obs in loop_ndays_to_process:
            sim_day = pd.to_datetime(starting_day_str) + pd.Timedelta(days=obs)
            sim_day_str = sim_day.strftime("%Y-%m-%d")

            # If the last file already exists, skip processing
            if os.path.isfile(output_dir + f"arc_n_start_{starting_day_str}_obs_{sim_day_str}_rc_2_sc_5.npz"):
                continue

            # First, get a list of trajectories that were valid in this obs
            valid_trajectory_mask = global_start_times <= obs
            valid_trajectory_ids = global_trajectory_id[valid_trajectory_mask]

            for release_i, release_class in enumerate(release_classes):
                # Get only the trajectories for this release type and plastic class
                release_type_trajectory_ids = release_class_traj[release_class]
                valid_release_trajectory_ids = np.intersect1d(valid_trajectory_ids, release_type_trajectory_ids)


                for plastic_i, plastic_size in enumerate(plastic_sizes):
                    if not os.path.isfile(output_dir + f"arc_n_start_{starting_day_str}_obs_{sim_day_str}_rc_{release_class}_sc_{plastic_i}.npz"):
                        # Get only the trajectories for this plastic size
                        plastic_size_trajectory_ids = plastic_class_traj[plastic_size]
                        valid_plastic_trajectory_ids = np.intersect1d(valid_release_trajectory_ids, plastic_size_trajectory_ids)

                        # Construct 2 dataarrays to select over - see example here: https://docs.xarray.dev/en/latest/user-guide/interpolation.html#advanced-interpolation
                        x = xr.DataArray(valid_plastic_trajectory_ids, dims="traj")
                        y = xr.DataArray(obs - global_start_times[valid_plastic_trajectory_ids], dims="traj")

                        object_ds = sim_ds.sel(trajectory=x, obs=y) # The object we want to compute the histogram for!
                        
                        #compute histograms here!
                        H_arc, yedges, xedges = np.histogram2d(object_ds.lat.values,
                                                        (object_ds.lon.values+180) % 360 - 180,
                                                        weights=object_ds.plastic_amount.values, # These need to be rescaled later!
                                                        bins=(arc_y_edges, arc_x_edges),
                                                        density=False)
                        H_arc_modify = H_arc.T
                        H_arc_modify[0,:] += H_arc_modify[-1,:] # Add the far RHS column to the LHS column for the periodicity
                        H_arc_modify = H_arc_modify[:-1, :]  # Remove the last column

                        
                        # Now compute surface only histograms
                        surface_object_ds = (object_ds.z <= surface_threshhold).rename("surface_mask").compute()
                        surface_object_ds = object_ds.sel(traj=surface_object_ds)
                        # Compute 2D histogram of particle counts in ARC
                        H_arc_surf, yedges, xedges = np.histogram2d(surface_object_ds.lat.values,
                                                        (surface_object_ds.lon.values+180) % 360 - 180,
                                                        weights=surface_object_ds.plastic_amount.values, # These need to be rescaled later!
                                                        bins=(arc_y_edges, arc_x_edges),
                                                        density=False)
                        H_arc_surf_modify = H_arc_surf.T
                        H_arc_surf_modify[0,:] += H_arc_surf_modify[-1,:] # Add the far RHS column to the LHS column for the periodicity
                        H_arc_surf_modify = H_arc_surf_modify[:-1, :]  # Remove the last column
                    

                        # Save to file
                        # filename structure {simulation}_{normal or restart}_start_{start day}_obs_{obs day}_rc_{release class}_sc_{size class}.npz
                        np.savez(output_dir + f"arc_n_start_{starting_day_str}_obs_{sim_day_str}_rc_{release_class}_sc_{plastic_i}.npz",
                                H_arc=H_arc_modify, H_arc_surf=H_arc_surf_modify) # Save the WC and surface histograms to the same file.
                        
                        # del surface_object_ds, H_arc, H_arc_modify, H_arc_surf, H_arc_surf_modify
                        # gc.collect()


        # Make the restart data the data we will look at
        sim_ds = sim_restart_ds
        # Reset the trajectory id's so that we can index correctly
        sim_ds = sim_ds.assign_coords(
            trajectory=('trajectory', np.arange(sim_ds.sizes['trajectory']))
        )
        restart_starting_day = pd.to_datetime(sim_ds.isel(obs=0, trajectory=0).time.values)

        # pre-processing to forward fill missing data (when particles get deleted). Assumption -> particles deleted remain at the last known position!
        sim_ds['lon'] = sim_ds.lon.ffill(dim='obs') # Forward fill longitude
        sim_ds['lat'] = sim_ds.lat.ffill(dim='obs') # Forward fill latitude
        sim_ds['z'] = sim_ds.z.ffill(dim='obs') # Forward fill depth

        # Construct a list of "starting times" for each trajectory, for when the day index is >=, can use the traj
        # These really should just be 0s, because all particles will have the same start day
        global_start_times = (sim_ds.isel(obs=0).time.values - sim_ds.isel(obs=0, trajectory=0).time.values).astype('timedelta64[D]').astype(int)

        # Construct a list of trajectory IDS
        global_trajectory_id = sim_ds.trajectory.values

        # List of plastic sizes we model
        plastic_sizes = np.unique(sim_ds.plastic_diameter.values)

        # List of release types (0=river, 1=coastal, 2=fisheries)
        release_classes = np.unique(sim_ds.release_class.values).astype(int)

        # List of trajectories for each plastic size
        plastic_class_traj = {}
        for p_size in plastic_sizes:
            mask = ((sim_ds.plastic_diameter == p_size)).rename("plastic_mask")
            plastic_class_traj[p_size] = sim_ds.sel(trajectory=mask).trajectory.values

        # List of trajectories for each release type
        release_class_traj = {}
        for r_class in release_classes:
            mask = ((sim_ds.release_class == r_class)).rename("release_mask")
            release_class_traj[r_class] = sim_ds.sel(trajectory=mask).trajectory.values


        # Big loop to construct histograms for each day
        loop_ndays_to_process = range(sim_ds.obs.size)
        if order == 'backward':
            loop_ndays_to_process = reversed(loop_ndays_to_process)

        for obs in loop_ndays_to_process:
            sim_day = restart_starting_day + pd.Timedelta(days=obs)
            sim_day_str = sim_day.strftime("%Y-%m-%d")

            # If the last file already exists, skip processing
            if os.path.isfile(output_dir + f"arc_r_start_{starting_day_str}_obs_{sim_day_str}_rc_2_sc_5.npz"):
                continue

            # First, get a list of trajectories that were valid in this obs
            valid_trajectory_mask = global_start_times <= obs
            valid_trajectory_ids = global_trajectory_id[valid_trajectory_mask]

            for release_i, release_class in enumerate(release_classes):
                # Get only the trajectories for this release type and plastic class
                release_type_trajectory_ids = release_class_traj[release_class]
                valid_release_trajectory_ids = np.intersect1d(valid_trajectory_ids, release_type_trajectory_ids)


                for plastic_i, plastic_size in enumerate(plastic_sizes):
                    if not os.path.isfile(output_dir + f"arc_r_start_{starting_day_str}_obs_{sim_day_str}_rc_{release_class}_sc_{plastic_i}.npz"):
                        # Get only the trajectories for this plastic size
                        plastic_size_trajectory_ids = plastic_class_traj[plastic_size]
                        valid_plastic_trajectory_ids = np.intersect1d(valid_release_trajectory_ids, plastic_size_trajectory_ids)

                        # Construct 2 dataarrays to select over - see example here: https://docs.xarray.dev/en/latest/user-guide/interpolation.html#advanced-interpolation
                        x = xr.DataArray(valid_plastic_trajectory_ids, dims="traj")
                        y = xr.DataArray(obs - global_start_times[valid_plastic_trajectory_ids], dims="traj")

                        object_ds = sim_ds.sel(trajectory=x, obs=y) # The object we want to compute the histogram for!
                        
                        #compute histograms here!
                        H_arc, yedges, xedges = np.histogram2d(object_ds.lat.values,
                                                        (object_ds.lon.values+180) % 360 - 180,
                                                        weights=object_ds.plastic_amount.values, # These need to be rescaled later!
                                                        bins=(arc_y_edges, arc_x_edges),
                                                        density=False)
                        H_arc_modify = H_arc.T
                        H_arc_modify[0,:] += H_arc_modify[-1,:] # Add the far RHS column to the LHS column for the periodicity
                        H_arc_modify = H_arc_modify[:-1, :]  # Remove the last column

                        
                        # Now compute surface only histograms
                        surface_object_ds = (object_ds.z <= surface_threshhold).rename("surface_mask").compute()
                        surface_object_ds = object_ds.sel(traj=surface_object_ds)
                        # Compute 2D histogram of particle counts in ARC
                        H_arc_surf, yedges, xedges = np.histogram2d(surface_object_ds.lat.values,
                                                        (surface_object_ds.lon.values+180) % 360 - 180,
                                                        weights=surface_object_ds.plastic_amount.values, # These need to be rescaled later!
                                                        bins=(arc_y_edges, arc_x_edges),
                                                        density=False)
                        H_arc_surf_modify = H_arc_surf.T
                        H_arc_surf_modify[0,:] += H_arc_surf_modify[-1,:] # Add the far RHS column to the LHS column for the periodicity
                        H_arc_surf_modify = H_arc_surf_modify[:-1, :]  # Remove the last column

                        # Save to file
                        # filename structure {simulation}_{normal or restart}_start_{start day}_obs_{obs day}_rc_{release class}_sc_{size class}.npz
                        np.savez(output_dir + f"arc_r_start_{starting_day_str}_obs_{sim_day_str}_rc_{release_class}_sc_{plastic_i}.npz",
                                H_arc=H_arc_modify, H_arc_surf=H_arc_surf_modify) # Save the WC and surface histograms to the same file.


                        # del surface_object_ds, H_arc, H_arc_modify, H_arc_surf, H_arc_surf_modify
                        # gc.collect()
                        
        print("Finished processing restart for starting day:", starting_day_str)

else:
    for starting_day_str in start_dates_to_process:
        print(f"Processing ARC starting day: {starting_day_str}")

        # Load the dataset
        sim_ds = xr.open_zarr(data_path + f"particles_{starting_day_str}.zarr")

        # pre-processing to forward fill missing data (when particles get deleted). Assumption -> particles deleted remain at the last known position!
        sim_ds['lon'] = sim_ds.lon.ffill(dim='obs') # Forward fill longitude
        sim_ds['lat'] = sim_ds.lat.ffill(dim='obs') # Forward fill latitude
        sim_ds['z'] = sim_ds.z.ffill(dim='obs') # Forward fill depth
        
        # Construct a list of "starting times" for each trajectory
        global_start_times = (sim_ds.isel(obs=0).time.values - sim_ds.isel(obs=0, trajectory=0).time.values).astype('timedelta64[D]').astype(int)

        # Construct a list of trajectory IDS
        global_trajectory_id = sim_ds.trajectory.values

        # List of plastic sizes we model
        plastic_sizes = np.unique(sim_ds.plastic_diameter.values)

        # List of release types (0=river, 1=coastal, 2=fisheries)
        release_classes = np.unique(sim_ds.release_class.values).astype(int)


        # List of trajectories for each plastic size
        plastic_class_traj = {}
        for p_size in plastic_sizes:
            mask = ((sim_ds.plastic_diameter == p_size)).rename("plastic_mask")
            plastic_class_traj[p_size] = sim_ds.sel(trajectory=mask).trajectory.values

        # List of trajectories for each release type
        release_class_traj = {}
        for r_class in release_classes:
            mask = ((sim_ds.release_class == r_class)).rename("release_mask")
            release_class_traj[r_class] = sim_ds.sel(trajectory=mask).trajectory.values


        # Big loop to construct histograms for each day
        loop_ndays_to_process = range(sim_ds.obs.size)
        if order == 'backward':
            loop_ndays_to_process = reversed(loop_ndays_to_process)

        for obs in loop_ndays_to_process:
            sim_day = pd.to_datetime(starting_day_str) + pd.Timedelta(days=obs)
            sim_day_str = sim_day.strftime("%Y-%m-%d")

            # If the last file already exists, skip processing
            if os.path.isfile(output_dir + f"arc_n_start_{starting_day_str}_obs_{sim_day_str}_rc_2_sc_5.npz"):
                continue

            # First, get a list of trajectories that were valid in this obs
            valid_trajectory_mask = global_start_times <= obs
            valid_trajectory_ids = global_trajectory_id[valid_trajectory_mask]

            for release_i, release_class in enumerate(release_classes):
                # Get only the trajectories for this release type and plastic class
                release_type_trajectory_ids = release_class_traj[release_class]
                valid_release_trajectory_ids = np.intersect1d(valid_trajectory_ids, release_type_trajectory_ids)


                for plastic_i, plastic_size in enumerate(plastic_sizes):
                    if not os.path.isfile(output_dir + f"arc_n_start_{starting_day_str}_obs_{sim_day_str}_rc_{release_class}_sc_{plastic_i}.npz"):
                        # Get only the trajectories for this plastic size
                        plastic_size_trajectory_ids = plastic_class_traj[plastic_size]
                        valid_plastic_trajectory_ids = np.intersect1d(valid_release_trajectory_ids, plastic_size_trajectory_ids)

                        # Construct 2 dataarrays to select over - see example here: https://docs.xarray.dev/en/latest/user-guide/interpolation.html#advanced-interpolation
                        x = xr.DataArray(valid_plastic_trajectory_ids, dims="traj")
                        y = xr.DataArray(obs - global_start_times[valid_plastic_trajectory_ids], dims="traj")

                        object_ds = sim_ds.sel(trajectory=x, obs=y) # The object we want to compute the histogram for!
                        
                        #compute histograms here!
                        H_arc, yedges, xedges = np.histogram2d(object_ds.lat.values,
                                                        (object_ds.lon.values+180) % 360 - 180,
                                                        weights=object_ds.plastic_amount.values, # These need to be rescaled later!
                                                        bins=(arc_y_edges, arc_x_edges),
                                                        density=False)
                        H_arc_modify = H_arc.T
                        H_arc_modify[0,:] += H_arc_modify[-1,:] # Add the far RHS column to the LHS column for the periodicity
                        H_arc_modify = H_arc_modify[:-1, :]  # Remove the last column

                        
                        # Now compute surface only histograms
                        surface_object_ds = (object_ds.z <= surface_threshhold).rename("surface_mask").compute()
                        surface_object_ds = object_ds.sel(traj=surface_object_ds)
                        # Compute 2D histogram of particle counts in ARC
                        H_arc_surf, yedges, xedges = np.histogram2d(surface_object_ds.lat.values,
                                                        (surface_object_ds.lon.values+180) % 360 - 180,
                                                        weights=surface_object_ds.plastic_amount.values, # These need to be rescaled later!
                                                        bins=(arc_y_edges, arc_x_edges),
                                                        density=False)
                        H_arc_surf_modify = H_arc_surf.T
                        H_arc_surf_modify[0,:] += H_arc_surf_modify[-1,:] # Add the far RHS column to the LHS column for the periodicity
                        H_arc_surf_modify = H_arc_surf_modify[:-1, :]  # Remove the last column



                        # Save to file
                        # filename structure {simulation}_{normal or restart}_start_{start day}_obs_{obs day}_rc_{release class}_sc_{size class}.npz
                        np.savez(output_dir + f"arc_n_start_{starting_day_str}_obs_{sim_day_str}_rc_{release_class}_sc_{plastic_i}.npz",
                                H_arc=H_arc_modify, H_arc_surf=H_arc_surf_modify) # Save the WC and surface histograms to the same file.
        
        print("Finished processing normal for starting day:", starting_day_str)

print("Post-processing complete.")