Changeset 4192 for trunk/LMDZ.COMMON/libf/evolution/deftank/visu_layering_evo.py

Timestamp:

Apr 20, 2026, 3:58:44 PM (4 weeks ago)

Author:

jbclement

Message:

PEM:
Codebase for pore ice representation and evolution in the layered deposits:

• replace single pore ice volume fraction with coefficients of polymial expansion to describe the stratum pore ice profile;
• add the formation date in the stratum information;
• make default initialization for 'stratum';
• centralize the definition of stratum types;
• merge of similar strata;
• add interpolation and polynomial fitting tools;
• remove subsurface strata.

JBC

File:

• trunk/LMDZ.COMMON/libf/evolution/deftank/visu_layering_evo.py (modified) (8 diffs)

trunk/LMDZ.COMMON/libf/evolution/deftank/visu_layering_evo.py

@@ -97 +97 @@
 
 
+def _get_axis(dim_names, aliases):
+    """Return axis index matching aliases in a dimension-name list."""
+    for i, name in enumerate(dim_names):
+        lname = name.lower()
+        for alias in aliases:
+            if lname == alias or alias in lname:
+                return i
+    return None
+
+
+def _var_to_tsgi(nc_var):
+    """
+    Convert a NetCDF stratification variable to shape (time, slope, strata, grid).
+    Handles both slope-specific variables and generic variables with an nslope dimension.
+    """
+    arr = np.asarray(nc_var[:], dtype=np.float32)
+    dims = [d.lower() for d in nc_var.dimensions]
+
+    ax_time = _get_axis(dims, ('time',))
+    ax_slope = _get_axis(dims, ('nslope', 'slope'))
+    ax_str = _get_axis(dims, ('nb_str_max', 'nb_str'))
+    ax_grid = _get_axis(dims, ('physical_points', 'ngrid', 'grid'))
+
+    if arr.ndim == 4 and None not in (ax_time, ax_slope, ax_str, ax_grid):
+        return np.transpose(arr, (ax_time, ax_slope, ax_str, ax_grid))
+
+    if arr.ndim == 3 and None not in (ax_time, ax_str, ax_grid):
+        arr_tsg = np.transpose(arr, (ax_time, ax_str, ax_grid))
+        return arr_tsg[:, np.newaxis, :, :]
+
+    if arr.ndim == 3 and None not in (ax_slope, ax_str, ax_grid):
+        arr_ssg = np.transpose(arr, (ax_slope, ax_str, ax_grid))
+        return arr_ssg[np.newaxis, :, :, :]
+
+    if arr.ndim == 3:
+        # Historical fallback used by this script: (time, strata, grid)
+        return arr[:, np.newaxis, :, :]
+
+    if arr.ndim == 2:
+        return arr[np.newaxis, np.newaxis, :, :]
+
+    raise ValueError(f"Unsupported variable shape {arr.shape} for '{nc_var.name}'")
+
+
+def _slope_index_from_var(vname):
+    """Return 0-based slope index from a slope-specific variable name, else None."""
+    if 'slope' not in vname:
+        return None
+    try:
+        return int(vname.split('slope')[1].split('_')[0]) - 1
+    except (ValueError, IndexError):
+        return None
+
+
 def collect_stratification_variables(files, base_name):
     """
@@ -114 +168 @@
 
     property_markers = {
-        'heights':   'stratif_slope',    # "..._top_elevation"
+        'heights':   'top_elevation',
         'co2_ice':   'h_co2ice',
         'h2o_ice':   'h_h2oice',
         'dust':      'h_dust',
         'pore':      'h_pore',
-        'pore_ice':  'poreice_volfrac'
+        'poreice_coef1': 'poreice_coef1',
+        'poreice_coef2': 'poreice_coef2',
+        'poreice_coef3': 'poreice_coef3',
+        'poreice_coef4': 'poreice_coef4'
     }
     var_info = {prop: set() for prop in property_markers}
@@ -128 +185 @@
                 max_nb_str = max(max_nb_str, ds.dimensions['nb_str_max'].size)
 
-            nslope = ds.dimensions['nslope'].size
-            for k in range(1, nslope + 1):
-                var_name = f"stratif_slope{k:02d}_top_elevation"
-                if var_name in ds.variables:
-                    arr = ds.variables[var_name][:]
-                    min_base_elev = min(min_base_elev, np.min(arr))
-                    max_top_elev = max(max_top_elev, np.max(arr))
-                    var_info['heights'].add(var_name)
-
             for full_var in ds.variables:
+                if 'stratif' not in full_var:
+                    continue
+
+                if property_markers['heights'] in full_var:
+                    arr = np.asarray(ds.variables[full_var][:])
+                    if arr.size > 0:
+                        min_base_elev = min(min_base_elev, float(np.nanmin(arr)))
+                        max_top_elev = max(max_top_elev, float(np.nanmax(arr)))
+                    var_info['heights'].add(full_var)
+
                 for prop, marker in property_markers.items():
-                    if (marker in full_var) and prop != 'heights':
+                    if prop != 'heights' and marker in full_var:
                         var_info[prop].add(full_var)
 
@@ -174 +232 @@
     for t_idx, ds in enumerate(datasets):
         for var_name in var_info['heights']:
-            slope_idx = int(var_name.split("slope")[1].split("_")[0]) - 1
-            if 0 <= slope_idx < nslope:
-                raw = ds.variables[var_name][0, :, :]  # (n_strata, ngrid)
-                heights_data[t_idx][slope_idx] = raw.T  # (ngrid, n_strata)
+            if var_name not in ds.variables:
+                continue
+            raw_tsgi = _var_to_tsgi(ds.variables[var_name])
+            raw_sgi = raw_tsgi[0]  # (nslope_var, nstrata, ngrid)
+
+            slope_idx = _slope_index_from_var(var_name)
+            if slope_idx is not None and raw_sgi.shape[0] == 1:
+                if 0 <= slope_idx < nslope:
+                    heights_data[t_idx][slope_idx] = raw_sgi[0].T
+                continue
+
+            nsl = min(nslope, raw_sgi.shape[0])
+            for isl in range(nsl):
+                heights_data[t_idx][isl] = raw_sgi[isl].T
 
     raw_prop_arrays = {}
@@ -185 +253 @@
         raw_prop_arrays[prop] = np.zeros((ngrid, ntime, nslope, max_nb_str), dtype=np.float32)
 
-    def slope_index_from_var(vname):
-        return int(vname.split("slope")[1].split("_")[0]) - 1
-
     for prop in raw_prop_arrays:
-        slope_map = {}
-        for vname in var_info[prop]:
-            isl = slope_index_from_var(vname)
-            if 0 <= isl < nslope:
-                slope_map[isl] = vname
-
         arr = raw_prop_arrays[prop]
         for t_idx, ds in enumerate(datasets):
-            for isl, var_name in slope_map.items():
-                raw = ds.variables[var_name][0, :, :]  # (n_strata, ngrid)
-                n_strata_current = raw.shape[0]
-                arr[:, t_idx, isl, :n_strata_current] = raw.T
+            for var_name in var_info[prop]:
+                if var_name not in ds.variables:
+                    continue
+
+                raw_tsgi = _var_to_tsgi(ds.variables[var_name])
+                raw_sgi = raw_tsgi[0]  # (nslope_var, nstrata, ngrid)
+                slope_idx = _slope_index_from_var(var_name)
+
+                if slope_idx is not None and raw_sgi.shape[0] == 1:
+                    if 0 <= slope_idx < nslope:
+                        n_strata_current = min(raw_sgi.shape[1], max_nb_str)
+                        arr[:, t_idx, slope_idx, :n_strata_current] = raw_sgi[0, :n_strata_current, :].T
+                    continue
+
+                nsl = min(nslope, raw_sgi.shape[0])
+                for isl in range(nsl):
+                    n_strata_current = min(raw_sgi.shape[1], max_nb_str)
+                    arr[:, t_idx, isl, :n_strata_current] = raw_sgi[isl, :n_strata_current, :].T
 
     return heights_data, raw_prop_arrays, ntime
@@ -209 +282 @@
     Given raw_prop_arrays for 'co2_ice', 'h2o_ice', 'dust', 'pore' (in meters),
     normalize each set of strata so that the sum of those four = 1 per cell.
-    Returns:
-      - frac_arrays: dict mapping same keys -> 4D arrays of fractions (0..1).
+        Returns:
+            - frac_arrays: dict mapping same keys -> 4D arrays of fractions (0..1).
+            - includes an additional 'pore_ice' fraction derived from pore-ice coefficients.
     """
     co2 = raw_prop_arrays['co2_ice']
@@ -224 +298 @@
     frac_dust = np.zeros_like(dust, dtype=np.float32)
     frac_pore = np.zeros_like(pore, dtype=np.float32)
+    frac_pore_ice = np.zeros_like(pore, dtype=np.float32)
 
     frac_co2[mask] = co2[mask] / total[mask]
@@ -230 +305 @@
     frac_pore[mask] = pore[mask] / total[mask]
 
+    if all(key in raw_prop_arrays for key in ('poreice_coef1', 'poreice_coef2', 'poreice_coef3', 'poreice_coef4')):
+        z_norm = np.linspace(0.0, 1.0, 33, dtype=np.float32)
+        c1 = raw_prop_arrays['poreice_coef1']
+        c2 = raw_prop_arrays['poreice_coef2']
+        c3 = raw_prop_arrays['poreice_coef3']
+        c4 = raw_prop_arrays['poreice_coef4']
+        profile = (c1[..., None] + c2[..., None] * z_norm + c3[..., None] * z_norm**2 + c4[..., None] * z_norm**3)
+        profile = np.clip(profile, 0.0, 1.0)
+        poreice_mean = np.clip(np.trapz(profile, z_norm, axis=-1), 0.0, 1.0).astype(np.float32)
+    else:
+        poreice_mean = np.zeros_like(pore, dtype=np.float32)
+
+    frac_pore_ice[mask] = pore[mask] * poreice_mean[mask] / total[mask]
+
     return {
         'co2_ice': frac_co2,
         'h2o_ice': frac_h2o,
         'dust':     frac_dust,
-        'pore':     frac_pore
+        'pore':     frac_pore,
+        'pore_ice': frac_pore_ice
     }