surface precip product compatibility in S&H example and netCDF exporter + fix in CI examples-setup editable instal (#1322)

slayoo · web-flow · commit 063feedf546d · 2024-05-12T15:17:53.000+02:00
diff --git a/.github/workflows/tests+artifacts+pypi.yml b/.github/workflows/tests+artifacts+pypi.yml
@@ -198,7 +198,7 @@ jobs:
       - if: steps.cache.outputs.cache-hit != 'true'
         run: |
           python -m pip uninstall -y jupyterlab-server # https://github.com/pypa/pip/issues/6275
-          python -m pip install -e . ./examples
+          python -m pip install -e . -e ./examples
           python -m pip install -r tests/devops_tests/requirements.txt
       - if: steps.cache.outputs.cache-hit != 'true'
         uses: actions/cache/save@v4
diff --git a/PySDM/exporters/netcdf_exporter_1d.py b/PySDM/exporters/netcdf_exporter_1d.py
@@ -75,7 +75,9 @@ def _create_variables(self, ncdf):
                 )
 
             n_dimensions = len(instance.shape)
-            if n_dimensions == 1:
+            if n_dimensions == 0:
+                dimensions = ("time",)
+            elif n_dimensions == 1:
                 dimensions = ("height", "time")
             elif n_dimensions == 2:
                 dim_name = name.replace(" ", "_") + "_bin_index"
@@ -94,7 +96,9 @@ def _create_variables(self, ncdf):
     def _write_variables(self):
         for var in self.simulator.particulator.products.keys():
             n_dimensions = len(self.simulator.particulator.products[var].shape)
-            if n_dimensions == 1:
+            if n_dimensions == 0:
+                self.vars[var][:] = self.data[var][:]
+            elif n_dimensions == 1:
                 self.vars[var][:, :] = self.data[var][-self.nz_export :, :]
             elif n_dimensions == 2:
                 if self.n_save_spec == 0:
diff --git a/examples/PySDM_examples/Shipway_and_Hill_2012/fig_1.ipynb b/examples/PySDM_examples/Shipway_and_Hill_2012/fig_1.ipynb
@@ -44,6 +44,15 @@
     "    pip_install_on_colab('PySDM-examples')"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 2,
@@ -85,6 +94,9 @@
     "    \"kappa\": .9,\n",
     "    \"particles_per_volume_STP\": 50 / si.cm**3\n",
     "}\n",
+    "if 'CI' in os.environ:\n",
+    "    common_params[\"t_max\"] = 10 * common_params[\"dt\"]\n",
+    "    common_params[\"n_sd_per_gridbox\"] = 16 \n",
     "\n",
     "output = {}\n",
     "settings = {}\n",
@@ -104,7 +116,7 @@
     "            precip=precip\n",
     "        )\n",
     "        simulation[key] = Simulation(settings[key])\n",
-    "        output[key] = simulation[key].run().products\n"
+    "        output[key] = simulation[key].run().products"
    ]
   },
   {
@@ -327,7 +339,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.9.4 ('edjPySDM')",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -341,7 +353,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.4"
+   "version": "3.9.2"
   },
   "vscode": {
    "interpreter": {
@@ -350,5 +362,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }
diff --git a/examples/PySDM_examples/Shipway_and_Hill_2012/simulation.py b/examples/PySDM_examples/Shipway_and_Hill_2012/simulation.py
@@ -168,6 +168,7 @@ def zZ_to_z_above_reservoir(zZ):
                     PySDM_products.CoalescenceRatePerGridbox(
                         name="coalescence_rate",
                     ),
+                    PySDM_products.SurfacePrecipitation(),
                 ]
             )
         self.particulator = self.builder.build(
@@ -182,7 +183,9 @@ def zZ_to_z_above_reservoir(zZ):
         }
         self.output_products = {}
         for k, v in self.particulator.products.items():
-            if len(v.shape) == 1:
+            if len(v.shape) == 0:
+                self.output_products[k] = np.zeros(self.nt + 1)
+            elif len(v.shape) == 1:
                 self.output_products[k] = np.zeros((self.mesh.grid[-1], self.nt + 1))
             elif len(v.shape) == 2:
                 number_of_time_sections = len(self.save_spec_and_attr_times)
@@ -203,7 +206,10 @@ def save_scalar(self, step):
         for k, v in self.particulator.products.items():
             if len(v.shape) > 1:
                 continue
-            self.output_products[k][:, step] = v.get()
+            if len(v.shape) == 1:
+                self.output_products[k][:, step] = v.get()
+            else:
+                self.output_products[k][step] = v.get()
 
     def save_spectrum(self, index):
         for k, v in self.particulator.products.items():
diff --git a/tests/smoke_tests/kinematic_1d/shipway_and_hill_2012/test_1d_exporters.py b/tests/smoke_tests/kinematic_1d/shipway_and_hill_2012/test_1d_exporters.py
@@ -1,6 +1,8 @@
 # pylint: disable = missing-module-docstring,missing-class-docstring,missing-function-docstring
 import os
 import platform
+from collections import namedtuple
+from pathlib import Path
 from tempfile import TemporaryDirectory
 
 import numpy as np
@@ -20,127 +22,138 @@ def simulation_1d_fixture():
         n_sd_per_gridbox=n_sd_per_gridbox,
         dt=60 * si.s,
         dz=200 * si.m,
-        precip=False,
+        precip=True,
         rho_times_w_1=2 * si.m / si.s * si.kg / si.m**3,
     )
     settings.t_max = 20 * settings.dt
     settings.save_spec_and_attr_times = [0 * si.min, 10 * si.min, 20 * si.min]
     simulation = Simulation(settings)
     results = simulation.run()
-    return results, settings, simulation
+    return namedtuple(
+        Path(__file__).stem + "_FixtureData", ("results, settings, simulation")
+    )(results=results, settings=settings, simulation=simulation)
 
 
-@pytest.mark.parametrize(
-    "exclude_particle_reservoir",
-    (
-        False,
-        True,
-    ),
-)
-def test_netcdf_exporter_1d(simulation_1d, exclude_particle_reservoir):
-    # Arrange
-    data = simulation_1d[0].products
-    settings = simulation_1d[1]
-    simulation = simulation_1d[2]
-    nz_export = (
-        int(settings.z_max / settings.dz) if exclude_particle_reservoir else settings.nz
+class Test2DExporters:
+    @staticmethod
+    @pytest.mark.parametrize(
+        "exclude_particle_reservoir",
+        (
+            False,
+            True,
+        ),
     )
-
-    # Act
-    file = TemporaryFile(".nc")
-    netcdf_exporter = NetCDFExporter_1d(
-        data,
-        settings,
-        simulation,
-        filename=file.absolute_path,
-        exclude_particle_reservoir=exclude_particle_reservoir,
-    )
-    netcdf_exporter.run()
-    data_from_file = readNetCDF_1d(file.absolute_path)
-
-    # Assert
-    assert data_from_file.products["time"].shape == data["t"].shape
-    assert data_from_file.products["height"].shape == data["z"][-nz_export:].shape
-    assert data_from_file.products["T"].shape == data["T"][-nz_export:, :].shape
-    assert (
-        data_from_file.products["dry spectrum"].shape
-        == data["dry spectrum"][-nz_export:, :, :].shape
-    )
-
-    assert np.amin(data_from_file.products["time"]) == np.amin(data["t"])
-    assert np.amax(data_from_file.products["time"]) == np.amax(data["t"])
-    assert np.amin(data_from_file.products["height"]) == np.amin(data["z"][-nz_export:])
-    assert np.amax(data_from_file.products["height"]) == np.amax(data["z"][-nz_export:])
-    assert data_from_file.products["rhod"].mean() == pytest.approx(
-        data["rhod"][-nz_export:, :].mean(), 1e-6
-    )
-    assert data_from_file.products["wet spectrum"].mean() == pytest.approx(
-        data["wet spectrum"][-nz_export:, :, :].mean(), 1e-6
-    )
-
-    assert data_from_file.settings["precip"] == settings.precip
-    assert data_from_file.settings["kappa"] == settings.kappa
-    assert data_from_file.settings["r_bins_edges"].size == settings.number_of_bins + 1
-
-
-@pytest.mark.skipif(
-    platform.architecture()[0] == "32bit", reason="Not available vtk module!"
-)
-@pytest.mark.parametrize(
-    "exclude_particle_reservoir",
-    (
-        False,
-        True,
-    ),
-)
-def test_vtk_exporter_1d(
-    simulation_1d, exclude_particle_reservoir
-):  # pylint: disable=too-many-locals
-    # Arrange
-    data = simulation_1d[0].attributes
-    settings = simulation_1d[1]
-    z0 = 0.0 if exclude_particle_reservoir else -settings.particle_reservoir_depth
-    exported_particles_indexes = {}
-    number_of_exported_particles = []
-    for i, t in enumerate(settings.save_spec_and_attr_times):
-        exported_particles_indexes[t] = np.where(
-            data["cell origin"][i][0]
-            >= int((z0 + settings.particle_reservoir_depth) / settings.dz)
+    def test_netcdf_exporter_1d(simulation_1d, exclude_particle_reservoir):
+        # Arrange
+        data = simulation_1d.results.products
+        settings = simulation_1d.settings
+        nz_export = (
+            int(settings.z_max / settings.dz)
+            if exclude_particle_reservoir
+            else settings.nz
         )
-        number_of_exported_particles.append(exported_particles_indexes[t][0].size)
 
-    # Act
-    with TemporaryDirectory() as tmpdir_:
-        tmpdir = tmpdir_ + "/"
-        vtk_exporter = VTKExporter_1d(
+        # Act
+        file = TemporaryFile(".nc")
+        netcdf_exporter = NetCDFExporter_1d(
             data,
             settings,
-            path=tmpdir,
+            simulation_1d.simulation,
+            filename=file.absolute_path,
             exclude_particle_reservoir=exclude_particle_reservoir,
         )
-        vtk_exporter.run()
-        written_files_list = os.listdir(tmpdir)
-        data_from_file = {}
-        for t in settings.save_spec_and_attr_times:
-            leading_zeros_in_filename = [
-                "0" for i in range(len(str(settings.t_max)) - len(str(t)))
-            ]
-            filename = "time" + "".join(leading_zeros_in_filename) + str(t) + ".vtu"
-            data_from_file[t] = readVTK_1d(tmpdir + filename)
+        netcdf_exporter.run()
+        data_from_file = readNetCDF_1d(file.absolute_path)
+
+        # Assert
+        assert data_from_file.products["time"].shape == data["t"].shape
+        assert data_from_file.products["height"].shape == data["z"][-nz_export:].shape
+        assert data_from_file.products["T"].shape == data["T"][-nz_export:, :].shape
+        assert (
+            data_from_file.products["dry spectrum"].shape
+            == data["dry spectrum"][-nz_export:, :, :].shape
+        )
 
-    # Assert
-    for i, t in enumerate(settings.save_spec_and_attr_times):
-        filename_leading_zeros = "".join(
-            ["0" for i in range(len(str(settings.t_max)) - len(str(t)))]
+        assert np.amin(data_from_file.products["time"]) == np.amin(data["t"])
+        assert np.amax(data_from_file.products["time"]) == np.amax(data["t"])
+        assert np.amin(data_from_file.products["height"]) == np.amin(
+            data["z"][-nz_export:]
+        )
+        assert np.amax(data_from_file.products["height"]) == np.amax(
+            data["z"][-nz_export:]
         )
-        filename = "time" + filename_leading_zeros + str(t) + ".vtu"
-        assert filename in written_files_list
-        assert z0 <= np.amin(data_from_file[t]["z"])
-        assert np.amax(data_from_file[t]["z"]) <= settings.z_max
-        assert data_from_file[t]["z"].size == number_of_exported_particles[i]
-        data_from_file[t].pop("z")
-        assert data.keys() == data_from_file[t].keys()
+        assert data_from_file.products["rhod"].mean() == pytest.approx(
+            data["rhod"][-nz_export:, :].mean(), 1e-6
+        )
+        assert data_from_file.products["wet spectrum"].mean() == pytest.approx(
+            data["wet spectrum"][-nz_export:, :, :].mean(), 1e-6
+        )
+
+        assert data_from_file.settings["precip"] == settings.precip
+        assert data_from_file.settings["kappa"] == settings.kappa
         assert (
-            data["radius"][i][exported_particles_indexes[t]].mean()
-            == data_from_file[t]["radius"].mean()
+            data_from_file.settings["r_bins_edges"].size == settings.number_of_bins + 1
         )
+
+    @staticmethod
+    @pytest.mark.skipif(
+        platform.architecture()[0] == "32bit", reason="Not available vtk module!"
+    )
+    @pytest.mark.parametrize(
+        "exclude_particle_reservoir",
+        (
+            False,
+            True,
+        ),
+    )
+    def test_vtk_exporter_1d(
+        simulation_1d, exclude_particle_reservoir
+    ):  # pylint: disable=too-many-locals
+        # Arrange
+        data = simulation_1d.results.attributes
+        settings = simulation_1d[1]
+        z0 = 0.0 if exclude_particle_reservoir else -settings.particle_reservoir_depth
+        exported_particles_indexes = {}
+        number_of_exported_particles = []
+        for i, t in enumerate(settings.save_spec_and_attr_times):
+            exported_particles_indexes[t] = np.where(
+                data["cell origin"][i][0]
+                >= int((z0 + settings.particle_reservoir_depth) / settings.dz)
+            )
+            number_of_exported_particles.append(exported_particles_indexes[t][0].size)
+
+        # Act
+        with TemporaryDirectory() as tmpdir_:
+            tmpdir = tmpdir_ + "/"
+            vtk_exporter = VTKExporter_1d(
+                data,
+                settings,
+                path=tmpdir,
+                exclude_particle_reservoir=exclude_particle_reservoir,
+            )
+            vtk_exporter.run()
+            written_files_list = os.listdir(tmpdir)
+            data_from_file = {}
+            for t in settings.save_spec_and_attr_times:
+                leading_zeros_in_filename = [
+                    "0" for i in range(len(str(settings.t_max)) - len(str(t)))
+                ]
+                filename = "time" + "".join(leading_zeros_in_filename) + str(t) + ".vtu"
+                data_from_file[t] = readVTK_1d(tmpdir + filename)
+
+        # Assert
+        for i, t in enumerate(settings.save_spec_and_attr_times):
+            filename_leading_zeros = "".join(
+                ["0" for i in range(len(str(settings.t_max)) - len(str(t)))]
+            )
+            filename = "time" + filename_leading_zeros + str(t) + ".vtu"
+            assert filename in written_files_list
+            assert z0 <= np.amin(data_from_file[t]["z"])
+            assert np.amax(data_from_file[t]["z"]) <= settings.z_max
+            assert data_from_file[t]["z"].size == number_of_exported_particles[i]
+            data_from_file[t].pop("z")
+            assert data.keys() == data_from_file[t].keys()
+            assert (
+                data["radius"][i][exported_particles_indexes[t]].mean()
+                == data_from_file[t]["radius"].mean()
+            )
diff --git a/tests/smoke_tests/kinematic_1d/shipway_and_hill_2012/test_few_steps.py b/tests/smoke_tests/kinematic_1d/shipway_and_hill_2012/test_few_steps.py
@@ -23,7 +23,7 @@ def test_few_steps_no_precip(particle_reservoir_depth, plot=False):
         n_sd_per_gridbox=n_sd_per_gridbox,
         dt=30 * si.s,
         dz=60 * si.m,
-        precip=False,
+        precip=True,
         rho_times_w_1=2 * si.m / si.s * si.kg / si.m**3,
     )
     settings.particle_reservoir_depth = particle_reservoir_depth
@@ -80,6 +80,7 @@ def mean_profile_over_last_steps(var, smooth=True):
     assert max(mean_profile_over_last_steps("activating")) == 0
     assert max(mean_profile_over_last_steps("ripening")) > 0
     assert max(mean_profile_over_last_steps("deactivating")) > 0
+    assert max(output["surface precipitation"]) == 0
 
 
 def test_fixed_thd():
