From f1b53ae8395becf85c61be8dfec8416a693cbb23 Mon Sep 17 00:00:00 2001 From: voigta80 <aiko.voigt@univie.ac.at> Date: Tue, 7 Jun 2022 00:31:37 +0200 Subject: [PATCH] Changed tricco function descriptions to docstrings --- tricco/compute.py | 73 ++++++++-------- tricco/cubulation_functions.py | 153 ++++++++++++++++++--------------- tricco/grid_functions.py | 16 ++-- tricco/prepare.py | 65 +++++++++----- 4 files changed, 176 insertions(+), 131 deletions(-) diff --git a/tricco/compute.py b/tricco/compute.py index e33c803..bd36264 100644 --- a/tricco/compute.py +++ b/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, diff --git a/tricco/cubulation_functions.py b/tricco/cubulation_functions.py index 344cd65..7375562 100644 --- a/tricco/cubulation_functions.py +++ b/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() diff --git a/tricco/grid_functions.py b/tricco/grid_functions.py index e9cc843..9a2b9eb 100644 --- a/tricco/grid_functions.py +++ b/tricco/grid_functions.py @@ -1,17 +1,21 @@ # 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 diff --git a/tricco/prepare.py b/tricco/prepare.py index 2f8a1c1..f04ac73 100644 --- a/tricco/prepare.py +++ b/tricco/prepare.py @@ -1,26 +1,30 @@ 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 -- GitLab