merge.py

Go to the documentation of this file.

# This file is part of the Peano project. For conditions of distribution and
# use, please see the copyright notice at www.peano-framework.org
"""!
 
Merge the partial patch files of a Peano simulation
 
A parallel run writes one patch file per spacetree (and rank) and snapshot.
The script merges the files of each snapshot into one file, and writes a new
meta file merged_meta_file.peano-patch-file next to the original one which
links to the merged files. It handles files with ASCII and with binary
quantities. Call it with
 
    python3 -m tarch.plotter.api.merge --meta-file solution-Euler.peano-patch-file
 
@ref documentation_tarch_visualisation_merging-data
 
"""
import os
import argparse
 
 
def merge_partial_files(partial_files, merged_file_path):
    """!
 
    Merge the patches of several patch files into one file
 
    The files are processed as bytes. A file with format Binary holds raw
    doubles between its text lines, and we copy those verbatim. The header
    of the merged file is the header of the first partial file, i.e. all
    partial files have to have the same format. Everything from the first
    "begin patch" in a file on is a patch, and everything before it is the
    header: A header holds text only, so the first occurrence is the right
    one, even in a binary file.
 
    @param partial_files: [str]
      Files to merge. Their patches end up in the merged file in this order.
 
    @param merged_file_path: str
      File to write.
 
    """
    print("Start to merge file:", merged_file_path, end="...")
    patch_marker = b"begin patch"
 
    with open(merged_file_path, "wb") as merged_patch_file:
        # First we need to write the header for our new patch
        # file. We can open any of the partial patch files
        # and pass this across
        with open(partial_files[0], "rb") as patch_file:
            content = patch_file.read()
            first_patch = content.find(patch_marker)
            merged_patch_file.write(content if first_patch < 0 else content[:first_patch])
 
        for infile in partial_files:
            with open(infile, "rb") as patch_file:
                content = patch_file.read()
            first_patch = content.find(patch_marker)
            if first_patch >= 0:
                merged_patch_file.write(b"\n")
                merged_patch_file.write(content[first_patch:])
    print("Merged.")
 
 
def read_meta_file(path_to_metafile, start, end):
    # merged_files = []  # Stores names of merged files
    partial_files = []  # Stores partial files that need to be merged
    # into the same file
    # timestamps = []
    merged_file_number = 0
 
    f = open(
        os.path.join(
            os.path.dirname(path_to_metafile), "merged_meta_file.peano-patch-file"
        ),
        "w",
    )
    f.write(
        "# \n" "# Peano patch file \n" "# Version 0.2 \n" "# \n" "format ASCII \n" "\n"
    )
 
    with open(path_to_metafile, "r") as metafile:
        for line in metafile.readlines():
            if "end dataset" in line and (
                merged_file_number < start or merged_file_number > end
            ):
                partial_files.clear()
                print("jump snapshot", merged_file_number)
                merged_file_number += 1
            elif "end dataset" in line:
                # Merged files will be put in the same dir as the original meta-file
                merged_file_name = (
                    "merged_patch_file_" + str(merged_file_number) + ".peano-patch-file"
                )
 
                if os.path.isabs(path_to_metafile):
                    merged_file_name = os.path.join(
                        os.path.dirname(path_to_metafile), merged_file_name
                    )
 
                merge_partial_files(partial_files, merged_file_name)
 
                partial_files.clear()
                merged_file_number += 1
                print()
 
                f.write(
                    "begin dataset\n" + timestamp_line + "\n"
                    f'  include "{merged_file_name}"\n'
                    "end dataset\n"
                    "\n"
                )
 
            if "timestamp" in line and merged_file_number >= start:
                timestamp_line = line
                print(line[0:-2])
 
            words = line.split()
 
            if "include" in words:
                file_path = words[1].strip('"')  # This is the input path
 
                # Where the filepath listed in the meta-file is not
                # an absolute path with assume the file lives in the
                # same directory as the meta-file itself:
                if not os.path.isabs(file_path):
                    file_path = os.path.join(
                        os.path.dirname(path_to_metafile), file_path
                    )
 
                partial_files.append(file_path)
 
    f.close()
 
 
def count_total_snapshot_number(path_to_metafile):
    count = 0
    with open(path_to_metafile, "r") as metafile:
        for line in metafile.readlines():
            if "timestamp" in line:
                count += 1
    return count
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Peano 4 - Merge partial patch files "
        "that are linked to by a meta file"
    )
    parser.add_argument(
        "--meta-file",
        dest="metafile",
        help="The main file that links to a set of Peano patch files",
        required=True,
    )
    parser.add_argument(
        "-start",
        "--start-snapshot-index",
        dest="start",
        type=int,
        help="Set the snapshot number where the merge starts (all snapshots before will be ignored), default is 0 (starts from beginning)",
        default=0,
    )
    parser.add_argument(
        "-end",
        "--end-snapshot-index",
        dest="end",
        type=int,
        help="Set the snapshot number where the merge ends (all snapshots after will be ignored), default is -1 (ends at the last)",
        default=-1,
    )
 
    args = parser.parse_args()
 
    if not os.path.exists(args.metafile):
        raise Exception(
            "specified input file '{}' does not exist, exiting...".format(args.metafile)
        )
 
    total_snapshot_number = count_total_snapshot_number(args.metafile)
    if (args.start > args.end and args.end != -1) or args.start < 0 or args.end < -1:
        raise Exception("Specified snapshot index is wrong, please check.")
    elif args.end == -1:
        start_snapshot = args.start
        end_snapshot = (
            total_snapshot_number - 1
        )  # Processing all snapshot if no end snapshot is specified.
    elif args.end > (total_snapshot_number - 1):
        start_snapshot = args.start
        end_snapshot = total_snapshot_number - 1
        print(
            "Warning: specified end snapshot index exceeds the total number, reset to the last snapshot."
        )
    else:
        start_snapshot = args.start
        end_snapshot = args.end
 
    print(
        "The meta file includes",
        total_snapshot_number,
        "snapshots, processing snapshot",
        start_snapshot,
        "to",
        end_snapshot,
    )
 
    read_meta_file(args.metafile, start_snapshot, end_snapshot)