Update profile_stratification.py

British-Oceanographic-Data-Centre · May 14, 2024 · 5227549 · 5227549
1 parent 0d46d1a
commit 5227549
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/coast/diagnostics/profile_stratification.py b/coast/diagnostics/profile_stratification.py
@@ -222,12 +222,15 @@ def calc_pea(
 
         if not "density" in profile.dataset:
             profile.construct_density(CT_AS=CT_AS, pot_dens=True)
+
+        # Update Zd_mask to exlude nan points
+        Zd_mask = Zd_mask.where(
+            np.isfinite(profile.dataset.variables["density"]), 0
+        )  
+
         if not "density_bar" in profile.dataset:
             profile.construct_density(CT_AS=CT_AS, rhobar=True, Zd_mask=Zd_mask, pot_dens=True)
         rho = profile.dataset.variables["density"].fillna(0)  # density
-        Zd_mask = Zd_mask.where(
-            np.isfinite(profile.dataset.variables["density"]), 0
-        )  # update Zd_mask to exclude nan pts
         rhobar = profile.dataset.variables["density_bar"]  # density with depth-mean T and S
 
         pot_energy_anom = (