NECCTONsimulations/post_process_nws.py at main · Parcels-code/NECCTONsimulations · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
# Script to post-process global simulation data and compute histograms for NWS
import numpy as np
import os
import xarray as xr
import pandas as pd
import os

from argparse import ArgumentParser

# Parse the arguments
p = ArgumentParser(description="""NWS 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/additional_NWS/"
output_dir = '/storage/shared/oceanparcels/output_data/data_Michael/NECCTONsimulations/data/histograms_nws/'

# Construct the NWS 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_NWS = bathy.sel(longitude=slice(-15.1,10.1), latitude=slice(44.9,61.1))

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

lon_bins = (bathy_NWS.longitude.values[:-1] + bathy_NWS.longitude.values[1:]) / 2
lon_centers = bathy_NWS.longitude.values[1:-1]

# Create a grid to compute the densitites over
nws_x_edges = lon_bins
nws_y_edges = lat_bins

surface_threshhold = 5  # meters

# Save to file
if not os.path.isfile(output_dir + "nws_grid.npz"):
    np.savez(output_dir + "nws_grid.npz", nws_x_edges=nws_x_edges, nws_y_edges=nws_y_edges, lon_centers=lon_centers, lat_centers=lat_centers)


# List of start times without and with restart files
# start_dates = ['2008-01-05',
#                '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]
# else:
#     raise ValueError("Invalid processtype argument. Choose from: normal_0, normal_1, normal_2, normal_3")

start_dates_to_process = [startdate]

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

# Loop over release dates without restarts
for starting_day_str in start_dates_to_process:
    print(f"Processing NWS 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
    #NOTE: No need to figure out future times, because we are handling times separately.

    # 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"nws_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"nws_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_nws, 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=(nws_y_edges, nws_x_edges),
                                                    density=False
                                                    )


                    # 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 NWS
                    H_nws_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=(nws_y_edges, nws_x_edges),
                                                    density=False
                                                    )
                    # 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"nws_n_start_{starting_day_str}_obs_{sim_day_str}_rc_{release_class}_sc_{plastic_i}.npz",
                            H_nws=H_nws.T, H_nws_surf=H_nws_surf.T) # Save the WC and surface histograms to the same file.
    print("Finished processing normal for starting day:", starting_day_str)

print("Post-processing complete.")