#!/usr/bin/env python
import os
import sys
import shutil
import random
import argparse
from tqdm import tqdm
from hurry.filesize import size
[docs]
def set_up_test(args):
"""Set up a testing directory in the requested directory.
Parameters (keys in `args`):
----------
test_dir : string
The directory where the test directory is to be set up.
dsr : float
Downsampling rate when creating testing directory.
"""
if not os.path.isdir(args.test_dir):
sys.exit(f"Directory {args.test_dir} does not exist.")
for folder in os.listdir(args.mesa_dir):
if os.path.isdir(os.path.join(args.mesa_dir, folder)):
is_mesa_run = False
sub_dir = os.listdir(os.path.join(args.mesa_dir, folder))
track_dirs = []
for _f in sub_dir:
if "_grid_index_" in _f:
is_mesa_run = True
track_dirs.append(os.path.join(args.mesa_dir, folder, _f))
if is_mesa_run: # checking if directory is a mesa run
os.mkdir(os.path.join(args.test_dir, folder))
# choosing which tracks to copy over
inds = random.sample(list(range(len(track_dirs))),
int(len(track_dirs) * args.dsr))
for ind in inds:
if os.path.isdir(track_dirs[ind]):
shutil.copytree(track_dirs[ind], os.path.join(
args.test_dir,
folder,
os.path.split(track_dirs[ind])[1]))
else:
shutil.copy(track_dirs[ind], os.path.join(
args.test_dir,
folder,
os.path.split(track_dirs[ind])[1]))
print(f"Created Test Directory at {args.test_dir}.")
[docs]
def compress_dir(args):
"""Compresses a directory containing tracks evolved with MESA.
Parameters (keys in `args`):
-----------
mesa_dir : string
The directory where the MESA tracks are stored.
"""
def get_size(start_path="."):
total_size = 0
for dirpath, _, filenames in os.walk(start_path):
for filename in filenames:
filepath = os.path.join(dirpath, filename)
# skip if it is symbolic link
if not os.path.islink(filepath):
total_size += os.path.getsize(filepath)
return total_size
if not os.path.isdir(args.mesa_dir):
sys.exit("The MESA directory does not exist.")
og_size = size(get_size(args.mesa_dir))
for folder in tqdm(os.listdir(args.mesa_dir)):
if os.path.isdir(os.path.join(args.mesa_dir, folder)):
is_mesa_run = False
sub_dir = os.listdir(os.path.join(args.mesa_dir, folder))
track_dirs = []
for _f in sub_dir:
if "_grid_index_" in _f: # checking if directory is mesa run
is_mesa_run = True
if (os.path.isdir(os.path.join(args.mesa_dir, folder, _f))
and _f not in ["star1", "star2", "binary"]):
track_dirs.append(os.path.join(args.mesa_dir, folder, _f))
if is_mesa_run:
# iterating over directory containing MESA runs
for track_obj in track_dirs:
for root, _dir, files in os.walk(track_obj):
# traversing over directory tree of copied mesa tracks
# and compressing .data, .mod, .txt files
for file in files:
ext = file.split(".")[-1]
if ext not in ["data", "mod", "txt"]:
continue
os.system(f"gzip -1 {os.path.join(root, file)}")
print("\nCompressed MESA tracks\nOriginal size {} | Compressed size {}\n".
format(og_size, size(get_size(args.mesa_dir))))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--function", type=str,
help="The name of the function to be called.")
parser.add_argument("-td", "--test_dir", type=str,
help="The path to where the testing directory should "
"be set up, either set_up_test or compress_dir.")
parser.add_argument("-dsr", "--dsr", type=str,
help="Downsampling rate when creating testing "
"directory", default=0.01)
parser.add_argument("-md", "--mesa_dir", type=str,
help="The path to the directory containing "
"MESA-generated data")
functions = {"set_up_test": set_up_test, "compress_dir": compress_dir}
arguments = parser.parse_args()
if arguments.function is None:
sys.exit("A function must be specified to use this script.")
functions[arguments.function](arguments)