02_generate_output_from_code.bash

#!/usr/bin/env bash

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "${SCRIPT_DIR}"

# Full output of every experiment script is always captured here, so a second
# terminal can `tail -f 02_generate_output_from_code.log` regardless of VERBOSE.
LOGFILE="02_generate_output_from_code.log"
LOGFILE_PATH="${SCRIPT_DIR}/${LOGFILE}"
: > "${LOGFILE_PATH}"

# VERBOSE!=0 also streams each script's output to this console.
# VERBOSE=0 (default) keeps the console clean and instead prints a
# per-script wall-time and peak-memory summary; the full output is still in the
# LOGFILE_PATH.
VERBOSE="${VERBOSE:-0}"

if [[ -t 1 ]]; then
    _green=$'\033[32m' _bold=$'\033[1m' _red=$'\033[31m' _blue=$'\033[34m' _reset=$'\033[0m'
else
    _green="" _bold="" _red="" _blue="" _reset=""
fi

log()  { echo "${_green}:: $*${_reset}"; }
raise()  { echo "${_red}:: [ERROR] $*${_reset}" >&2;}

if [[ "${VERBOSE}" == "0" ]]; then
    log "Suppressing experiment output: tail -f ${LOGFILE} to follow, or run with VERBOSE=1."
    step()  { echo -n "${_blue}:: $* ... ${_reset}"; }
else
    step()  { echo "${_blue}:: $*:${_reset}"; }
fi

if [[ "${QUICKTEST:-0}" == "1" ]]; then
    log "QUICKTEST=1 detected: only running coarse experiments."
fi

PYTHON_PATH="$(command -v python 2>/dev/null || true)"
EXPECTED_SUFFIX=".conda_env/bin/python"

if [[ "$PYTHON_PATH" == "/opt/env/${EXPECTED_SUFFIX}" ]]; then
    log "Running inside container."
elif [[ "$PYTHON_PATH" == "${SCRIPT_DIR}/${EXPECTED_SUFFIX}" ]]; then
    log "Running with local conda env."
else
    raise "Wrong Python active: '${PYTHON_PATH:-none found}'."
    raise "  Run: source 01_activate_env.bash"
    exit 1
fi

# Memory is measured via Python's resource.getrusage, which is portable and needs
# no /usr/bin/time.
run_experiment() {
    local exp="$1"
    VERBOSE="${VERBOSE}" LOGFILE_PATH="${LOGFILE_PATH}" python - "${exp}" <<'PY'
import os, resource, subprocess, sys, time

script = sys.argv[1]
verbose = os.environ.get('VERBOSE', '0') != '0'
LOGFILE_PATH = os.environ['LOGFILE_PATH']

start = time.perf_counter()
with open(LOGFILE_PATH, 'a') as log:
    log.write(f'\n===== {script} =====\n')
    log.flush()
    if verbose:
        # Stream to both console and LOGFILE_PATH.
        proc = subprocess.Popen([sys.executable, script],
                                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                text=True, bufsize=1)
        for line in proc.stdout:
            sys.stdout.write(line)
            log.write(line)
        proc.wait()
    else:
        # LOGFILE_PATH only; console stays clean.
        proc = subprocess.run([sys.executable, script], stdout=log, stderr=log)
wall = time.perf_counter() - start

# ru_maxrss is reported in KiB on Linux (the container platform).
maxrss_kib = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
if sys.platform == 'darwin':  # macOS reports bytes
    maxrss_kib /= 1024
mem_mib = maxrss_kib / 1024
h, rem = divmod(int(wall), 3600)
m, s = divmod(rem, 60)
time_str=f"{s:02d}s"
if m > 0:
    time_str = f"{m:02d}m{time_str}"
if h > 0:
    time_str = f"{h:d}h{time_str}"
mem_str = f"{mem_mib/1024:.1f}GiB" if mem_mib > 1024 else f"{mem_mib:.0f}MiB"
print(f" (took {time_str} and {mem_str})")
sys.exit(proc.returncode)
PY
}

log "Setting environment variables for reproducability ..."
export OMP_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export MKL_NUM_THREADS=1
export MKL_CBWR=COMPATIBLE
export MPLBACKEND=Agg

log "Running Mass-spring-damper experiments:"
cd "${SCRIPT_DIR}/code/msd_chain"
for exp in run_*.py
do
    step "Running msd_chain/${exp}"
    run_experiment "${exp}"
done
cd "${SCRIPT_DIR}"

log "Running Navier-Stokes experiments:"
cd "${SCRIPT_DIR}/code/navier_stokes"
for exp in run_*.py
do
    step "Running navier_stokes/${exp}"
    run_experiment "${exp}"
done
cd "${SCRIPT_DIR}"

log "All experiments finished."