Changed tricco function descriptions to docstrings

tricco/compute.py

@@ -3,19 +3,21 @@ import numpy as np
 from .grid_functions import get_edges_of_cell
 from .cubulation_functions import EnhancedTri, cubing_next_round, shift_coordinates, compute_connected_components, compute_list_of_components 
 
-# compute_cubulation
-#   Assign to each triangle a 3d-coordinate called cube coordinate.
-# INPUT: 
-#   start_triangle      int         the triangle index where the cubulatiion algorithm starts
-#   radius              int         the number of "outward" steps
-#   print_progress      Boolean     print progress of iteration (default False)
-# OUTPUT: 
-#   Returns a list of arrays of the following form:
-#   np.array([<triangle index>, <cube array>])
-#   where <cube array> is an array with three entries that encode the cube coordinate of the triangle with <triangle index>. 
 
 def compute_cubulation(start_triangle, radius, print_progress=False): 
-    
+    """"
+    Assign to each triangle a 3d-coordinate called cube coordinate.
+       
+        Input: 
+            start_triangle      int         the triangle index where the cubulatiion algorithm starts
+            radius              int         the number of "outward" steps
+            print_progress      Boolean     print progress of iteration (default False)
+     
+        Output: 
+            Returns a list of arrays of the following form:
+            np.array([<triangle index>, <cube array>])
+            where <cube array> is an array with three entries that encode the cube coordinate of the triangle with <triangle index>. 
+    """
     # initialization
     init = EnhancedTri(start_triangle, np.array([0,0,0])) # assign coordinates (0,0,0) to start_triangle
     cube_coordinates = [np.array([init.triangle, init.cube], dtype=object)]
@@ -37,19 +39,19 @@ def compute_cubulation(start_triangle, radius, print_progress=False):
     return shift_coordinates(cube_coordinates, radius)
     
     
-# compute_connected_2d_components:
-#   From given cubuluation, compute connected components and output them as a list of lists
-# INPUT:
-#   cubulation          list of arrays      output of compute_cubulation 
-#   field_cube          numpy array         input field in cubulated 3d-coordinates
-#   connectivity        string              use 'vertex' (default) or 'edge' connectivity
-# OUTPUT: 
-#   - Returns a list of lists. 
-#   - Each inner list corresponds to a connected component and contains 
-#     all triangle indices that belong to this connected component
-
 def compute_connected_components_2d(cubulation, field_cube, connectivity = 'vertex'): 
-    
+    """
+    From given cubuluation, compute connected components and output them as a list of lists
+
+        Input:
+            cubulation          list of arrays      output of compute_cubulation 
+            field_cube          numpy array         input field in cubulated 3d-coordinates
+            connectivity        string              use 'vertex' (default) or 'edge' connectivity
+        Output: 
+            - Returns a list of lists. 
+            - Each inner list corresponds to a connected component and contains 
+              all triangle indices that belong to this connected component
+    """    
     # compute connected components on field_cube
     # components_cube is a 3-dimensional array that contains
     # the connected component index of a triangle with cube coordinates (x,y,z) at position component_array[x][y][z]
@@ -62,19 +64,20 @@ def compute_connected_components_2d(cubulation, field_cube, connectivity = 'vert
     return components
 
 
-# compute_connected_3d_components:
-#   From given cubuluation, compute connected components and output them as a list of lists
-# INPUT:
-#   cubulation          list of arrays      output of compute_cubulation 
-#   field_cube          numpy array         input field in dimensions of (lev x cubulated coordinates)
-#   connectivity        string              use 'vertex' (default) or 'edge' connectivity
-# OUTPUT: 
-#   - Returns a list of lists. 
-#   - Each inner list corresponds to a connected component and contains all tuples of (level, triangle index) that
-#     belong to this connected component
-
 def compute_connected_components_3d(cubulation, field_cube, connectivity = 'vertex'):
-    
+    """
+    From given cubuluation, compute connected components and output them as a list of lists
+   
+        Input:
+            cubulation          list of arrays      output of compute_cubulation 
+            field_cube          numpy array         input field in dimensions of (lev x cubulated coordinates)
+            connectivity        string              use 'vertex' (default) or 'edge' connectivity
+ 
+        Output: 
+            - Returns a list of lists. 
+            - Each inner list corresponds to a connected component and contains all tuples of (level, triangle index) that
+              belong to this connected component
+    """    
     nlev = np.shape(field_cube)[0] # number of vertical levels
     
     # Step 1: compute connected components for each level separately, 

tricco/cubulation_functions.py

@@ -3,9 +3,12 @@ import cc3d # external library for connected component labeling on cubic grids
 
 from .grid_functions import get_vertices_of_edge, get_edges_of_cell, get_neighbors_of_cell
 
-# class to enhance a triangle with its cube coordinate
-#   for a given triangle create an object consisting of the triangle index and the triangle's cube coordinate
+
 class EnhancedTri:
+    """
+    class to enhance a triangle with its cube coordinate
+    for a given triangle create an object consisting of the triangle index and the triangle's cube coordinate
+    """
     def __init__(self, triangle, cube):
         self.triangle = triangle
         self.cube = cube
@@ -13,19 +16,21 @@ class EnhancedTri:
         print("Triangle index:", self.triangle, " Cube coordinate:", self.cube)
         
 
-# colouring edges of a triangle
-#   situation:  the colours of two edges are already known, 
-#               the third edge must be coloured
-#   called from colour_new_edges in this special situation
-# INPUT:
-#   first_known_edge      int     index of a already coloured edge
-#   second_known_edge     int     index of a already coloured edge
-#   new_edge              int     index of the not yet coloured edge
-#   edge_colours          dict    dictionary storing to each edge index the colour
-# OUTPUT:
-#   None, but updates edge_colours
 def colour_exactly_one_new_edge(first_known_edge, second_known_edge, new_edge, edge_colours):
-    
+    """
+    Colouring edges of a triangle
+    Situation:  the colours of two edges are already known, 
+                the third edge must be coloured
+    called from colour_new_edges in this special situation
+   
+        Input:
+            first_known_edge      int     index of a already coloured edge
+            second_known_edge     int     index of a already coloured edge
+            new_edge              int     index of the not yet coloured edge
+            edge_colours          dict    dictionary storing to each edge index the colour
+        Output:
+            none, but updates edge_colours
+    """    
     # get the colours of the already coloured edges
     first_colour = edge_colours[first_known_edge]
     second_colour = edge_colours[second_known_edge]
@@ -37,17 +42,19 @@ def colour_exactly_one_new_edge(first_known_edge, second_known_edge, new_edge, e
             return
         
         
-# colouring edges of a triangle
-#   from old_triangle with already coloured edges determine the edge_colours of the adjacent new_triangle
-# INPUT:
-#   old_triangle_edges      set     contains the edge indices of old_triangle's edges
-#   new_triangle_edges      set     contains the edge indices of new_triangle's edges
-#   joint_edge              int     index of the shared edge from old_triangle and new_triangle
-#   edge_colours            dict    dictionary storing to each edge index the colour
-# OUTPUT:
-#   None, but updates edge_colours
 def colour_new_edges(old_triangle_edges, new_triangle_edges, joint_edge, edge_colours):
-    
+    """
+    Colouring edges of a triangle
+    From old_triangle with already coloured edges determine the edge_colours of the adjacent new_triangle
+
+        Input:
+            old_triangle_edges      set     contains the edge indices of old_triangle's edges
+            new_triangle_edges      set     contains the edge indices of new_triangle's edges
+            joint_edge              int     index of the shared edge from old_triangle and new_triangle
+            edge_colours            dict    dictionary storing to each edge index the colour
+        Output:
+            none, but updates edge_colours
+    """    
     # determine the edges of new_triangle that are not shared with old_triangle
     new_edges   = new_triangle_edges.difference( old_triangle_edges ) 
     first_edge  = new_edges.pop()
@@ -95,17 +102,20 @@ def colour_new_edges(old_triangle_edges, new_triangle_edges, joint_edge, edge_co
             edge_colours[ second_edge ] = edge_colours[ first_comparison_edge ]
             
         
-# compute cube coordinate of a new triangle
-#   given the old_triangle's cube_coordinate derive the cube coordinate of the adjacent new_triangle 
-#   by using the colour of their joint_edge 
-# INPUT:
-#   old                     EnhancedTri     the old EnhancedTri with known cube_coordinates
-#   joint_edge              int             index of the shared edge from old_triangle and new_triangle
-#   edge_colours            dict            dictionary storing to each edge index the colour
-# OUTPUT:
-#   new_cube_coordinate, the cube_coordinate of new_triangle
 def determine_cube_coordinates (old, old_triangle_edges, new_triangle_edges, joint_edge, edge_colours):
-    
+    """
+    Compute cube coordinate of a new triangle
+    Given the old_triangle's cube_coordinate derive the cube coordinate of the adjacent new_triangle 
+    by using the colour of their joint_edge 
+
+        Input:
+            old                     EnhancedTri     the old EnhancedTri with known cube_coordinates
+            joint_edge              int             index of the shared edge from old_triangle and new_triangle
+            edge_colours            dict            dictionary storing to each edge index the colour
+
+        Output:
+            new_cube_coordinate, the cube_coordinate of new_triangle
+    """    
     # get colour of the joint_edge and call it direction
     direction = edge_colours[joint_edge]
     # the direction encodes the parallel class of the joint edge
@@ -130,16 +140,18 @@ def determine_cube_coordinates (old, old_triangle_edges, new_triangle_edges, joi
     return new_cube_coordinate
 
 
-# iterate over the outmost triangles and compute the cube_coordinates of their adjacent triangles
-# INPUT:
-#   cube_coordinates        array       EnhancedTri's of the triangles whose cube_coordinates are already computed
-#   visited_triangles       list        indices of the triangles whose cube_coordinates are already computed
-#   outmost                 list        EnhancedTri's of the triangles considered in the round before   
-#   edge_colours            dict        dictionary storing to each edge index the colour
-# OUTPUT:
-#   updated outmost, EnhancedTri's of the new triangles considered this round
 def cubing_next_round(cube_coordinates, visited_triangles, outmost, edge_colours):
+    """
+    Iterate over the outmost triangles and compute the cube_coordinates of their adjacent triangles
     
+        Input:
+            cube_coordinates        array       EnhancedTri's of the triangles whose cube_coordinates are already computed
+            visited_triangles       list        indices of the triangles whose cube_coordinates are already computed
+            outmost                 list        EnhancedTri's of the triangles considered in the round before   
+            edge_colours            dict        dictionary storing to each edge index the colour
+        Output:
+            updated outmost, EnhancedTri's of the new triangles considered this round
+    """    
     # list for triangles that will be outmost in next iteration
     new_outmost = []
    
@@ -169,36 +181,38 @@ def cubing_next_round(cube_coordinates, visited_triangles, outmost, edge_colours
     return new_outmost
 
 
-# shift cube coordinates such that they are all positive
-# INPUT:
-#   cube_coordinates    array   stores triangle - cube_coordinate pairs
-#   radius              int     same as in compute_cube_coordiantes
-# OUTPUT:
-#   shifted cube_coordinates
 def shift_coordinates (cube_coordinates, radius):
+    """
+    Shift cube coordinates such that they are all positive
+
+        Input:
+            cube_coordinates    array   stores triangle - cube_coordinate pairs
+            radius              int     same as in compute_cube_coordiantes
+        
+        Output:
+            shifted cube_coordinates
+    """
     shift = int(radius/2)
-    
     # update all coordinates
     for entry in cube_coordinates:
         cube = entry[1] # get cube_coordinate 
-        
         # consider all three directions
         for direction in (0,1,2):
             cube[direction] += shift # shift coordinate
-    
     return cube_coordinates
 
 
-# compute connected components
-# INPUT:
-#   field_cube     array    obtained as part of field prepaation
-#   connectivity   string   use 'vertex' (default) or 'edge' connectivity
-# OUTPUT:
-#   component_cube, 3d array containing the component indices
 def compute_connected_components(field_cube, connectivity = 'vertex'):
+    """"
+    Compute connected components
     
+        Input:
+            field_cube     array    obtained as part of field prepaation
+            connectivity   string   use 'vertex' (default) or 'edge' connectivity
+        Output:
+            component_cube, 3d array containing the component indices
+    """"    
     import sys
-    
     # translate connectivity - string into input value for 3d connected component labeling
     if connectivity == 'edge':
         connectivity_value = 6
@@ -206,23 +220,24 @@ def compute_connected_components(field_cube, connectivity = 'vertex'):
         connectivity_value = 26     
         if connectivity != 'vertex':    # use default 'vertex' connectivity for invalid input
             sys.exit("Invalid input: only 'vertex' or 'edge' connectivity are allowed.")
-            
     # calling external 3d connected component labeling
     component_cube = cc3d.connected_components(field_cube, connectivity = connectivity_value)
-    
     return component_cube
 
 
-# which triangle belongs to which component?
-# INPUT:
-#   cubulation         array       output of compute_cube_coordinates   
-#   component_cube     3d arry     output of compute_connected_components
-# OUTPUT:
-#   Returns a list of lists. 
-#   Each inner list corresponds to one connected component and contains the indices
-#   of all triangles that belong to it
 def compute_list_of_components(cubulation, component_cube):
-    
+    """
+    Which triangle belongs to which component?
+        
+        Input:
+            cubulation         array       output of compute_cube_coordinates   
+            component_cube     3d arry     output of compute_connected_components
+        
+        Output:
+            Returns a list of lists. 
+            Each inner list corresponds to one connected component and contains the indices
+            of all triangles that belong to it.
+    """
     # number of found components on cubical grid
     ncomp = component_cube.max()
     

tricco/grid_functions.py

 # xarray dataset that holds the grid
 grid = None
 
-# get_neighbors_of_cell:
-#   returns edge neighbors of a given triangle
 def get_neighbors_of_cell(triangle):
+    """
+    Returns edge neighbors of a given triangle
+    """
     return grid.neighbor_cell_index[:,triangle].values
    
-# get_edges_of_cell:
-#   returns the three edges of a given triangle
 def get_edges_of_cell(triangle):
+    """
+    Returns the three edges of a given triangle
+    """
     return grid.edge_of_cell[:, triangle].values
 
-# get_vertices_of_edge:
-#   returns the two vertices of a given edge
+
 def get_vertices_of_edge(edge):
+    """
+    Returns the two vertices of a given edge
+    """
     return grid.edge_vertices[:, edge].values

tricco/prepare.py

 import numpy as np
 import xarray as xr
 
-# path_input:
-#   create globale variables to open the xArray dataset from gridfile and datafile
-# INPUT:
-#   path        string      location where gridfile and datafile are stored
-#   gridfile    string      name of the file containing the grid
-#   datafile    string      name of the file containing the field
-
-# we only need the following information about the horizontal triangular grid:
-# - the edge neighbors of a given triangle
-# - the three edges of a given triangle
-# - the two vertices of a given edge 
 
 def prepare_grid(model, path, file):
+    """
+    Returns grid as xarray dataset.
+    
+    We only need the following information about the horizontal triangular grid:
+        - the edge neighbors of a given triangle
+        - the three edges of a given triangle
+        - the two vertices of a given edge 
+    
+        Input:
+            model string      name of considered model
+            path  string      location where gridfile is stored
+            file  string      name of the file containing the grid
+
+        Output:
+            grid as xarray dataset
+    """
     if model == 'ICON':
         _ds_grid = _grid_icon(path+file)
     return _ds_grid
 
 
 # implementation for ICON model
-
 def _grid_icon(_gridfile):
     
     ds_grid = xr.open_dataset(_gridfile)
@@ -40,22 +44,41 @@ def _grid_icon(_gridfile):
     return ds_grid
 
 
-# prepare_field
-#   returns the data field as a numpy array, also applies threshold-based mask
-# INPUT:
-#   file: file containing the data field
-#   var:  variabe name
-#   threshold: set field to one if field>=threshold, and to zero otherwise
-# OUTPUT:
-#   field both on the triangular grid as well as on the cubulated grid
-
 
 def prepare_field(model, path, file, var, threshold, cubulation):
+    """
+    Returns the one-level data field as a numpy array, also applies threshold-based mask.
+    
+        Input:
+            model     : name of considered model
+            path      : path to file
+            file      : file containing the data field
+            var       :  variabe name
+            threshold : set field to one if field>=threshold, and to zero otherwise
+            cubulation: cubulation of horizontal grid
+
+        Output:
+            field both on the triangular grid as well as on the cubulated grid
+    """
     if model == 'ICON':
         _field, _field_cube = _field_icon(path, file, var, threshold, cubulation)
     return _field, _field_cube
 
 def prepare_field_lev(model, path, file, var, threshold, cubulation):
+    """
+    Returns the many-level data field as a numpy array, also applies threshold-based mask.
+    
+        Input:
+            model     : name of considered model
+            path      : path to file
+            file      : file containing the data field
+            var       :  variabe name
+            threshold : set field to one if field>=threshold, and to zero otherwise
+            cubulation: cubulation of horizontal grid
+            
+        Output:
+            field both on the triangular grid as well as on the cubulated grid
+    """
     if model == 'ICON':
         _field, _field_cube = _field_icon_lev(path, file, var, threshold, cubulation)
     return _field, _field_cube