Reorganize for multiple dates

Separate the definition of the fields of shapefiles and the output of
the namelist from the output of the shapes. This will allow us to loop
on dates only for the output of the shapes. Create two sets of writers
so that we will be able to alternate between the two sets when we loop
on multiple dates. So we create a function `define_fields`, we replace
the argument `SHP_triplet` of function write by argument writers and
we create the dictionary factories.

Analysis/mat_to_shapefiles.py

@@ -14,83 +14,98 @@ import scipy.io as sio
 import datetime
 import f90nml
 
-def write(cell, cyclone, SHP_triplet):
-    """cell is a numpy.ndarray with shape (number of eddies, 21).
+def define_fields(writers):
+    """writers is a dictionary of shapefile.Writer objects."""
+    
+    writers["extr"].field("ssh", "N", 13, 6)
+    writers["extr"].field("date_index", "N", 5)
+    writers["extr"].field("eddy_index", "N", 5)
+    writers["extr"].field("interpolat", "N", 1)
+    writers["extr"].field("cyclone", "N", 1)
+    writers["extr"].field("valid", "N", 1)
+    writers["extr"].field("speed", "N", 13, 6)
+
+    writers["outer"].field("r_eq_area", "N", 10, 4)
+    writers["outer"].field("ssh", "N", 13, 6)
+    writers["outer"].field("date_index", "N", 5)
+    writers["outer"].field("eddy_index", "N", 5)
+    writers["outer"].field("radius4", "N", 2)
+
+    writers["max_speed"].field("r_eq_area", "N", 10, 4)
+    writers["max_speed"].field("ssh", "N", 13, 6)
+    writers["max_speed"].field("date_index", "N", 5)
+    writers["max_speed"].field("eddy_index", "N", 5)
+
+def write(writers, cell, cyclone):
+    """writers is a dictionary. cell is a numpy.ndarray with shape (number
+    of eddies, 21).
 
     """
-
-    with shapefile.Writer(path.join(SHP_triplet, "extremum"),
-                          shapeType = shapefile.POINT) as w_extr, \
-      shapefile.Writer(path.join(SHP_triplet, "outermost_contour"),
-                       shapeType = shapefile.POLYGON) as w_outer, \
-      shapefile.Writer(path.join(SHP_triplet, "max_speed_contour"),
-                       shapeType = shapefile.POLYGON) as w_max_speed:
-        w_extr.field("ssh", "N", 13, 6)
-        w_extr.field("date_index", "N", 5)
-        w_extr.field("eddy_index", "N", 5)
-        w_extr.field("interpolat", "N", 1)
-        w_extr.field("cyclone", "N", 1)
-        w_extr.field("valid", "N", 1)
-        w_extr.field("speed", "N", 13, 6)
-
-        w_outer.field("r_eq_area", "N", 10, 4)
-        w_outer.field("ssh", "N", 13, 6)
-        w_outer.field("date_index", "N", 5)
-        w_outer.field("eddy_index", "N", 5)
-        w_outer.field("radius4", "N", 2)
-
-        w_max_speed.field("r_eq_area", "N", 10, 4)
-        w_max_speed.field("ssh", "N", 13, 6)
-        w_max_speed.field("date_index", "N", 5)
-        w_max_speed.field("eddy_index", "N", 5)
-
-        eddy_index = 0
-
-        for eddy in cell:
-            # eddy is a numpy.ndarray with shape (21,). The
-            # description of each element of eddy is given by
-            # matlab_data["Fields"]. Note that amplitudes in
-            # eddy[7] and eddy[11] are positive regardless of eddy
-            # orientation.
-            eddy_index += 1
-
-            w_extr.point(eddy[0], eddy[1])
-            i = np.argwhere(matlab_data["X"] == eddy[0]).item()
-            j = np.argwhere(matlab_data["Y"] == eddy[1]).item()
-            speed = 1e4 if np.isnan(eddy[12]) else eddy[12]
-            w_extr.record(ssh = matlab_data["ADT"][i, j],
-                          date_index = date_index, eddy_index = eddy_index,
-                          interpolat = 0, cyclone = cyclone, valid = 1,
-                          speed = speed)
-
-            if cyclone:
-                ssh = matlab_data["ADT"][i, j] + eddy[7]
-            else:
-                ssh =  matlab_data["ADT"][i, j] - eddy[7]
-
-            w_outer.record(r_eq_area = eddy[6], ssh = ssh, 
-                           date_index = date_index, eddy_index = eddy_index,
-                           radius4 = - 1)
-            polyline = np.stack((eddy[4], eddy[5]), axis = 1)
-            w_outer.poly([polyline])
-
-            if cyclone:
-                ssh = matlab_data["ADT"][i, j] + eddy[11]
-            else:
-                ssh =  matlab_data["ADT"][i, j] - eddy[11]
-
-            w_max_speed.record(r_eq_area = eddy[10], ssh = ssh, 
-                               date_index = date_index, eddy_index = eddy_index)
-            polyline = np.stack((eddy[8], eddy[9]), axis = 1)
-            w_max_speed.poly([polyline])
-
-    nml.write(path.join(SHP_triplet, "grid_nml.txt"))
+    
+    eddy_index = 0
+
+    for eddy in cell:
+        # eddy is a numpy.ndarray with shape (21,). The
+        # description of each element of eddy is given by
+        # matlab_data["Fields"]. Note that amplitudes in
+        # eddy[7] and eddy[11] are positive regardless of eddy
+        # orientation.
+        eddy_index += 1
+
+        writers["extr"].point(eddy[0], eddy[1])
+        i = np.argwhere(matlab_data["X"] == eddy[0]).item()
+        j = np.argwhere(matlab_data["Y"] == eddy[1]).item()
+        speed = 1e4 if np.isnan(eddy[12]) else eddy[12]
+        writers["extr"].record(ssh = matlab_data["ADT"][i, j],
+                      date_index = date_index, eddy_index = eddy_index,
+                      interpolat = 0, cyclone = cyclone, valid = 1,
+                      speed = speed)
+
+        if cyclone:
+            ssh = matlab_data["ADT"][i, j] + eddy[7]
+        else:
+            ssh =  matlab_data["ADT"][i, j] - eddy[7]
+
+        writers["outer"].record(r_eq_area = eddy[6], ssh = ssh, 
+                       date_index = date_index, eddy_index = eddy_index,
+                       radius4 = - 1)
+        polyline = np.stack((eddy[4], eddy[5]), axis = 1)
+        writers["outer"].poly([polyline])
+
+        if cyclone:
+            ssh = matlab_data["ADT"][i, j] + eddy[11]
+        else:
+            ssh =  matlab_data["ADT"][i, j] - eddy[11]
+
+        writers["max_speed"].record(r_eq_area = eddy[10], ssh = ssh, 
+                           date_index = date_index, eddy_index = eddy_index)
+        polyline = np.stack((eddy[8], eddy[9]), axis = 1)
+        writers["max_speed"].poly([polyline])
 
 if len(sys.argv) != 2: sys.exit("Required argument: .mat file")
 
 matlab_data = sio.loadmat(sys.argv[1], squeeze_me = True)
 # matlab_data is a dictionary.
 
+factories = [dict(writers = {}, cell_name = "Anticyclonic_Cell",
+                  cyclone = False, SHP_triplet = "Snapshot_anti"),
+             dict(writers = {}, cell_name = "Cyclonic_Cell", cyclone = True,
+                  SHP_triplet = "Snapshot_cyclo")]
+
+for f in factories:
+    file = path.join(f["SHP_triplet"], "extremum")
+    f["writers"]["extr"] = shapefile.Writer(file, shapeType = shapefile.POINT)
+    
+    file = path.join(f["SHP_triplet"], "outermost_contour")
+    f["writers"]["outer"] = shapefile.Writer(file,
+                                             shapeType = shapefile.POLYGON)
+    
+    file = path.join(f["SHP_triplet"], "max_speed_contour")
+    f["writers"]["max_speed"] = shapefile.Writer(file,
+                                                 shapeType = shapefile.POLYGON)
+    
+    define_fields(f["writers"])
+    
 date_index = datetime.date.fromordinal(matlab_data["date_num"] - 366) \
     - datetime.date(1950, 1, 1)
 date_index = date_index.days
@@ -98,7 +113,10 @@ nml = f90nml.Namelist()
 nml["grid_nml"] = dict(corner_deg = [matlab_data["X"][0], matlab_data["Y"][0]],
                        nlon = matlab_data["X"].size,
                        nlat = matlab_data["Y"].size)
-write(matlab_data["Anticyclonic_Cell"], cyclone = False,
-      SHP_triplet = "Snapshot_anti")
-write(matlab_data["Cyclonic_Cell"], cyclone = True,
-      SHP_triplet = "Snapshot_cyclo")
+
+for f in factories:
+    write(f["writers"], matlab_data[f["cell_name"]], f["cyclone"])
+
+for f in factories:
+    nml.write(path.join(f["SHP_triplet"], "grid_nml.txt"))
+    for k, v in f["writers"].items(): v.close()