docstring cleanup

Author: lazyoracle

src/isotropic/e2.py

@@ -1,4 +1,4 @@
-"""This module contains functions for generating the vector e_2."""
+"""This module contains functions for generating the vector $e_2$."""
 
 from typing import Callable, Tuple
 
@@ -13,7 +13,7 @@
 
 def F_j(theta_j: float, j: int, d: int) -> Array:
     """
-    Calculate the function F_j for the given angle theta_j and index j in dimension d.
+    Calculate the function $F_j$ for the given angle $\\theta_j$ and index $j$ in dimension $d$.
 
     Parameters
     ----------
@@ -27,7 +27,7 @@ def F_j(theta_j: float, j: int, d: int) -> Array:
     Returns
     -------
     Array
-        The value of the function F_j evaluated at theta_j.
+        The value of the function $F_j$ evaluated at $\\theta_j$.
     """
     dj = d - j
     numoverden = double_factorial_ratio_scipy(dj - 2, dj - 1)
@@ -84,23 +84,24 @@ def get_e2_coeffs(
     d: int, F_j: Callable, key: ArrayLike = random.PRNGKey(0)
 ) -> Tuple[Array, Array]:
     """
-    Generate the vector e_2 in R^d.
+    Generate the coefficients of the vector $e_2$.
 
     Parameters
     ----------
     d : int
         Dimension of the space.
     F_j : Callable
-        Function to compute F_j for the given angle, dimension and index.
+        Function to compute $F_j$ for the given angle, dimension and index.
     key : ArrayLike, optional
         Random key for reproducibility, by default random.PRNGKey(0).
 
     Returns
     -------
     Tuple[Array, Array]
         A tuple containing:
-        - theta: Array of angles used to construct e_2.
-        - e2: Array representing the vector e_2 in R^d.
+
+        - theta: Array of angles used to construct $e_2$.
+        - e2: Array representing the coefficients of the vector $e_2$.
     """
     theta: Array = jnp.zeros(d - 1)
 

src/isotropic/orthonormal.py

@@ -9,10 +9,10 @@
 
 def get_orthonormal_basis(Phi: ArrayLike) -> Array:
     """
-    Construct an orthonormal basis given a point Φ on a unit sphere.
+    Construct an orthonormal basis given a point $\\Phi$ on a unit sphere.
 
-    The point Φ is given by a d+1 dimensional vector and the orthonormal basis consists of d vectors
-    each of dimension d+1, which are orthogonal to Φ and to each other.
+    The point $\\Phi$ is given by a d+1 dimensional vector and the orthonormal basis consists of d vectors
+    each of dimension d+1, which are orthogonal to $\\Phi$ and to each other.
 
     Parameters
     ----------

src/isotropic/thetazero.py

@@ -12,22 +12,22 @@
 
 def get_theta_zero(x: ArrayLike, g: Callable) -> float:
     """
-    Calculate the inverse angle theta_0 with a normal distribution given a value x.
+    Calculate the inverse angle $\\theta_0$ with a normal distribution given a value x.
 
-    This function finds the angle theta_0 such that the integral of g from 0 to theta_0 equals x.
+    This function finds the angle $\\theta_0$ such that the integral of g from 0 to $\\theta_0$ equals x.
     It uses Simpson's rule for numerical integration and a bisection method to find the root.
 
     Parameters
     ----------
     x : ArrayLike
-        Value for which to find the inverse, should be uniformly distributed in [0, 1].
+        Value for which to find the inverse, should be uniformly distributed in $[0, 1]$.
     g : Callable
-        Function g(theta) that is integrated to calculate F(theta).
+        Function $g(\\theta)$ that is integrated to calculate $F(\\theta)$.
 
     Returns
     -------
     float
-        Value of theta_0.
+        Value of $\\theta_0$.
     """
 
     # We wrap the function g into a callable F that integrates g from 0 to theta.

src/isotropic/utils/bisection.py

@@ -1,4 +1,4 @@
-"""This module contains functions for the bisection algorithm to calculate F inverse"""
+"""This module contains functions for the bisection algorithm to calculate $F^{-1}$"""
 
 from typing import Callable
 
@@ -9,23 +9,13 @@ def get_theta(
     F: Callable, a: float, b: float, x: float | ArrayLike, eps: float
 ) -> float:
     """
-    Finds the value of theta such that F(theta) = x using the bisection method.
-    This function assumes that F is an increasing function in the interval [a, b]
-    and that F(a) ≤ x ≤ F(b).
+    Finds the value of theta such that $F(\\theta) = x$ using the bisection method.
+    This function assumes that $F$ is an increasing function in the interval $[a, b]$
+    and that $F(a) \\leq x \\leq F(b)$.
 
     The bisection method is a root-finding method that repeatedly bisects an interval
     and then selects a subinterval in which a root exists.
 
-    Bisection algorithm:
-    Input: Function F increasing in the interval [a, b], a value x such that
-    F(a) ≤ x ≤ F(b) and the error bound ε.
-    Output: Value of theta such that |theta - theta^*| < ε where theta^* is the solution
-    (F(theta^*) = x).
-    1. Step 1: Calculate the midpoint of the interval c = a + b.
-    2. Step 2: If F(c)≤ x, update a=c and if F(c)>x update b=c.
-    5. Step 3: If b-a<ε, return theta_0 =c, and finish.
-    6. Step 4: Repeat the process again.
-
     Parameters
     ----------
     F : Callable
@@ -42,7 +32,7 @@ def get_theta(
     Returns
     -------
     float
-        The value of theta such that F(theta) = x.
+        The value of $theta$ such that $F(\\theta) = x$.
     """
     while b - a > eps:
         c = (a + b) / 2.0

src/isotropic/utils/distribution.py

@@ -77,7 +77,7 @@ def double_factorial_ratio_jax(num: int, den: int) -> Array:
 
     Notes
     -----
-    For very large numbers, this is numerically stable only when |num - den| is ~5.
+    For very large numbers, this is numerically stable only when |num - den| ~5.
     """
     warnings.warn(
         "This is an experimental implementation. There are known issues with using this for numbers larger than 2**8",
@@ -115,7 +115,7 @@ def normal_integrand(theta: float, d: int, sigma: float) -> Array:
     Computes the function g(θ) that is integrated to calculate F(θ) which is the
     distribution function for the angle θ in a normal distribution:
 
-        g(θ) = [(d-1)!! * (1-σ²) * sin^(d-1)(θ)] / [π * (d-2)!! * (1+σ²-2σcos(θ))^((d+1)/2)]
+    $$g(\\theta) = \\frac{(d-1)!! \\times (1-\\sigma^2) \\times \\sin^{d-1}(\\theta)}{\\pi \\times (d-2)!! \\times (1+\\sigma^2-2\\sigma\\cos(\\theta))^{(d+1)/2}}$$
 
     Parameters
     ----------

src/isotropic/utils/linalg.py

@@ -6,7 +6,7 @@
 
 def jax_null_space(A: ArrayLike) -> Array:
     """
-    Compute the null space of a matrix A using JAX.
+    Compute the null space of a matrix $A$ using JAX.
 
     Parameters
     ----------

src/isotropic/utils/state_transforms.py

@@ -15,17 +15,17 @@
 
 def statevector_to_hypersphere(Phi: Array) -> Array:
     """
-    Generate the hypersphere Phi from statevector Psi
+    Generate the hypersphere from statevector $\\Phi$
 
     Parameters
     ----------
-    psi: ArrayLike
-        statevector as a complex JAX array of dimension 2^n, for n-qubits
+    Phi: ArrayLike
+        statevector as a complex JAX array of dimension $2^n$, for n-qubits
 
     Returns
     -------
     Array
-        hypersphere as a real JAX array of dimension 2^{n+1}
+        hypersphere as a real JAX array of dimension $2^{n+1}$
     """
     S = jnp.zeros(int(2 ** (log(Phi.shape[0], 2) + 1)), dtype=float)
     for x in range(S.shape[0] // 2):
@@ -36,17 +36,17 @@ def statevector_to_hypersphere(Phi: Array) -> Array:
 
 def hypersphere_to_statevector(S: Array) -> Array:
     """
-    Generate the statevector Psi from hypersphere Phi
+    Generate the statevector $\\Phi$ from hypersphere $S$
 
     Parameters
     ----------
     S: ArrayLike
-        hypersphere as a real JAX array of dimension 2^{n+1} for n qubits
+        hypersphere as a real JAX array of dimension $2^{n+1}$ for n qubits
 
     Returns
     -------
     Array
-        statevector as a complex JAX array of dimension 2^n
+        statevector as a complex JAX array of dimension $2^n$
     """
     Phi = jnp.zeros(int(2 ** (log(S.shape[0], 2) - 1)), dtype=complex)
     for x in range(Phi.shape[0]):
@@ -56,16 +56,16 @@ def hypersphere_to_statevector(S: Array) -> Array:
 
 def add_isotropic_error(Phi_sp: Array, e2: Array, theta_zero: float) -> Array:
     """
-    Add isotropic error to state Phi given e2 and theta_zero
+    Add isotropic error to state $\\Phi$ given $e_2$ and $\\theta_0$
 
     Parameters
     ----------
     Phi_sp : ArrayLike
         state to which isotropic error is added (in spherical form)
     e2 : ArrayLike
-        vector e2 in S_{d-1} with uniform distribution
+        vector $e_2$ in $S_{d-1}$ with uniform distribution
     theta_zero : float
-        angle θ_0 in [0,π] with density function f(θ_0)
+        angle $\\theta_0$ in $[0,\\pi]$ with density function $f(\\theta_0)$
 
     Returns
     -------
@@ -87,7 +87,7 @@ def generate_and_add_isotropic_error(
     Parameters
     ----------
     Phi : ArrayLike
-        The input statevector as a complex JAX array of dimension 2^n, for n-qubits.
+        The input statevector as a complex JAX array of dimension $2^n$, for n-qubits.
     sigma : float, optional
         The standard deviation for the isotropic error, by default 0.9.
     key : ArrayLike, optional