Added units to the docstrings (#68)

* Added units to the docstrings * style: pre-commit fixes * Adding another file * Made changes to the docstrings * style: pre-commit fixes * Made the changes, fixed indentation * style: pre-commit fixes * Added spacing, fixed indentation, added unitless * Made changes and adjusted the spacing * style: pre-commit fixes * docs: Updated units and formatting including cosmology, lenses, and sources * fix: Removed tuple and moved returns * fix: Remove tuple from return and moved separated returns * style: pre-commit fixes --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: uwcdc <[email protected]>
uw-ssec · Feb 15, 2024 · 1dde0c9 · 1dde0c9
1 parent 61c4c51
commit 1dde0c9
Show file tree

Hide file tree

Showing 18 changed files with 1,602 additions and 144 deletions.
diff --git a/src/caustics/cosmology/base.py b/src/caustics/cosmology/base.py
@@ -17,12 +17,13 @@ class Cosmology(Parametrized):
     This class provides an interface for cosmological computations used in lensing
     such as comoving distance and critical surface density.
 
-    Units
-    -----
     Distance
-        Mpc
+
+        *Unit: megaparsec*
+
     Mass
-        solar mass
+
+        *Unit: solMass*
 
     Attributes
     ----------
@@ -50,13 +51,19 @@ def critical_density(self, z: Tensor, params: Optional["Packed"] = None) -> Tens
         ----------
         z: Tensor
             The redshifts.
+
+            *Unit: unitless*
+
         params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The critical density at each redshift.
+
+            *Unit: solMass/megaparsec^3*
+
         """
         ...
 
@@ -72,13 +79,19 @@ def comoving_distance(
         ----------
         z: Tensor
             The redshifts.
-        params: (Packed, optional0
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The comoving distance to each redshift.
+
+            *Unit: megaparsec*
+
         """
         ...
 
@@ -94,13 +107,19 @@ def transverse_comoving_distance(
         ----------
         z: Tensor
             The redshifts.
+
+            *Unit: unitless*
+
         params: (Packed, optional)
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The transverse comoving distance to each redshift in Mpc.
+
+            *Unit: megaparsec*
+
         """
         ...
 
@@ -115,15 +134,24 @@ def comoving_distance_z1z2(
         ----------
         z1: Tensor
             The starting redshifts.
+
+            *Unit: unitless*
+
         z2: Tensor
             The ending redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The comoving distance between each pair of redshifts.
+
+            *Unit: megaparsec*
+
         """
         return self.comoving_distance(z2, params) - self.comoving_distance(z1, params)
 
@@ -138,15 +166,24 @@ def transverse_comoving_distance_z1z2(
         ----------
         z1: Tensor
             The starting redshifts.
+
+            *Unit: unitless*
+
         z2: Tensor
             The ending redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The transverse comoving distance between each pair of redshifts in Mpc.
+
+            *Unit: megaparsec*
+
         """
         return self.transverse_comoving_distance(
             z2, params
@@ -163,13 +200,19 @@ def angular_diameter_distance(
         -----------
         z: Tensor
             The redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The angular diameter distance to each redshift.
+
+            *Unit: megaparsec*
+
         """
         return self.comoving_distance(z, params, **kwargs) / (1 + z)
 
@@ -184,15 +227,24 @@ def angular_diameter_distance_z1z2(
         ----------
         z1: Tensor
             The starting redshifts.
+
+            *Unit: unitless*
+
         z2: Tensor
             The ending redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The angular diameter distance between each pair of redshifts.
+
+            *Unit: megaparsec*
+
         """
         return self.comoving_distance_z1z2(z1, z2, params, **kwargs) / (1 + z2)
 
@@ -212,15 +264,24 @@ def time_delay_distance(
         ----------
         z_l: Tensor
             The lens redshifts.
+
+            *Unit: unitless*
+
         z_s: Tensor
             The source redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The time delay distance for each pair of lens and source redshifts.
+
+            *Unit: megaparsec*
+
         """
         d_l = self.angular_diameter_distance(z_l, params)
         d_s = self.angular_diameter_distance(z_s, params)
@@ -243,15 +304,24 @@ def critical_surface_density(
         ----------
         z_l: Tensor
             The lens redshifts.
+
+            *Unit: unitless*
+
         z_s: Tensor
             The source redshifts.
-        params: (Packed, optional)
+
+            *Unit: unitless*
+
+        params: Packed, optional
             Dynamic parameter container for the computation.
 
         Returns
         -------
         Tensor
             The critical surface density for each pair of lens and source redshifts.
+
+            *Unit: solMass/megaparsec^2*
+
         """
         d_l = self.angular_diameter_distance(z_l, params)
         d_s = self.angular_diameter_distance(z_s, params)