Databook default all column

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 This file records changes to the codebase grouped by version release. Unreleased changes are generally only present during development (relevant parts of the changelog can be written and saved in that section before a version number has been assigned)
 
+## [1.31.0] - 2025-02-18
+
+- Framework definitions of compartments, characteristics, and parameters support a new column 'databook default all'. If set to 'y', then when a databook is produced, the data entry table will only contain a record for 'All' instead of having population-specific rows. Further manual editing of the tables is supported as normal. 
+
 ## [1.30.0] - 2025-02-17
 
 - Automatic calibration can now selectively weight parts of the time series to select or prioritise a subset of time points.

README.md

@@ -2,7 +2,7 @@
 
 [![Build Status](https://dev.azure.com/AtomicaTeam/Atomica/_apis/build/status/atomicateam.atomica?branchName=master)](https://dev.azure.com/AtomicaTeam/Atomica/_build/latest?definitionId=1&branchName=master)
 
-[![PyPi version](https://badgen.net/pypi/v/atomica/)](https://pypi.com/project/atomica)
+[![PyPi version](https://badgen.net/pypi/v/atomica/)](https://pypi.org/project/atomica)
 
 Atomica is a simulation engine for compartmental models. It can be used to simulate disease epidemics, health care cascades, and many other things.
 

atomica/data.py

@@ -21,6 +21,7 @@
 from collections import defaultdict
 import pandas as pd
 import itertools
+from .version import version, gitinfo
 
 __all__ = ["InvalidDatabook", "ProjectData"]
 
@@ -58,11 +59,21 @@ def __init__(self, framework):
         self.tdve_pages = sc.odict()  #: This is an odict mapping worksheet name to an (ordered) list of TDVE code names appearing on that sheet
 
         # Internal storage used with methods while writing
-        self._pop_types = list(framework.pop_types.keys())  # : Store set of valid population types from framework
+        self._pop_types = list(framework.pop_types.keys())  #: Store set of valid population types from framework
         self._formats = None  #: Temporary storage for the Excel formatting while writing a databook
         self._book = None  #: Temporary storage for the workbook while writing a databook
         self._references = None  #: Temporary storage for cell references while writing a databook
 
+        self.version = version  #: Current Atomica version
+        self.gitinfo = sc.dcp(gitinfo)  #: Atomica Git version information, if being run in a Git repository
+
+    def __setstate__(self, d):
+        from .migration import migrate
+
+        self.__dict__ = d
+        projectdata = migrate(self)
+        self.__dict__ = projectdata.__dict__
+
     def tables(self):
         """
         Return iterator over all TDVE and TDC tables
@@ -261,13 +272,15 @@ class instance (e.g. if creating a new databook).
                     pop_type = spec.get("population type")
                     databook_order = spec.get("databook order")
                     full_name = spec["display name"]
+                    default_all = spec["databook default all"] == "y"
+                    allowed_units = [framework.get_databook_units(full_name)]
 
                     if pd.isna(databook_order):
                         order = np.inf
                     else:
                         order = databook_order
                     pages[databook_page].append((spec.name, order))
-                    data.tdve[spec.name] = TimeDependentValuesEntry(full_name, data.tvec, allowed_units=[framework.get_databook_units(full_name)], comment=spec["guidance"], pop_type=pop_type)
+                    data.tdve[spec.name] = TimeDependentValuesEntry(full_name, data.tvec, allowed_units=allowed_units, comment=spec["guidance"], pop_type=pop_type, default_all=default_all)
                     data.tdve[spec.name].write_units = True
                     data.tdve[spec.name].write_uncertainty = True
                     if obj_type == "pars":
@@ -277,6 +290,10 @@ class instance (e.g. if creating a new databook).
                             data.tdve[spec.name].write_uncertainty = False  # Don't show uncertainty for timed parameters. In theory users could manually add the column and sample over it, but because the duration is rounded to the timestep, it's likely to have confusing stepped effects
                     data.tdve[spec.name].pop_type = pop_type
 
+                    if default_all:
+                        # add_pop normally adds TDVE rows, but it won't operate on any TDVEs that default to 'All' so we need to add the 'All' rows here
+                        data.tdve[spec.name].ts["All"] = TimeSeries(units=allowed_units[0])
+
         # Now convert pages to full names and sort them into the correct order
         for _, spec in framework.sheets["databook pages"][0].iterrows():
 
@@ -627,7 +644,7 @@ def add_pop(self, code_name: str, full_name: str, pop_type: str = None) -> None:
 
         :param code_name: The code name for the new population
         :param full_name: The full name/label for the new population
-        :param pop_type: String with the population type code name
+        :param pop_type: String with the population type code name (optional) - default is the type of the first population
 
         """
 
@@ -653,8 +670,11 @@ def add_pop(self, code_name: str, full_name: str, pop_type: str = None) -> None:
         for tdve in self.tdve.values():
             # Since TDVEs in databooks must have the unit set in the framework, all ts objects must share the same units
             # And, there is only supposed to be one type of unit allowed for TDVE tables (if the unit is empty, it will be 'N.A.')
-            # so can just pick the first of the allowed units
-            if tdve.pop_type == pop_type:
+            # so can just pick the first of the allowed units. We will add the population row if the pop type matches and if
+            # the TDVE is either not a 'default_all' or if the user has removed the 'All' row from the TDVE despite it being default_all
+            if tdve.pop_type != pop_type or (tdve.default_all and ("All" in tdve.ts or "all" in tdve.ts)):
+                continue
+            else:
                 tdve.ts[code_name] = TimeSeries(units=tdve.allowed_units[0])
 
     def rename_pop(self, existing_code_name: str, new_code_name: str, new_full_name: str) -> None:

atomica/excel.py

@@ -884,7 +884,7 @@ class TimeDependentValuesEntry:
 
     """
 
-    def __init__(self, name, tvec: np.array = None, ts=None, allowed_units: list = None, comment: str = None, pop_type: str = None):
+    def __init__(self, name, tvec: np.array = None, ts=None, allowed_units: list = None, comment: str = None, pop_type: str = None, default_all: bool = False):
 
         if ts is None:
             ts = sc.odict()
@@ -907,6 +907,8 @@ def __init__(self, name, tvec: np.array = None, ts=None, allowed_units: list = N
         self.write_uncertainty = None  #: Write a column for uncertainty (if None, uncertainty will be written if any of the TimeSeries have uncertainty)
         self.write_assumption = None  #: Write a column for assumption/constant (if None, assumption will be written if any of the TimeSeries have an assumption)
 
+        self.default_all = default_all  #: Record whether the framework specifies that this TDVE should default to having an 'All' row instead of population-specific rows (the user can manually modify further)
+
     def __repr__(self):
         output = sc.prepr(self)
         return output

atomica/framework.py

@@ -632,6 +632,7 @@ def _sanitize_compartments(self) -> None:
             "is source": "n",
             "is junction": "n",
             "databook page": None,
+            "databook default all": "n",
             "default value": None,
             "population type": None,
             "databook order": None,  # Default is for it to be randomly ordered if the databook page is not None
@@ -643,6 +644,7 @@ def _sanitize_compartments(self) -> None:
             "is sink": {"y", "n"},
             "is source": {"y", "n"},
             "is junction": {"y", "n"},
+            "databook default all": {"y", "n"},
         }
         numeric_columns = ["databook order", "default value"]
 
@@ -723,13 +725,15 @@ def _sanitize_characteristics(self) -> None:
             "default value": None,
             "databook page": None,
             "databook order": None,
+            "databook default all": "n",
             "guidance": None,
             "population type": None,
             "provenance": FS.DEFAULT_PROVENANCE,
         }
         valid_content = {
             "display name": None,
             "components": None,
+            "databook default all": {"y", "n"},
         }
         numeric_columns = ["databook order", "default value"]
 
@@ -861,6 +865,7 @@ def _sanitize_parameters(self) -> None:
             "function": None,
             "databook page": None,
             "databook order": None,
+            "databook default all": "n",
             "targetable": "n",
             "guidance": None,
             "timescale": None,
@@ -874,6 +879,7 @@ def _sanitize_parameters(self) -> None:
             "targetable": {"y", "n"},
             "is derivative": {"y", "n"},
             "timed": {"y", "n"},
+            "databook default all": {"y", "n"},
         }
         numeric_columns = ["databook order", "default value", "minimum value", "maximum value", "timescale"]
 

atomica/migration.py

@@ -812,3 +812,44 @@ def _convert_framework_columns(framework):
 def _parset_add_initialization(parset):
     parset.initialization = None
     return parset
+
+
+@migration("ProjectData", "1.30.0", "1.31.0", "Add pop types attribute and default_all")
+def _projectdata_add_types_default(D):
+
+    if not hasattr(D, "_pop_types"):
+        # We have to check for existence because
+        # Are any population types present?
+        pop_types = list(x["type"] for x in D.pops.values() if "type" in x)
+        if not pop_types:
+            pop_types = [FS.DEFAULT_POP_TYPE]
+
+        for pop, spec in D.pops.items():
+            if "type" not in spec:
+                spec["type"] = pop_types[0]
+        D._pop_types = list({x["type"]: None for x in D.pops.values()}.keys())  # Using a dictionary here allows for order-preserving unique
+
+    for tdve in D.tdve.values():
+
+        # Fix the '_add_pop_type' migration which added tdve.type instead of tdve.pop_type
+        if hasattr(tdve, "type"):
+            tdve.pop_type = tdve.type
+            delattr(tdve, "type")
+        elif not hasattr(tdve, "pop_type"):
+            tdve.pop_type = D._pop_types[0]  # The majority of the time, if the pop_type is missing, the default needs to be added
+
+        # Also add the default_all attribute
+        tdve.default_all = False
+
+        # Some old TDVE tables have a lowercase 'n.a.' instead of the correct default value
+        for ts in tdve.ts.values():
+            if ts.units == FS.DEFAULT_SYMBOL_INAPPLICABLE.lower():
+                ts.units = FS.DEFAULT_SYMBOL_INAPPLICABLE
+
+        # A TDVE should always have some allowed units, if the allowed units have not been populated, then draw
+        # them from the ts entries. There are some older saved files that may have no allowed units even though the ts
+        # entries themselves have units.
+        if not hasattr(tdve, "allowed_units") or not tdve.allowed_units:
+            tdve.allowed_units = list({x.units: None for x in tdve.ts.values()}.keys())
+
+    return D

atomica/plotting.py

@@ -1250,31 +1250,32 @@ def _get_legend_handles(ax, handles, labels):
     if handles is None:
         if ax is None:
             ax = plt.gca()
-        elif isinstance(ax, plt.Figure): # Allows an argument of a figure instead of an axes # pragma: no cover
+        elif isinstance(ax, plt.Figure):  # Allows an argument of a figure instead of an axes # pragma: no cover
             ax = ax.axes[-1]
         handles, labels = ax.get_legend_handles_labels()
-    else: # pragma: no cover
+    else:  # pragma: no cover
         if labels is None:
             labels = [h.get_label() for h in handles]
         else:
             assert len(handles) == len(labels), f"Number of handles ({len(handles)}) and labels ({len(labels)}) must match"
     return ax, handles, labels
 
+
 # Temporary copy of function from Sciris to remove after Sciris update
 def separatelegend(ax=None, handles=None, labels=None, reverse=False, figsettings=None, legendsettings=None):
-    """ Allows the legend of a figure to be rendered in a separate window instead """
+    """Allows the legend of a figure to be rendered in a separate window instead"""
 
     # Handle settings
-    f_settings = sc.mergedicts({'figsize':(4.0,4.8)}, figsettings) # (6.4,4.8) is the default, so make it a bit narrower
-    l_settings = sc.mergedicts({'loc': 'center', 'bbox_to_anchor': None, 'frameon': False}, legendsettings)
+    f_settings = sc.mergedicts({"figsize": (4.0, 4.8)}, figsettings)  # (6.4,4.8) is the default, so make it a bit narrower
+    l_settings = sc.mergedicts({"loc": "center", "bbox_to_anchor": None, "frameon": False}, legendsettings)
 
     # Get handles and labels
     _, handles, labels = _get_legend_handles(ax, handles, labels)
 
     # Set up new plot
     fig = plt.figure(**f_settings)
     ax = fig.add_subplot(111)
-    ax.set_position([-0.05,-0.05,1.1,1.1]) # This cuts off the axis labels, ha-ha
+    ax.set_position([-0.05, -0.05, 1.1, 1.1])  # This cuts off the axis labels, ha-ha
     ax.set_axis_off()  # Hide axis lines
 
     # A legend renders the line/patch based on the object handle. However, an object
@@ -1290,17 +1291,16 @@ def separatelegend(ax=None, handles=None, labels=None, reverse=False, figsetting
         handles2.append(h2)
 
     # Reverse order, e.g. for stacked plots
-    if reverse: # pragma: no cover
+    if reverse:  # pragma: no cover
         handles2 = handles2[::-1]
-        labels   = labels[::-1]
+        labels = labels[::-1]
 
     # Plot the new legend
     ax.legend(handles=handles2, labels=labels, **l_settings)
 
     return fig
 
 
-
 def plot_bars(plotdata, stack_pops=None, stack_outputs=None, outer=None, legend_mode=None, show_all_labels=False, orientation="vertical") -> list:
     """
     Produce a bar plot

atomica/version.py

@@ -6,6 +6,6 @@
 
 import sciris as sc
 
-version = "1.30.0"
+version = "1.31.0"
 versiondate = "2025-02-17"
 gitinfo = sc.gitinfo(__file__)

tests/test_tox_databooks.py

@@ -181,9 +181,27 @@ def test_databook_all():
     at.plot_series(d, axis="pops", data=P.data)
 
 
+def test_databook_default_all():
+    F = at.ProjectFramework(testdir / "sir_framework_default_all.xlsx")
+    D = at.ProjectData.new(F, np.arange(2000, 2005), pops=4, transfers=1)
+    D.save(tmpdir / "databook_default_all_test.xlsx")  # Test saving it back
+
+    D.add_pop("new", "new")
+    assert list(D.tdve["sus"].ts.keys()) == ["pop_0", "pop_1", "pop_2", "pop_3", "new"]
+    assert list(D.tdve["ch_prev"].ts.keys()) == ["All"]
+
+    D.tdve["contacts"].ts["pop_0"] = D.tdve["contacts"].ts["All"]
+    del D.tdve["contacts"].ts["All"]
+    D.add_pop("new2", "new2")
+    assert list(D.tdve["ch_prev"].ts.keys()) == ["All"]
+    assert list(D.tdve["contacts"].ts.keys()) == ["pop_0", "new2"]
+    D.save(tmpdir / "databook_default_all_test_2.xlsx")  # Test saving it back
+
+
 if __name__ == "__main__":
     # test_mixed_years_2()
     # test_mixed_years_1()
     # test_databooks()
     # test_databook_comments()
-    test_databook_all()
+    # test_databook_all()
+    test_databook_default_all()

tests/test_tox_migration.py

@@ -33,6 +33,11 @@ def test_migration():
     P.databook.save(tmpdir / "migration_test_databook_save")  # Save original databook
     P.data.save(tmpdir / "migration_test_data_save")  # Re-convert data to spreadsheet and save
 
+    # Test databook operations
+    P.data.add_pop("test_pop", "test_pop")  # Requires migration
+    P.data.add_transfer("test_transfer", "test_transfer")
+    P.data.add_interaction("test_interaction", "test_interaction")
+
 
 if __name__ == "__main__":
     test_migration()