# |##########################################|
# | Base C++ Extension Publication |
# |##########################################|
from ._core import ( # type: ignore[import]
Agent,
DataBase,
DiagramType,
Entity,
Model,
ModelDiagram,
NativeUpdateFun,
Tool,
UpdateFun,
Virus,
__epiworld_version__,
)
from . import epimodels
__all__ = [
"Agent",
"DataBase",
"DiagramType",
"Entity",
"Model",
"ModelDiagram",
"NativeUpdateFun",
"Tool",
"UpdateFun",
"Virus",
"epimodels",
]
# |##########################################|
# | Platform-independent saver module |
# |##########################################|
import os
import csv
import json
import numpy as np # type: ignore[import]
from pathlib import Path
from typing import Any
DatabaseResults = dict[str, Any]
[docs]
def write_db_results_multiple_csv(
results: DatabaseResults, base: str = "results", directory: Path | str = Path(".")
) -> None:
"""Write each result to its own CSV file.
Parameters
----------
results : dict
Results from :func:`extract_database_results`.
base : str, default "results"
Prefix of the file names; each file is named ``{base}-{key}.csv``.
directory : str or pathlib.Path, default "."
Directory for the files, created if needed.
"""
if isinstance(directory, str):
directory = Path(directory)
os.makedirs(directory, exist_ok=True)
assert directory.exists() and directory.is_dir()
for key, value in results.items():
with open(directory / f"{base}-{key}.csv", "w", newline="") as f:
writer = csv.writer(f)
if isinstance(value, dict):
writer.writerow(value.keys())
rows = zip(*value.values())
writer.writerows(rows)
else:
writer.writerows(value)
[docs]
class NumpyJSONEncoder(json.JSONEncoder):
"""JSON encoder that stores NumPy scalars and arrays with their types.
Use it with :func:`json.dump` (``cls=NumpyJSONEncoder``), and
:func:`numpy_json_decoder` to read the values back.
"""
[docs]
def default(self, o):
if isinstance(o, np.integer):
return {"__numpy__": True, "dtype": str(o.dtype), "value": int(o)}
elif isinstance(o, np.floating):
return {"__numpy__": True, "dtype": str(o.dtype), "value": float(o)}
elif isinstance(o, np.ndarray):
return {
"__numpy__": True,
"dtype": str(o.dtype),
"shape": o.shape,
"value": o.tolist(),
}
return super().default(o)
[docs]
def numpy_json_decoder(dct):
"""Restore the NumPy values written by :class:`NumpyJSONEncoder`.
Use it as the ``object_hook`` of :func:`json.load`.
Parameters
----------
dct : dict
A decoded JSON object.
Returns
-------
object
The NumPy scalar or array ``dct`` encodes, or ``dct`` unchanged.
"""
if "__numpy__" in dct:
dtype = np.dtype(dct["dtype"])
if "shape" in dct:
return np.array(dct["value"], dtype=dtype).reshape(dct["shape"])
else:
return np.array(dct["value"], dtype=dtype).item()
return dct
[docs]
def write_db_results_json(
results: DatabaseResults, filename: Path = Path("results.json")
) -> None:
"""Write all the results to one JSON file.
Parameters
----------
results : dict
Results from :func:`extract_database_results`.
filename : str or pathlib.Path, default "results.json"
Path of the file.
"""
if isinstance(filename, str):
filename = Path(filename)
with open(filename, "w") as f:
json.dump(results, f, cls=NumpyJSONEncoder)
[docs]
def write_db_results_multiple_json(
results: DatabaseResults,
base: str = "results",
directory: Path = Path("."),
) -> None:
"""Write each result to its own JSON file.
Parameters
----------
results : dict
Results from :func:`extract_database_results`.
base : str, default "results"
Prefix of the file names; each file is named ``{base}-{key}.json``.
directory : str or pathlib.Path, default "."
Directory for the files, created if needed.
"""
if isinstance(directory, str):
directory = Path(directory)
os.makedirs(directory, exist_ok=True)
assert directory.exists() and directory.is_dir()
for key, value in results.items():
with open(directory / f"{base}-{key}.json", "w") as f:
json.dump(value, f, cls=NumpyJSONEncoder)
try:
import h5py # type: ignore[import]
def write_db_results_hdf5( # type: ignore
results: DatabaseResults, filename: Path = Path("results.h5"), **hdf5_kwargs
) -> None:
"""Write the results to an HDF5 file, one group per result.
Requires the ``h5py`` package.
Parameters
----------
results : dict
Results from :func:`extract_database_results`.
filename : str or pathlib.Path, default "results.h5"
Path of the file.
**hdf5_kwargs
Unused; kept for compatibility.
"""
def write_group(grp: h5py.Group, value):
if (
isinstance(value, dict)
and len(value) == 2
and "values" in value
and "indexes" in value
):
arr = np.asarray(value["indexes"])
if arr.dtype == object:
arr = arr.astype("U")
dset = grp.create_dataset("data", data=arr)
v = value["values"]
if isinstance(v, np.ndarray):
v = v.tolist()
dset.attrs["values"] = v
elif isinstance(value, dict):
for subkey, subval in value.items():
subgrp = grp.create_group(subkey)
write_group(subgrp, subval)
else:
arr = np.asarray(value)
if arr.dtype == object:
arr = arr.astype("U")
grp.create_dataset("data", data=arr)
if isinstance(filename, str):
filename = Path(filename)
with h5py.File(filename, "w") as root:
for key, value in results.items():
subgrp = root.create_group(key)
write_group(subgrp, value)
except ImportError:
[docs]
def write_db_results_hdf5(results, filename=None):
"""Write the results to an HDF5 file (requires ``h5py``, not installed).
Parameters
----------
results : dict
Results from :func:`extract_database_results`.
filename : str or pathlib.Path, optional
Path of the file.
Raises
------
ImportError
Always, because ``h5py`` is not installed.
"""
raise ImportError(
"The `h5py' package is not installed, cannot write HDF5 files."
)
try:
import zarr # type: ignore[import]
def write_db_results_zarr( # type: ignore
results: DatabaseResults, filename: Path = Path("results.zarr")
) -> None:
"""Write the results to a Zarr store, one group per result.
Requires the ``zarr`` package.
Parameters
----------
results : dict
Results from :func:`extract_database_results`.
filename : str or pathlib.Path, default "results.zarr"
Path of the store.
"""
def write_group(grp: zarr.Group, value):
if (
isinstance(value, dict)
and len(value) == 2
and "values" in value
and "indexes" in value
):
arr = np.asarray(value["indexes"])
if arr.dtype == object:
arr = arr.astype("U")
dset = grp.create_array("data", data=arr)
v = value["values"]
if isinstance(v, np.ndarray):
v = v.tolist()
dset.attrs["values"] = v
elif isinstance(value, dict):
for subkey, subval in value.items():
subgrp = grp.create_group(subkey)
write_group(subgrp, subval)
else:
arr = np.asarray(value)
if arr.dtype == object:
arr = arr.astype("U")
grp.create_array("data", data=arr)
if isinstance(filename, str):
filename = Path(filename)
root = zarr.open_group(filename, mode="w")
for key, value in results.items():
subgrp = root.create_group(key)
write_group(subgrp, value)
except ImportError:
[docs]
def write_db_results_zarr(results, filename=None):
"""Write the results to a Zarr store (requires ``zarr``, not installed).
Parameters
----------
results : dict
Results from :func:`extract_database_results`.
filename : str or pathlib.Path, optional
Path of the store.
Raises
------
ImportError
Always, because ``zarr`` is not installed.
"""
raise ImportError(
"The `zarr' package is not installed, cannot write Zarr files."
)
__all__ += [
"DatabaseResults",
"extract_database_results",
"write_db_results_multiple_csv",
"write_db_results_json",
"write_db_results_multiple_json",
]
if write_db_results_hdf5 is not None:
__all__.append("write_db_results_hdf5")
if write_db_results_zarr is not None:
__all__.append("write_db_results_zarr")