Fixes to allow non-lazy compenents.

Author: pp-mo

lib/iris/aux_factory.py

@@ -11,7 +11,7 @@
 import dask.array as da
 import numpy as np
 
-from iris._lazy_data import concatenate
+from iris._lazy_data import concatenate, is_lazy_data
 from iris.common import CFVariableMixin, CoordMetadata, metadata_manager_factory
 import iris.coords
 from iris.warnings import IrisIgnoringBoundsWarning
@@ -119,17 +119,6 @@ def _derive_array(self, *dep_arrays, **other_args):
         """
         result = self._calculate_array(*dep_arrays, **other_args)
 
-        # The dims of all the given components should be the same and, **presumably**,
-        #  the same as the result ??
-        for i_dep, (dep, name) in enumerate(zip(dep_arrays, self.dependencies.keys())):
-            if dep.ndim != result.ndim:
-                msg = (
-                    f"Dependency #{i_dep}, '{name}' has ndims={dep.ndim}, "
-                    "not matching result {result.ndim!r}"
-                    " (shapes {dep.shape}/{result.shape})."
-                )
-                raise ValueError(msg)
-
         # See if we need to improve on the chunking of the result
         from iris._lazy_data import _optimum_chunksize
         adjusted_chunks = _optimum_chunksize(
@@ -138,26 +127,45 @@ def _derive_array(self, *dep_arrays, **other_args):
             dtype=result.dtype,
         )
 
-        if adjusted_chunks != result.chunksize:
+        # Does optimum_chunksize say we should have smaller chunks in some dimensions?
+        if np.any(adjusted_chunks < result.chunksize):
+            # co-broadcast all the deps to get same dimensions for each
+            dep_arrays = np.broadcast_arrays(*dep_arrays)
+
+            # The dims of all the given components should now be the same and, *presumably*,
+            #  the same as the result ??
+            for i_dep, (dep, name) in enumerate(zip(dep_arrays, self.dependencies.keys())):
+                if dep.ndim != result.ndim:
+                    msg = (
+                        f"Dependency #{i_dep}, '{name}' has ndims={dep.ndim}, "
+                        f"not matching result {result.ndim!r}"
+                        f" : respective shapes {dep.shape}, {result.shape}."
+                    )
+                    raise ValueError(msg)
+
             # Re-do the result calculation, re-chunking the inputs along dimensions
             #  which it is suggested to reduce.
             # First make a (writable) copy of the inputs.....
             new_deps = []
             for i_dep, dep in enumerate(dep_arrays):
                 # Reduce each dependency chunksize to the result chunksize if smaller.
-                dep_chunks = dep.chunksize
+                dep_chunks = dep.chunksize if is_lazy_data(dep) else dep.shape
                 new_chunks = tuple([
                     min(dep_chunk, adj_chunk)
                     for dep_chunk, adj_chunk in zip(dep_chunks, adjusted_chunks)
                 ])
                 # If the dep chunksize was reduced, replace with a rechunked version.
                 if new_chunks != dep_chunks:
+                    if not is_lazy_data(dep):
+                        # I guess this is possible ?
+                        # TODO: needs a test
+                        dep = da.from_array(dep)
                     dep = dep.rechunk(new_chunks)
                 new_deps.append(dep)
 
             # Finally, re-do the calculation, which hopefully results in a better
             #  overall chunking for the result
-            result = self._calculate_array(*new_deps)
+            result = self._calculate_array(*new_deps, **other_args)
 
         return result