Mercurial > repos > pointsource_pde / file revision

#!.venv/bin/python3
import argparse
import csv
from pathlib import Path

import numpy as np

from plot import find_files


def extract_aux(prefix):

    iter_files = find_files(prefix)

    init_filename = "%s/init_aux.csv" % Path(prefix).parent
    filename = "%s/res_aux.csv" % prefix

    # Do true and noisy values

    initdata = [
        [
            "k_true",
            "c1_true",
            "c2_true",
            "k_noisy",
            "c1_noisy",
            "c2_noisy",
        ]
    ]

    data0 = np.load("%s/omega_0.npz" % Path(prefix).parent)
    k_true = data0["k_true"]
    c1_true = data0["c1_true"]
    c2_true = data0["c2_true"]
    initdata.append(
        [
            k_true,
            c1_true,
            c2_true,
            data0["k_noisy"],
            data0["c1_noisy"],
            data0["c2_noisy"],
        ]
    )
    with open(init_filename, mode="w", newline="") as file:
        # Create a csv.writer object
        writer = csv.writer(file, delimiter="\t")
        # Write data to the CSV file
        writer.writerows(initdata)

    # Do reconstructed values

    auxdata = [
        [
            "iter",
            "k",
            "c1",
            "c2",
            "k_dist",
            "c_dist",
        ]
    ]

    for iter, file in iter_files:
        d = np.load(file)
        k = d["k"]
        c1 = d["c1"]
        c2 = d["c2"]
        k_dist = np.abs(k - k_true)
        c_dist = np.sqrt((c1 - c1_true) ** 2 + (c2 - c2_true) ** 2)
        auxdata.append([iter, k, c1, c2, k_dist, c_dist])

    # Open the file in write mode
    with open(filename, mode="w", newline="") as file:
        # Create a csv.writer object
        writer = csv.writer(file, delimiter="\t")
        # Write data to the CSV file
        writer.writerows(auxdata)


def main():
    parser = argparse.ArgumentParser(
        prog="extract_aux.py",
        description="Extract auxiliary parameters from pointsource_pde",
    )

    parser.add_argument("filename")

    args = parser.parse_args()

    extract_aux(args.filename)


if __name__ == "__main__":
    main()