Funcs and tuple for pv

PySDM/exporters/parcel_vtk_exporter.py

@@ -0,0 +1,127 @@
+"""
+VTK Exporter for Pyrcel PySDM simulations.
+
+This module defines `VTKExporterPyrcel`, a subclass of `PySDM.exporters.VTKExporter`,
+that writes simulation outputs to VTK format using `pyevtk`. It exports product
+profiles (e.g., relative humidity) as unstructured grids and particle attributes
+as point clouds, along with `.pvd` collection files for time-series visualization
+in ParaView.
+"""
+
+from pyevtk.hl import unstructuredGridToVTK, pointsToVTK
+from pyevtk.vtk import VtkHexahedron, VtkGroup
+import numpy as np
+
+from PySDM.exporters import VTKExporter
+
+
+class VTKExporterPyrcel(VTKExporter):
+    """
+    Custom VTK exporter for Pyrcel PySDM, exporting products as grids
+    and attributes as point clouds for ParaView visualization.
+    """
+
+    def __init__(self, n_sd, output, mass_of_dry_air):
+        super().__init__()
+        self.output = output
+        self.x = np.random.random(n_sd)  # pylint: disable=invalid-name
+        self.y = np.random.random(n_sd)  # pylint: disable=invalid-name
+        self.z = np.random.random(n_sd)  # pylint: disable=invalid-name
+        self.half_diagonal = []
+        self.n_levels = len(self.output["products"]["z"])
+
+        _volume = mass_of_dry_air / output["products"]["rhod"][0]
+        _dz = output["products"]["z"][1] - output["products"]["z"][0]
+        for level in range(self.n_levels):
+            if level > 0:
+                prev_to_curr_density_ratio = (
+                    output["products"]["rhod"][level - 1]
+                    / output["products"]["rhod"][level]
+                )
+                _volume *= prev_to_curr_density_ratio
+                _dz = (
+                    output["products"]["z"][level] - output["products"]["z"][level - 1]
+                )
+            _area = _volume / _dz
+            self.half_diagonal.append((2 * _area) ** 0.5)
+
+    def write_pvd(self):
+        pvd_attributes = VtkGroup(self.attributes_file_path)
+        for key, value in self.exported_times["attributes"].items():
+            pvd_attributes.addFile(key + ".vtu", sim_time=value)
+        pvd_attributes.save()
+
+        pvd_products = VtkGroup(self.products_file_path)
+        for key, value in self.exported_times["products"].items():
+            pvd_products.addFile(key + ".vtu", sim_time=value)
+        pvd_products.save()
+
+    def export_products(
+        self, step, simulation
+    ):  # pylint: disable=arguments-differ, too-many-locals
+        path = self.products_file_path + "_num" + self.add_leading_zeros(step)
+        self.exported_times["products"][path] = step * simulation.particulator.dt
+
+        n_vertices = 4 * self.n_levels
+        x = np.zeros(n_vertices)  # pylint: disable=invalid-name
+        y = np.zeros(n_vertices)  # pylint: disable=invalid-name
+        z = np.zeros(n_vertices)  # pylint: disable=invalid-name
+        conn = [0, 1, 2, 3]
+        ctype = np.zeros(self.n_levels - 1)
+        ctype[:] = VtkHexahedron.tid
+
+        for level in range(self.n_levels):
+            hd = self.half_diagonal[level]
+            _z = self.output["products"]["z"][level]
+            x[i := level * 4], y[i], z[i] = -hd, -hd, _z
+            x[i := i + 1], y[i], z[i] = -hd, hd, _z
+            x[i := i + 1], y[i], z[i] = hd, hd, _z
+            x[i := i + 1], y[i], z[i] = hd, -hd, _z
+            conn += [*range(4 * (level + 1), 4 * (level + 2))] * 2
+        conn = np.asarray(conn[:-4])
+        offset = np.asarray(range(8, 8 * self.n_levels, 8))
+
+        _ = {"test_pd": np.array([44] * n_vertices)}  # pointData
+
+        _RH = self.output["products"]["S_max_percent"]  # pylint: disable=invalid-name
+        cell_data = {"RH": np.full(shape=(len(_RH) - 1,), fill_value=np.nan)}
+        cell_data["RH"][:step] = (np.array(_RH[:-1] + np.diff(_RH) / 2))[:step]
+        unstructuredGridToVTK(
+            path,
+            x,
+            y,
+            z,
+            connectivity=conn,
+            offsets=offset,
+            cell_types=ctype,
+            cellData=cell_data,
+            # pointData=point_data,
+            # fieldData=field_data,
+        )
+
+    def export_attributes(self, step, simulation):  # pylint: disable=arguments-differ
+        path = self.attributes_file_path + "_num" + self.add_leading_zeros(step)
+        self.exported_times["attributes"][path] = step * simulation.particulator.dt
+
+        payload = {}
+
+        for k in self.output["attributes"].keys():
+            payload[k] = np.asarray(self.output["attributes"][k])[:, step].copy()
+        # payload["size"] = np.full(simulation.particulator.n_sd, 100.0)
+        if step != 0:
+            dz = (
+                self.output["products"]["z"][step]
+                - self.output["products"]["z"][step - 1]
+            )  # pylint: disable=invalid-name
+            if step == 1:
+                self.z *= dz
+            else:
+                self.z += dz
+
+        pointsToVTK(
+            path,
+            2 * (self.x - 0.5) * self.half_diagonal[step],
+            2 * (self.y - 0.5) * self.half_diagonal[step],
+            self.z,
+            data=payload,
+        )

examples/PySDM_examples/Pyrcel/paraview.ipynb

@@ -0,0 +1,140 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "6dcbcc6d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "c:\\Users\\strza\\Desktop\\PySDM\\examples\\PySDM_examples\\Pyrcel\\../../../PySDM/exporters/\n"
+     ]
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "from PySDM import Formulae\n",
+    "from PySDM.physics import si\n",
+    "from PySDM.initialisation.spectra import Lognormal\n",
+    "from PySDM.products import (\n",
+    "    ParcelDisplacement,\n",
+    "    AmbientTemperature,\n",
+    "    AmbientDryAirDensity,\n",
+    "    AmbientRelativeHumidity,\n",
+    "    ParticleSizeSpectrumPerVolume,\n",
+    "    ParticleVolumeVersusRadiusLogarithmSpectrum,\n",
+    ")\n",
+    "\n",
+    "from PySDM_examples.Pyrcel import Settings, Simulation\n",
+    "import sys\n",
+    "import os\n",
+    "\n",
+    "notebook_dir = os.path.dirname(os.path.abspath(\"__file__\"))\n",
+    "exporters_path = os.path.join(notebook_dir, \"../../../PySDM/exporters/\")\n",
+    "print(exporters_path)\n",
+    "sys.path.append(exporters_path)\n",
+    "\n",
+    "from parcel_vtk_exporter import VTKExporterPyrcel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "cd6288b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "settings = Settings(\n",
+    "    dz=10 * si.m,\n",
+    "    n_sd_per_mode=(25, 25),\n",
+    "    aerosol_modes_by_kappa={\n",
+    "        0.54: Lognormal(norm_factor=850 / si.cm**3, m_mode=15 * si.nm, s_geom=1.6),\n",
+    "        1.2: Lognormal(norm_factor=10 / si.cm**3, m_mode=850 * si.nm, s_geom=1.2),\n",
+    "    },\n",
+    "    vertical_velocity=1.0 * si.m / si.s,\n",
+    "    initial_pressure=775 * si.mbar,\n",
+    "    initial_temperature=274 * si.K,\n",
+    "    initial_relative_humidity=0.90,\n",
+    "    displacement=1000 * si.m,\n",
+    "    formulae=Formulae(constants={\"MAC\": 0.3}),\n",
+    ")\n",
+    "\n",
+    "dry_radius_bin_edges = np.logspace(\n",
+    "    np.log10(1e-3 * si.um), np.log10(5e0 * si.um), 33, endpoint=False\n",
+    ")\n",
+    "\n",
+    "simulation = Simulation(\n",
+    "    settings,\n",
+    "    products=(\n",
+    "        ParcelDisplacement(name=\"z\"),\n",
+    "        AmbientRelativeHumidity(name=\"S_max_percent\", unit=\"%\", var=\"RH\"),\n",
+    "        AmbientTemperature(name=\"T\"),\n",
+    "        ParticleSizeSpectrumPerVolume(\n",
+    "            name=\"dry:dN/dR\", radius_bins_edges=dry_radius_bin_edges, dry=True\n",
+    "        ),\n",
+    "        ParticleVolumeVersusRadiusLogarithmSpectrum(\n",
+    "            name=\"dry:dV/dlnR\", radius_bins_edges=dry_radius_bin_edges, dry=True\n",
+    "        ),\n",
+    "        AmbientDryAirDensity(),\n",
+    "    ),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "3ff50a25",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "output = simulation.run()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "15d59657",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "e = VTKExporterPyrcel(n_sd=simulation.particulator.n_sd, output=output, mass_of_dry_air=simulation.particulator.environment.mass_of_dry_air)\n",
+    "for step in settings.output_steps:\n",
+    "    e.export_products(step, simulation)\n",
+    "    e.export_attributes(step, simulation)\n",
+    "e.write_pvd()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4a6ea0e6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}