feat: new mass balance

examples/example1.py

@@ -7,7 +7,7 @@
 #       format_version: '1.3'
 #       jupytext_version: 1.18.1
 #   kernelspec:
-#     display_name: .venv
+#     display_name: flodym
 #     language: python
 #     name: python3
 # ---
@@ -58,9 +58,7 @@
 elements = Dimension(
     name="Elements",
     letter="e",
-    items=[
-        "single material",
-    ],
+    items=["single material"],
 )
 dimensions = DimensionSet(dim_list=[time, elements])
 
@@ -116,9 +114,8 @@
 # %%
 class SimpleMFA(MFASystem):
     def compute(self):
-        self.flows["sysenv => process 1"][...] = self.parameters[
-            "D"
-        ]  # the elipsis slice [...] ensures the dimensionality of the flow is not changed
+        # the elipsis slice [...] ensures the dimensionality of the flow is not changed
+        self.flows["sysenv => process 1"][...] = self.parameters["D"]
         self.flows["process 1 => process 2"][...] = (
             1 / (1 - self.parameters["alpha"]) * self.parameters["D"]
         )

examples/example2.ipynb

@@ -57,6 +57,7 @@
     "    ParameterDefinition,\n",
     "    FlowDefinition,\n",
     "    StockDefinition,\n",
+    "    ProcessDefinition,\n",
     "    MFASystem,\n",
     "    SimpleFlowDrivenStock,\n",
     ")\n",
@@ -108,12 +109,12 @@
    "outputs": [],
    "source": [
     "process_names = [\n",
-    "    \"sysenv\",\n",
-    "    \"shredder\",\n",
-    "    \"demolition\",\n",
-    "    \"remelting\",\n",
-    "    \"landfills\",\n",
-    "    \"slag piles\",\n",
+    "    ProcessDefinition(id=0, name=\"sysenv\"),\n",
+    "    ProcessDefinition(id=1, name=\"shredder\"),\n",
+    "    ProcessDefinition(id=2, name=\"demolition\"),\n",
+    "    ProcessDefinition(id=3, name=\"remelting\"),\n",
+    "    ProcessDefinition(id=4, name=\"landfills\"),\n",
+    "    ProcessDefinition(id=5, name=\"slag piles\"),\n",
     "]"
    ]
   },
@@ -275,6 +276,7 @@
     "    dimension_sheets={d.name: d.name for d in dimension_definitions},\n",
     "    parameter_sheets={p.name: p.name for p in parameter_definitions},\n",
     ")\n",
+    "\n",
     "mfa_example.compute()"
    ]
   },
@@ -306,7 +308,7 @@
     "    array=remelted,\n",
     "    intra_line_dim=\"Time\",\n",
     "    linecolor_dim=\"Material\",\n",
-    "    title=\"GDP-per-capita\",\n",
+    "    title=\"Total remelted material\",\n",
     ")\n",
     "fig = plotter.plot(do_show=True)"
    ]
@@ -457,7 +459,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "flodym",
    "language": "python",
    "name": "python3"
   },

examples/example2.py

@@ -7,7 +7,7 @@
 #       format_version: '1.3'
 #       jupytext_version: 1.18.1
 #   kernelspec:
-#     display_name: .venv
+#     display_name: flodym
 #     language: python
 #     name: python3
 # ---
@@ -53,6 +53,7 @@
     ParameterDefinition,
     FlowDefinition,
     StockDefinition,
+    ProcessDefinition,
     MFASystem,
     SimpleFlowDrivenStock,
 )
@@ -85,12 +86,12 @@
 
 # %%
 process_names = [
-    "sysenv",
-    "shredder",
-    "demolition",
-    "remelting",
-    "landfills",
-    "slag piles",
+    ProcessDefinition(id=0, name="sysenv"),
+    ProcessDefinition(id=1, name="shredder"),
+    ProcessDefinition(id=2, name="demolition"),
+    ProcessDefinition(id=3, name="remelting"),
+    ProcessDefinition(id=4, name="landfills"),
+    ProcessDefinition(id=5, name="slag piles"),
 ]
 
 # %%
@@ -207,6 +208,7 @@ def compute(self):
     dimension_sheets={d.name: d.name for d in dimension_definitions},
     parameter_sheets={p.name: p.name for p in parameter_definitions},
 )
+
 mfa_example.compute()
 
 # %% [markdown]
@@ -224,7 +226,7 @@ def compute(self):
     array=remelted,
     intra_line_dim="Time",
     linecolor_dim="Material",
-    title="GDP-per-capita",
+    title="Total remelted material",
 )
 fig = plotter.plot(do_show=True)
 

examples/example3.py

@@ -7,7 +7,7 @@
 #       format_version: '1.3'
 #       jupytext_version: 1.18.1
 #   kernelspec:
-#     display_name: .venv
+#     display_name: flodym
 #     language: python
 #     name: python3
 # ---

examples/example5.py

@@ -7,7 +7,7 @@
 #       format_version: '1.3'
 #       jupytext_version: 1.18.1
 #   kernelspec:
-#     display_name: .venv
+#     display_name: flodym
 #     language: python
 #     name: python3
 # ---

flodym/__init__.py

@@ -4,6 +4,7 @@
     DimensionDefinition as DimensionDefinition,
     StockDefinition as StockDefinition,
     ParameterDefinition as ParameterDefinition,
+    ProcessDefinition as ProcessDefinition,
 )
 from flodym.mfa_system import MFASystem as MFASystem
 from flodym.dimensions import DimensionSet as DimensionSet, Dimension as Dimension
@@ -20,6 +21,7 @@
     DynamicStockModel as DynamicStockModel,
     InflowDrivenDSM as InflowDrivenDSM,
     StockDrivenDSM as StockDrivenDSM,
+    FlexibleDSM as FlexibleDSM,
 )
 from flodym.lifetime_models import (
     LifetimeModel as LifetimeModel,
@@ -41,3 +43,12 @@
     ExcelParameterReader as ExcelParameterReader,
     CompoundDataReader as CompoundDataReader,
 )
+from flodym.config import config as config
+
+Flow.model_rebuild()
+Stock.model_rebuild()
+SimpleFlowDrivenStock.model_rebuild()
+DynamicStockModel.model_rebuild()
+InflowDrivenDSM.model_rebuild()
+StockDrivenDSM.model_rebuild()
+FlexibleDSM.model_rebuild()

flodym/config.py

@@ -0,0 +1,75 @@
+from pydantic import BaseModel, confloat
+from typing import Optional
+from enum import Enum
+import logging
+
+
+class ErrorBehavior(str, Enum):
+    ERROR = "error"
+    WARN = "warn"
+    INFO = "info"
+    IGNORE = "ignore"
+
+
+def handle_error(behavior: ErrorBehavior, message: str = None, error: Exception = None):
+
+    # ensures valid ErrorBehavior
+    behavior = ErrorBehavior(behavior)
+
+    if message is None == error is None:
+        raise ValueError("Exactly one of message or error must be provided.")
+
+    if behavior == ErrorBehavior.ERROR:
+        raise error or ValueError(message)
+    elif behavior == ErrorBehavior.WARN:
+        logging.warning(message or str(error))
+    elif behavior == ErrorBehavior.INFO:
+        logging.info(message or str(error))
+    elif behavior == ErrorBehavior.IGNORE:
+        pass
+
+
+class Checks(BaseModel, validate_assignment=True):
+    """
+    Which checks to perform during computation.
+    """
+
+    mass_balance_stocks: bool = True
+    """Whether to check mass balance after each `stock.compute()` call."""
+
+
+class ErrorBehaviors(BaseModel, validate_assignment=True):
+    """
+    What to do if a check fails.
+    """
+
+    mass_balance: ErrorBehavior = ErrorBehavior.ERROR
+    """what to do if mass balance is not satisfied."""
+    check_flows: ErrorBehavior = ErrorBehavior.WARN
+    """What to do if a flow has negative or NaN value."""
+
+
+class Config(BaseModel, validate_assignment=True):
+    """
+    Configuration class for Flodym.
+    """
+
+    checks: Checks = Checks()
+    """Which checks to perform during computation."""
+
+    error_behaviors: ErrorBehaviors = ErrorBehaviors()
+    """What to do if a check fails."""
+
+    relative_tolerance: Optional[confloat(ge=0)] = 10.0
+    """Tolerance relative to a float precision metric used for mass balance and other checks.
+    Increase if your mass balance checks fail due to rounding errors.
+    Overridden by `absolute_tolerance` if set.
+    """
+    absolute_tolerance: Optional[confloat(ge=0)] = None
+    """Absolute tolerance used for mass balance and other checks.
+    Overrides `relative_tolerance` if set.
+    If None, `relative_tolerance` is used.
+    """
+
+
+config = Config()

flodym/dimensions.py

@@ -195,11 +195,6 @@ def shape(self) -> tuple[int]:
         """shape of the array that would be created with the dimensions in the set"""
         return tuple(self.size(dim) for dim in self.letters)
 
-    @property
-    def ndim(self):
-        """the number of dimensions in the set"""
-        return len(self.dim_list)
-
     def get_subset(self, dims: tuple = None) -> "DimensionSet":
         """Selects :py:class:`Dimension` objects from the object attribute dim_list,
         according to the dims passed, which can be either letters or names.
@@ -273,6 +268,10 @@ def intersect_with(self, other: "DimensionSet") -> "DimensionSet":
         intersection_letters = [dim.letter for dim in self.dim_list if dim.letter in other.letters]
         return self.get_subset(intersection_letters)
 
+    def __and__(self, other: "DimensionSet") -> "DimensionSet":
+        """Intersection operator for two DimensionSets."""
+        return self.intersect_with(other)
+
     def union_with(self, other: "DimensionSet") -> "DimensionSet":
         """Get the union of two DimensionSets.
 
@@ -285,6 +284,10 @@ def union_with(self, other: "DimensionSet") -> "DimensionSet":
         added_dims = [dim for dim in other.dim_list if dim.letter not in self.letters]
         return self.expand_by(added_dims)
 
+    def __or__(self, other: "DimensionSet") -> "DimensionSet":
+        """Union operator for two DimensionSets."""
+        return self.union_with(other)
+
     def difference_with(self, other: "DimensionSet") -> "DimensionSet":
         """Get the set difference of two DimensionSets.
 
@@ -299,6 +302,27 @@ def difference_with(self, other: "DimensionSet") -> "DimensionSet":
         ]
         return self.get_subset(difference_letters)
 
+    def __sub__(self, other: "DimensionSet") -> "DimensionSet":
+        """Difference operator for two DimensionSets."""
+        return self.difference_with(other)
+
+    def __xor__(self, other: "DimensionSet") -> "DimensionSet":
+        """Symmetric difference operator for two DimensionSets."""
+        return (self - other) | (other - self)
+
+    @property
+    def ndim(self):
+        """the number of dimensions in the set"""
+        return len(self.dim_list)
+
+    def __len__(self) -> int:
+        """Return the number of dimensions in the set."""
+        return len(self.dim_list)
+
+    def __bool__(self) -> bool:
+        """Return True if the set is not empty."""
+        return len(self.dim_list) > 0
+
     @property
     def names(self):
         """A tuple of the names of the dimensions in the set."""