|
1 | | -""" |
2 | | -This file is part of the openPMD-api. |
3 | | -
|
4 | | -This module provides functions that are wrapped into sys.exit(...()) calls by |
5 | | -the setuptools (setup.py) "entry_points" -> "console_scripts" generator. |
6 | | -
|
7 | | -Copyright 2021 openPMD contributors |
8 | | -Authors: Franz Poeschel |
9 | | -License: LGPLv3+ |
10 | | -""" |
11 | | -from .. import openpmd_api_cxx as io |
12 | | -import argparse |
13 | | -import sys # sys.stderr.write |
14 | | - |
15 | | -# MPI is an optional dependency |
16 | | -try: |
17 | | - from mpi4py import MPI |
18 | | - HAVE_MPI = True |
19 | | -except ImportError: |
20 | | - HAVE_MPI = False |
21 | | - |
22 | | -debug = False |
23 | | - |
24 | | - |
25 | | -class FallbackMPICommunicator: |
26 | | - def __init__(self): |
27 | | - self.size = 1 |
28 | | - self.rank = 0 |
29 | | - |
30 | | - |
31 | | -def parse_args(): |
32 | | - parser = argparse.ArgumentParser(description=""" |
33 | | -openPMD Pipe. |
34 | | -
|
35 | | -This tool connects an openPMD-based data source with an openPMD-based data sink |
36 | | -and forwards all data from source to sink. |
37 | | -Possible uses include conversion of data from one backend to another one |
38 | | -or multiplexing the data path in streaming setups. |
39 | | -Parallelization with MPI is optionally possible and is done automatically |
40 | | -as soon as the mpi4py package is found and this tool is called in an MPI |
41 | | -context. In that case, each dataset will be equally sliced along the dimension |
42 | | -with the largest extent.""") |
43 | | - |
44 | | - parser.add_argument('--infile', type=str, help='In file') |
45 | | - parser.add_argument('--outfile', type=str, help='Out file') |
46 | | - parser.add_argument('--inconfig', |
47 | | - type=str, |
48 | | - default='{}', |
49 | | - help='JSON config for the in file') |
50 | | - parser.add_argument('--outconfig', |
51 | | - type=str, |
52 | | - default='{}', |
53 | | - help='JSON config for the out file') |
54 | | - |
55 | | - return parser.parse_args() |
56 | | - |
57 | | - |
58 | | -class Chunk: |
59 | | - """ |
60 | | - A Chunk is an n-dimensional hypercube, defined by an offset and an extent. |
61 | | - Offset and extent must be of the same dimensionality (Chunk.__len__). |
62 | | - """ |
63 | | - def __init__(self, offset, extent): |
64 | | - assert (len(offset) == len(extent)) |
65 | | - self.offset = offset |
66 | | - self.extent = extent |
67 | | - |
68 | | - def __len__(self): |
69 | | - return len(self.offset) |
70 | | - |
71 | | - def slice1D(self, mpi_rank, mpi_size, dimension=None): |
72 | | - """ |
73 | | - Slice this chunk into mpi_size hypercubes along one of its |
74 | | - n dimensions. The dimension is given through the 'dimension' |
75 | | - parameter. If None, the dimension with the largest extent on |
76 | | - this hypercube is automatically picked. |
77 | | - Returns the mpi_rank'th of the sliced chunks. |
78 | | - """ |
79 | | - if dimension is None: |
80 | | - # pick that dimension which has the highest count of items |
81 | | - dimension = 0 |
82 | | - maximum = self.extent[0] |
83 | | - for k, v in enumerate(self.extent): |
84 | | - if v > maximum: |
85 | | - dimension = k |
86 | | - assert (dimension < len(self)) |
87 | | - # no offset |
88 | | - assert (self.offset == [0 for _ in range(len(self))]) |
89 | | - offset = [0 for _ in range(len(self))] |
90 | | - stride = self.extent[dimension] // mpi_size |
91 | | - rest = self.extent[dimension] % mpi_size |
92 | | - |
93 | | - # local function f computes the offset of a rank |
94 | | - # for more equal balancing, we want the start index |
95 | | - # at the upper gaussian bracket of (N/n*rank) |
96 | | - # where N the size of the dataset in dimension dim |
97 | | - # and n the MPI size |
98 | | - # for avoiding integer overflow, this is the same as: |
99 | | - # (N div n)*rank + round((N%n)/n*rank) |
100 | | - def f(rank): |
101 | | - res = stride * rank |
102 | | - padDivident = rest * rank |
103 | | - pad = padDivident // mpi_size |
104 | | - if pad * mpi_size < padDivident: |
105 | | - pad += 1 |
106 | | - return res + pad |
107 | | - |
108 | | - offset[dimension] = f(mpi_rank) |
109 | | - extent = self.extent.copy() |
110 | | - if mpi_rank >= mpi_size - 1: |
111 | | - extent[dimension] -= offset[dimension] |
112 | | - else: |
113 | | - extent[dimension] = f(mpi_rank + 1) - offset[dimension] |
114 | | - return Chunk(offset, extent) |
115 | | - |
116 | | - |
117 | | -class deferred_load: |
118 | | - def __init__(self, source, dynamicView, offset, extent): |
119 | | - self.source = source |
120 | | - self.dynamicView = dynamicView |
121 | | - self.offset = offset |
122 | | - self.extent = extent |
123 | | - |
124 | | - |
125 | | -class particle_patch_load: |
126 | | - """ |
127 | | - A deferred load/store operation for a particle patch. |
128 | | - Our particle-patch API requires that users pass a concrete value for |
129 | | - storing, even if the actual write operation occurs much later at |
130 | | - series.flush(). |
131 | | - So, unlike other record components, we cannot call .store_chunk() with |
132 | | - a buffer that has not yet been filled, but must wait until the point where |
133 | | - we actual have the data at hand already. |
134 | | - In short: calling .store() must be deferred, until the data has been fully |
135 | | - read from the sink. |
136 | | - This class stores the needed parameters to .store(). |
137 | | - """ |
138 | | - def __init__(self, data, dest): |
139 | | - self.data = data |
140 | | - self.dest = dest |
141 | | - |
142 | | - def run(self): |
143 | | - for index, item in enumerate(self.data): |
144 | | - self.dest.store(index, item) |
145 | | - |
146 | | - |
147 | | -class pipe: |
148 | | - """ |
149 | | - Represents the configuration of one "pipe" pass. |
150 | | - """ |
151 | | - def __init__(self, infile, outfile, inconfig, outconfig, comm): |
152 | | - self.infile = infile |
153 | | - self.outfile = outfile |
154 | | - self.inconfig = inconfig |
155 | | - self.outconfig = outconfig |
156 | | - self.loads = [] |
157 | | - self.comm = comm |
158 | | - |
159 | | - def run(self): |
160 | | - if self.comm.size == 1: |
161 | | - print("Opening data source") |
162 | | - sys.stdout.flush() |
163 | | - inseries = io.Series(self.infile, io.Access.read_only, |
164 | | - self.inconfig) |
165 | | - print("Opening data sink") |
166 | | - sys.stdout.flush() |
167 | | - outseries = io.Series(self.outfile, io.Access.create, |
168 | | - self.outconfig) |
169 | | - print("Opened input and output") |
170 | | - sys.stdout.flush() |
171 | | - else: |
172 | | - print("Opening data source on rank {}.".format(self.comm.rank)) |
173 | | - sys.stdout.flush() |
174 | | - inseries = io.Series(self.infile, io.Access.read_only, self.comm, |
175 | | - self.inconfig) |
176 | | - print("Opening data sink on rank {}.".format(self.comm.rank)) |
177 | | - sys.stdout.flush() |
178 | | - outseries = io.Series(self.outfile, io.Access.create, self.comm, |
179 | | - self.outconfig) |
180 | | - print("Opened input and output on rank {}.".format(self.comm.rank)) |
181 | | - sys.stdout.flush() |
182 | | - self.__copy(inseries, outseries) |
183 | | - |
184 | | - def __copy(self, src, dest, current_path="/data/"): |
185 | | - """ |
186 | | - Worker method. |
187 | | - Copies data from src to dest. May represent any point in the openPMD |
188 | | - hierarchy, but src and dest must both represent the same layer. |
189 | | - """ |
190 | | - if (type(src) != type(dest) |
191 | | - and not isinstance(src, io.IndexedIteration) |
192 | | - and not isinstance(dest, io.Iteration)): |
193 | | - raise RuntimeError( |
194 | | - "Internal error: Trying to copy mismatching types") |
195 | | - attribute_dtypes = src.attribute_dtypes |
196 | | - for key in src.attributes: |
197 | | - attr = src.get_attribute(key) |
198 | | - attr_type = attribute_dtypes[key] |
199 | | - dest.set_attribute(key, attr, attr_type) |
200 | | - container_types = [ |
201 | | - io.Mesh_Container, io.Particle_Container, io.ParticleSpecies, |
202 | | - io.Record, io.Mesh, io.Particle_Patches, io.Patch_Record |
203 | | - ] |
204 | | - if isinstance(src, io.Series): |
205 | | - # main loop: read iterations of src, write to dest |
206 | | - write_iterations = dest.write_iterations() |
207 | | - for in_iteration in src.read_iterations(): |
208 | | - if self.comm.rank == 0: |
209 | | - print("Iteration {0} contains {1} meshes:".format( |
210 | | - in_iteration.iteration_index, |
211 | | - len(in_iteration.meshes))) |
212 | | - for m in in_iteration.meshes: |
213 | | - print("\t {0}".format(m)) |
214 | | - print("") |
215 | | - print( |
216 | | - "Iteration {0} contains {1} particle species:".format( |
217 | | - in_iteration.iteration_index, |
218 | | - len(in_iteration.particles))) |
219 | | - for ps in in_iteration.particles: |
220 | | - print("\t {0}".format(ps)) |
221 | | - print("With records:") |
222 | | - for r in in_iteration.particles[ps]: |
223 | | - print("\t {0}".format(r)) |
224 | | - out_iteration = write_iterations[in_iteration.iteration_index] |
225 | | - sys.stdout.flush() |
226 | | - self.__particle_patches = [] |
227 | | - self.__copy( |
228 | | - in_iteration, out_iteration, |
229 | | - current_path + str(in_iteration.iteration_index) + "/") |
230 | | - for deferred in self.loads: |
231 | | - deferred.source.load_chunk( |
232 | | - deferred.dynamicView.current_buffer(), deferred.offset, |
233 | | - deferred.extent) |
234 | | - in_iteration.close() |
235 | | - for patch_load in self.__particle_patches: |
236 | | - patch_load.run() |
237 | | - out_iteration.close() |
238 | | - self.__particle_patches.clear() |
239 | | - self.loads.clear() |
240 | | - sys.stdout.flush() |
241 | | - elif isinstance(src, io.Record_Component): |
242 | | - shape = src.shape |
243 | | - offset = [0 for _ in shape] |
244 | | - dtype = src.dtype |
245 | | - dest.reset_dataset(io.Dataset(dtype, shape)) |
246 | | - if src.empty: |
247 | | - # empty record component automatically created by |
248 | | - # dest.reset_dataset() |
249 | | - pass |
250 | | - elif src.constant: |
251 | | - dest.make_constant(src.get_attribute("value")) |
252 | | - else: |
253 | | - chunk = Chunk(offset, shape) |
254 | | - local_chunk = chunk.slice1D(self.comm.rank, self.comm.size) |
255 | | - if debug: |
256 | | - end = local_chunk.offset.copy() |
257 | | - for i in range(len(end)): |
258 | | - end[i] += local_chunk.extent[i] |
259 | | - print("{}\t{}/{}:\t{} -- {}".format( |
260 | | - current_path, self.comm.rank, self.comm.size, |
261 | | - local_chunk.offset, end)) |
262 | | - span = dest.store_chunk(local_chunk.offset, local_chunk.extent) |
263 | | - self.loads.append( |
264 | | - deferred_load(src, span, local_chunk.offset, |
265 | | - local_chunk.extent)) |
266 | | - elif isinstance(src, io.Patch_Record_Component): |
267 | | - dest.reset_dataset(io.Dataset(src.dtype, src.shape)) |
268 | | - if self.comm.rank == 0: |
269 | | - self.__particle_patches.append( |
270 | | - particle_patch_load(src.load(), dest)) |
271 | | - elif isinstance(src, io.Iteration): |
272 | | - self.__copy(src.meshes, dest.meshes, current_path + "meshes/") |
273 | | - self.__copy(src.particles, dest.particles, |
274 | | - current_path + "particles/") |
275 | | - elif any([ |
276 | | - isinstance(src, container_type) |
277 | | - for container_type in container_types |
278 | | - ]): |
279 | | - for key in src: |
280 | | - self.__copy(src[key], dest[key], current_path + key + "/") |
281 | | - if isinstance(src, io.ParticleSpecies): |
282 | | - self.__copy(src.particle_patches, dest.particle_patches) |
283 | | - else: |
284 | | - raise RuntimeError("Unknown openPMD class: " + str(src)) |
285 | | - |
286 | | - |
287 | | -def main(): |
288 | | - args = parse_args() |
289 | | - if not args.infile or not args.outfile: |
290 | | - print("Please specify parameters --infile and --outfile.") |
291 | | - sys.exit(1) |
292 | | - if (HAVE_MPI): |
293 | | - run_pipe = pipe(args.infile, args.outfile, args.inconfig, |
294 | | - args.outconfig, MPI.COMM_WORLD) |
295 | | - else: |
296 | | - run_pipe = pipe(args.infile, args.outfile, args.inconfig, |
297 | | - args.outconfig, FallbackMPICommunicator()) |
298 | | - |
299 | | - run_pipe.run() |
300 | | - |
301 | | - |
302 | | -if __name__ == "__main__": |
303 | | - main() |
304 | | - sys.exit() |
| 1 | +__doc__ = "for usage documentation, call this module's main with --help" |
0 commit comments