[com1]: add option to only save certain particle props

avaframe/com1DFA/com1DFA.py

@@ -10,7 +10,6 @@
 import os
 import pathlib
 import pickle
-import platform
 import re
 import time
 from datetime import datetime
@@ -2159,7 +2158,7 @@ def DFAIterate(cfg, particles, fields, dem, inputSimLines, outDir, cuSimName, si
         exportFields(cfg, t, fields, dem, outDir, cuSimName, TSave="initial")
 
         if "particles" in resTypes:
-            savePartToPickle(particles, outDirData, cuSimName)
+            savePartToPickle(particles, outDirData, cuSimName, cfg=cfg)
 
         # Update dtSave to remove the initial timestep we just saved
         dtSave = updateSavingTimeStep(dtSaveOriginal, cfgGen, t)
@@ -2284,7 +2283,7 @@ def DFAIterate(cfg, particles, fields, dem, inputSimLines, outDir, cuSimName, si
 
                 # export particles dictionaries of saving time steps
                 if "particles" in resTypes:
-                    savePartToPickle(particles, outDirData, cuSimName)
+                    savePartToPickle(particles, outDirData, cuSimName, cfg=cfg)
 
             # export particles properties for visulation
             if cfg["VISUALISATION"].getboolean("writePartToCSV"):
@@ -2416,7 +2415,7 @@ def DFAIterate(cfg, particles, fields, dem, inputSimLines, outDir, cuSimName, si
 
         # export particles dictionaries of saving time steps
         if "particles" in resTypes:
-            savePartToPickle(particles, outDirData, cuSimName)
+            savePartToPickle(particles, outDirData, cuSimName, cfg=cfg)
 
     # save contour line for each sim only if the field is properly computed (not a dummy array)
     contourResType = cfg["VISUALISATION"]["contourResType"]
@@ -2834,7 +2833,42 @@ def releaseSecRelArea(cfg, particles, fields, dem, zPartArray0, reportAreaInfo):
     return particles, zPartArray0, reportAreaInfo
 
 
-def savePartToPickle(dictList, outDir, logName):
+def _buildParticlePropertiesFilter(cfg, dictKeys):
+    """Build a set of particle property keys to keep when saving.
+
+    Returns None if all properties should be saved (no filtering).
+    Validates both export and track particle properties against dictKeys.
+    """
+    if not isinstance(cfg, configparser.ConfigParser):
+        return None
+
+    exportProperties = cfg["EXPORTS"]["exportParticleProperties"]
+    if exportProperties == "":
+        return None
+
+    exportKeys = exportProperties.split("|")
+    nonExisting = [k for k in exportKeys if k not in dictKeys]
+    if nonExisting:
+        message = "These particle properties are not available %s" % nonExisting
+        log.error(message)
+        raise KeyError(message)
+
+    propertiesFilter = {"t"} | set(exportKeys)
+
+    if cfg["TRACKPARTICLES"].getboolean("trackParticles"):
+        trackProperties = cfg["TRACKPARTICLES"]["particleProperties"]
+        if trackProperties != "":
+            # Track properties are validated upstream by trackParticles();
+            # no validation against dictKeys here since custom keys may be present.
+            trackKeys = trackProperties.split("|")
+            propertiesFilter |= set(trackKeys)
+        # Core properties always needed for particle tracking
+        propertiesFilter |= {"x", "y", "z", "ux", "uy", "uz", "m", "h"}
+
+    return propertiesFilter
+
+
+def savePartToPickle(dictList, outDir, logName, cfg=""):
     """Save each dictionary from a list to a pickle in outDir; works also for one dictionary instead of list
     Note: particle coordinates are still in com1DFA reference system with origin 0,0
 
@@ -2846,17 +2880,72 @@ def savePartToPickle(dictList, outDir, logName):
         path to output directory
     logName : str
         simulation Id
+    cfg: str or configparser object
+        ['EXPORTS'] and ['TRACKPARTICLES'] settings to provide particle properties to be saved,
+        if empty str all particle properties are saved, t (time info) always appended
     """
 
-    if isinstance(dictList, list):
-        for dict in dictList:
-            fi = open(outDir / ("particles_%s_%09.4f.pickle" % (logName, dict["t"])), "wb")
-            pickle.dump(dict, fi)
-            fi.close()
-    else:
-        fi = open(outDir / ("particles_%s_%09.4f.pickle" % (logName, dictList["t"])), "wb")
-        pickle.dump(dictList, fi)
-        fi.close()
+    dictKeys = [
+        "nPart",
+        "x",
+        "y",
+        "trajectoryLengthXY",
+        "trajectoryLengthXYCor",
+        "trajectoryLengthXYZ",
+        "z",
+        "m",
+        "dmDet",
+        "massPerPart",
+        "nPPK",
+        "mTot",
+        "h",
+        "ux",
+        "uy",
+        "uz",
+        "uAcc",
+        "stoppCriteria",
+        "kineticEne",
+        "trajectoryAngle",
+        "potentialEne",
+        "peakKinEne",
+        "peakMassFlowing",
+        "simName",
+        "xllcenter",
+        "yllcenter",
+        "ID",
+        "nID",
+        "parentID",
+        "t",
+        "inCellDEM",
+        "indXDEM",
+        "indYDEM",
+        "indPartInCell",
+        "partInCell",
+        "secondaryReleaseInfo",
+        "iterate",
+        "idFixed",
+        "peakForceSPH",
+        "forceSPHIni",
+        "totalEnthalpy",
+        "velocityMag",
+        "nExitedParticles",
+        "tPlot",
+        "dmEnt",
+        "stoppedParticles",
+        "massInitialized",
+        "massEntrained",
+        "massDetrained",
+        "massStopped",
+    ]
+
+    propertiesFilter = _buildParticlePropertiesFilter(cfg, dictKeys)
+
+    dicts = dictList if isinstance(dictList, list) else [dictList]
+    for d in dicts:
+        if propertiesFilter is not None:
+            d = {key: d[key] for key in propertiesFilter}
+        with open(outDir / ("particles_%s_%09.4f.pickle" % (logName, d["t"])), "wb") as fi:
+            pickle.dump(d, fi)
 
 
 def trackParticles(cfgTrackPart, dem, particlesList):

avaframe/com1DFA/com1DFACfg.ini

@@ -524,7 +524,7 @@ thresholdPointInPoly = 0.001
 
 [TRACKPARTICLES]
 # if particles should be tracked - don't forget to specify the "tSteps" you want to
-# save further up (for example tStep = 0:1 will lead to tracking patiles every 1 second)
+# save further up (for example tStep = 0:1 will lead to tracking particles every 1 second)
 trackParticles = False
 # centerTrackPartPoint of the location of the particles to track (x|y coordinates)
 centerTrackPartPoint = 2933|-4010
@@ -570,4 +570,7 @@ exportData = True
 exportRasters = False
 # use LZW compression when writing TIFF raster files
 useCompression = True
+# particle properties - list all properties that shall be saved, t is always added
+exportParticleProperties =
+
 

avaframe/tests/test_com1DFA.py

@@ -1821,6 +1821,75 @@ def test_savePartToPickle(tmp_path):
     assert np.array_equal(particlesRead3["m"], particles1["m"])
     assert particlesRead3["t"] == 0.0
 
+    # call function to be tested
+    logName = "simNameTest4"
+    cfg = configparser.ConfigParser()
+    cfg["EXPORTS"] = {"exportParticleProperties": "x|m"}
+    cfg["TRACKPARTICLES"] = {"trackParticles": False}
+    com1DFA.savePartToPickle(particles1, outDir, logName, cfg=cfg)
+
+    # read pickle
+    picklePath4 = outDir / "particles_simNameTest4_0000.0000.pickle"
+    particlesRead4 = pickle.load(open(picklePath4, "rb"))
+
+    assert np.array_equal(particlesRead4["x"], particles1["x"])
+    assert "y" not in particlesRead4.keys()
+    assert np.array_equal(particlesRead4["m"], particles1["m"])
+    assert particlesRead4["t"] == 0.0
+
+    # call function to be tested
+    logName = "simNameTest5"
+    cfg = configparser.ConfigParser()
+    cfg["EXPORTS"] = {"exportParticleProperties": "x|m"}
+    cfg["TRACKPARTICLES"] = {"trackParticles": True, "particleProperties": "iCell"}
+    particles1["ux"] = np.asarray([1.0, 2.0, 3.0])
+    particles1["uy"] = np.asarray([1.0, 4.0, 5.0])
+    particles1["uz"] = np.asarray([10.0, 11.0, 11.0])
+    particles1["iCell"] = np.asarray([10.0, 11.0, 11.0])
+    particles2["ux"] = np.asarray([1.0, 2.0, 3.0])
+    particles2["uy"] = np.asarray([1.0, 4.0, 5.0])
+    particles2["uz"] = np.asarray([10.0, 11.0, 11.0])
+    particles2["iCell"] = np.asarray([10.0, 11.0, 11.0])
+    particles1["z"] = np.asarray([1.0, 2.0, 3.0])
+    particles2["z"] = np.asarray([1.0, 4.0, 5.0])
+    particles1["h"] = np.asarray([1.0, 2.0, 3.0])
+    particles2["h"] = np.asarray([1.0, 4.0, 5.0])
+    com1DFA.savePartToPickle(particles1, outDir, logName, cfg=cfg)
+
+    # read pickle
+    picklePath5 = outDir / "particles_simNameTest5_0000.0000.pickle"
+    particlesRead5 = pickle.load(open(picklePath5, "rb"))
+
+    assert np.array_equal(particlesRead5["x"], particles1["x"])
+    assert "y" in particlesRead5.keys()
+    assert "ux" in particlesRead5.keys()
+    assert np.array_equal(particlesRead5["iCell"], particles1["iCell"])
+    assert np.array_equal(particlesRead5["m"], particles1["m"])
+    assert particlesRead5["t"] == 0.0
+
+    # call function to be tested
+    logName = "simNameTest6"
+    cfg = configparser.ConfigParser()
+    cfg["EXPORTS"] = {"exportParticleProperties": "x|m|hallo"}
+    cfg["TRACKPARTICLES"] = {"trackParticles": False}
+
+    with pytest.raises(KeyError) as e:
+        com1DFA.savePartToPickle(particles1, outDir, logName, cfg=cfg)
+    assert ("These particle properties are not available") in str(e.value)
+
+    # call function to be tested
+    logName = "simNameTest7"
+    cfg = configparser.ConfigParser()
+    cfg["EXPORTS"] = {"exportParticleProperties": ""}
+    cfg["TRACKPARTICLES"] = {"trackParticles": False}
+    com1DFA.savePartToPickle(particles1, outDir, logName, cfg=cfg)
+    # read pickle
+    picklePath7 = outDir / "particles_simNameTest7_0000.0000.pickle"
+    particlesRead7 = pickle.load(open(picklePath7, "rb"))
+
+    for pProp in particlesRead7:
+        assert pProp in ['ux', 'uy', 'uz', 'iCell', 'z', 'x', 'y', 'm', 'h', 't']
+
 
 def test_exportFields(tmp_path):
     """test exporting fields to ascii files"""