Error handling - IPM Simulation (#104)

Implemented error handling for common simulation errors, namely:
1. Incorrect parameters causing a negative or zero rate
2. Divide by zero errors, mainly occurring when the population of a node is 0
3. Incorrect parameters causing a negative probability in fork definitions
4. Incompatible clause/parameter values for MM simulations

For more details, check devlog 2024-04-16.ipynb

doc/devlog/2024-04-04-draw-demo.ipynb

@@ -7,6 +7,8 @@
                 "# 2024/04/04 Draw Module Demo\n",
                 "_Author: Izaac Molina_\n",
                 "\n",
+                "Provides demos on using the `draw` module in epymorph\n",
+                "\n",
                 "## Goal and Purpose\n",
                 "The traditional way to design a model in epymorph is to first draw a ipm model, then translate that model into an ipm object via code. This project seeks to go in the other direction: turning an ipm object into a model. This is great both for educational and debugging purposes as it allows users to present their model more clearly as well as check to ensure they have implemented the model correctly in epymorph\n",
                 "\n",

doc/devlog/README.md

@@ -47,6 +47,7 @@ This folder is a handy place to put Jupyter notebooks or other documents which h
 | 2024-03-01.ipynb | Tyler | | Getting the indices of IPM events and compartments by name with wildcard support. |
 | 2024-03-13.ipynb | Tyler | | Showing off movement data collection (NEW!) |
 | 2024-04-04-draw-demo.ipynb | Izaac | | Showing the new draw module for visualising IPMs (NEW!) |
+| 2024-04-16.ipynb | Izaac | | Showing error handling for common ipm errors (NEW!)|
 
 ## Contributing
 

epymorph/__init__.py

@@ -1,5 +1,7 @@
 """epymorph's main package and main exports"""
 
+from numpy import seterr
+
 import epymorph.compartment_model as IPM
 from epymorph.data import geo_library, ipm_library, mm_library
 from epymorph.data_shape import Shapes
@@ -10,6 +12,10 @@
 from epymorph.proxy import dim, geo
 from epymorph.simulation import SimDType, TimeFrame, default_rng
 
+# set numpy errors to raise exceptions instead of warnings, useful for catching
+# simulation errrors
+seterr(all='raise')
+
 __all__ = [
     'IPM',
     'ipm_library',

epymorph/engine/ipm_exec.py

@@ -12,6 +12,8 @@
                                         TransitionDef, exogenous_states)
 from epymorph.engine.context import RumeContext, Tick
 from epymorph.engine.world import World
+from epymorph.error import (IpmSimInvalidProbsException,
+                            IpmSimLessThanZeroException, IpmSimNaNException)
 from epymorph.simulation import SimDType
 from epymorph.sympy_shim import SympyLambda, lambdify, lambdify_list
 from epymorph.util import index_of
@@ -160,12 +162,41 @@ def _events(self, node: int, tick: Tick, effective_pop: NDArray[SimDType]) -> ND
         for t in self._trxs:
             match t:
                 case _IndependentTrx(rate_lambda):
-                    rate = rate_lambda(rate_args)
+                    # get rate from lambda expression, catch divide by zero error
+                    try:
+                        rate = rate_lambda(rate_args)
+                    except (ZeroDivisionError, FloatingPointError):
+                        raise IpmSimNaNException(
+                            self._get_zero_division_args(
+                                rate_args, node, tick, t)
+                        )
+                    # check for < 0 rate, throw error in this case
+                    if rate < 0:
+                        raise IpmSimLessThanZeroException(
+                            self._get_default_error_args(rate_args, node, tick)
+                        )
                     occur[index] = self._ctx.rng.poisson(rate * tick.tau)
                 case _ForkedTrx(size, rate_lambda, prob_lambda):
-                    rate = rate_lambda(rate_args)
+                    # get rate from lambda expression, catch divide by zero error
+                    try:
+                        rate = rate_lambda(rate_args)
+                    except (ZeroDivisionError, FloatingPointError):
+                        raise IpmSimNaNException(
+                            self._get_zero_division_args(
+                                rate_args, node, tick, t)
+                        )
+                    # check for < 0 base, throw error in this case
+                    if rate < 0:
+                        raise IpmSimLessThanZeroException(
+                            self._get_default_error_args(rate_args, node, tick)
+                        )
                     base = self._ctx.rng.poisson(rate * tick.tau)
                     prob = prob_lambda(rate_args)
+                    # check for negative probs
+                    if any(n < 0 for n in prob):
+                        raise IpmSimInvalidProbsException(
+                            self._get_invalid_prob_args(rate_args, node, tick, t)
+                        )
                     stop = index + size
                     occur[index:stop] = self._ctx.rng.multinomial(
                         base, prob)
@@ -199,6 +230,56 @@ def _events(self, node: int, tick: Tick, effective_pop: NDArray[SimDType]) -> ND
                         desired, available)
         return occur
 
+    def _get_default_error_args(self, rate_attrs: list, node: int, tick: Tick) -> list[tuple[str, dict]]:
+        arg_list = []
+        arg_list.append(("Node : Timestep", {node: tick.step}))
+        arg_list.append(("compartment values", {
+            name: value for (name, value) in zip(self._ctx.ipm.compartment_names,
+                                                 rate_attrs[:self._ctx.dim.compartments])
+        }))
+        arg_list.append(("ipm params", {
+            attribute.name: value for (attribute, value) in zip(self._ctx.ipm.attributes,
+                                                                rate_attrs[self._ctx.dim.compartments:])
+        }))
+
+        return arg_list
+
+    def _get_invalid_prob_args(self, rate_attrs: list, node: int, tick: Tick,
+                               transition: _ForkedTrx) -> list[tuple[str, dict]]:
+        arg_list = self._get_default_error_args(rate_attrs, node, tick)
+
+        transition_index = self._trxs.index(transition)
+        corr_transition = self._ctx.ipm.transitions[transition_index]
+        if isinstance(corr_transition, ForkDef):
+            to_compartments = ", ".join([str(edge.compartment_to)
+                                        for edge in corr_transition.edges])
+            from_compartment = corr_transition.edges[0].compartment_from
+            arg_list.append(("corresponding fork transition and probabilities",
+                             {
+                                 f"{from_compartment}->({to_compartments})": corr_transition.rate,
+                                 f"Probabilities": ', '.join([str(expr) for expr in corr_transition.probs]),
+                             }))
+
+        return arg_list
+
+    def _get_zero_division_args(self, rate_attrs: list, node: int, tick: Tick,
+                                transition: _IndependentTrx | _ForkedTrx) -> list[tuple[str, dict]]:
+        arg_list = self._get_default_error_args(rate_attrs, node, tick)
+
+        transition_index = self._trxs.index(transition)
+        corr_transition = self._ctx.ipm.transitions[transition_index]
+        if isinstance(corr_transition, EdgeDef):
+            arg_list.append(("corresponding transition", {
+                            f"{corr_transition.compartment_from}->{corr_transition.compartment_to}": corr_transition.rate}))
+        if isinstance(corr_transition, ForkDef):
+            to_compartments = ", ".join([str(edge.compartment_to)
+                                        for edge in corr_transition.edges])
+            from_compartment = corr_transition.edges[0].compartment_from
+            arg_list.append(("corresponding fork transition", {
+                            f"{from_compartment}->({to_compartments})": corr_transition.rate}))
+
+        return arg_list
+
     def _distribute(self, cohorts: NDArray[SimDType], events: NDArray[SimDType]) -> NDArray[SimDType]:
         """Distribute all events across a location's cohorts and return the compartment deltas for each."""
         x = cohorts.shape[0]

epymorph/error.py

@@ -2,6 +2,7 @@
 A common exception framework for epymorph.
 """
 from contextlib import contextmanager
+from textwrap import dedent
 from typing import Self
 
 
@@ -99,6 +100,74 @@ class IpmSimException(SimulationException):
     """Exception during IPM processing."""
 
 
+class IpmSimExceptionWithFields(IpmSimException):
+    """
+    Exception during IPM processing where it is appropriate to show specific
+    fields within the simulation.
+    To create a new error with fields, create a subclass of this and set the 
+    displayed error message along with the fields to print.
+    See 'IpmSimNaNException' for an example.
+    """
+
+    display_fields: tuple[str, dict] | list[tuple[str, dict]]
+
+    def __init__(self, message: str, display_fields: tuple[str, dict] | list[tuple[str, dict]]):
+        super().__init__(message)
+        if isinstance(display_fields, tuple):
+            display_fields = [display_fields]
+        self.display_fields = display_fields
+
+    def __str__(self):
+        msg = super().__str__()
+        fields = ""
+        for name, values in self.display_fields:
+            fields += f"Showing current {name}\n"
+            for key, value in values.items():
+                fields += f"{key}: {value}\n"
+            fields += "\n"
+        return f"{msg}\n{fields}"
+
+
+class IpmSimNaNException(IpmSimExceptionWithFields):
+    """Exception for handling NaN (not a number) rate values"""
+
+    def __init__(self, display_fields: tuple[str, dict] | list[tuple[str, dict]]):
+        msg = '''
+              NaN (not a number) rate detected. This is often the result of a divide by zero error.
+              When constructing the IPM, ensure that no edge transitions can result in division by zero
+              This commonly occurs when defining an S->I edge that is (some rate / sum of the compartments)
+              To fix this, change the edge to define the S->I edge as (some rate / Max(1/sum of the the compartments))
+              See examples of this in the provided example ipm definitions in the data/ipms folder.
+              '''
+        msg = dedent(msg)
+        super().__init__(msg, display_fields)
+
+
+class IpmSimLessThanZeroException(IpmSimExceptionWithFields):
+    """ Exception for handling less than 0 rate values """
+
+    def __init__(self, display_fields: tuple[str, dict] | list[tuple[str, dict]]):
+        msg = '''
+              Less than zero rate detected. When providing or defining ipm parameters, ensure that
+              they will not result in a negative rate. Note: this can often happen unintentionally
+              if a function is given as a parameter.
+              '''
+        msg = dedent(msg)
+        super().__init__(msg, display_fields)
+
+
+class IpmSimInvalidProbsException(IpmSimExceptionWithFields):
+    """ Exception for handling invalid probability values """
+
+    def __init__(self, display_fields: tuple[str, dict] | list[tuple[str, dict]]):
+        msg = '''
+              Invalid probabilities for fork definition detected. Probabilities for a 
+              given tick should always be nonnegative and sum to 1
+              '''
+        msg = dedent(msg)
+        super().__init__(msg, display_fields)
+
+
 class MmSimException(SimulationException):
     """Exception during MM processing."""
 

epymorph/movement/movement_model.py

@@ -11,7 +11,8 @@
 from numpy.typing import NDArray
 
 from epymorph.compartment_model import CompartmentModel
-from epymorph.error import AttributeException, MmValidationException
+from epymorph.error import (AttributeException, MmSimException,
+                            MmValidationException)
 from epymorph.geo.geo import Geo
 from epymorph.movement.parser import Attribute as MmAttribute
 from epymorph.params import ContextParams
@@ -115,6 +116,11 @@ def requested(self, ctx: MovementContext, predef: PredefParams, tick: Tick) -> N
             # until we can properly validate the MM clauses.
             msg = f"Missing attribute {e} required by movement model clause '{self.name}'."
             raise AttributeException(msg) from None
+        except Exception as e:
+            # NOTE: catching exceptions here is necessary to get nice error messages
+            # for some value error cause by incorrect parameter and/or clause definition
+            msg = f"Error from applying clause '{self.name}': see exception trace"
+            raise MmSimException(msg) from e
 
     def returns(self, ctx: MovementContext, tick: Tick) -> TickDelta:
         return self._returns(ctx, tick)