fix exception messages and more fixtures in tests

Author: Hevagog

tests/examples_tests/Loftus_and_Wordsworth_2021/accuracy_test.py

@@ -1,7 +1,11 @@
+from __future__ import annotations
+
 import pytest
 import os
 import numpy as np
 from scipy.optimize import fsolve
+import warnings
+from typing import Tuple, List, Generator
 
 from PySDM import Formulae
 from PySDM.physics import si
@@ -10,13 +14,15 @@
 
 
 class GroundTruthLoader:
-    def __init__(self, groundtruth_dir_path, n_samples=2, random_seed=2137):
+    def __init__(
+        self, groundtruth_dir_path: str, n_samples: int = 2, random_seed: int = 2137
+    ):
         self.dir_path = groundtruth_dir_path
         self.RHs = None
         self.r0grid = None
         self.m_frac_evap = None
-        self.n_samples = n_samples  # Number of random samples to test
-        np.random.seed(random_seed)  # reproducible random samples during debugging
+        self.n_samples = n_samples
+        np.random.seed(random_seed)
 
     def __enter__(self):
         try:
@@ -25,16 +31,64 @@ def __enter__(self):
             self.m_frac_evap = np.load(os.path.join(self.dir_path, "m_frac_evap.npy"))
             return self
         except FileNotFoundError as e:
-            print(f"Error loading ground truth files: {e}")
-            raise
+            pytest.fail(f"Error loading ground truth files: {e}")
         except Exception as e:
-            print(f"An unexpected error occurred while loading ground truth data: {e}")
-            raise
+            pytest.fail(f"Unexpected error loading ground truth data: {e}")
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         pass
 
 
+@pytest.fixture(scope="module")
+def ground_truth_data(request) -> Generator[GroundTruthLoader, None, None]:
+    current_dir = os.path.dirname(os.path.abspath(request.fspath))
+    groundtruth_dir = os.path.abspath(os.path.join(current_dir, "ground_truth"))
+    if not os.path.isdir(groundtruth_dir):
+        pytest.fail(f"Groundtruth directory not found at {groundtruth_dir}")
+    with GroundTruthLoader(groundtruth_dir) as gt:
+        yield gt
+
+
+@pytest.fixture
+def ground_truth_sample(ground_truth_data: GroundTruthLoader) -> List[dict]:
+    gt = ground_truth_data
+    n_rh_values = len(gt.RHs)
+    n_radius_values = gt.r0grid.shape[1]
+    total_points = n_rh_values * n_radius_values
+    n_samples = min(gt.n_samples, total_points)
+    if n_samples == 0:
+        pytest.skip("No data points available to sample.")
+    all_indices = np.array(
+        [(i, j) for i in range(n_rh_values) for j in range(n_radius_values)]
+    )
+    sampled_indices_flat = np.random.choice(len(all_indices), n_samples, replace=False)
+    sampled_ij_pairs = all_indices[sampled_indices_flat]
+    return [
+        {
+            "rh": gt.RHs[i_rh],
+            "r_m": gt.r0grid[0, j_r],
+            "expected_m_frac_evap": gt.m_frac_evap[i_rh, j_r],
+            "i_rh": i_rh,
+            "j_r": j_r,
+        }
+        for i_rh, j_r in sampled_ij_pairs
+    ]
+
+
+@pytest.fixture(scope="module")
+def static_arrays() -> Tuple[np.ndarray, np.ndarray, np.ndarray, object]:
+    formulae = Formulae(
+        ventilation="PruppacherAndRasmussen1979",
+        saturation_vapour_pressure="AugustRocheMagnus",
+        diffusion_coordinate="WaterMassLogarithm",
+    )
+    radius_array = np.logspace(-4.5, -2.5, 50) * si.m
+    RH_array = np.linspace(0.25, 0.99, 50)
+    output_matrix = np.full((len(RH_array), len(radius_array)), np.nan)
+    const = formulae.constants
+    return radius_array, RH_array, output_matrix, const
+
+
 class TestNPYComparison:
     @staticmethod
     def _mix(dry_prop, vap_prop, ratio):
@@ -44,15 +98,11 @@ def _calculate_cloud_properties(
         self, planet: EarthLike, surface_RH: float, formulae_instance: Formulae
     ):
         const = formulae_instance.constants
-
         planet.RH_zref = surface_RH
-
         pvs_stp = formulae_instance.saturation_vapour_pressure.pvs_water(planet.T_STP)
-
         initial_water_vapour_mixing_ratio = const.eps / (
             planet.p_STP / planet.RH_zref / pvs_stp - 1
         )
-
         R_air_mix = self._mix(const.Rd, const.Rv, initial_water_vapour_mixing_ratio)
         cp_mix = self._mix(const.c_pd, const.c_pv, initial_water_vapour_mixing_ratio)
 
@@ -78,154 +128,113 @@ def solve_Tcloud(T_candidate):
         )
         return initial_water_vapour_mixing_ratio, Tcloud, Zcloud, pcloud
 
-    def test_figure_2_replication_accuracy(self):
-        current_dir = os.path.dirname(os.path.abspath(__file__))
-        groundtruth_dir = os.path.abspath(os.path.join(current_dir, "ground_truth"))
-        if not os.path.isdir(groundtruth_dir):
-            pytest.fail(f"Groundtruth directory not found at {groundtruth_dir}")
-
+    def test_figure_2_replication_accuracy(self, ground_truth_sample, static_arrays):
         formulae = Formulae(
             ventilation="PruppacherAndRasmussen1979",
             saturation_vapour_pressure="AugustRocheMagnus",
             diffusion_coordinate="WaterMassLogarithm",
         )
-
-        with GroundTruthLoader(groundtruth_dir) as gt:
-            if gt.RHs is None or gt.r0grid is None or gt.m_frac_evap is None:
-                pytest.fail("Ground truth data (.npy files) not loaded properly.")
-
-            n_rh_values = len(gt.RHs)
-            n_radius_values = gt.r0grid.shape[1]
-
-            total_points = n_rh_values * n_radius_values
-            n_samples = min(gt.n_samples, total_points)
-
-            if n_samples == 0:
-                pytest.skip("No data points available to sample.")
-
-            all_indices = np.array(
-                [(i, j) for i in range(n_rh_values) for j in range(n_radius_values)]
+        for sample in ground_truth_sample:
+            planet = EarthLike()
+            rh = sample["rh"]
+            r_m = sample["r_m"]
+            expected = sample["expected_m_frac_evap"]
+            i_rh = sample["i_rh"]
+            j_r = sample["j_r"]
+            try:
+                iwvmr, Tcloud, Zcloud, pcloud = self._calculate_cloud_properties(
+                    planet, rh, formulae
+                )
+                simulated = self.calc_simulated_m_frac_evap_point(
+                    planet,
+                    formulae,
+                    i_rh,
+                    j_r,
+                    rh,
+                    r_m,
+                    expected,
+                    iwvmr,
+                    Tcloud,
+                    Zcloud,
+                    pcloud,
+                )
+            except Exception as e:
+                pytest.fail(
+                    f"Error in _calculate_cloud_properties for RH={rh} (sample idx {i_rh},{j_r}): {e}."
+                )
+            error_context = (
+                f"Sample (RH_idx={i_rh}, R_idx={j_r}), RH={rh:.4f}, R_m={r_m:.3e}. "
+                f"Expected: {expected}, Got: {simulated}"
             )
-
-            sampled_indices_flat = np.random.choice(len(all_indices), n_samples, replace=False)
-            sampled_ij_pairs = all_indices[sampled_indices_flat]
-
-            for i_rh, j_r in sampled_ij_pairs:
-                current_planet_state = EarthLike()
-
-                current_rh = gt.RHs[i_rh]
-                current_r_m = gt.r0grid[0, j_r]
-                expected_m_frac_evap = gt.m_frac_evap[i_rh, j_r]
-
-                try:
-                    iwvmr, Tcloud, Zcloud, pcloud = self._calculate_cloud_properties(
-                        current_planet_state, current_rh, formulae
-                    )
-                    simulated_m_frac_evap_point = self.calc_simulated_m_frac_evap_point(
-                        current_planet_state,
-                        formulae,
-                        i_rh,
-                        j_r,
-                        current_rh,
-                        current_r_m,
-                        expected_m_frac_evap,
-                        iwvmr,
-                        Tcloud,
-                        Zcloud,
-                        pcloud,
-                    )
-                except Exception as e:
-                    print(
-                        f"Warning: Error in _calculate_cloud_properties for RH={current_rh} (sample idx {i_rh},{j_r}): {e}."
-                    )
-
-                error_context = (
-                    f"Sample (RH_idx={i_rh}, R_idx={j_r}), "
-                    f"RH={current_rh:.4f}, R_m={current_r_m:.3e}. "
-                    f"Expected: {expected_m_frac_evap}, Got: {simulated_m_frac_evap_point}"
+            if np.isnan(expected):
+                assert np.isnan(
+                    simulated
+                ), f"NaN Mismatch. {error_context} (Expected NaN, got non-NaN)"
+            else:
+                assert not np.isnan(
+                    simulated
+                ), f"NaN Mismatch. {error_context} (Expected non-NaN, got NaN)"
+                np.testing.assert_allclose(
+                    simulated,
+                    expected,
+                    rtol=1e-1,
+                    atol=1e-1,
+                    err_msg=f"Value Mismatch. {error_context}",
                 )
 
-                if np.isnan(expected_m_frac_evap):
-                    assert np.isnan(
-                        simulated_m_frac_evap_point
-                    ), f"NaN Mismatch. {error_context} (Expected NaN, got non-NaN)"
-                else:
-                    assert not np.isnan(
-                        simulated_m_frac_evap_point
-                    ), f"NaN Mismatch. {error_context} (Expected non-NaN, got NaN)"
-                    np.testing.assert_allclose(
-                        simulated_m_frac_evap_point,
-                        expected_m_frac_evap,
-                        rtol=1e-1,  # Relative tolerance
-                        atol=1e-1,  # Absolute tolerance
-                        err_msg=f"Value Mismatch. {error_context}",
-                    )
-
     def calc_simulated_m_frac_evap_point(
         self,
-        current_planet_state,
+        planet,
         formulae,
         i_rh,
         j_r,
-        current_rh,
-        current_r_m,
-        expected_m_frac_evap,
+        rh,
+        r_m,
+        expected,
         iwvmr,
         Tcloud,
         Zcloud,
         pcloud,
     ):
-
-        simulated_m_frac_evap_point = np.nan
-
-        if np.isnan(current_r_m) or current_r_m <= 0:
-            print(f"Warning: Invalid radius current_r_m={current_r_m} for sample idx {i_rh},{j_r}.")
-        else:
-            settings = Settings(
-                planet=current_planet_state,
-                r_wet=current_r_m,
-                mass_of_dry_air=1e5 * si.kg,
-                initial_water_vapour_mixing_ratio=iwvmr,
-                pcloud=pcloud,
-                Zcloud=Zcloud,
-                Tcloud=Tcloud,
-                formulae=formulae,
-            )
-            simulation = Simulation(settings)
-
-            try:
-                output = simulation.run()
-                if (
-                    output
-                    and "r" in output
-                    and len(output["r"]) > 0
-                    and "z" in output
-                    and len(output["z"]) > 0
-                ):
-                    final_radius_um = output["r"][-1]
-                    if np.isnan(final_radius_um) or final_radius_um < 0:
-                        final_radius_m = final_radius_um * 1e-6
-                        if final_radius_m < 0:  # Non-physical radius
-                            simulated_m_frac_evap_point = 1.0  # 1.0 means fully evaporated
-                        else:
-                            simulated_m_frac_evap_point = np.nan
-                    else:
-                        final_radius_m = final_radius_um * 1e-6
-                        if current_r_m == 0:
-                            frac_evap = 1.0
-                        else:
-                            frac_evap = 1.0 - (final_radius_m / current_r_m) ** 3
-                        frac_evap = np.clip(frac_evap, 0.0, 1.0)
-                        simulated_m_frac_evap_point = frac_evap
-                else:
-                    simulated_m_frac_evap_point = np.nan
-            except Exception as e:
-                print(
-                    f"Warning: Simulation run failed for RH={current_rh}, r={current_r_m} (sample idx {i_rh},{j_r}): {e}."
-                )
-                if np.isclose(expected_m_frac_evap, 1.0, atol=1e-6):
-                    simulated_m_frac_evap_point = 1.0
+        if np.isnan(r_m) or r_m <= 0:
+            pytest.fail(f"Invalid radius r_m={r_m} for sample idx {i_rh},{j_r}.")
+        settings = Settings(
+            planet=planet,
+            r_wet=r_m,
+            mass_of_dry_air=1e5 * si.kg,
+            initial_water_vapour_mixing_ratio=iwvmr,
+            pcloud=pcloud,
+            Zcloud=Zcloud,
+            Tcloud=Tcloud,
+            formulae=formulae,
+        )
+        simulation = Simulation(settings)
+        try:
+            output = simulation.run()
+            if (
+                output
+                and "r" in output
+                and len(output["r"]) > 0
+                and "z" in output
+                and len(output["z"]) > 0
+            ):
+                final_radius_um = output["r"][-1]
+                if np.isnan(final_radius_um) or final_radius_um < 0:
+                    final_radius_m = final_radius_um * 1e-6
+                    if final_radius_m < 0:
+                        return 1.0
+                    return np.nan
+                final_radius_m = final_radius_um * 1e-6
+                if r_m == 0:
+                    frac_evap = 1.0
                 else:
-                    simulated_m_frac_evap_point = np.nan
-
-        return simulated_m_frac_evap_point
+                    frac_evap = 1.0 - (final_radius_m / r_m) ** 3
+                return np.clip(frac_evap, 0.0, 1.0)
+            return np.nan
+        except Exception as e:
+            warnings.warn(
+                f"Simulation run failed for RH={rh:.4f}, r={r_m:.3e} (sample idx {i_rh},{j_r}): {type(e).__name__}: {e}"
+            )
+            if np.isclose(expected, 1.0, atol=1e-6):
+                return 1.0
+            return np.nan