1 | ############################################################################################## |
---|
2 | ### Python script to output the stratification data over time from the "startpem.nc" files ### |
---|
3 | ############################################################################################## |
---|
4 | |
---|
5 | import os |
---|
6 | import numpy as np |
---|
7 | from netCDF4 import Dataset |
---|
8 | import matplotlib.pyplot as plt |
---|
9 | from scipy import interpolate |
---|
10 | |
---|
11 | ######################### |
---|
12 | ### Parameters to fill in |
---|
13 | folder_path = "starts" |
---|
14 | base_name = "restartpem" |
---|
15 | dz = 0.1 # Discrization step of the reference grid for the elevation [m] |
---|
16 | infofile = 'info_PEM.txt' |
---|
17 | ######################### |
---|
18 | |
---|
19 | |
---|
20 | ############################################################################################## |
---|
21 | ### Function to read the "startpem.nc" files and process their stratification data |
---|
22 | def process_files(folder_path,base_name): |
---|
23 | # Find all files in the directory with the pattern {base_name}{num}.nc |
---|
24 | nfile = 0 |
---|
25 | for file_name in sorted(os.listdir(folder_path)): |
---|
26 | if file_name.startswith(base_name) and file_name.endswith('.nc'): |
---|
27 | file_number = file_name[len(base_name):-3] |
---|
28 | if file_number.isdigit(): |
---|
29 | nfile += 1 |
---|
30 | |
---|
31 | if nfile == 0: |
---|
32 | print("No files found. Exiting...") |
---|
33 | return |
---|
34 | |
---|
35 | # Process each file and collect data |
---|
36 | datasets = [] |
---|
37 | max_top_elevation = 0 |
---|
38 | max_nb_str = 0 |
---|
39 | ngrid = Dataset(os.path.join(folder_path,base_name + "1.nc"),'r').dimensions['physical_points'].size # ngrid is the same for all files |
---|
40 | nslope = Dataset(os.path.join(folder_path,base_name + "1.nc"),'r').dimensions['nslope'].size # nslope is the same for all files |
---|
41 | longitude = Dataset(os.path.join(folder_path,base_name + "1.nc"),'r').variables['longitude'][:] |
---|
42 | latitude = Dataset(os.path.join(folder_path,base_name + "1.nc"),'r').variables['latitude'][:] |
---|
43 | for i in range(1,nfile + 1): |
---|
44 | file_path = os.path.join(folder_path,base_name + str(i) + ".nc") |
---|
45 | #print(f"Processing file: {file_path}") |
---|
46 | ds = Dataset(file_path,'r') |
---|
47 | datasets.append(ds) |
---|
48 | |
---|
49 | # Track max of nb_str_max |
---|
50 | max_nb_str = max(ds.dimensions['nb_str_max'].size,max_nb_str) |
---|
51 | |
---|
52 | # Track max of top_elevation across all slopes |
---|
53 | for k in range(1,nslope + 1): |
---|
54 | slope_var_name = f"stratif_slope{k:02d}_top_elevation" |
---|
55 | max_top_elevation = max(max_top_elevation,np.max(ds.variables[slope_var_name][:])) |
---|
56 | |
---|
57 | print(f"> number of files = {nfile}") |
---|
58 | print(f"> ngrid = {ngrid}") |
---|
59 | print(f"> nslope = {nslope}") |
---|
60 | print(f"> max(nb_str_max) = {max_nb_str}") |
---|
61 | print(f"> max(top_elevation) = {max_top_elevation}") |
---|
62 | |
---|
63 | # Concatenate stratif variables with dimension 'nb_str_max' along the "Time" dimension |
---|
64 | stratif_data = [] |
---|
65 | stratif_heights = np.zeros((ngrid,nfile,nslope,max_nb_str)) |
---|
66 | stratif_co2ice = np.zeros((ngrid,nfile,nslope,max_nb_str)) |
---|
67 | stratif_h2oice = np.zeros((ngrid,nfile,nslope,max_nb_str)) |
---|
68 | stratif_dust = np.zeros((ngrid,nfile,nslope,max_nb_str)) |
---|
69 | stratif_air = np.zeros((ngrid,nfile,nslope,max_nb_str)) |
---|
70 | for var_name in datasets[0].variables: |
---|
71 | if 'top_elevation' in var_name: |
---|
72 | for i in range(0,ngrid): |
---|
73 | for j in range(0,nfile): |
---|
74 | for k in range(0,nslope): |
---|
75 | if f'slope{k + 1:02d}' in var_name: |
---|
76 | stratif_heights[i,j,k,:datasets[j].variables[var_name].shape[1]] = datasets[j].variables[var_name][0,:,i] |
---|
77 | print(f"Processed variable: {var_name}") |
---|
78 | elif 'co2ice_volfrac' in var_name: |
---|
79 | for i in range(0,ngrid): |
---|
80 | for j in range(0,nfile): |
---|
81 | for k in range(0,nslope): |
---|
82 | if f'slope{k + 1:02d}' in var_name: |
---|
83 | stratif_co2ice[i,j,k,:datasets[j].variables[var_name].shape[1]] = datasets[j].variables[var_name][0,:,i] |
---|
84 | print(f"Processed variable: {var_name}") |
---|
85 | elif 'h2oice_volfrac' in var_name: |
---|
86 | for i in range(0,ngrid): |
---|
87 | for j in range(0,nfile): |
---|
88 | for k in range(0,nslope): |
---|
89 | if f'slope{k + 1:02d}' in var_name: |
---|
90 | stratif_h2oice[i,j,k,:datasets[j].variables[var_name].shape[1]] = datasets[j].variables[var_name][0,:,i] |
---|
91 | print(f"Processed variable: {var_name}") |
---|
92 | elif 'dust_volfrac' in var_name: |
---|
93 | for i in range(0,ngrid): |
---|
94 | for j in range(0,nfile): |
---|
95 | for k in range(0,nslope): |
---|
96 | if f'slope{k + 1:02d}' in var_name: |
---|
97 | stratif_dust[i,j,k,:datasets[j].variables[var_name].shape[1]] = datasets[j].variables[var_name][0,:,i] |
---|
98 | print(f"Processed variable: {var_name}") |
---|
99 | elif 'air_volfrac' in var_name: |
---|
100 | for i in range(0,ngrid): |
---|
101 | for j in range(0,nfile): |
---|
102 | for k in range(0,nslope): |
---|
103 | if f'slope{k + 1:02d}' in var_name: |
---|
104 | stratif_air[i,j,k,:datasets[j].variables[var_name].shape[1]] = datasets[j].variables[var_name][0,:,i] |
---|
105 | print(f"Processed variable: {var_name}") |
---|
106 | |
---|
107 | # Close the datasets |
---|
108 | for ds in datasets: |
---|
109 | ds.close() |
---|
110 | |
---|
111 | stratif_data = [stratif_heights,stratif_co2ice,stratif_h2oice,stratif_dust,stratif_air] |
---|
112 | return stratif_data, max_top_elevation, longitude, latitude |
---|
113 | |
---|
114 | ### Function to interpolate the stratification data on a reference grid |
---|
115 | def interpolate_data(stratif_data,max_top_elevation,dz): |
---|
116 | # Define the reference ref_grid |
---|
117 | ref_grid = np.arange(0,max_top_elevation,dz) |
---|
118 | print(f"> Number of ref_grid points = {len(ref_grid)}") |
---|
119 | |
---|
120 | # Interpolate the strata properties on the ref_grid |
---|
121 | gridded_stratif_data = -1.*np.ones((np.shape(stratif_data)[0] - 1,np.shape(stratif_data)[1],np.shape(stratif_data)[2],np.shape(stratif_data)[3],len(ref_grid))) |
---|
122 | for iprop in range(1,np.shape(stratif_data)[0]): |
---|
123 | for i in range(np.shape(stratif_data)[1]): |
---|
124 | for j in range(np.shape(stratif_data)[2]): |
---|
125 | for k in range(np.shape(stratif_data)[3]): |
---|
126 | i_hmax = np.max(np.nonzero(stratif_data[0][i,j,k,:])) |
---|
127 | hmax = stratif_data[0][i,j,k,i_hmax] |
---|
128 | i_zmax = np.searchsorted(ref_grid,hmax,'left') |
---|
129 | f = interpolate.interp1d(stratif_data[0][i,j,k,:i_hmax + 1],stratif_data[iprop][i,j,k,:i_hmax + 1],kind = 'next')#,fill_value = "extrapolate") |
---|
130 | gridded_stratif_data[iprop - 1,i,j,k,:i_zmax] = f(ref_grid[:i_zmax]) |
---|
131 | |
---|
132 | return ref_grid, gridded_stratif_data |
---|
133 | |
---|
134 | ### Function to read the "info_PEM.txt" file |
---|
135 | def read_infofile(file_name): |
---|
136 | with open(file_name,'r') as file: |
---|
137 | # Read the first line to get the parameters |
---|
138 | first_line = file.readline().strip() |
---|
139 | parameters = list(map(float,first_line.split())) |
---|
140 | |
---|
141 | # Read the following lines |
---|
142 | data_lines = [] |
---|
143 | date_time = [] |
---|
144 | for line in file: |
---|
145 | data = list(map(int,line.split())) |
---|
146 | data_lines.append(data) |
---|
147 | date_time.append(data[0]) |
---|
148 | |
---|
149 | return date_time |
---|
150 | |
---|
151 | ### Processing |
---|
152 | stratif_data, max_top_elevation, longitude, latitude = process_files(folder_path,base_name) |
---|
153 | ref_grid, gridded_stratif_data = interpolate_data(stratif_data,max_top_elevation,dz) |
---|
154 | date_time = read_infofile(infofile) |
---|
155 | |
---|
156 | ### Figures plotting |
---|
157 | subtitle = ['CO2 ice','H2O ice','Dust','Air'] |
---|
158 | cmap = plt.get_cmap('viridis').copy() |
---|
159 | cmap.set_under('white') |
---|
160 | for igr in range(np.shape(gridded_stratif_data)[1]): |
---|
161 | for isl in range(np.shape(gridded_stratif_data)[3]): |
---|
162 | fig, axes = plt.subplots(2,2,figsize = (10,8)) |
---|
163 | fig.suptitle(f'Contents variation over time in the layered-deposit of grid point {igr + 1} and slope {isl + 1}') |
---|
164 | iprop = 0 |
---|
165 | for ax in axes.flat: |
---|
166 | time_mesh, elevation_mesh = np.meshgrid(date_time,ref_grid) |
---|
167 | #im = ax.imshow(np.transpose(gridded_stratif_data[iprop][0,:,0,:]),aspect = 'auto',cmap = 'viridis',origin = 'lower',extent = [date_time[0],date_time[-1],ref_grid[0],ref_grid[-1]],vmin = 0,vmax = 1) |
---|
168 | im = ax.pcolormesh(time_mesh,elevation_mesh,np.transpose(gridded_stratif_data[iprop][igr,:,isl,:]),cmap = cmap,shading = 'auto',vmin = 0,vmax = 1) |
---|
169 | ax.set_title(subtitle[iprop]) |
---|
170 | ax.set(xlabel = 'Time (y)',ylabel = 'Elevation (m)') |
---|
171 | #ax.label_outer() |
---|
172 | iprop += 1 |
---|
173 | cbar = fig.colorbar(im,ax = axes.ravel().tolist(),label = 'Content value') |
---|
174 | plt.savefig(f"layering_evolution_ig{igr + 1}_is{isl + 1}.png") |
---|
175 | plt.show() |
---|