vicsek-simulation/SimulationManager.py at master · Airfooox/vicsek-simulation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import os
import json
import numpy as np
import multiprocessing as mp
from pathos.multiprocessing import ProcessingPool

from Simulation import Simulation
from util import printProgressBar
from tqdm import tqdm

class SimulationManager:
    def __init__(self, dataDir):
        self.dataDir = dataDir
        self.simulationGroups = []

        if not (os.path.exists(self.dataDir) and os.path.isdir(self.dataDir)):
            os.mkdir(self.dataDir)

    def appendGroup(self, simulationGroup):
        self.simulationGroups.append(simulationGroup)

    def simulate(self):
        simulationScenarios = []
        for simulationGroup in self.simulationGroups:
            simulationScenarios = simulationScenarios + simulationGroup.simulationScenarios

        pool = ProcessingPool(nodes=(mp.cpu_count() - 1))
        for _ in tqdm(pool.imap(SimulationManager.runSimulation, simulationScenarios), total=len(simulationScenarios)):
            pass
        pool.close()
        pool.join()

    @staticmethod
    def runSimulation(simulationScenario):
        simulationIndex = simulationScenario.simulationIndex
        numSimulation = simulationScenario.numSimulation
        scenarioDataDir = simulationScenario.scenarioDataDir
        scenarioConfig = simulationScenario.scenarioConfig
        timePercentageUsedForMean = simulationScenario.timePercentageUsedForMean
        saveTrajectoryData = simulationScenario.saveTrajectoryData

        simulation = Simulation(simulationIndex, numSimulation, scenarioConfig, timePercentageUsedForMean)
        simulation.simulate()
        absoluteVelocities = simulation.absoluteVelocities
        totalAbsoluteVelocity = simulation.totalAbsoluteVelocity
        totalAbsoluteGroupVelocities = simulation.totalAbsoluteGroupVelocities
        vectorialVelocities = simulation.vectorialVelocities
        totalVectorialVelocity = simulation.totalVectorialVelocity
        vectorialGroupVelocities = simulation.vectorialGroupVelocities
        totalVectorialGroupVelocities = simulation.totalVectorialGroupVelocities
        totalNematicOrderParameter = simulation.totalNematicOrderParameter
        totalNematicOrderParameterGroups = simulation.totalNematicOrderParameterGroups

        if not (os.path.exists(scenarioDataDir) and os.path.isdir(scenarioDataDir)):
            os.mkdir(scenarioDataDir)

        with open(os.path.join(scenarioDataDir, 'config.json'), 'w') as configFile:
            json.dump(scenarioConfig, configFile)

        with open(os.path.join(scenarioDataDir, 'absoluteVelocities.npy'), 'wb') as absoluteVelocitiesFile:
            np.save(absoluteVelocitiesFile, absoluteVelocities)

        with open(os.path.join(scenarioDataDir, 'vectorialVelocities.npy'), 'wb') as vectorialVelocitiesFile:
            np.save(vectorialVelocitiesFile, vectorialVelocities)

        with open(os.path.join(scenarioDataDir, 'vectorialGroupVelocities.npy'), 'wb') as vectorialGroupVelocitiesFile:
            np.save(vectorialGroupVelocitiesFile, vectorialGroupVelocities)

        totalVelocities = {
            'totalAbsoluteVelocity': totalAbsoluteVelocity,
            'totalAbsoluteGroupVelocities': totalAbsoluteGroupVelocities.tolist(),
            'totalVectorialVelocity': totalVectorialVelocity.tolist(),
            'totalVectorialGroupVelocities': totalVectorialGroupVelocities.tolist(),
            'totalNematicOrderParameter': totalNematicOrderParameter,
            'totalNematicOrderParameterGroups': totalNematicOrderParameterGroups.tolist()
        }

        with open(os.path.join(scenarioDataDir, 'totalVelocities.json'), 'w') as totalVelocitiesFile:
            json.dump(totalVelocities, totalVelocitiesFile)

        if saveTrajectoryData or totalAbsoluteVelocity <= 0:
            statesData = simulation.states
            print(totalAbsoluteVelocity, os.path.join(scenarioDataDir, 'statesData.npy'))
            with open(os.path.join(scenarioDataDir, 'statesData.npy'), 'wb') as statesFile:
                np.save(statesFile, statesData)


class SimulationGroup:
    def __init__(self, simulationDataDir, configFunction, numSimulation, repeatNum, timePercentageUsedForMean, saveTrajectoryData = False):
        self.simulationDataDir = simulationDataDir
        self.constantsFunc = configFunction
        self.numSimulation = numSimulation
        self.repeatNum = repeatNum
        self.timePercentageUsedForMean = timePercentageUsedForMean
        self.saveTrajectoryData = saveTrajectoryData

        if not (os.path.exists(self.simulationDataDir) and os.path.isdir(self.simulationDataDir)):
            os.mkdir(self.simulationDataDir)

        simulationGroupConfig = {
            'numSimulation': self.numSimulation,
            'repeatNum': self.repeatNum,
            'timePercentageUsedForMean': self.timePercentageUsedForMean,
            'saveTrajectoryData': self.saveTrajectoryData,
        }
        with open(os.path.join(simulationDataDir, "simulationGroupConfig.json"), "w") as simulationGroupConfigFile:
            json.dump(simulationGroupConfig, simulationGroupConfigFile)

        self.simulationScenarios = []
        for simulationIndex in range(numSimulation):
            for subSimulationIndex in range(repeatNum):
                scenarioDataDir = self.simulationDataDir + '/' + str(simulationIndex) + '_' + str(subSimulationIndex)
                self.simulationScenarios.append(SimulationScenario(simulationIndex, numSimulation, scenarioDataDir, self.constantsFunc(simulationIndex, numSimulation), self.timePercentageUsedForMean, self.saveTrajectoryData))


class SimulationScenario:
    def __init__(self, simulationIndex, numSimulation, scenarioDataDir, scenarioConfig, timePercentageUsedForMean, saveTrajectoryData):
        self.simulationIndex = simulationIndex
        self.numSimulation = numSimulation
        self.scenarioDataDir = scenarioDataDir
        self.scenarioConfig = scenarioConfig
        self.timePercentageUsedForMean = timePercentageUsedForMean
        self.saveTrajectoryData = saveTrajectoryData