Generalise VectorField to N components

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Features
 
 - Added unstructured mesh support (`UnstructuredSubMesh`, generators, and interface coupling) for arbitrary 2D/3D domains. Hexahedra must have planar faces (warped hexes raise a `GeometryError`), and `UserSuppliedUnstructuredMesh` accepts tetrahedral, triangular, and quadrilateral cells only. ([#5687](https://github.com/pybamm-team/PyBaMM/pull/5687))
+- Generalised `VectorField` to N components and added `Component`/`Norm` operators for multi-dimensional vector fields. ([#5686](https://github.com/pybamm-team/PyBaMM/pull/5686))
 
 ## Bug fixes
 

docs/source/api/expression_tree/unary_operator.rst

@@ -73,6 +73,12 @@ Unary Operators
 .. autoclass:: pybamm.Downwind
   :members:
 
+.. autoclass:: pybamm.Component
+  :members:
+
+.. autoclass:: pybamm.Norm
+  :members:
+
 .. autofunction:: pybamm.grad
 
 .. autofunction:: pybamm.div

packages/pybamm/src/pybamm/discretisations/discretisation.py

@@ -882,7 +882,11 @@ def process_symbol(self, symbol):
 
         # Assign mesh as an attribute to the processed variable
         if symbol.domain != []:
-            discretised_symbol.mesh = self.mesh[symbol.domain]
+            mesh_for_symbol = self.mesh[symbol.domain]
+            discretised_symbol.mesh = mesh_for_symbol
+            if isinstance(discretised_symbol, pybamm.VectorField):
+                for comp in discretised_symbol.components:
+                    comp.mesh = mesh_for_symbol
         else:
             discretised_symbol.mesh = None
 
@@ -900,6 +904,38 @@ def process_symbol(self, symbol):
 
         return discretised_symbol
 
+    def _process_vector_field_binary(self, symbol, disc_left, disc_right):
+        """Broadcast a scalar side, then apply ``symbol`` component-wise."""
+        left_is_vf = isinstance(disc_left, pybamm.VectorField)
+        right_is_vf = isinstance(disc_right, pybamm.VectorField)
+        if left_is_vf and right_is_vf:
+            if disc_left.n_components != disc_right.n_components:
+                raise pybamm.DiscretisationError(
+                    f"Cannot combine VectorFields with {disc_left.n_components} and "
+                    f"{disc_right.n_components} components"
+                )
+            n = disc_left.n_components
+        elif left_is_vf:
+            n = disc_left.n_components
+            disc_right = pybamm.VectorField(*[disc_right] * n)
+        else:
+            n = disc_right.n_components
+            disc_left = pybamm.VectorField(*[disc_left] * n)
+        new_comps = [
+            pybamm.simplify_if_constant(
+                symbol.create_copy(
+                    new_children=[disc_left.components[k], disc_right.components[k]]
+                )
+            )
+            for k in range(n)
+        ]
+        result = pybamm.VectorField(*new_comps)
+        if disc_left._disc_state_vector is not None:
+            result._disc_state_vector = disc_left._disc_state_vector
+        else:
+            result._disc_state_vector = disc_right._disc_state_vector
+        return result
+
     def _process_symbol(self, symbol):
         """See :meth:`Discretisation.process_symbol()`."""
 
@@ -934,23 +970,9 @@ def _process_symbol(self, symbol):
                 if isinstance(disc_left, pybamm.VectorField) or isinstance(
                     disc_right, pybamm.VectorField
                 ):
-                    if not isinstance(disc_right, pybamm.VectorField):
-                        disc_right = pybamm.VectorField(disc_right, disc_right)
-                    if not isinstance(disc_left, pybamm.VectorField):
-                        disc_left = pybamm.VectorField(disc_left, disc_left)
-                    else:  # both are vector fields already
-                        pass
-                    disc_lr = pybamm.simplify_if_constant(
-                        symbol.create_copy(
-                            new_children=[disc_left.lr_field, disc_right.lr_field]
-                        )
+                    return self._process_vector_field_binary(
+                        symbol, disc_left, disc_right
                     )
-                    disc_tb = pybamm.simplify_if_constant(
-                        symbol.create_copy(
-                            new_children=[disc_left.tb_field, disc_right.tb_field]
-                        )
-                    )
-                    return pybamm.VectorField(disc_lr, disc_tb)
 
                 return pybamm.simplify_if_constant(
                     symbol.create_copy(new_children=[disc_left, disc_right])
@@ -1092,6 +1114,18 @@ def _process_symbol(self, symbol):
             elif isinstance(symbol, pybamm.NotConstant):
                 # After discretisation, we can make the symbol constant
                 return disc_child
+            elif isinstance(symbol, pybamm.Component):
+                if not isinstance(disc_child, pybamm.VectorField):
+                    raise pybamm.DiscretisationError(
+                        "Component can only be applied to a VectorField"
+                    )
+                return disc_child.components[symbol.index]
+            elif isinstance(symbol, pybamm.Norm):
+                if not isinstance(disc_child, pybamm.VectorField):
+                    raise pybamm.DiscretisationError(
+                        "Norm can only be applied to a VectorField"
+                    )
+                return sum(c**2 for c in disc_child.components) ** 0.5
             elif isinstance(symbol, pybamm.Magnitude):
                 if not isinstance(disc_child, pybamm.VectorField):
                     raise ValueError("Magnitude can only be applied to a vector field")
@@ -1104,10 +1138,13 @@ def _process_symbol(self, symbol):
                     raise ValueError("Invalid direction")
             else:
                 if isinstance(disc_child, pybamm.VectorField):
-                    return pybamm.VectorField(
-                        symbol.create_copy(new_children=[disc_child.lr_field]),
-                        symbol.create_copy(new_children=[disc_child.tb_field]),
-                    )
+                    new_comps = [
+                        symbol.create_copy(new_children=[c])
+                        for c in disc_child.components
+                    ]
+                    result = pybamm.VectorField(*new_comps)
+                    result._disc_state_vector = disc_child._disc_state_vector
+                    return result
                 else:
                     return symbol.create_copy(new_children=[disc_child])
 
@@ -1181,10 +1218,8 @@ def _process_symbol(self, symbol):
             )
 
         elif isinstance(symbol, pybamm.VectorField):
-            # VectorField is a subclass of TensorField, handle it first for specificity
-            left_symbol = self.process_symbol(symbol.lr_field)
-            right_symbol = self.process_symbol(symbol.tb_field)
-            return symbol.create_copy(new_children=[left_symbol, right_symbol])
+            processed = [self.process_symbol(c) for c in symbol.components]
+            return symbol.create_copy(new_children=processed)
 
         elif isinstance(symbol, pybamm.TensorField):
             # General TensorField handling (rank-2 tensors)

packages/pybamm/src/pybamm/expression_tree/unary_operators.py

@@ -1509,6 +1509,54 @@ def _unary_new_copy(self, child, perform_simplifications=True):
         return self.__class__(child, self.direction)
 
 
+class Component(UnaryOperator):
+    """
+    Extract component *index* from a VectorField.
+
+    Parameters
+    ----------
+    child : :class:`pybamm.Symbol`
+        A VectorField symbol.
+    index : int
+        Zero-based component index.
+    """
+
+    def __init__(self, child, index):
+        super().__init__(f"component({index})", child)
+        self.index = index
+
+    def to_json(self):
+        return {
+            "name": self.name,
+            "domains": self.domains,
+            "index": self.index,
+        }
+
+    @classmethod
+    def _from_json(cls, snippet):
+        return cls(snippet["children"][0], snippet["index"])
+
+    def _unary_new_copy(self, child, perform_simplifications=True):
+        return self.__class__(child, self.index)
+
+
+class Norm(UnaryOperator):
+    """
+    Euclidean norm of a VectorField: ``sqrt(sum(comp_i ** 2))``.
+
+    Parameters
+    ----------
+    child : :class:`pybamm.Symbol`
+        A VectorField symbol.
+    """
+
+    def __init__(self, child):
+        super().__init__("norm", child)
+
+    def _unary_new_copy(self, child, perform_simplifications=True):
+        return self.__class__(child)
+
+
 class Upwind(UpwindDownwind):
     """
     Upwinding operator. To be used if flow velocity is positive (left to right).

packages/pybamm/src/pybamm/expression_tree/vector_field.py

@@ -1,9 +1,11 @@
 """
-VectorField class - a rank-1 tensor field for 2D simulations.
+VectorField class - a rank-1 tensor field with N components.
 """
 
 from __future__ import annotations
 
+import casadi
+
 import pybamm
 from pybamm.expression_tree.tensor_field import TensorField
 
@@ -13,67 +15,80 @@ class VectorField(TensorField):
     A node in the expression tree representing a vector field.
 
     VectorField is a convenience subclass of TensorField for rank-1 tensors
-    with two components (lr and tb directions in 2D).
+    with N >= 2 components.  Components are stored by integer index; the
+    properties ``lr_field`` and ``tb_field`` are aliases for ``[0]`` and ``[1]``.
 
     Parameters
     ----------
-    lr_field : pybamm.Symbol
-        The left-right (x) component of the vector field.
-    tb_field : pybamm.Symbol
-        The top-bottom (y) component of the vector field.
+    *components : pybamm.Symbol
+        Two or more component symbols, all sharing the same domain.
     """
 
-    def __init__(self, lr_field, tb_field):
-        if lr_field.domain != tb_field.domain:
-            raise ValueError("lr_field and tb_field must have the same domain")
-        # Initialize as a rank-1 TensorField with two components
-        super().__init__([lr_field, tb_field], domain=lr_field.domain)
-        # Override the name to maintain backward compatibility
+    # Set by discretisation for unstructured FV edge-averaging; None until then.
+    _disc_state_vector = None
+
+    def __init__(self, *components):
+        if len(components) < 2:
+            raise ValueError(
+                f"VectorField requires at least 2 components, got {len(components)}"
+            )
+        ref_domain = components[0].domain
+        for i, c in enumerate(components[1:], start=1):
+            if c.domain != ref_domain:
+                raise ValueError(
+                    f"All components must have the same domain: "
+                    f"component {i} has {c.domain}, expected {ref_domain}"
+                )
+        super().__init__(list(components), domain=ref_domain)
         self.name = "vector_field"
 
     @classmethod
     def _from_json(cls, snippet):
         # Two positional args, not a single list -- override TensorField._from_json.
         return cls(snippet["children"][0], snippet["children"][1])
 
+    @property
+    def n_components(self):
+        """Number of vector components."""
+        return len(self.components)
+
+    # ---- aliases for structured-grid directions ----
+
     @property
     def lr_field(self):
-        """The left-right (x) component of the vector field."""
-        return self._components[0]
+        """Component 0 (left-right / x)."""
+        return self.components[0]
 
     @property
     def tb_field(self):
-        """The top-bottom (y) component of the vector field."""
-        return self._components[1]
+        """Component 1 (top-bottom / y)."""
+        return self.components[1]
 
     def create_copy(
         self,
         new_children: list[pybamm.Symbol] | None = None,
         perform_simplifications: bool = True,
     ):
-        """Create a copy of this vector field with optional new children."""
         if new_children is None:
             new_children = [
-                self.lr_field.create_copy(
-                    perform_simplifications=perform_simplifications
-                ),
-                self.tb_field.create_copy(
-                    perform_simplifications=perform_simplifications
-                ),
+                c.create_copy(perform_simplifications=perform_simplifications)
+                for c in self.components
             ]
         return VectorField(*new_children)
 
+    def _to_casadi(self, t, y, y_dot, inputs, casadi_symbols):
+        """See :meth:`pybamm.Symbol._to_casadi()`."""
+        return casadi.vertcat(
+            *self._children_to_casadi(t, y, y_dot, inputs, casadi_symbols)
+        )
+
     def evaluates_on_edges(self, dimension: str) -> bool:
-        """Check if components evaluate on edges.
-
-        Overrides TensorField to provide more specific error message.
-        """
-        left_evaluates_on_edges = self.lr_field.evaluates_on_edges(dimension)
-        right_evaluates_on_edges = self.tb_field.evaluates_on_edges(dimension)
-        if left_evaluates_on_edges == right_evaluates_on_edges:
-            return left_evaluates_on_edges
-        else:
-            raise ValueError(
-                "lr_field and tb_field must either both evaluate on edges "
-                "or both not evaluate on edges"
-            )
+        statuses = [c.evaluates_on_edges(dimension) for c in self.components]
+        if all(statuses):
+            return True
+        if not any(statuses):
+            return False
+        raise ValueError(
+            "All VectorField components must either all evaluate on edges "
+            "or none evaluate on edges"
+        )

packages/pybamm/src/pybamm/solvers/processed_variable.py

@@ -1387,6 +1387,12 @@ def process_variable(name: str, base_variables, *args, **kwargs):
         return ProcessedVariable3DSciKitFEM(name, base_variables, *args, **kwargs)
 
     if mesh and hasattr(mesh, "edges_lr") and hasattr(mesh, "edges_tb"):
+        if isinstance(base_variables[0], pybamm.VectorField):
+            raise NotImplementedError(
+                "Reading VectorField variables from a Solution is not supported "
+                "on structured 2D finite-volume meshes. Use pybamm.Component to "
+                "extract a scalar component first."
+            )
         return ProcessedVariable2DFVM(name, base_variables, *args, **kwargs)
 
     # check variable shape

packages/pybamm/src/pybamm/solvers/solution.py

@@ -748,16 +748,30 @@ def _update_variable(self, name: str):
                     "solve. Please re-run the solve with `output_variables` set to "
                     "include this variable."
                 )
-            var_casadi, var_pybamm, time_integral = self._update_model_variable(
-                model,
-                _var_pybamm,
-                inputs=inputs,
-                ys_shape=ys.shape,
-                time_integral=time_integral,
-                cache_key=name,
-            )
-            vars_pybamm[i] = var_pybamm
-            vars_casadi[i] = var_casadi
+            if isinstance(_var_pybamm, pybamm.VectorField):
+                comp_casadi = []
+                for k, comp in enumerate(_var_pybamm.components):
+                    cc, _, _ = self._update_model_variable(
+                        model,
+                        comp,
+                        inputs=inputs,
+                        ys_shape=ys.shape,
+                        time_integral=None,
+                        cache_key=f"{name}[{k}]",
+                    )
+                    comp_casadi.append(cc)
+                vars_casadi[i] = comp_casadi
+            else:
+                var_casadi, var_pybamm, time_integral = self._update_model_variable(
+                    model,
+                    _var_pybamm,
+                    inputs=inputs,
+                    ys_shape=ys.shape,
+                    time_integral=time_integral,
+                    cache_key=name,
+                )
+                vars_pybamm[i] = var_pybamm
+                vars_casadi[i] = var_casadi
         var = pybamm.process_variable(
             name, vars_pybamm, vars_casadi, self, time_integral=time_integral
         )