From 2da4bd6e0729d360dd8b0059c28467c9b2e87f9b Mon Sep 17 00:00:00 2001 From: Anna Wirbel Date: Fri, 8 May 2026 14:59:26 +0200 Subject: [PATCH 1/2] add option to only save certain particle props check for particle properties, rename key add pytest --- avaframe/com1DFA/com1DFA.py | 104 +++++++++++++++++++++++++++++--- avaframe/com1DFA/com1DFACfg.ini | 5 +- avaframe/tests/test_com1DFA.py | 69 +++++++++++++++++++++ 3 files changed, 170 insertions(+), 8 deletions(-) diff --git a/avaframe/com1DFA/com1DFA.py b/avaframe/com1DFA/com1DFA.py index 5b60a9911..694548625 100644 --- a/avaframe/com1DFA/com1DFA.py +++ b/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,7 @@ def releaseSecRelArea(cfg, particles, fields, dem, zPartArray0, reportAreaInfo): return particles, zPartArray0, reportAreaInfo -def savePartToPickle(dictList, outDir, logName): +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,16 +2845,107 @@ def savePartToPickle(dictList, outDir, logName): path to output directory logName : str simulation Id + cfg: str or configparser object + ['EXPORTS'] and ['GENERAL'] settings to provide particle properties to be saved, + if empty str all particle properties are saved, t (time info) always appended """ + 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", + ] + + # create list of particle properties and append t (time info) + if isinstance(cfg, configparser.ConfigParser): + if cfg["EXPORTS"]["exportParticleProperties"] == "": + particleProperties = "" + else: + # first check if particle properties are valid + nonExisting = [ + item + for item in cfg["EXPORTS"]["exportParticleProperties"].split("|") + if item not in dictKeys + ] + if len(nonExisting) > 0: + message = "These particle properties are not available %s" % nonExisting + log.error(message) + raise AttributeError(message) + + particleProperties = list(set(["t"] + cfg["EXPORTS"]["exportParticleProperties"].split("|"))) + if cfg["TRACKPARTICLES"].getboolean("trackParticles"): + trackParticleProperties = cfg["TRACKPARTICLES"]["particleProperties"].split("|") + particleProperties = set( + ["x", "y", "z", "ux", "uy", "uz", "m", "h"] + + particleProperties + + trackParticleProperties + ) + else: + particleProperties = "" + if isinstance(dictList, list): for dict in dictList: + if particleProperties != "": + particlesToSave = {key: dict[key] for key in particleProperties} + else: + particlesToSave = dict fi = open(outDir / ("particles_%s_%09.4f.pickle" % (logName, dict["t"])), "wb") - pickle.dump(dict, fi) + pickle.dump(particlesToSave, fi) fi.close() else: + if particleProperties != "": + particlesToSave = {key: dictList[key] for key in particleProperties} + else: + particlesToSave = dictList fi = open(outDir / ("particles_%s_%09.4f.pickle" % (logName, dictList["t"])), "wb") - pickle.dump(dictList, fi) + pickle.dump(particlesToSave, fi) fi.close() diff --git a/avaframe/com1DFA/com1DFACfg.ini b/avaframe/com1DFA/com1DFACfg.ini index 6c305d86b..0927cca80 100644 --- a/avaframe/com1DFA/com1DFACfg.ini +++ b/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 = + diff --git a/avaframe/tests/test_com1DFA.py b/avaframe/tests/test_com1DFA.py index 63bb20ee6..abf772987 100644 --- a/avaframe/tests/test_com1DFA.py +++ b/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(AttributeError) 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""" From 41adeb3d746ac52f5126587dfb100b807c1bc229 Mon Sep 17 00:00:00 2001 From: Felix Oesterle <6945681+fso42@users.noreply.github.com> Date: Mon, 18 May 2026 09:59:03 +0200 Subject: [PATCH 2/2] refactor(com1DFA): extract particle properties filtering into helper function - Introduced `_buildParticlePropertiesFilter` to centralize key validation and filtering logic. - Simplified `savePartToPickle` by reusing the new helper function, improving readability. - Replaced `AttributeError` with `KeyError` for invalid properties to better signify the issue. --- avaframe/com1DFA/com1DFA.py | 87 +++++++++++++++++----------------- avaframe/tests/test_com1DFA.py | 2 +- 2 files changed, 44 insertions(+), 45 deletions(-) diff --git a/avaframe/com1DFA/com1DFA.py b/avaframe/com1DFA/com1DFA.py index 694548625..85a6cf81d 100644 --- a/avaframe/com1DFA/com1DFA.py +++ b/avaframe/com1DFA/com1DFA.py @@ -2833,6 +2833,41 @@ def releaseSecRelArea(cfg, particles, fields, dem, zPartArray0, reportAreaInfo): return particles, zPartArray0, reportAreaInfo +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,7 +2881,7 @@ def savePartToPickle(dictList, outDir, logName, cfg=""): logName : str simulation Id cfg: str or configparser object - ['EXPORTS'] and ['GENERAL'] settings to provide particle properties to be saved, + ['EXPORTS'] and ['TRACKPARTICLES'] settings to provide particle properties to be saved, if empty str all particle properties are saved, t (time info) always appended """ @@ -2903,50 +2938,14 @@ def savePartToPickle(dictList, outDir, logName, cfg=""): "massStopped", ] - # create list of particle properties and append t (time info) - if isinstance(cfg, configparser.ConfigParser): - if cfg["EXPORTS"]["exportParticleProperties"] == "": - particleProperties = "" - else: - # first check if particle properties are valid - nonExisting = [ - item - for item in cfg["EXPORTS"]["exportParticleProperties"].split("|") - if item not in dictKeys - ] - if len(nonExisting) > 0: - message = "These particle properties are not available %s" % nonExisting - log.error(message) - raise AttributeError(message) - - particleProperties = list(set(["t"] + cfg["EXPORTS"]["exportParticleProperties"].split("|"))) - if cfg["TRACKPARTICLES"].getboolean("trackParticles"): - trackParticleProperties = cfg["TRACKPARTICLES"]["particleProperties"].split("|") - particleProperties = set( - ["x", "y", "z", "ux", "uy", "uz", "m", "h"] - + particleProperties - + trackParticleProperties - ) - else: - particleProperties = "" + propertiesFilter = _buildParticlePropertiesFilter(cfg, dictKeys) - if isinstance(dictList, list): - for dict in dictList: - if particleProperties != "": - particlesToSave = {key: dict[key] for key in particleProperties} - else: - particlesToSave = dict - fi = open(outDir / ("particles_%s_%09.4f.pickle" % (logName, dict["t"])), "wb") - pickle.dump(particlesToSave, fi) - fi.close() - else: - if particleProperties != "": - particlesToSave = {key: dictList[key] for key in particleProperties} - else: - particlesToSave = dictList - fi = open(outDir / ("particles_%s_%09.4f.pickle" % (logName, dictList["t"])), "wb") - pickle.dump(particlesToSave, fi) - fi.close() + 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): diff --git a/avaframe/tests/test_com1DFA.py b/avaframe/tests/test_com1DFA.py index abf772987..49854b070 100644 --- a/avaframe/tests/test_com1DFA.py +++ b/avaframe/tests/test_com1DFA.py @@ -1873,7 +1873,7 @@ def test_savePartToPickle(tmp_path): cfg["EXPORTS"] = {"exportParticleProperties": "x|m|hallo"} cfg["TRACKPARTICLES"] = {"trackParticles": False} - with pytest.raises(AttributeError) as e: + with pytest.raises(KeyError) as e: com1DFA.savePartToPickle(particles1, outDir, logName, cfg=cfg) assert ("These particle properties are not available") in str(e.value)