Issue 1582 concentration

pybamm/__init__.py

@@ -35,7 +35,7 @@
 #
 # Utility classes and methods
 #
-from .util import Timer, TimerTime, FuzzyDict
+from .util import Timer, TimerTime, FuzzyDict, DomainDict
 from .util import root_dir, load_function, rmse, get_infinite_nested_dict, load
 from .util import get_parameters_filepath, have_jax, install_jax, is_jax_compatible
 from .logger import logger, set_logging_level

pybamm/discretisations/discretisation.py

@@ -236,6 +236,16 @@ def process_model(self, model, inplace=True, check_model=True):
             model_disc
         )
 
+        # Process length scales
+        new_length_scales = {}
+        for domain, scale in model.length_scales.items():
+            new_scale = self.process_symbol(scale)
+            if isinstance(new_scale, pybamm.Array):
+                # Convert possible arrays of length 1 to scalars
+                new_scale = pybamm.Scalar(float(new_scale.evaluate()))
+            new_length_scales[domain] = new_scale
+        model_disc._length_scales = new_length_scales
+
         # Check that resulting model makes sense
         if check_model:
             pybamm.logger.verbose("Performing model checks for {}".format(model.name))
@@ -317,9 +327,7 @@ def _get_variable_size(self, variable):
         else:
             size = 0
             spatial_method = self.spatial_methods[variable.domain[0]]
-            repeats = spatial_method._get_auxiliary_domain_repeats(
-                variable.auxiliary_domains
-            )
+            repeats = spatial_method._get_auxiliary_domain_repeats(variable.domains)
             for dom in variable.domain:
                 size += spatial_method.mesh[dom].npts_for_broadcast_to_nodes * repeats
             return size
@@ -782,9 +790,7 @@ def process_dict(self, var_eqn_dict):
         for eqn_key, eqn in var_eqn_dict.items():
             # Broadcast if the equation evaluates to a number (e.g. Scalar)
             if np.prod(eqn.shape_for_testing) == 1 and not isinstance(eqn_key, str):
-                eqn = pybamm.FullBroadcast(
-                    eqn, eqn_key.domain, eqn_key.auxiliary_domains
-                )
+                eqn = pybamm.FullBroadcast(eqn, broadcast_domains=eqn_key.domains)
 
             # note we are sending in the key.id here so we don't have to
             # keep calling .id
@@ -824,9 +830,9 @@ def process_symbol(self, symbol):
             else:
                 discretised_symbol.mesh = None
             # Assign secondary mesh
-            if "secondary" in symbol.auxiliary_domains:
+            if symbol.domains["secondary"] != []:
                 discretised_symbol.secondary_mesh = self.mesh.combine_submeshes(
-                    *symbol.auxiliary_domains["secondary"]
+                    *symbol.domains["secondary"]
                 )
             else:
                 discretised_symbol.secondary_mesh = None
@@ -937,10 +943,7 @@ def _process_symbol(self, symbol):
                     out = disc_child * pybamm.Vector([1])
                 else:
                     out = spatial_method.broadcast(
-                        disc_child,
-                        symbol.domain,
-                        symbol.auxiliary_domains,
-                        symbol.broadcast_type,
+                        disc_child, symbol.domains, symbol.broadcast_type
                     )
                 return out
 
@@ -976,8 +979,7 @@ def _process_symbol(self, symbol):
         elif isinstance(symbol, pybamm.VariableDot):
             return pybamm.StateVectorDot(
                 *self.y_slices[symbol.get_variable().id],
-                domain=symbol.domain,
-                auxiliary_domains=symbol.auxiliary_domains,
+                domains=symbol.domains,
             )
 
         elif isinstance(symbol, pybamm.Variable):
@@ -992,8 +994,7 @@ def _process_symbol(self, symbol):
                     return pybamm.ExternalVariable(
                         symbol.name,
                         size=self._get_variable_size(symbol),
-                        domain=symbol.domain,
-                        auxiliary_domains=symbol.auxiliary_domains,
+                        domains=symbol.domains,
                     )
                 else:
                     # We have to use a special name since the concatenation doesn't have
@@ -1002,8 +1003,7 @@ def _process_symbol(self, symbol):
                     ext = pybamm.ExternalVariable(
                         name,
                         size=self._get_variable_size(parent),
-                        domain=parent.domain,
-                        auxiliary_domains=parent.auxiliary_domains,
+                        domains=parent.domains,
                     )
                     out = pybamm.Index(ext, slice(start, end))
                     out.domain = symbol.domain
@@ -1025,11 +1025,7 @@ def _process_symbol(self, symbol):
                             symbol.name
                         )
                     )
-                return pybamm.StateVector(
-                    *y_slices,
-                    domain=symbol.domain,
-                    auxiliary_domains=symbol.auxiliary_domains,
-                )
+                return pybamm.StateVector(*y_slices, domains=symbol.domains)
 
         elif isinstance(symbol, pybamm.SpatialVariable):
             return spatial_method.spatial_variable(symbol)

pybamm/expression_tree/array.py

@@ -25,6 +25,11 @@ class Array(pybamm.Symbol):
         list of domains the parameter is valid over, defaults to empty list
     auxiliary_domainds : dict, optional
         dictionary of auxiliary domains, defaults to empty dict
+    domains : dict
+        A dictionary equivalent to {'primary': domain, auxiliary_domains}. Either
+        'domain' and 'auxiliary_domains', or just 'domains', should be provided
+        (not both). In future, the 'domain' and 'auxiliary_domains' arguments may be
+        deprecated.
     entries_string : str
         String representing the entries (slow to recalculate when copying)
 
@@ -37,6 +42,7 @@ def __init__(
         name=None,
         domain=None,
         auxiliary_domains=None,
+        domains=None,
         entries_string=None,
     ):
         # if
@@ -49,7 +55,9 @@ def __init__(
         self._entries = entries.astype(float)
         # Use known entries string to avoid re-hashing, where possible
         self.entries_string = entries_string
-        super().__init__(name, domain=domain, auxiliary_domains=auxiliary_domains)
+        super().__init__(
+            name, domain=domain, auxiliary_domains=auxiliary_domains, domains=domains
+        )
 
     @property
     def entries(self):
@@ -105,9 +113,8 @@ def create_copy(self):
         return self.__class__(
             self.entries,
             self.name,
-            self.domain,
-            self.auxiliary_domains,
-            self.entries_string,
+            domains=self.domains,
+            entries_string=self.entries_string,
         )
 
     def _base_evaluate(self, t=None, y=None, y_dot=None, inputs=None):

pybamm/expression_tree/averages.py

@@ -75,12 +75,7 @@ def _unary_new_copy(self, child):
 
 class SizeAverage(_BaseAverage):
     def __init__(self, child, f_a_dist):
-        R = pybamm.SpatialVariable(
-            "R",
-            domain=child.domain,
-            auxiliary_domains=child.auxiliary_domains,
-            coord_sys="cartesian",
-        )
+        R = pybamm.SpatialVariable("R", domains=child.domains, coord_sys="cartesian")
         integration_variable = [R]
         super().__init__(child, "size-average", integration_variable)
         self.f_a_dist = f_a_dist
@@ -134,22 +129,22 @@ def x_average(symbol):
             dom in ["negative electrode", "separator", "positive electrode"]
             for dom in symbol.secondary_domain
         ):
-            aux = {}
-            if "tertiary" in symbol.auxiliary_domains:
-                aux["secondary"] = symbol.auxiliary_domains["tertiary"]
-            return pybamm.FullBroadcast(symbol.orphans[0], symbol.broadcast_domain, aux)
-        elif (
-            isinstance(symbol, pybamm.FullBroadcast)
-            and "tertiary" in symbol.auxiliary_domains
-            and all(
-                dom in ["negative electrode", "separator", "positive electrode"]
-                for dom in symbol.tertiary_domain
-            )
+            domains = {
+                "primary": symbol.domains["primary"],
+                "secondary": symbol.domains["tertiary"],
+                "tertiary": symbol.domains["quaternary"],
+            }
+            return pybamm.FullBroadcast(symbol.orphans[0], broadcast_domains=domains)
+        elif isinstance(symbol, pybamm.FullBroadcast) and all(
+            dom in ["negative electrode", "separator", "positive electrode"]
+            for dom in symbol.tertiary_domain
         ):
-            aux = {"secondary": symbol.auxiliary_domains["secondary"]}
-            if "quaternary" in symbol.auxiliary_domains:
-                aux["tertiary"] = symbol.auxiliary_domains["quaternary"]
-            return pybamm.FullBroadcast(symbol.orphans[0], symbol.broadcast_domain, aux)
+            domains = {
+                "primary": symbol.domains["primary"],
+                "secondary": symbol.domains["secondary"],
+                "tertiary": symbol.domains["quaternary"],
+            }
+            return pybamm.FullBroadcast(symbol.orphans[0], broadcast_domains=domains)
         else:  # pragma: no cover
             # It should be impossible to get here
             raise NotImplementedError
@@ -174,14 +169,14 @@ def x_average(symbol):
         if out.domain != []:
             return out
         # Otherwise we may need to broadcast it
-        elif child.auxiliary_domains == {}:
+        elif child.domains["secondary"] == []:
             return out
         else:
-            domain = child.auxiliary_domains["secondary"]
-            if "tertiary" not in child.auxiliary_domains:
+            domain = child.domains["secondary"]
+            if child.domains["tertiary"] == []:
                 return pybamm.PrimaryBroadcast(out, domain)
             else:
-                auxiliary_domains = {"secondary": child.auxiliary_domains["tertiary"]}
+                auxiliary_domains = {"secondary": child.domains["tertiary"]}
                 return pybamm.FullBroadcast(out, domain, auxiliary_domains)
     # Otherwise, use Integral to calculate average value
     else:
@@ -218,7 +213,7 @@ def z_average(symbol):
         return symbol
     # If symbol is a Broadcast, its average value is its child
     elif isinstance(symbol, pybamm.Broadcast):
-        return symbol.orphans[0]
+        return symbol.reduce_one_dimension()
     # Otherwise, define a ZAverage
     else:
         return ZAverage(symbol)
@@ -251,12 +246,16 @@ def yz_average(symbol):
         return symbol
     # If symbol is a Broadcast, its average value is its child
     elif isinstance(symbol, pybamm.Broadcast):
-        return symbol.orphans[0]
+        return symbol.reduce_one_dimension()
     # Otherwise, define a YZAverage
     else:
         return YZAverage(symbol)
 
 
+def xyz_average(symbol):
+    return yz_average(x_average(symbol))
+
+
 def r_average(symbol):
     """
     Convenience function for creating an average in the r-direction.
@@ -276,11 +275,7 @@ def r_average(symbol):
         raise ValueError("Can't take the r-average of a symbol that evaluates on edges")
     # Otherwise, if symbol doesn't have a particle domain,
     # its r-averaged value is itself
-    elif symbol.domain not in [
-        ["positive particle"],
-        ["negative particle"],
-        ["working particle"],
-    ]:
+    elif symbol.domain not in [["positive particle"], ["negative particle"]]:
         return symbol
     # If symbol is a secondary broadcast onto "negative electrode" or
     # "positive electrode", take the r-average of the child then broadcast back
@@ -291,12 +286,10 @@ def r_average(symbol):
         child_av = pybamm.r_average(child)
         return pybamm.PrimaryBroadcast(child_av, symbol.domains["secondary"])
     # If symbol is a Broadcast onto a particle domain, its average value is its child
-    elif isinstance(symbol, pybamm.PrimaryBroadcast) and symbol.domain in [
-        ["positive particle"],
-        ["negative particle"],
-        ["working particle"],
-    ]:
-        return symbol.orphans[0]
+    elif isinstance(
+        symbol, (pybamm.PrimaryBroadcast, pybamm.FullBroadcast)
+    ) and symbol.domain in [["positive particle"], ["negative particle"]]:
+        return symbol.reduce_one_dimension()
     else:
         return RAverage(symbol)
 
@@ -345,10 +338,7 @@ def size_average(symbol, f_a_dist=None):
         if f_a_dist is None:
             geo = pybamm.geometric_parameters
             R = pybamm.SpatialVariable(
-                "R",
-                domain=symbol.domain,
-                auxiliary_domains=symbol.auxiliary_domains,
-                coord_sys="cartesian",
+                "R", domains=symbol.domains, coord_sys="cartesian"
             )
             if ["negative particle size"] in symbol.domains.values():
                 f_a_dist = geo.f_a_dist_n(R)

pybamm/expression_tree/binary_operators.py

@@ -25,42 +25,15 @@ def preprocess_binary(left, right):
         )
 
     # Do some broadcasting in special cases, to avoid having to do this manually
-    if left.domain != [] and right.domain != []:
-        if (
-            left.domain != right.domain
-            and "secondary" in right.auxiliary_domains
-            and left.domain == right.auxiliary_domains["secondary"]
-        ):
+    if left.domain != [] and right.domain != [] and left.domain != right.domain:
+        if left.domain == right.secondary_domain:
             left = pybamm.PrimaryBroadcast(left, right.domain)
-        if (
-            right.domain != left.domain
-            and "secondary" in left.auxiliary_domains
-            and right.domain == left.auxiliary_domains["secondary"]
-        ):
+        elif right.domain == left.secondary_domain:
             right = pybamm.PrimaryBroadcast(right, left.domain)
 
     return left, right
 
 
-def get_binary_children_domains(ldomain, rdomain):
-    """Combine domains from children in appropriate way."""
-    if ldomain == rdomain:
-        return ldomain
-    elif ldomain == []:
-        return rdomain
-    elif rdomain == []:
-        return ldomain
-    else:
-        raise pybamm.DomainError(
-            """
-            children must have same (or empty) domains, but left.domain is '{}'
-            and right.domain is '{}'
-            """.format(
-                ldomain, rdomain
-            )
-        )
-
-
 class BinaryOperator(pybamm.Symbol):
     """
     A node in the expression tree representing a binary operator (e.g. `+`, `*`)
@@ -83,14 +56,8 @@ class BinaryOperator(pybamm.Symbol):
     def __init__(self, name, left, right):
         left, right = preprocess_binary(left, right)
 
-        domain = get_binary_children_domains(left.domain, right.domain)
-        auxiliary_domains = self.get_children_auxiliary_domains([left, right])
-        super().__init__(
-            name,
-            children=[left, right],
-            domain=domain,
-            auxiliary_domains=auxiliary_domains,
-        )
+        domains = self.get_children_domains([left, right])
+        super().__init__(name, children=[left, right], domains=domains)
         self.left = self.children[0]
         self.right = self.children[1]
 
@@ -413,12 +380,7 @@ def _binary_evaluate(self, left, right):
         if issparse(left):
             return csr_matrix(left.multiply(1 / right))
         else:
-            if isinstance(right, numbers.Number) and right == 0:
-                # don't raise RuntimeWarning for NaNs
-                with np.errstate(invalid="ignore"):
-                    return left * np.inf
-            else:
-                return left / right
+            return left / right
 
 
 class Inner(BinaryOperator):
@@ -988,6 +950,10 @@ def simplified_subtraction(left, right):
         ):
             return left
 
+    # Return constant if both sides are constant
+    if left.is_constant() and right.is_constant():
+        return pybamm.simplify_if_constant(Subtraction(left, right))
+
     # a symbol minus itself is 0s of the same shape
     if left.id == right.id:
         return pybamm.zeros_like(left)