emod-api 3.0.2__py3-none-any.whl

emod_api/serialization/dtk_file_utility.py

@@ -0,0 +1,141 @@
+#!/usr/bin/python
+
+import argparse
+import emod_api.serialization.dtk_file_tools as dft
+import json
+import os
+import sys
+
+
+def __do_read__(args):
+
+    if args.output is not None:
+        prefix = args.output
+    else:
+        root, _ = os.path.splitext(args.filename)
+        prefix = root
+
+    if args.raw:
+        extension = 'bin'
+    else:
+        extension = 'json'
+
+    dtk_file = dft.read(args.filename)
+
+    if args.header:
+        with open(args.header, 'w') as handle:
+            json.dump(dtk_file.header, handle, indent=2, separators=(',', ':'))
+
+    print(f'File header: {dtk_file.header}')
+
+    for index in range(len(dtk_file.chunks)):
+        if args.raw:
+            # Write raw chunks to disk
+            output = dtk_file.chunks[index]
+        else:
+            if args.unformatted:
+                # Expand compressed contents, but don't serialize and format
+                output = dtk_file.contents[index]
+            else:
+                # Expand compressed contents, serialize, write out formatted
+                obj = dtk_file.objects[index]
+                print(f'Formatting chunk {index + 1} of {len(dtk_file.chunks)}... ', end='')
+                output = json.dumps(obj, indent=2, separators=(',', ':'))
+
+        if index == 0:
+            output_filename = '.'.join([prefix, 'simulation', extension])
+        else:
+            output_filename = '.'.join([prefix, f'node-{index:0>5}', extension])
+
+        print(f"Writing file '{output_filename}'")
+        if args.raw:
+            with open(output_filename, 'wb') as handle:
+                handle.write(output)
+        else:
+            with open(output_filename, 'wt', encoding='utf-8') as handle:
+                handle.write(output)
+
+    return
+
+
+def __do_write__(args):
+
+    print(f"Writing file '{args.filename}'", file=sys.stderr)
+    print(f"Reading simulation data from '{args.simulation}'", file=sys.stderr)
+    print(f"Reading node data from {args.nodes}", file=sys.stderr)
+    print(f"Author = {args.author}", file=sys.stderr)
+    print(f"Tool = {args.tool}", file=sys.stderr)
+    print(f"{'Compressing' if args.compress else 'Not compressing'} contents", file=sys.stderr)
+    print(f"{'Verifying' if args.verify else 'Not verifying'} contents", file=sys.stderr)
+    print(f"Using compression engine '{args.engine}'", file=sys.stderr)
+
+    dtk_file = dft.DtkFileV3()
+    dtk_file.author = args.author
+    dtk_file.tool = args.tool
+    dtk_file.compression = args.engine
+
+    _prepare_simulation_data(args.simulation, dtk_file)
+    _prepare_node_data(args.nodes, dtk_file)
+
+    dft.write(dtk_file, args.filename)
+
+    return
+
+
+def _prepare_simulation_data(filename, dtk_file):
+
+    with open(filename, 'rb') as handle:
+        data = handle.read()
+        # Do not use dtk_file.simulation property because this is text rather than a Python object
+        dtk_file.contents.append(data)
+
+    return
+
+
+# noinspection SpellCheckingInspection
+def _prepare_node_data(filenames, dtk_file):
+
+    for filename in filenames:
+        with open(filename, 'rb') as handle:
+            data = handle.read()
+            # Do not use dtk_file.nodes here because this is text rather than a Python object
+            dtk_file.contents.append(data)
+
+    return
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(help='add_subparsers help')
+
+    read_parser = subparsers.add_parser('read', help='read help')
+    read_parser.add_argument('filename')
+    read_parser.add_argument('--header', default=None, help='Write header to file', metavar='<filename>')
+    read_parser.add_argument('-r', '--raw', default=False, action='store_true',
+                             help='Write raw contents of chunks to disk')
+    # noinspection SpellCheckingInspection
+    read_parser.add_argument('-u', '--unformatted', default=False, action='store_true',
+                             help='Write unformatted (compact) JSON to disk')
+    read_parser.add_argument('-o', '--output', default=None,
+                             help='Output filename prefix, defaults to input filename with .json extension')
+    read_parser.set_defaults(func=__do_read__)
+
+    username = os.environ['USERNAME'] if 'USERNAME' in os.environ else os.environ['USER']
+    tool_name = os.path.basename(__file__)
+
+    write_parser = subparsers.add_parser('write', help='write help')
+    write_parser.add_argument('filename', help='Output .dtk filename')
+    write_parser.add_argument('simulation', help='Filename for simulation JSON')
+    write_parser.add_argument('nodes', nargs='+', help='Filename(s) for node JSON')
+    write_parser.add_argument('-a', '--author', default=username, help=f'Author name for header [{username}]')
+    write_parser.add_argument('-t', '--tool', default=tool_name, help=f'Tool name for header [{tool_name}]')
+    write_parser.add_argument('-u', '--uncompressed', default=True, action='store_false', dest='compress',
+                              help='Do not compress contents of new .dtk file')
+    write_parser.add_argument('-v', '--verify', default=False, action='store_true',
+                              help='Verify JSON in simulation and nodes (could be slow).')
+    write_parser.add_argument('-e', '--engine', default='LZ4', help='Compression engine {NONE|LZ4|SNAPPY} [LZ4]')
+    write_parser.set_defaults(func=__do_write__)
+
+    commandline_args = parser.parse_args()
+    commandline_args.func(commandline_args)

emod_api/serialization/serialized_population.py

@@ -0,0 +1,205 @@
+"""Class to load and manipulate a saved population."""
+import difflib
+import emod_api.serialization.dtk_file_tools as dft
+
+from collections.abc import Iterable
+from typing import Union
+
+COUNTER = 0
+
+
+class SerializedPopulation:
+    """Opens the passed file and reads in all the nodes.
+
+    Args:
+        file: serialized population file
+
+    Examples:
+        Create an instance of SerializedPopulation::
+            import emod_api.serialization.SerializedPopulation as SerPop
+            ser_pop = SerPop.SerializedPopulation('state-00001.dtk')
+
+     """
+
+    def __init__(self, file: str):
+        self.next_infection_suid = None
+        self.next_infection_suid_initialized = False
+        self.dtk = dft.read(file)
+
+    @property
+    def nodes(self):
+        """All nodes.
+
+        Examples:
+            Delete number_of_ind individuals from node 0::
+                node = ser_pop.nodes[0]
+                del node.individualHumans[0:number_of_ind]
+
+            Only keep individuals with a certain condition::
+                node.individualHumans = [ind for ind in node.individualHumans if keep_fct(ind)]
+
+            Change susceptibility of an individual::
+                print(node.individualHumans[0].susceptibility)
+                new_susceptibility = {"age": 101.01, "mod_acquire": 0}
+                node.individualHumans[0].susceptibility.update(new_susceptibility)
+
+            Copy individual[0] from node 0, change properties and add individual as new individual::
+                import copy
+                individual_properties = {"m_age": 1234}
+                individual = copy.deepcopy(node.individualHumans[0])
+                individual["suid"] = ser_pop.get_next_individual_suid(0)
+                individual.update(individual_properties)
+                ser_pop.nodes[0].individualHumans.append(individual)
+
+            Infect an individual with an infection copied from another individual::
+                import copy
+                individual_0 = node.individualHumans[0]
+                individual_1 = node.individualHumans[1]
+                new_infection = copy.deepcopy(individual_0.infections[0])
+                new_infection["suid"] = ser_pop.get_next_infection_suid()
+                individual_1.infections.append(new_infection)
+                individual_1.m_is_infected = True
+
+        """
+        return self.dtk.nodes
+
+    def flush(self):
+        """Save all made changes to the node(s)."""
+        for idx in range(len(self.dtk.nodes)):
+            self.dtk.nodes[idx] = self.dtk.nodes[idx]
+
+    def write(self, output_file: str = "my_sp_file.dtk"):
+        """Write the population to a file.
+
+        Args:
+            output_file: output file
+        """
+        self.flush()
+        sim = self.dtk.simulation
+        sim["infectionSuidGenerator"]["next_suid"] = self.get_next_infection_suid()
+        self.dtk.simulation = sim
+
+        print(f"Saving file {output_file}.")
+        dft.write(self.dtk, output_file)
+
+    def get_next_infection_suid(self):
+        """Each infection needs a unique identifier, this function returns one."""
+        sim = self.dtk.simulation
+        if not self.next_infection_suid_initialized:
+            self.next_infection_suid = sim["infectionSuidGenerator"]["next_suid"]
+            self.next_infection_suid_initialized = True
+        else:
+            self.next_infection_suid["id"] = (
+                self.next_infection_suid["id"]
+                + sim["infectionSuidGenerator"]["numtasks"]
+            )
+
+        return dict(self.next_infection_suid)
+
+    def get_next_individual_suid(self, node_id: int) -> dict:
+        """Each individual needs a unique identifier, this function returns one.
+
+        Args:
+            node_id: The first parameter.
+
+        Returns:
+            The return value. True for success, False otherwise.
+
+        Examples:
+            To get a unique id for an individual::
+
+                print(sp.get_next_individual_suid(0))
+                {'id': 2}
+        """
+        suid = self.dtk.nodes[node_id]["m_IndividualHumanSuidGenerator"]["next_suid"]
+        self.dtk.nodes[node_id]["m_IndividualHumanSuidGenerator"]["id"] = (
+            suid["id"]
+            + self.dtk.nodes[node_id]["m_IndividualHumanSuidGenerator"]["numtasks"]
+        )
+        return dict(suid)
+
+
+# Some useful functions
+def find(name: str,
+         handle: Union[str, Iterable],
+         currentlevel: str = "dtk.nodes"):
+    """Recursively searches for a paramters that matches or is close to name and prints out where to find it in the file.
+
+    Args:
+        name: the paramter you are looking for e.g. "age", "gender".
+        handle: some iterable data structure, can be a list of
+                nodes, a node, list of individuals, etc
+        currentlevel: just a string to print out where the found item
+                is located e.g. "dtk.nodes" or "dtk.node.individuals"
+
+    Examples:
+        What is the exact paramteter name used for the age of an individual?::
+
+            SerPop.find("age", node)
+            ...
+            1998   Found in:  dtk.nodes.individualHumans[999].m_age
+            1999   Found in:  dtk.nodes.individualHumans[999].susceptibility.age
+            2000   Found in:  dtk.nodes.m_vectorpopulations[0].EggQueues[0].age
+            2001   Found in:  dtk.nodes.m_vectorpopulations[0].EggQueues[1].age
+            ...
+    """
+    global COUNTER
+    if isinstance(handle, str) and difflib.get_close_matches(name, [handle], cutoff=0.6):
+        print(COUNTER, "  Found in: ", currentlevel)
+        COUNTER += 1
+        return
+
+    if isinstance(handle, str) or not isinstance(handle, Iterable):
+        return
+
+    # key can be a string or on dict/list/..
+    for idx, key in enumerate(handle):
+        level = (
+            currentlevel + "." + key
+            if isinstance(key, str)
+            else currentlevel + "[" + str(idx) + "]"
+        )
+        try:
+            tmp = handle[key]
+            if isinstance(tmp, Iterable):
+                find(name, key, level + "[]")
+            else:
+                find(name, key, level)
+        except BaseException:
+            # list or keys of a dict, works in all cases but misses objects in
+            # dicts
+            find(name, key, level)
+        if isinstance(handle, dict):
+            find(name, handle[key], level)  # check if string is key for a dict
+
+
+def get_parameters(handle: Union[str, Iterable],
+                   currentlevel: str = "dtk.nodes"):
+    """Return a set of all parameters in the serialized population file. Helpful to get an overview about what is in the serialized population file.
+
+    Args:
+        handle: some iterable data structure, can be a list of
+                nodes, a node, list of individuals, etc
+        currentlevel: just a string to print out where the found item
+                is located e.g. "dtk.nodes" or "dtk.node.individuals
+    Examples:
+        Print all parameters in serialized population file::
+            for n in sorted(SerPop.get_parameters(node)):
+                print(n)
+    """
+    param = set()
+
+    if isinstance(handle, str):
+        param.add(currentlevel)
+        return param
+
+    if not isinstance(handle, Iterable):
+        return param
+
+    for _, d in enumerate(handle):
+        level = currentlevel + " " + d if isinstance(d, str) else currentlevel
+        param.update(get_parameters(d, level))
+        if isinstance(handle, dict):
+            param.update(get_parameters(handle[d], level))
+
+    return param

emod_api/spatialreports/__main__.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+
+import os
+from argparse import ArgumentParser
+from emod_api.spatialreports.spatial import SpatialReport
+
+MATPLOTLIB = True
+try:
+    import matplotlib
+    matplotlib.use('TkAgg')
+    import matplotlib.pyplot as plt
+except ModuleNotFoundError:
+    print("This example requires the matplotlib package.")
+    MATPLOTLIB = False
+
+SCRIPT_PATH = os.path.realpath(__file__)
+WORKING_DIRECTORY = os.path.dirname(SCRIPT_PATH)
+
+
+def main(filename: str):
+
+    report = SpatialReport(filename)
+
+    plt.xkcd()
+
+    y_axis_guess = os.path.basename(filename).strip("SpatialReport_").strip(".bin")
+    # Show node-wise time series data
+    for node_id in report.node_ids:
+        plt.plot(report[node_id].data, label=f"Node {node_id}")
+    plt.title(f"{y_axis_guess} for Nodes")
+    plt.legend()
+    plt.xlabel("Time Step")
+    plt.ylabel(f"{y_axis_guess}")
+    plt.show()
+
+    return
+
+
+def dump_source():
+
+    with open(SCRIPT_PATH, "r") as file:
+        print(file.read())
+
+    return
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument(
+        "-f", "--filename", default=None, help="spatial report filename"
+    )
+    parser.add_argument(
+        "-g",
+        "--get",
+        default=False,
+        action="store_true",
+        help="Write source code for this example to stdout.",
+    )
+
+    args = parser.parse_args()
+
+    if args.get:
+        dump_source()
+    elif args.filename and MATPLOTLIB:
+        main(args.filename)
+    else:
+        parser.print_help()

emod_api/spatialreports/plot_spat_means.py

@@ -0,0 +1,99 @@
+"""
+This script assumes local SpatialReport_XXX.bin files which have been downloaded (from COMPS)
+using pyCOMPS getexpout function or equivalent. Note that this interacts with files on an experiment
+basis, not simulation basis. It assumes the files are in a subdirectory named after the experiment id,
+and then in subdirectories of that named after the simulation id.
+
+<exp_id>/
+  <sim1_id>/
+      SpatialReport_XXX.bin
+  <sim2_id>/
+      SpatialReport_XXX.bin
+
+The idea is that the data is most interesting not on a simulation basis, but for an experiment,
+especially aggregated on a certain sweep param and value. This plot calculates means and plots those.
+
+Option 1: For each node, plot the mean of the specified channel for all files (values) found in experiment.
+Option 2: For each node, plot the mean of the specified channel for all files (values) found in experiment _limited_ by specified tag key-value pair.
+There is little to no assistance here so you need to specify a valid key and value.
+"""
+
+import os
+import emod_api.spatialreports.spatial as sr
+import matplotlib.pyplot as plt
+import numpy as np
+import sqlite3
+
+
+def collect(exp_id, chan="Prevalence", tag=None):
+    node_chan_data = {}
+    node_chan_means = {}
+    groupby_values = {}
+    if tag:
+        if len(tag.split("=")) == 1:
+            raise ValueError("When passing tag, has to have key=value format.")
+
+        groupby_key = tag.split("=")[0]
+        groupby_value = tag.split("=")[1]
+        db = os.path.join("latest_experiment", "results.db")
+        con = sqlite3.connect(db)
+        cur = con.cursor()
+
+        query = f"SELECT sim_id FROM results where CAST({groupby_key} AS DECIMAL)-CAST({groupby_value} AS DECIMAL)<0.0001"
+        all_results = cur.execute(query)
+        groupby_values["ref"] = list()
+        for result in all_results:
+            sim_id = result[0]
+            groupby_values["ref"].append(sim_id)
+    else:
+        groupby_values["ref"] = os.listdir(exp_id)
+        if "results.db" in groupby_values["ref"]:
+            groupby_values["ref"].remove("results.db")
+
+    for sim_id in groupby_values["ref"]:
+        report_path = os.path.join(str(exp_id), sim_id, "SpatialReport_" + chan + ".bin")
+        data = sr.SpatialReport(report_path)
+        for node_id in data.node_ids:
+            chan_data = data[node_id].data
+            if node_id not in node_chan_data:
+                node_chan_data[node_id] = list()
+                node_chan_means[node_id] = list()
+            node_chan_data[node_id].append(chan_data)
+            node_chan_means[node_id] = np.zeros(len(node_chan_data[node_id]))
+    for node in node_chan_data.keys():
+        node_chan_means[node] = np.mean(np.array(node_chan_data[node]), axis=0)
+
+    return node_chan_means
+
+
+def plot(exp_id, chan="Prevalence", tag=None):
+
+    node_chan_means = collect(exp_id, chan, tag)
+
+    for node in node_chan_means.keys():
+        plt.plot(node_chan_means[node], label=f"node={node}")
+    plt.xlabel("Timestep")
+    plt.ylabel(chan)
+    plt.title(f"Mean values of {chan} over time by node.")
+    plt.legend()
+    plt.show()
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description='Spatial Report Plotting')
+    parser.add_argument('-c', '--channel', action='store', default="Prevalence", help='channel(s) to display [Prevalence]')
+    parser.add_argument('-e', '--experiment_id', action='store', default=None, help='experiment id to plot, data assumed to be local')
+    parser.add_argument('-t', '--tag', action='store', default="", help='tag constraint')
+    args = parser.parse_args()
+    if not args.experiment_id:
+        with open("COMPS_ID", "r") as fp:
+            args.experiment_id = fp.read()
+    if not args.channel:
+        args.channel = 'Prevalence' # should not be necessary
+
+    # check that folder with name experiment_id exists
+    if not os.path.exists(str(args.experiment_id)):
+        raise ValueError(f"Don't see folder for {args.experiment_id}.")
+
+    plot(str(args.experiment_id), args.channel, tag=args.tag)