コード例 #1
0
ファイル: test_outputfiles.py プロジェクト: p-robot/MetaWards
def test_openfiles(prompt=None):
    outdir = os.path.join(script_dir, "test_openfiles_output")

    if os.path.exists(outdir):
        OutputFiles.remove(outdir, prompt=prompt)

    of = OutputFiles(outdir)

    assert of.is_open()
    assert not of.is_closed()

    with pytest.raises(ValueError):
        of.open("../test.txt")

    with pytest.raises(ValueError):
        of.open("/test.txt")

    FILE = of.open("test.txt")

    assert of.is_open()

    FILE.write("hello\n")

    of.close()

    assert of.is_closed()

    assert open(os.path.join(outdir, "test.txt")).readline() == "hello\n"

    with pytest.raises(FileExistsError):
        OutputFiles(outdir, prompt=None)

    with OutputFiles(outdir, force_empty=True, prompt=None) as of:
        FILE = of.open("test.txt")
        assert of.is_open()

        assert of.get_filename("test.txt").endswith("test.txt")

        FILE.write("goodbye\n")

        FILE = of.open("test2.txt", auto_bzip=True)
        FILE.write("hello ")

        assert of.get_filename("test2.txt").endswith("test2.txt.bz2")

        FILE = of.open("test2.txt")
        FILE.write("world\n")

    assert open(os.path.join(outdir, "test.txt")).readline() == "goodbye\n"

    import bz2
    line = bz2.open(os.path.join(outdir, "test2.txt.bz2"), "rt").readline()
    assert line == "hello world\n"

    OutputFiles.remove(outdir, prompt=None)
コード例 #2
0
def output_wards_i_serial(network: metawards.Network,
                          population: metawards.Population,
                          output_dir: metawards.OutputFiles,
                          workspace: metawards.Workspace, **kwargs):

    if network.name is None:
        name = ""
    else:
        name = "_" + network.name.replace(" ", "_")

    d_vars = network.params.disease_params
    u_vars = network.params.user_params
    infect_file = output_dir.open(f"wards_trajectory{name}_I.csv")
    rclass_file = output_dir.open(f"wards_trajectory{name}_R.csv")

    # Just print the design parameters:
    head_str = "beta[2],beta[3],progress[1],progress[2],progress[3],.scale_rate[1],.scale_rate[2]"

    design_out = [
        d_vars.beta[2], d_vars.beta[3], d_vars.progress[1], d_vars.progress[2],
        d_vars.progress[3]
    ]
    design_out += [u_vars["lock_1_restrict"], u_vars["lock_2_release"]]

    if population.day == 0:
        # Print header
        ident: List[str] = ["day", "date"]
        ident += [head_str]
        ident += [
            "ward[" + str(i) + "]" for i, _ in enumerate(workspace.I_in_wards)
        ]
        header = ','.join(ident)
        infect_file.write(header)
        infect_file.write("\n")
        rclass_file.write(header)
        rclass_file.write("\n")

    day = str(population.day) + ","
    date = str(population.date) + ","

    # NOTE: The standard extractors write .dat tables, but R doesn't load these properly (and pandas sometimes)
    # Write a CSV instead
    infect_file.write(day)
    infect_file.write(date)
    infect_file.write(",".join(str(x) for x in design_out))
    infect_file.write(",")
    infect_file.write(",".join([str(x) for x in workspace.I_in_wards]))
    infect_file.write("\n")

    # NOTE: The standard extractors write .dat tables, but R doesn't load these properly (and pandas sometimes)
    # Write a CSV instead
    rclass_file.write(day)
    rclass_file.write(date)
    rclass_file.write(",".join(str(x) for x in design_out))
    rclass_file.write(",")
    rclass_file.write(",".join([str(x) for x in workspace.R_in_wards]))
    rclass_file.write("\n")