diff --git a/avaframe/com1DFA/com1DFA.py b/avaframe/com1DFA/com1DFA.py index 5b60a9911..85a6cf81d 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,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): 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..49854b070 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(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"""