Skip to content

Main

Synthesize data.

main()

Main command-line driver.

Source code in src/snailz/main.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def main():
    """Main command-line driver."""

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

    params = _initialize(args)
    data = _synthesize(params)
    data["changes"] = do_all_effects(params, data)
    data["tidy_grids"] = Grid.tidy(data["grids"])

    if args.outdir is not None:
        _save_params(args.outdir, params)
        _save_csv(args.outdir, data)
        if args.outdir != "-":
            _save_db(args.outdir, data)

    return 0

_ensure_dir(dirname)

Ensure directory exists.

Source code in src/snailz/main.py 
46
47
48
49
50
51
def _ensure_dir(dirname):
    """Ensure directory exists."""

    dirpath = Path(dirname)
    if not dirpath.is_dir():
        dirpath.mkdir(exist_ok=True)

_initialize(args)

Initialize for data synthesis.

Source code in src/snailz/main.py 
54
55
56
57
58
59
60
61
62
63
64
65
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 src/snailz/main.py 
68
69
70
71
72
73
74
75
76
77
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_csv(outdir, data)

Save synthesized data as CSV.

Source code in src/snailz/main.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def _save_csv(outdir, data):
    """Save synthesized data as CSV."""

    if outdir == "-":
        outdir = None
    else:
        _ensure_dir(Path(outdir))

    persist.grids_to_csv(outdir, data["grids"], data["tidy_grids"])
    for name, cls in (
        ("machines", Machine),
        ("persons", Person),
        ("ratings", Rating),
        ("samples", Sample),
    ):
        with utils.file_or_std(outdir, f"{name}.csv", "w") as writer:
            persist.objects_to_csv(writer, data[name])

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

_save_db(outdir, data)

Save synthesized data as CSV.

Source code in src/snailz/main.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def _save_db(outdir, data):
    """Save synthesized data as CSV."""

    _ensure_dir(outdir)
    dbpath = outdir / DB_FILE
    dbpath.unlink(missing_ok=True)

    cnx = sqlite3.connect(dbpath)

    for table, name in (
        ("machine", "machines"),
        ("person", "persons"),
        ("rating", "ratings"),
        ("sample", "samples"),
    ):
        persist.objects_to_db(cnx, table, data[name])

    persist.grids_to_db(cnx, data["tidy_grids"])

    cnx.close()

_save_params(outdir, params)

Save parameters.

Source code in src/snailz/main.py 
124
125
126
127
128
129
130
131
132
def _save_params(outdir, params):
    """Save parameters."""

    if outdir == "-":
        sys.stdout.write(utils.json_dump(params))
    else:
        _ensure_dir(Path(outdir))
        with open(Path(outdir, "params.json"), "w") as writer:
            writer.write(utils.json_dump(params))

_synthesize(params)

Synthesize data.

Source code in src/snailz/main.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def _synthesize(params):
    """Synthesize data."""

    grids = Grid.make(params)
    persons = Person.make(params)
    samples = Sample.make(params, grids, persons)
    machines = Machine.make(params)
    ratings = Rating.make(persons, machines)
    return {
        "grids": grids,
        "persons": persons,
        "samples": samples,
        "machines": machines,
        "ratings": ratings,
    }