DeepJetCore/MultiDataCollection.py at master · DL4Jets/DeepJetCore · GitHub

History

192 lines (162 loc) · 6.46 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

from DataCollection import DataCollection

from multiprocessing import cpu_count

from itertools import izip

from pdb import set_trace

import numpy as np

from copy import deepcopy

class MultiDataCollection(object):

'''This class allows the simultaneous use of multiple

DataCollections for a training, it provides the same interface

and adds the functionality of adding Y targets as well as flags returned instead of the weights.

In case of weights the flag is multiplied by the weight value

Constructor ([infiles = None[, nprocs = -1[, add_ys = [][, flags=[]]]]])

optional parameters:

infiles: list of input dataCollection files to be opened

nprocs: number of processors to use

add_ys: list of additional Y targets to be added at generator time, must be the same length of the input collections. The list content must be iterable, each iteration produces a new target, only scalar type supprted for now

flags: like add_ys, same rules apply. The flags gets multiplied by the event weight is case weights are used. The lenght of the flags must be the same as the TOTAL number of Y targets. Flags are returned instead of the event weights

'''

def __init__(self, infiles = None, nprocs = -1, add_ys = [] ,flags=[]):

'''Constructor'''

self.collections = []

self.nprocs = nprocs

self.meansnormslimit=500000

self.flags = []

self.generator_modifier = lambda x: x

self.additional_ys = []

if infiles:

self.collections = [

DataCollection(

cpu_count()/len(infiles) if nprocs == -1 else nprocs/len(infiles)

) for i in infiles]

if flags:

self.setFlags(flags)

if add_ys:

self.addYs(add_ys)

@property

def useweights(self):

return all(i.useweights for i in self.collections)

@useweights.setter

def useweights(self, val):

for i in self.collections:

i.useweights = val

def addYs(self, add_ys):

'adds Ys that will be appended on the fly to the generator, Ys are a list of iterables'

if len(add_ys) != len(self.collections):

raise ValueError('The Ys must be the same lenght of the input collections')

self.additional_ys = add_ys

def readFromFile(self, infiles):

self.collections = [

DataCollection(

i, cpu_count()/len(infiles) if self.nprocs == -1 else self.nprocs/len(infiles)

) for i in infiles]

def setFlags(self, flags):

'adds flags that will be added on the fly to the generator, flags are a list of iterables'

if len(flags) != len(self.collections):

raise ValueError('The flags must be the same lenght of the input collections')

self.flags = flags

def getInputShapes(self):

'Gets the input shapes from the data class description'

shapes = [i.getInputShapes() for i in self.collections]