Parcourir la source

improved docstring info

gnidovec il y a 2 ans
Parent
commit
4992059aff
1 fichiers modifiés avec 31 ajouts et 10 suppressions
  1. 31 10
      contact_function.py

+ 31 - 10
contact_function.py

@@ -39,6 +39,7 @@ clib.contact_function_multi.restype = ctypes.c_void_p
 
 
 def minimizer_map(minimizer: str) -> int:
+    """Maps minimizer name to an integer for use in C functions from "overlap_algorithm.c"."""
     if minimizer == 'brent':
         return 0
     elif minimizer == 'brent_early':
@@ -53,17 +54,22 @@ def minimizer_map(minimizer: str) -> int:
 
 @dataclass
 class Contact:
+    """Stores contact data."""
     contact_f: float
     min_eig_T: float
     med_eig_T: float
-    min_vec_scalar: float
-    med_vec_scalar: float
-    feval_min: int
-    feval_med: int
-    branch: int
+    min_vec_scalar: float  # result of scalar product test fot the eigenvector corresponding to minimal eigenvalue
+    med_vec_scalar: float  # result of scalar product test fot the eigenvector corresponding to median eigenvalue
+    feval_min: int  # number of minimal eigenvalue evaluations
+    feval_med: int  # number of median eigenvalue evaluations
+    branch: int  # contact function solution branch (0: min eig, 1: med eig, 2: Omega)
 
 
-def contact_function(a0, a1, b0, b1, coord0, orient0, coord1, orient1, minimizer='brent') -> Contact:
+def contact_function(a0: float, a1: float, b0: float, b1: float,
+                    coord0: np.ndarray, orient0: np.ndarray, coord1: np.ndarray, orient1: np.ndarray,
+                    minimizer: str = 'brent') -> Contact:
+    """Calculates contact data for a pair of spherical ellipses."""
+
     other_results = np.zeros(4, dtype=np.float64)
     fevals = np.zeros(2, dtype=np.intc)
     branch = np.zeros(1, dtype=np.intc)
@@ -77,13 +83,21 @@ def contact_function(a0, a1, b0, b1, coord0, orient0, coord1, orient1, minimizer
 
 @dataclass
 class ContactMulti:
+    """
+    Stores the contact function for many pairs of ellipses, along with average numbers of eigenvalue evaluations
+    and total calculation time.
+    """
     contact_f: np.ndarray
     avg_evals_mineig: float
     avg_evals_medeig: float
     time: float
 
 
-def contact_function_multi(a0, a1, b0, b1, coord0, orient0, coord1, orient1, minimizer='brent') -> ContactMulti:
+def contact_function_multi(a0: float, a1: float, b0: float, b1: float,
+                           coord0: np.ndarray, orient0: np.ndarray, coord1: np.ndarray, orient1: np.ndarray,
+                           minimizer: str = 'brent') -> ContactMulti:
+    """Calculates contact function for multiple configurations of spherical ellipses."""
+
     n = len(coord0)
     results = np.zeros(n, dtype=np.float64)
     all_evals = np.zeros(2, dtype=np.intc)
@@ -98,9 +112,9 @@ def contact_function_multi(a0, a1, b0, b1, coord0, orient0, coord1, orient1, min
 
 class EllipsePair:
 
-    def __init__(self, a0, a1, epsilon):
+    def __init__(self, a0: float, a1: float, epsilon: float):
         """
-        Class in which optimization of ellipse configurations on a sphere is performed.
+        Class to calculate the contact function between  two spherical ellipses.
         :param a0: semi-major axis for the first elliptical cylinder
         :param a1: semi-major axis for the second elliptical cylinder
         :param epsilon: aspect ratio
@@ -114,10 +128,16 @@ class EllipsePair:
 
     def contact(self, coord0, orient0, coord1=np.array([0., 0., 1.]), orient1=np.array([1., 0., 0.]),
                 minimizer='brent') -> Contact:
+        """Evaluates the contact function along with oth contact data."""
         return contact_function(self.a0, self.a1, self.b0, self.b1, coord0, orient0, coord1, orient1, minimizer)
 
     def contact_multi_dist(self, distances, num_ang, axis=np.array([1, 0, 0]), minimizer='brent')\
             -> (np.ndarray, np.ndarray, np.ndarray):
+        """
+        Evaluates the contact function at different distances between two ellipses as well as multiple
+        mutual orientations. Returns mean calculation time, mean number of minimum eigenvalue evaluations and
+        mean number of median eigenvalue evaluations at each distance.
+        """
 
         timings = np.zeros(len(distances))
         avg_evals_mineig = np.zeros(len(distances))
@@ -136,7 +156,8 @@ class EllipsePair:
         return timings, avg_evals_mineig, avg_evals_medeig
 
 
-def generate_confgs(dist, n, axis=np.array([1, 0, 0])):
+def generate_confgs(dist: np.ndarray, n: int, axis: np.ndarray = np.array([1, 0, 0])) \
+        -> (np.ndarray, np.ndarray, np.ndarray, np.ndarray):
     """
     At a given distance, generate all possible orientational configurations of two spherical ellipses,
     with n steps in a half-circle rotation. This results in n x n configurations.