Coverage for dragonfly_sim/utils/checkpoint.py: 89%

1""" 

2Partition-independent checkpoints for time-dependent runs. 

3 

4A checkpoint is one .npz written on rank 0. Every DG field is stored per cell, 

5keyed by the cell's index in the INPUT mesh (topology.original_cell_index) 

6and the cell-local dof order, so a run may be restarted on any number of 

7ranks. The dof coordinates are stored as well and checked on load. 

8""" 

9import json 

10 

11import numpy as np 

12from mpi4py import MPI 

13 

14 

15def _owned_cell_dofs(V): 

16 """(owned cells' input-mesh indices, their dof blocks, the dof coordinates).""" 

17 mesh = V.mesh 

18 n_owned = mesh.topology.index_map(mesh.topology.dim).size_local 

19 dofs = V.dofmap.list[:n_owned] 

20 orig = np.asarray(mesh.topology.original_cell_index[:n_owned], dtype=np.int64) 

21 return orig, dofs, V.tabulate_dof_coordinates()[dofs] 

22 

23 

24def save_checkpoint(path, functions, meta): 

25 """ 

26 Write `functions` (DG Functions on one mesh, distinct names) and the JSON 

27 serializable dict `meta` to `path`. The dof coordinates are stored as 

28 "coords" for the first function's space and "<name>__coords" for any 

29 function on another space. Collective. 

30 """ 

31 comm = functions[0].function_space.mesh.comm 

32 V0 = functions[0].function_space 

33 local = {} 

34 for f in functions: 

35 V = f.function_space 

36 orig, dofs, coords = _owned_cell_dofs(V) 

37 local["orig"] = orig # the same owned cells for every space 

38 local["coords" if V == V0 else f.name + "__coords"] = coords 

39 local[f.name] = f.x.array.reshape(-1, V.dofmap.index_map_bs)[dofs] 

40 gathered = comm.gather(local, root=0) 

41 if comm.rank == 0: 

42 orig_all = np.concatenate([g["orig"] for g in gathered]) 

43 order = np.argsort(orig_all) 

44 assert np.array_equal(orig_all[order], np.arange(orig_all.size)), \ 

45 "owned cells do not tile the input mesh exactly once" 

46 np.savez(path, meta=json.dumps(meta), 

47 **{k: np.concatenate([g[k] for g in gathered])[order] 

48 for k in local if k != "orig"}) 

49 comm.Barrier() 

50 

51 

52def load_checkpoint(path, functions, coord_tol=1e-10, optional=False): 

53 """ 

54 Fill `functions` from a save_checkpoint file, on any rank count, and 

55 return its meta dict. With optional=True a field missing from the file is 

56 left untouched instead of raising (e.g. far-field history when restarting 

57 a riemann2 run from a plain riemann checkpoint). 

58 """ 

59 comm = functions[0].function_space.mesh.comm 

60 data = np.load(path) 

61 for f in functions: 

62 if f.name not in data.files: 

63 if optional: 

64 continue 

65 raise KeyError("{} has no field {!r}".format(path, f.name)) 

66 V = f.function_space 

67 orig, dofs, coords = _owned_cell_dofs(V) 

68 ckey = f.name + "__coords" if f.name + "__coords" in data.files else "coords" 

69 ref = data[ckey] 

70 if ref.shape[0] != comm.allreduce(orig.size, MPI.SUM) or ref.shape[1:] != coords.shape[1:]: 

71 raise ValueError("checkpoint field {} does not match this mesh/space".format(f.name)) 

72 err = comm.allreduce(float(np.abs(ref[orig] - coords).max(initial=0.0)), MPI.MAX) 

73 if err > coord_tol: 

74 raise ValueError("checkpoint dof coordinates of {} are off by {:.2e}".format(f.name, err)) 

75 f.x.array.reshape(-1, V.dofmap.index_map_bs)[dofs] = data[f.name][orig] 

76 f.x.scatter_forward() 

77 return json.loads(str(data["meta"]))