Skip to content

Main

Synthesize data.

main()

Main command-line driver.

Source code in snailz/main.py
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def main():
    """Main command-line driver."""

    args = _parse_args()

    if args.defaults:
        print(utils.json_dump(Parameters()))
        return 0

    params = _initialize(args)
    grids, persons, samples = _synthesize(params)
    changes = do_all_effects(params, grids, persons, samples)
    if args.outdir:
        _save(args, grids, persons, samples, changes)

    return 0

_initialize(args)

Initialize for data synthesis.

Source code in snailz/main.py
35
36
37
38
39
40
41
42
43
44
45
46
def _initialize(args):
    """Initialize for data synthesis."""

    if args.params:
        with open(args.params, "r") as reader:
            params = Parameters.model_validate(json.load(reader))
    else:
        params = Parameters()

    random.seed(params.seed)

    return params

_parse_args()

Parse command-line arguments.

Source code in snailz/main.py
49
50
51
52
53
54
55
56
57
58
def _parse_args():
    """Parse command-line arguments."""

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--defaults", action="store_true", help="show default parameters"
    )
    parser.add_argument("--outdir", default=None, help="output directory")
    parser.add_argument("--params", default=None, help="JSON parameter file")
    return parser.parse_args()

_save(args, grids, persons, samples, changes)

Save synthesized data.

Source code in snailz/main.py
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def _save(args, grids, persons, samples, changes):
    """Save synthesized data."""

    if args.outdir == "-":
        outdir = None
    else:
        outdir = Path(args.outdir)
        if not outdir.is_dir():
            outdir.mkdir(exist_ok=True)

    for g in grids:
        with utils.file_or_std(outdir, f"{g.id}.csv", "w") as writer:
            print(g, file=writer)

    with utils.file_or_std(outdir, "grids.csv", "w") as writer:
        print(Grid.tidy(grids), file=writer)

    for name, cls, data in (("persons", Person, persons), ("samples", Sample, samples)):
        with utils.file_or_std(outdir, f"{name}.csv", "w") as writer:
            print(cls.csv_header(), file=writer)
            for record in data:
                print(record, file=writer)

    with utils.file_or_std(outdir, "changes.json", "w") as writer:
        json.dump(changes, writer)

_synthesize(params)

Synthesize data.

Source code in snailz/main.py
88
89
90
91
92
93
94
def _synthesize(params):
    """Synthesize data."""

    grids = [Grid.make(params) for _ in range(params.num_grids)]
    persons = [Person.make(params) for _ in range(params.num_persons)]
    samples = [Sample.make(params, grids, persons) for _ in range(params.num_samples)]
    return grids, persons, samples