Changeset 3809 for trunk/LMDZ.COMMON/libf/evolution/deftank
- Timestamp:
- Jun 16, 2025, 5:11:50 PM (7 weeks ago)
- File:
-
- 1 edited
Legend:
- Unmodified
- Added
- Removed
-
trunk/LMDZ.COMMON/libf/evolution/deftank/visu_evol_layering.py
r3792 r3809 1 1 #!/usr/bin/env python3 2 2 ####################################################################################################### 3 ### Python script to output the stratification data over time from the "restartpem#.nc" files files ### 3 ### Python script to output stratification data over time from "restartpem#.nc" files ### 4 ### and to plot orbital parameters from "obl_ecc_lsp.asc" ### 4 5 ####################################################################################################### 5 6 6 7 7 import os … … 47 47 ).strip() or "info_PEM.txt" 48 48 49 return folder_path, base_name, infofile 49 orbfile = input( 50 "Enter the name of the orbital parameters ASCII file " 51 "(press Enter for default [obl_ecc_lsp.asc]): " 52 ).strip() or "obl_ecc_lsp.asc" 53 while not os.path.isfile(orbfile): 54 print(f" » \"{orbfile}\" does not exist or is not a file.") 55 orbfile = input( 56 "Enter a valid orbital parameters ASCII filename (press Enter for default [obl_ecc_lsp.asc]): " 57 ).strip() or "info_PEM.txt" 58 59 return folder_path, base_name, infofile, orbfile 50 60 51 61 … … 195 205 """ 196 206 Given raw_prop_arrays for 'co2_ice', 'h2o_ice', 'dust', 'pore' (in meters), 197 normalize each set of strata so that the sum of those four = 1 per strata.207 normalize each set of strata so that the sum of those four = 1 per cell. 198 208 Returns: 199 209 - frac_arrays: dict mapping same keys -> 4D arrays of fractions (0..1). … … 228 238 """ 229 239 Reads "info_PEM.txt". Expects: 230 - First line: parameters (ignored). 231 - Each subsequent line: floats where first value is timestamp. 232 Returns: 1D numpy array of timestamps. 240 - First line: parameters where the 3rd value is martian_to_earth conversion factor. 241 - Each subsequent line: floats where first value is simulation timestamp (in Mars years). 242 Returns: 243 - date_time: 1D numpy array of timestamps (Mars years) 244 - martian_to_earth: float conversion factor 233 245 """ 234 246 date_time = [] 235 247 with open(file_name, 'r') as fp: 236 fp.readline() 248 first = fp.readline().split() 249 martian_to_earth = float(first[2]) 237 250 for line in fp: 238 251 parts = line.strip().split() … … 243 256 except ValueError: 244 257 continue 245 return np.array(date_time, dtype=np.float64) 258 return np.array(date_time, dtype=np.float64), martian_to_earth 246 259 247 260 … … 293 306 ): 294 307 """ 295 Build a reference grid and interpolate strata fractions (0..1) onto it. 296 297 Also returns a 'top_index' array of shape (ngrid, ntime, nslope) that 298 indicates, for each (ig, t_idx, isl), the number of ref_grid levels 299 covered by the topmost valid stratum. 300 301 Args: 302 - heights_data: list of lists where heights_data[t][isl] is a 2D array 303 (ngrid, n_strata_current) of top_elevation values. 304 - prop_arrays: dict mapping each property_name to a 4D array of shape 305 (ngrid, ntime, nslope, max_nb_str) holding fractions [0..1]. 306 - min_base_for_interp: float; if exclude_sub=True, this is 0.0. 307 - max_top_elev: float 308 - dz: float 309 - exclude_sub: bool. If True, ignore strata with elevation < 0. 308 Build a reference elevation grid and interpolate strata fractions onto it. 310 309 311 310 Returns: 312 311 - ref_grid: 1D array of elevations (nz,) 313 - gridded_data: dict mapping each property_name to a 4D array of shape 314 (ngrid, ntime, nslope, nz) with interpolated fractions. 315 - top_index: 3D array (ngrid, ntime, nslope) of ints: number of levels 316 of ref_grid covered by the topmost stratum. 317 """ 318 # Build ref_grid, ensuring at least two points if surface-only and dz > max_top_elev 312 - gridded_data: dict mapping each property_name to 4D array 313 (ngrid, ntime, nslope, nz) with interpolated fractions. 314 - top_index: 3D array (ngrid, ntime, nslope) of ints: 315 number of levels covered by the topmost stratum. 316 """ 319 317 if exclude_sub and (dz > max_top_elev): 320 318 ref_grid = np.array([0.0, max_top_elev], dtype=np.float32) 321 319 else: 322 ref_grid = np.arange(min_base_for_interp, max_top_elev + dz /2, dz)320 ref_grid = np.arange(min_base_for_interp, max_top_elev + dz/2, dz) 323 321 nz = len(ref_grid) 324 322 print(f"> Number of reference grid points = {nz}") 325 323 326 # Dimensions327 324 sample_prop = next(iter(prop_arrays.values())) 328 ngrid, ntime, nslope, max_nb_str = sample_prop.shape[0:4] 329 330 # Prepare outputs 325 ngrid, ntime, nslope, max_nb_str = sample_prop.shape 326 331 327 gridded_data = { 332 328 prop: np.full((ngrid, ntime, nslope, nz), -1.0, dtype=np.float32) … … 342 338 continue 343 339 344 raw_h = h_mat[ig, :] # (n_strata_current,) 345 # Create h_all of length max_nb_str, fill with NaN 340 raw_h = h_mat[ig, :] 346 341 h_all = np.full((max_nb_str,), np.nan, dtype=np.float32) 347 342 n_strata_current = raw_h.shape[0] … … 359 354 h_valid = h_all[valid_mask] 360 355 top_h = np.max(h_valid) 361 362 # Find i_zmax = number of ref_grid levels z <= top_h363 356 i_zmax = np.searchsorted(ref_grid, top_h, side='right') 364 357 top_index[ig, t_idx, isl] = i_zmax 365 366 358 if i_zmax == 0: 367 # top_h < ref_grid[0], skip interpolation368 359 continue 369 360 370 361 for prop, arr in prop_arrays.items(): 371 prop_profile_all = arr[ig, t_idx, isl, :] # (max_nb_str,)372 prop_profile = prop_profile_all[valid_mask] # (n_valid_strata,)362 prop_profile_all = arr[ig, t_idx, isl, :] 363 prop_profile = prop_profile_all[valid_mask] 373 364 if prop_profile.size == 0: 374 365 continue 375 366 376 # Step‐wise interpolation (kind='next')377 367 f_interp = interp1d( 378 368 h_valid, … … 382 372 fill_value=-1.0 383 373 ) 384 385 # Evaluate for ref_grid[0:i_zmax]386 374 gridded_data[prop][ig, t_idx, isl, :i_zmax] = f_interp(ref_grid[:i_zmax]) 387 375 … … 399 387 ): 400 388 """ 401 For each grid point (ig) and slope (isl), generate a 2×2 figure:389 For each grid point and slope, generate a 2×2 figure of: 402 390 - CO2 ice fraction 403 391 - H2O ice fraction 404 392 - Dust fraction 405 393 - Pore fraction 406 407 Fractions are in [0..1]. Values < 0 (fill) are masked. 408 Using top_index, any elevation above the last stratum is forced to NaN (white). 409 410 Additionally, draw horizontal violet line segments at each stratum top elevation 411 only over the interval [date_time[t_idx], date_time[t_idx+1]] where that stratum 412 exists at time t_idx. This way, boundaries appear only where the strata exist. 413 """ 414 import numpy as np 415 import matplotlib.pyplot as plt 416 394 """ 417 395 prop_names = ['co2_ice', 'h2o_ice', 'dust', 'pore'] 418 396 titles = ["CO2 ice", "H2O ice", "Dust", "Pore"] … … 426 404 if exclude_sub: 427 405 positive_indices = np.where(ref_grid >= 0.0)[0] 428 if positive_indices.size == 0:429 print("Warning: no positive elevations in ref_grid → nothing to display.")430 return431 406 sub_ref_grid = ref_grid[positive_indices] 432 407 else: … … 438 413 fig, axes = plt.subplots(2, 2, figsize=(10, 8)) 439 414 fig.suptitle( 440 f"Content variation over time for (Grid Point {ig+1}, Slope {isl+1})",415 f"Content variation over time for (Grid point {ig+1}, Slope {isl+1})", 441 416 fontsize=14 442 417 ) 443 418 444 # For each time step t_idx, gather this stratum's valid tops 445 # and draw a line segment from date_time[t_idx] to date_time[t_idx+1]. 446 # We'll skip t_idx = ntime - 1 since no next point. 447 # Precompute, for each t_idx, the array of valid top elevations: 419 # Precompute valid stratum tops per time 448 420 valid_tops_per_time = [] 449 421 for t_idx in range(ntime): 450 raw_h = heights_data[t_idx][isl][ig, :] # (n_strata_current,) 451 # Exclude NaNs or zeros 422 raw_h = heights_data[t_idx][isl][ig, :] 452 423 h_all = raw_h[~np.isnan(raw_h)] 453 424 if exclude_sub: … … 457 428 for idx, prop in enumerate(prop_names): 458 429 ax = axes.flat[idx] 459 data_3d = gridded_data[prop][ig, :, isl, :] # shape (ntime, nz) 460 mat_full = data_3d.T # shape (nz, ntime) 461 mat = mat_full[positive_indices, :].copy() # (nz_pos, ntime) 462 463 # Mask fill values (< 0) as NaN 430 data_3d = gridded_data[prop][ig, :, isl, :] 431 mat_full = data_3d.T 432 mat = mat_full[positive_indices, :].copy() 464 433 mat[mat < 0.0] = np.nan 465 434 466 # Mask everything above the top stratum using top_index435 # Mask above top stratum 467 436 for t_idx in range(ntime): 468 437 i_zmax = top_index[ig, t_idx, isl] … … 473 442 mat[count_z:, t_idx] = np.nan 474 443 475 # Draw pcolormesh476 444 im = ax.pcolormesh( 477 445 date_time, … … 484 452 ) 485 453 ax.set_title(titles[idx], fontsize=12) 486 ax.set_xlabel("Time ( y)")454 ax.set_xlabel("Time (Mars years)") 487 455 ax.set_ylabel("Elevation (m)") 488 456 489 # Draw horizontal violet segments only where strata exist490 for t_idx in range(ntime - 1):491 h_vals = valid_tops_per_time[t_idx]492 if h_vals.size == 0:493 continue494 t_left = date_time[t_idx]495 t_right = date_time[t_idx + 1]496 for h in h_vals:497 # Only draw if h falls within sub_ref_grid498 if h < sub_ref_grid[0] or h > sub_ref_grid[-1]:499 continue500 ax.hlines(501 y=h,502 xmin=t_left,503 xmax=t_right,504 color='violet',505 linewidth=0.7,506 linestyle='-'507 )508 509 # Reserve extra space on the right for the colorbar510 457 fig.subplots_adjust(right=0.88) 511 512 # Place a single shared colorbar in its own axes513 458 cbar_ax = fig.add_axes([0.90, 0.15, 0.02, 0.7]) 514 fig.colorbar( 515 im, 516 cax=cbar_ax, 517 orientation='vertical', 518 label="Content" 519 ) 520 521 # Tight layout excluding the region we reserved (0.88) 459 fig.colorbar(im, cax=cbar_ax, orientation='vertical', label="Content") 522 460 fig.tight_layout(rect=[0, 0, 0.88, 1.0]) 523 461 … … 526 464 ) 527 465 fig.savefig(fname, dpi=150) 528 plt.show() 529 plt.close(fig) 466 467 468 def plot_strata_count_and_total_height(heights_data, date_time, output_folder="."): 469 """ 470 For each grid point and slope, plot: 471 - Number of strata vs time 472 - Total deposit height vs time 473 """ 474 ntime = len(heights_data) 475 nslope = len(heights_data[0]) 476 ngrid = heights_data[0][0].shape[0] 477 478 for ig in range(ngrid): 479 for isl in range(nslope): 480 n_strata_t = np.zeros(ntime, dtype=int) 481 total_height_t = np.zeros(ntime, dtype=float) 482 483 for t_idx in range(ntime): 484 h_mat = heights_data[t_idx][isl] 485 raw_h = h_mat[ig, :] 486 valid_mask = (~np.isnan(raw_h)) & (raw_h != 0.0) 487 if np.any(valid_mask): 488 h_valid = raw_h[valid_mask] 489 n_strata_t[t_idx] = h_valid.size 490 total_height_t[t_idx] = np.max(h_valid) 491 492 fig, axes = plt.subplots(2, 1, figsize=(8, 6), sharex=True) 493 fig.suptitle( 494 f"Strata count & total height over time for (Grid point {ig+1}, Slope {isl+1})", 495 fontsize=14 496 ) 497 498 axes[0].plot(date_time, n_strata_t, marker='+', linestyle='-') 499 axes[0].set_ylabel("Number of strata") 500 axes[0].grid(True) 501 502 axes[1].plot(date_time, total_height_t, marker='+', linestyle='-') 503 axes[1].set_xlabel("Time (Mars years)") 504 axes[1].set_ylabel("Total height (m)") 505 axes[1].grid(True) 506 507 fig.tight_layout(rect=[0, 0, 1, 0.95]) 508 fname = os.path.join( 509 output_folder, f"strata_count_height_ig{ig+1}_is{isl+1}.png" 510 ) 511 fig.savefig(fname, dpi=150) 512 513 514 def read_orbital_data(orb_file, martian_to_earth): 515 """ 516 Read the .asc file containing obliquity, eccentricity and Ls p. 517 Columns: 518 0 = time in thousand Martian years 519 1 = obliquity (deg) 520 2 = eccentricity 521 3 = Ls p (deg) 522 Converts times to Earth years. 523 """ 524 data = np.loadtxt(orb_file) 525 dates_mka = data[:, 0] 526 dates_yr = dates_mka * 1e3 / martian_to_earth 527 obliquity = data[:, 1] 528 eccentricity = data[:, 2] 529 lsp = data[:, 3] 530 return dates_yr, obliquity, eccentricity, lsp 531 532 533 def plot_orbital_parameters(infofile, orb_file, date_time, output_folder="."): 534 """ 535 Plot the evolution of obliquity, eccentricity and Ls p 536 versus simulated time. 537 """ 538 # Read conversion factor from infofile 539 _, martian_to_earth = read_infofile(infofile) 540 541 # Read orbital data 542 dates_yr, obl, ecc, lsp = read_orbital_data(orb_file, martian_to_earth) 543 544 # Interpolate orbital parameters at simulation dates (date_time) 545 obl_interp = interp1d(dates_yr, obl, kind='linear', bounds_error=False, fill_value="extrapolate")(date_time) 546 ecc_interp = interp1d(dates_yr, ecc, kind='linear', bounds_error=False, fill_value="extrapolate")(date_time) 547 lsp_interp = interp1d(dates_yr, lsp, kind='linear', bounds_error=False, fill_value="extrapolate")(date_time) 548 549 # Plot 550 fig, axes = plt.subplots(3, 1, figsize=(8, 10), sharex=True) 551 fig.suptitle("Orbital Parameters vs Simulated Time", fontsize=14) 552 553 axes[0].plot(date_time, obl_interp, 'r+', linestyle='-') 554 axes[0].set_ylabel("Obliquity (°)") 555 axes[0].grid(True) 556 557 axes[1].plot(date_time, ecc_interp, 'b+', linestyle='-') 558 axes[1].set_ylabel("Eccentricity") 559 axes[1].grid(True) 560 561 axes[2].plot(date_time, lsp_interp, 'g+', linestyle='-') 562 axes[2].set_ylabel("Ls p (°)") 563 axes[2].set_xlabel("Time (Mars years)") 564 axes[2].grid(True) 565 566 plt.tight_layout(rect=[0, 0, 1, 0.96]) 567 fname = os.path.join(output_folder, "orbital_parameters.png") 568 fig.savefig(fname, dpi=150) 530 569 531 570 532 571 def main(): 533 572 # 1) Get user inputs 534 folder_path, base_name, infofile = get_user_inputs()573 folder_path, base_name, infofile, orbfile = get_user_inputs() 535 574 536 575 # 2) List and verify NetCDF files 537 576 files = list_netcdf_files(folder_path, base_name) 538 577 if not files: 539 print(f"No NetCDF files named \"{base_name}#.nc\" found in \"{folder_path}\". Exiting.")578 print(f"No NetCDF files named \"{base_name}#.nc\" found in \"{folder_path}\".") 540 579 sys.exit(1) 541 nfile = len(files) 542 print(f"> Found {nfile} NetCDF file(s) matching \"{base_name}#.nc\".") 543 544 # 3) Open one sample to get ngrid, nslope, lon/lat 580 print(f"> Found {len(files)} NetCDF file(s).") 581 582 # 3) Open one sample to get grid dimensions & coordinates 545 583 sample_file = files[0] 546 584 ngrid, nslope, longitude, latitude = open_sample_dataset(sample_file) 547 print(f"> ngrid = {ngrid}") 548 print(f"> nslope = {nslope}") 549 550 # 4) Scan all files to collect variable info + global min/max elevations 551 var_info, max_nb_str, min_base_elev, max_top_elev = collect_stratification_variables( 552 files, base_name 553 ) 554 print(f"> max(nb_str_max) = {max_nb_str}") 555 print(f"> min(base_elevation) = {min_base_elev:.3f}") 556 print(f"> max(top_elevation) = {max_top_elev:.3f}") 557 558 # 5) Open all datasets for extraction 585 print(f"> ngrid = {ngrid}, nslope = {nslope}") 586 587 # 4) Collect variable info + global min/max elevations 588 var_info, max_nb_str, min_base_elev, max_top_elev = collect_stratification_variables(files, base_name) 589 print(f"> max strata per slope = {max_nb_str}") 590 print(f"> min base elev = {min_base_elev:.3f} m, max top elev = {max_top_elev:.3f} m") 591 592 # 5) Load full datasets 559 593 datasets = load_full_datasets(files) 560 594 561 # 6) Extract raw stratification data 562 heights_data, raw_prop_arrays, ntime = extract_stratification_data( 563 datasets, var_info, ngrid, nslope, max_nb_str 564 ) 565 566 # 7) Close all datasets 595 # 6) Extract stratification data 596 heights_data, raw_prop_arrays, ntime = extract_stratification_data(datasets, var_info, ngrid, nslope, max_nb_str) 597 598 # 7) Close datasets 567 599 for ds in datasets: 568 600 ds.close() 569 601 570 # 8) Normalize raw prop arrays to volumefractions602 # 8) Normalize to fractions 571 603 frac_arrays = normalize_to_fractions(raw_prop_arrays) 572 604 573 # 9) Ask whether to showsubsurface605 # 9) Ask whether to include subsurface 574 606 show_subsurface = get_yes_no_input("Show subsurface layers?") 575 607 exclude_sub = not show_subsurface 576 577 608 if exclude_sub: 578 609 min_base_for_interp = 0.0 579 print("> Will interpolateonly elevations >= 0 m (surface strata).")610 print("> Interpolating only elevations >= 0 m (surface strata).") 580 611 else: 581 612 min_base_for_interp = min_base_elev 582 print(f"> Will interpolate full depth (min base = {min_base_elev:.3f} m).")583 584 # 10) Prompt fordiscretization step613 print(f"> Interpolating full depth down to {min_base_elev:.3f} m.") 614 615 # 10) Prompt discretization step 585 616 dz = prompt_discretization_step(max_top_elev) 586 617 587 # 11) Build reference grid and interpolate (returns top_index as well)618 # 11) Build reference grid and interpolate 588 619 ref_grid, gridded_data, top_index = interpolate_data_on_refgrid( 589 heights_data, 590 frac_arrays, 591 min_base_for_interp, 592 max_top_elev, 593 dz, 594 exclude_sub=exclude_sub 620 heights_data, frac_arrays, min_base_for_interp, max_top_elev, dz, exclude_sub=exclude_sub 595 621 ) 596 622 597 # 12) Read time stamps from "info_PEM.txt"598 date_time = read_infofile(infofile)623 # 12) Read timestamps and conversion factor from infofile 624 date_time, martian_to_earth = read_infofile(infofile) 599 625 if date_time.size != ntime: 600 print( 601 "Warning: number of timestamps does not match number of NetCDF files " 602 f"({date_time.size} vs {ntime})." 603 ) 604 605 # 13) Plot and save figures (passing top_index and heights_data) 626 print(f"Warning: {date_time.size} timestamps vs {ntime} NetCDF files.") 627 628 # 13) Plot stratification over time 606 629 plot_stratification_over_time( 607 gridded_data, 608 ref_grid, 609 top_index, 610 heights_data, 611 date_time, 612 exclude_sub=exclude_sub, 613 output_folder="." 630 gridded_data, ref_grid, top_index, heights_data, date_time, 631 exclude_sub=exclude_sub, output_folder="." 614 632 ) 633 634 # 14) Plot strata count & total height vs time 635 plot_strata_count_and_total_height(heights_data, date_time, output_folder=".") 636 637 # 15) Plot orbital parameters 638 plot_orbital_parameters(infofile, orbfile, date_time, output_folder=".") 639 640 # 16) Show all figures 641 plt.show() 615 642 616 643
Note: See TracChangeset
for help on using the changeset viewer.
