Added strict keyword argument to zip

Author: brownbaerchen

pySDC/helpers/blocks.py

@@ -102,7 +102,7 @@ def ranks(self):
     @property
     def localBounds(self):
         iLocList, nLocList = [], []
-        for rank, nPoints, nBlocks in zip(self.ranks, self.gridSizes, self.nBlocks):
+        for rank, nPoints, nBlocks in zip(self.ranks, self.gridSizes, self.nBlocks, strict=True):
             n0 = nPoints // nBlocks
             nRest = nPoints - nBlocks * n0
             nLoc = n0 + 1 * (rank < nRest)

pySDC/helpers/fieldsIO.py

@@ -714,7 +714,7 @@ def writeFields_MPI(fileName, dtypeIdx, algo, nSteps, nVar, gridSizes):
 
     iLoc, nLoc = blocks.localBounds
     Rectilinear.setupMPI(comm, iLoc, nLoc)
-    s = [slice(i, i + n) for i, n in zip(iLoc, nLoc)]
+    s = [slice(i, i + n) for i, n in zip(iLoc, nLoc, strict=True)]
     u0 = u0[(slice(None), *s)]
 
     f1 = Rectilinear(DTYPES[dtypeIdx], fileName)

pySDC/helpers/transfer_helper.py

@@ -27,7 +27,7 @@ def next_neighbors_periodic(p, ps, k):
 
     # zip it
     value_index = []
-    for d, i in zip(distance_to_p, range(distance_to_p.size)):
+    for d, i in zip(distance_to_p, range(distance_to_p.size, strict=True)):
         value_index.append((d, i))
     # sort by distance
     value_index_sorted = sorted(value_index, key=lambda s: s[0])
@@ -53,7 +53,7 @@ def next_neighbors(p, ps, k):
     distance_to_p = np.abs(ps - p)
     # zip it
     value_index = []
-    for d, i in zip(distance_to_p, range(distance_to_p.size)):
+    for d, i in zip(distance_to_p, range(distance_to_p.size, strict=True)):
         value_index.append((d, i))
     # sort by distance
     value_index_sorted = sorted(value_index, key=lambda s: s[0])
@@ -80,7 +80,7 @@ def continue_periodic_array(arr, nn):
     else:
         cont_arr = [arr[nn[0]]]
         shift = 0.0
-        for n, d in zip(nn[1:], d_nn):
+        for n, d in zip(nn[1:], d_nn, strict=True):
             if d != 1:
                 shift = -1
             cont_arr.append(arr[n] + shift)
@@ -107,7 +107,7 @@ def restriction_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1):
 
     if periodic:
         M = np.zeros((coarse_grid.size, fine_grid.size))
-        for i, p in zip(range(n_g), coarse_grid):
+        for i, p in zip(range(n_g), coarse_grid, strict=True):
             nn = next_neighbors_periodic(p, fine_grid, k)
             circulating_one = np.asarray([1.0] + [0.0] * (k - 1))
             cont_arr = continue_periodic_array(fine_grid, nn)
@@ -120,7 +120,7 @@ def restriction_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1):
                 M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol)))
     else:
         M = np.zeros((coarse_grid.size, fine_grid.size + 2 * pad))
-        for i, p in zip(range(n_g), coarse_grid):
+        for i, p in zip(range(n_g), coarse_grid, strict=True):
             padded_f_grid = border_padding(fine_grid, pad, pad)
             nn = next_neighbors(p, padded_f_grid, k)
             # construct the lagrange polynomials for the k neighbors
@@ -158,7 +158,7 @@ def interpolation_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1,
         M = np.zeros((fine_grid.size, coarse_grid.size))
 
         if equidist_nested:
-            for i, p in zip(range(n_f), fine_grid):
+            for i, p in zip(range(n_f), fine_grid, strict=True):
                 if i % 2 == 0:
                     M[i, int(i / 2)] = 1.0
                 else:
@@ -189,7 +189,7 @@ def interpolation_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1,
                         M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol)))
 
         else:
-            for i, p in zip(range(n_f), fine_grid):
+            for i, p in zip(range(n_f), fine_grid, strict=True):
                 nn = next_neighbors_periodic(p, coarse_grid, k)
                 circulating_one = np.asarray([1.0] + [0.0] * (k - 1))
                 cont_arr = continue_periodic_array(coarse_grid, nn)
@@ -208,7 +208,7 @@ def interpolation_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1,
         padded_c_grid = border_padding(coarse_grid, pad, pad)
 
         if equidist_nested:
-            for i, p in zip(range(n_f), fine_grid):
+            for i, p in zip(range(n_f), fine_grid, strict=True):
                 if i % 2 != 0:
                     M[i, int((i - 1) / 2) + 1] = 1.0
                 else:
@@ -231,7 +231,7 @@ def interpolation_matrix_1d(fine_grid, coarse_grid, k=2, periodic=False, pad=1,
                         M[i, nn] = np.asarray(list(map(lambda x: x(p), bary_pol)))
 
         else:
-            for i, p in zip(range(n_f), fine_grid):
+            for i, p in zip(range(n_f), fine_grid, strict=True):
                 nn = next_neighbors(p, padded_c_grid, k)
                 # construct the lagrange polynomials for the k neighbors
                 circulating_one = np.asarray([1.0] + [0.0] * (k - 1))

pySDC/helpers/vtkIO.py

@@ -52,7 +52,7 @@ def field(u):
         return numpy_support.numpy_to_vtk(num_array=u.ravel(order='F'), deep=True, array_type=vtk.VTK_FLOAT)
 
     pointData = vtr.GetPointData()
-    for name, u in zip(varNames, data):
+    for name, u in zip(varNames, data, strict=True):
         uVTK = field(u)
         uVTK.SetName(name)
         pointData.AddArray(uVTK)

pySDC/implementations/problem_classes/AllenCahn_Temp_MPIFFT.py

@@ -134,7 +134,7 @@ def __init__(
         N = self.fft.global_shape()
         k = [np.fft.fftfreq(n, 1.0 / n).astype(int) for n in N[:-1]]
         k.append(np.fft.rfftfreq(N[-1], 1.0 / N[-1]).astype(int))
-        K = [ki[si] for ki, si in zip(k, s)]
+        K = [ki[si] for ki, si in zip(k, s, strict=True)]
         Ks = list(np.meshgrid(*K, indexing='ij', sparse=True))
         Lp = 2 * np.pi / L
         for i in range(ndim):

pySDC/implementations/problem_classes/Brusselator.py

@@ -178,7 +178,7 @@ def plot(self, u, t=None, fig=None):  # pragma: no cover
 
         vmin = u.min()
         vmax = u.max()
-        for i, label in zip([self.iU, self.iV], [r'$u$', r'$v$']):
+        for i, label in zip([self.iU, self.iV], [r'$u$', r'$v$'], strict=True):
             im = axs[i].pcolormesh(self.X[0], self.X[1], u[i], vmin=vmin, vmax=vmax)
             axs[i].set_aspect(1)
             axs[i].set_title(label)

pySDC/implementations/problem_classes/Burgers.py

@@ -324,7 +324,7 @@ def plot(self, u, t=None, fig=None, vmin=None, vmax=None):  # pragma: no cover
         imV = axs[1].pcolormesh(self.X, self.Z, u[iv].real, vmin=vmin, vmax=vmax)
         imVort = axs[2].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real)
 
-        for i, label in zip([0, 1, 2], [r'$u$', '$v$', 'vorticity']):
+        for i, label in zip([0, 1, 2], [r'$u$', '$v$', 'vorticity'], strict=True):
             axs[i].set_aspect(1)
             axs[i].set_title(label)
 

pySDC/implementations/problem_classes/GrayScott_MPIFFT.py

@@ -343,7 +343,7 @@ def plot(self, u, t=None, fig=None):  # pragma: no cover
 
         vmin = u.min()
         vmax = u.max()
-        for i, label in zip([self.iU, self.iV], [r'$u$', r'$v$']):
+        for i, label in zip([self.iU, self.iV], [r'$u$', r'$v$'], strict=True):
             im = axs[i].pcolormesh(self.X[0], self.X[1], u[i], vmin=vmin, vmax=vmax)
             axs[i].set_aspect(1)
             axs[i].set_title(label)

pySDC/implementations/problem_classes/RayleighBenard.py

@@ -363,7 +363,7 @@ def plot(self, u, t=None, fig=None, quantity='T'):  # pragma: no cover
 
         imT = axs[0].pcolormesh(self.X, self.Z, u[self.index(quantity)].real)
 
-        for i, label in zip([0, 1], [rf'${quantity}$', 'vorticity']):
+        for i, label in zip([0, 1], [rf'${quantity}$', 'vorticity'], strict=True):
             axs[i].set_aspect(1)
             axs[i].set_title(label)
 
@@ -630,6 +630,7 @@ def post_step(self, step, level_number):
         for key, value in zip(
             ['Nusselt', 'buoyancy_production', 'viscous_dissipation'],
             [Nusselt, buoyancy_production, viscous_dissipation],
+            strict=True,
         ):
             self.add_to_stats(
                 process=step.status.slot,

pySDC/implementations/problem_classes/RayleighBenard3D.py

@@ -255,7 +255,7 @@ def eval_f(self, u, *args, **kwargs):
 
         fexpl_pad = self.xp.zeros_like(u_pad)
         for i in derivative_indices:
-            for i_vel, iD in zip([iu, iv, iw], range(self.ndim)):
+            for i_vel, iD in zip([iu, iv, iw], range(self.ndim), strict=True):
                 fexpl_pad[i] -= u_pad[i_vel] * derivatives[iD][i]
 
         if self.spectral_space:
@@ -398,7 +398,7 @@ def get_frequency_spectrum(self, u):
 
         # compute local spectrum
         local_spectrum = self.xp.empty(shape=(2, energy.shape[3], n_k))
-        for i, k in zip(range(n_k), unique_k):
+        for i, k in zip(range(n_k), unique_k, strict=True):
             mask = xp.logical_or(abs_kx == k, abs_ky == k)
             local_spectrum[..., i] = xp.sum(energy[indices, mask, :], axis=1)
 
@@ -411,7 +411,7 @@ def get_frequency_spectrum(self, u):
 
         spectra = self.comm.allgather(local_spectrum)
         spectrum = self.xp.zeros(shape=(2, self.axes[2].N, n_k_all))
-        for ks, _spectrum in zip(k_all, spectra):
+        for ks, _spectrum in zip(k_all, spectra, strict=True):
             ks = list(ks)
             unique_k_all = list(unique_k_all)
             for k in ks: