Maintenance in porepy.ad (#1166)

* ENH: Introducing TimeDependentOperator and IterativeOperator, a generalized handling of arbitrary time step and iterate indices. * MOD: Maintenance on ad.Operator parsing, and operator function framework Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/equation_system.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/equation_system.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operator_functions.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operators.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operator_functions.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operators.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operators.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operators.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operators.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operators.py Co-authored-by: Ivar Stefansson <[email protected]> * Update src/porepy/numerics/ad/operators.py Co-authored-by: Ivar Stefansson <[email protected]> * MOD: Applying some changes from review. * Update src/porepy/numerics/ad/operators.py Co-authored-by: Eirik Keilegavlen <[email protected]> * DOC: Added comments regarding Variable ids * STY: black & flake8 * REF: EquationSystem.dofs_of again returns indices corresponding to order of input arguments. * MOD: Changing time step indexation: Previous time starts with index 1. --------- Co-authored-by: Ivar Stefansson <[email protected]> Co-authored-by: Eirik Keilegavlen <[email protected]>
pmgbergen · May 13, 2024 · 0e94301 · 0e94301
1 parent 7bf8c53
commit 0e94301
Show file tree

Hide file tree

Showing 38 changed files with 1,555 additions and 1,274 deletions.
diff --git a/src/porepy/applications/test_utils/models.py b/src/porepy/applications/test_utils/models.py
@@ -222,10 +222,10 @@ def compare_scaled_primary_variables(
     for var_name, var_unit in zip(variable_names, variable_units):
         # Obtain scaled values.
         scaled_values_0 = setup_0.equation_system.get_variable_values(
-            variables=[var_name], time_step_index=0
+            variables=[var_name], time_step_index=1
         )
         scaled_values_1 = setup_1.equation_system.get_variable_values(
-            variables=[var_name], time_step_index=0
+            variables=[var_name], time_step_index=1
         )
         # Convert back to SI units.
         values_0 = setup_0.fluid.convert_units(scaled_values_0, var_unit, to_si=True)

diff --git a/src/porepy/examples/mandel_biot.py b/src/porepy/examples/mandel_biot.py
@@ -1435,7 +1435,7 @@ def initial_condition(self) -> None:
             values=self.exact_sol.pressure(sd, 0),
             data=data,
             iterate_index=0,
-            time_step_index=0,
+            time_step_index=1,
         )
 
         # Set initial displacement
@@ -1444,7 +1444,7 @@ def initial_condition(self) -> None:
             values=self.exact_sol.displacement(sd, 0),
             data=data,
             iterate_index=0,
-            time_step_index=0,
+            time_step_index=1,
         )
 
     def after_simulation(self) -> None:

diff --git a/src/porepy/examples/terzaghi_biot.py b/src/porepy/examples/terzaghi_biot.py
@@ -699,7 +699,7 @@ def initial_condition(self) -> None:
             values=initial_p,
             data=data,
             iterate_index=0,
-            time_step_index=0,
+            time_step_index=1,
         )
 
     def after_simulation(self) -> None:

diff --git a/src/porepy/models/boundary_condition.py b/src/porepy/models/boundary_condition.py
@@ -92,7 +92,7 @@ def update_boundary_condition(
                 # No previous time step exists. The method was called during
                 # the initialization.
                 vals = function(bg)
-            pp.set_solution_values(name=name, values=vals, data=data, time_step_index=0)
+            pp.set_solution_values(name=name, values=vals, data=data, time_step_index=1)
 
             # Set the unknown time step values.
             vals = function(bg)

diff --git a/src/porepy/models/momentum_balance.py b/src/porepy/models/momentum_balance.py
@@ -790,7 +790,7 @@ def initial_condition(self) -> None:
         self.equation_system.set_variable_values(
             traction_vals.ravel("F"),
             [self.contact_traction_variable],
-            time_step_index=0,
+            time_step_index=1,
             iterate_index=0,
         )
 

diff --git a/src/porepy/models/solution_strategy.py b/src/porepy/models/solution_strategy.py
@@ -257,11 +257,14 @@ def domain(self) -> pp.Domain:
     def time_step_indices(self) -> np.ndarray:
         """Indices for storing time step solutions.
 
+        Note:
+            (Previous) Time step indices should start with 1.
+
         Returns:
             An array of the indices of which time step solutions will be stored.
 
         """
-        return np.array([0])
+        return np.array([1])
 
     @property
     def iterate_indices(self) -> np.ndarray:
@@ -299,9 +302,9 @@ def reset_state_from_file(self) -> None:
                     time_index,
                     times_file,
                 )
-            vals = self.equation_system.get_variable_values(time_step_index=0)
+            vals = self.equation_system.get_variable_values(time_step_index=1)
             self.equation_system.set_variable_values(
-                vals, iterate_index=0, time_step_index=0
+                vals, iterate_index=0, time_step_index=1
             )
             # Update the boundary conditions to both the time step and iterate solution.
             self.update_time_dependent_ad_arrays()
@@ -439,7 +442,7 @@ def after_nonlinear_convergence(self) -> None:
         solution = self.equation_system.get_variable_values(iterate_index=0)
         self.equation_system.shift_time_step_values()
         self.equation_system.set_variable_values(
-            values=solution, time_step_index=0, additive=False
+            values=solution, time_step_index=1, additive=False
         )
         self.convergence_status = True
         self.save_data_time_step()

diff --git a/src/porepy/numerics/ad/_ad_utils.py b/src/porepy/numerics/ad/_ad_utils.py
@@ -28,7 +28,7 @@
 from __future__ import annotations
 
 from abc import ABCMeta
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 import scipy.sparse as sps
@@ -294,6 +294,60 @@ def discretize_from_list(
                     pass
 
 
+def _validate_indices(
+    time_step_index: Optional[int] = None,
+    iterate_index: Optional[int] = None,
+) -> list[tuple[Any, int]]:
+    """Helper method to validate the indexation of getter and setter methods for
+    values in a grid's data dictionary.
+
+    See :func:`set_solution_values` and :func:`get_solution_values`.
+
+    """
+    if time_step_index is None and iterate_index is None:
+        raise ValueError(
+            "At least one of time_step_index and iterate_index needs to be different"
+            " from None."
+        )
+
+    out = []
+
+    if iterate_index is not None:
+        # Some iterate values of the current time
+        if iterate_index >= 0:
+            out.append((pp.ITERATE_SOLUTIONS, iterate_index))
+        # Negative iterate indices are not supported
+        else:
+            raise ValueError(
+                "Use increasing, non-negative integers for (previous) iterate values."
+            )
+
+    if time_step_index is not None:
+        # Some previous time.
+        if time_step_index > 0:
+            out.append((pp.TIME_STEP_SOLUTIONS, time_step_index))
+        # Current time. NOTE this is ambigous since the current time is an unknown and
+        # has multiple iterate values.
+        # Alternatively, we could associate time_step_index = 0 with iterate_index = 0
+        # the below elif branch introduces the convention that
+        # time step = iterate step = 0 are equivalent.
+        elif time_step_index == 0:
+            # if (pp.ITERATE_SOLUTIONS, 0) not in out:
+            #     out.append((pp.ITERATE_SOLUTIONS, 0))
+            raise ValueError(
+                "Using time_step_index = 0 (current time) is ambiguous."
+                + " Specify iterate_index instead."
+                + " First previous time step value is time_step_index = 1."
+            )
+        # Negative time step indices are not supported
+        else:
+            raise ValueError(
+                "Use increasing, non-negative integers for previous time step values."
+            )
+
+    return out
+
+
 def set_solution_values(
     name: str,
     values: np.ndarray,
@@ -302,7 +356,8 @@ def set_solution_values(
     iterate_index: Optional[int] = None,
     additive: bool = False,
 ) -> None:
-    """Function for setting values in the data dictionary.
+    """Function for setting values in the data dictionary, for some time-dependent or
+    iterative term.
 
     Parameters:
         name: Name of the quantity that is to be assigned values.
@@ -320,36 +375,29 @@ def set_solution_values(
             dictionary should be added to or overwritten.
 
     Raises:
-        ValueError: If neither of `time_step_index` or `iterate_index` have been
-            assigned a non-None value.
+        ValueError: In the case of inconsistent usage of indices
+            (both None, or negative values).
+        ValueError: If the user attempts to set values additively at an index where no
+            values were set before.
 
     """
-    if time_step_index is None and iterate_index is None:
-        raise ValueError(
-            "At least one of time_step_index and iterate_index needs to be different"
-            " from None."
-        )
-
-    if not additive:
-        if time_step_index is not None:
-            if pp.TIME_STEP_SOLUTIONS not in data:
-                data[pp.TIME_STEP_SOLUTIONS] = {}
-            if name not in data[pp.TIME_STEP_SOLUTIONS]:
-                data[pp.TIME_STEP_SOLUTIONS][name] = {}
-            data[pp.TIME_STEP_SOLUTIONS][name][time_step_index] = values.copy()
-
-        if iterate_index is not None:
-            if pp.ITERATE_SOLUTIONS not in data:
-                data[pp.ITERATE_SOLUTIONS] = {}
-            if name not in data[pp.ITERATE_SOLUTIONS]:
-                data[pp.ITERATE_SOLUTIONS][name] = {}
-            data[pp.ITERATE_SOLUTIONS][name][iterate_index] = values.copy()
-    else:
-        if time_step_index is not None:
-            data[pp.TIME_STEP_SOLUTIONS][name][time_step_index] += values
-
-        if iterate_index is not None:
-            data[pp.ITERATE_SOLUTIONS][name][iterate_index] += values
+    loc_index = _validate_indices(time_step_index, iterate_index)
+
+    for loc, index in loc_index:
+        if loc not in data:
+            data[loc] = {}
+        if name not in data[loc]:
+            data[loc][name] = {}
+
+        if additive:
+            if index not in data[loc][name]:
+                raise ValueError(
+                    f"Cannot set value additively for {name} at {(loc, index)}:"
+                    + " No values stored to add to."
+                )
+            data[loc][name][index] += values
+        else:
+            data[loc][name][index] = values.copy()
 
 
 def get_solution_values(
@@ -358,11 +406,8 @@ def get_solution_values(
     time_step_index: Optional[int] = None,
     iterate_index: Optional[int] = None,
 ) -> np.ndarray:
-    """Function for fetching values stored in the data dictionary.
-
-    This function should be used for obtaining solution values that are not related to a
-    variable. This is to avoid the cumbersome alternative of writing e.g.:
-    `data["solution_name"][pp.TIME_STEP_SOLUTION/pp.ITERATE_SOLUTION][0]`.
+    """Function for fetching values stored in the data dictionary, for some
+    time-dependent or iterative term.
 
     Parameters:
         name: Name of the parameter whose values we are interested in.
@@ -375,45 +420,31 @@ def get_solution_values(
             from before.
 
     Raises:
-        ValueError: If both time_step_index and iterate_index are None.
-
-        ValueErorr: If both time_step_index and iterate_index are assigned a value.
-
-        KeyError: If there are no data values assigned to the provided name.
-
-        KeyError: If there are no data values assigned to the time step/iterate index.
+        ValueError: In the case of inconsistent usage of indices
+            (both None or negative values).
+        AssertionError: If the user attempts to get iterate and time step values
+            simultanously. Only 1 index is permitted in getter
+        KeyError: If no values are stored for the passed index.
 
     Returns:
-        An array containing the solution values.
+        A copy of the values stored at the passed index.
 
     """
-    if time_step_index is None and iterate_index is None:
-        raise ValueError("Both time_step_index and iterate_index cannot be None.")
+    loc_index = _validate_indices(time_step_index, iterate_index)
+    assert (
+        len(loc_index) == 1
+    ), "Cannot get value from both iterate and time step at once. Call separately."
 
-    if time_step_index is not None and iterate_index is not None:
-        raise ValueError(
-            "Both time_step_index and iterate_index cannot be assigned a value."
-        )
+    loc, index = loc_index[0]
 
-    if time_step_index is not None:
-        if name not in data[pp.TIME_STEP_SOLUTIONS].keys():
-            raise KeyError(f"There are no values related the parameter name {name}.")
+    try:
+        value = data[loc][name][index].copy()
+    except KeyError as err:
+        raise KeyError(
+            f"No values stored for {name} at {(loc, index)}: {str(err)}."
+        ) from err
 
-        if time_step_index not in data[pp.TIME_STEP_SOLUTIONS][name].keys():
-            raise KeyError(
-                f"There are no values stored for time step index {time_step_index}."
-            )
-        return data[pp.TIME_STEP_SOLUTIONS][name][time_step_index].copy()
-
-    else:
-        if name not in data[pp.ITERATE_SOLUTIONS].keys():
-            raise KeyError(f"There are no values related the parameter name {name}.")
-
-        if iterate_index not in data[pp.ITERATE_SOLUTIONS][name].keys():
-            raise KeyError(
-                f"There are no values stored for iterate index {iterate_index}."
-            )
-        return data[pp.ITERATE_SOLUTIONS][name][iterate_index].copy()
+    return value
 
 
 class MergedOperator(operators.Operator):