Changeset 1492 in lmdz_wrf for trunk/tools/generic_tools.py

Timestamp:

Apr 4, 2017, 6:51:11 PM (8 years ago)

Author:

lfita

Message:

Adding:

`provide_slices': Function to provide a list of slices for a matrix giving a sub-section of running dimensions

File:

• trunk/tools/generic_tools.py (modified) (2 diffs)

trunk/tools/generic_tools.py

@@ -110 +110 @@
 # prime_numbers: Function to find all the prime numbers up to a given value above 17
 # printing_dictionary: Function to print the content of a dictionary
+# provide_slices: Function to provide a list of slices for a matrix giving a sub-section of running dimensions
 # radius_angle: Function to generate a matrix with the angle at a given point
 # radius_dist: Function to generate a matrix with the distance at a given point
@@ -11527 +11528 @@
     return np.array([1, b, c])
 
+def provide_slices(dimns, dimzs, rundims):
+    """ Function to provide a list of slices for a matrix giving a sub-section of running dimensions
+        dimns: list of names of dimensions
+        dimzs: list of sizes of dimensions in the same order as in [dimns]
+        rundims: name of dimensions to allow to run
+    >>> provide_slices(['z', 't', 'y', 'x'], [2, 3, 10, 20], ['t', 'z', 'l'])
+    [[1, 2, slice(0, 10, None), slice(0, 20, None)], 
+     [0, 1, slice(0, 10, None), slice(0, 20, None)], 
+     [1, 1, slice(0, 10, None), slice(0, 20, None)], 
+     [0, 0, slice(0, 10, None), slice(0, 20, None)], 
+     [1, 0, slice(0, 10, None), slice(0, 20, None)], 
+     [0, 2, slice(0, 10, None), slice(0, 20, None)]]
+    """
+    fname = 'provide_slices'
+
+    Ndims = len(dimns)
+
+    # Checking presence of dimension in list of dimensions
+    Tslicesize = 1
+    origslicesize = {}
+    runslicesize = {}
+    for dimn in rundims:
+        if not searchInlist(dimns, dimn):
+            #print warnmsg
+            #print '  ' + fname + ": dimension '" + dimn + "' not in list !!"
+            #print '    removing it'
+            rundims.remove(dimn)
+
+    for dimn in dimns:
+        idim = index_vec(dimns, dimn)
+        origslicesize[dimn] = dimzs[idim]
+        if searchInlist(rundims, dimn):
+            runslicesize[dimn] = dimzs[idim] - 1
+            Tslicesize = Tslicesize*dimzs[idim]
+
+    #print '  ' + fname + ': Total number of slices to provide:', Tslicesize,         \
+    #  'along:', runslicesize
+
+    slices = []
+    for il in range(Tslicesize):
+        islice = []
+        alreadychanged = False
+        for idim in range(Ndims):
+            dimn = dimns[idim]
+            if searchInlist(rundims,dimn):
+                islice.append(runslicesize[dimn])
+                # Only running value for a given dimension for the next
+                if not alreadychanged:
+                    runslicesize[dimn] = runslicesize[dimn] - 1
+                    alreadychanged = True
+                    if runslicesize[dimn] < 0:
+                        runslicesize[dimn] = origslicesize[dimn] - 1
+                        for rdimn in rundims:
+                            if rdimn != dimn:
+                                runslicesize[rdimn] = runslicesize[rdimn] - 1
+                                if runslicesize[rdimn] >= 0:
+                                    break
+                                else:
+                                    runslicesize[rdimn] = origslicesize[rdimn] - 1
+            else:
+                islice.append(slice(0,origslicesize[dimn]))
+
+        slices.append(islice)
+
+    return slices
+
 #quit()