From 0bbd45ded1878e8c5b22a49cd3826f33312a82ae Mon Sep 17 00:00:00 2001
From: jmoo2880 <jmoo2880@uni.sydney.edu.au>
Date: Tue, 24 Dec 2024 13:00:37 +1100
Subject: [PATCH] update docstrings for BEE/SEE

---
 docs/src/tools.md       |  3 ++-
 src/Analysis/analyse.jl | 37 ++++++++++++++++++++++---------------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/docs/src/tools.md b/docs/src/tools.md
index 21ff4d0..67cd050 100644
--- a/docs/src/tools.md
+++ b/docs/src/tools.md
@@ -21,7 +21,7 @@ A vector is returned where each entry contains the BEE spectrum for the class-sp
 For example, in the case of a two class problem, we obtain the individual BEE spectrums for the class 0 MPS and the class 1 MPS. 
 For an unsupervised problem with only a single class, there is only a single BEE spectrum. 
 #### Example
-To illustrate how we might use the BEE in a typical analysis, consider an example involving real world time series from the [ItalyPowerDemand](https://www.timeseriesclassification.com/description.php?Dataset=ItalyPowerDemand) UCR dataset. 
+To illustrate how we might use the BEE in a typical analysis, consider an example involving real world time series from the [ItalyPowerDemand](https://www.timeseriesclassification.com/description.php?Dataset=ItalyPowerDemand) (IPD) UCR dataset. 
 There are two classes corresponding to the power demand during: __(i)__ the winter months; __(ii)__ the summer months. 
 For this example, we will train an MPS to classify between summer and winter time-series data:
 ```Julia
@@ -74,4 +74,5 @@ MPSTime.bipartite_spectrum
 ## Internal Methods
 ```@docs
 MPSTime.von_neumann_entropy
+MPSTime.one_site_rdm
 ```
diff --git a/src/Analysis/analyse.jl b/src/Analysis/analyse.jl
index 869c24c..b5cba3f 100644
--- a/src/Analysis/analyse.jl
+++ b/src/Analysis/analyse.jl
@@ -1,21 +1,19 @@
 """
-```Julia
-von_neumann_entropy(mps::MPS; logfn::Function=log) -> Vector{Float64}
-```
-Compute the [von Neumann entanglement entropy](https://en.wikipedia.org/wiki/Entropy_of_entanglement) for each site in a Matrix Product State (MPS).
+   single_site_spectrum(mps::TrainedMPS) -> Vector{Vector{Float64}}
+
+Compute the single-site entanglement entropy (SEE) spectrum of a trained MPS.
 
-The von Neumann entropy quantifies the entanglement at each bond of the MPS by computing the entropy of the singular value spectrum obtained from a singular value decomposition (SVD). The entropy is computed as:
+The single-site entanglement entropy (SEE) quantifies the entanglement between each site in the MPS and all other sites. It is computed as:
 
-[ S = -sum_{i} p_i log(p_i) ]
+``\\text{SEE} = -\\text{tr}(\\rho \\log(\\rho))``
 
-where ( p_i ) are the squared singular values normalized to sum to 1.
+where ``\\rho`` is the single-site reduced density matrix (RDM).
 
 # Arguments
-- `mps::MPS`: The Matrix Product State (MPS) whose entanglement entropy is to be computed.
-- `logfn::Function`: (Optional) The logarithm function to use (`log`, `log2`, or `log10`). Defaults to the natural logarithm (`log`).
+- `mps::TrainedMPS`: A trained Matrix Product State (MPS) object, which includes the MPS and associated labels.
 
 # Returns
-A vector of `Float64` values where the i-th element represents the von Neumann entropy at site i of the MPS.
+A vector of vectors, where the outer vector corresponds to each label in the expanded MPS, and the inner vectors contain the SEE values for the respective sites.
 """
 function von_neumann_entropy(mps::MPS, logfn::Function=log)
     # adapted from http://itensor.org/docs.cgi?page=formulas/entanglement_mps
@@ -51,8 +49,8 @@ bipartite_spectrum(mps::TrainedMPS; logfn::Function=log) -> Vector{Vector{Float6
 Compute the bipartite entanglement entropy (BEE) of a trained MPS across each bond.
 Given a single unlabeled MPS the BEE is defined as:
 
-∑ α^2 log(α^2)
-where α are the eigenvalues obtained from the shmidt decomposition. 
+``\\sum_i \\alpha_i^2 \\log(\\alpha_i^2)``
+where ``\\alpha_i`` are the eigenvalues obtained from the shmidt decomposition. 
     
 Compute the bipartite entanglement entropy (BEE) of a trained MPS.
 """
@@ -92,6 +90,15 @@ function rho_correct(rho::Matrix, eigentol::Float64=sqrt(eps()))
     return rho_corrected
 end
 
+"""
+```Julia
+one_site_rdm(mps::MPS, site::Int) -> Matrix
+```
+Compute the single-site reduced density matrix (RDM) of the 
+MPS at a given site. 
+If the RDM is not positive semidefinite, clamp the negative eigenvalues
+(if within the tolerance) and reconstruct the rdm.
+"""
 function one_site_rdm(mps::MPS, site::Int)
     s = siteinds(mps)
     orthogonalize!(mps, site)
@@ -120,12 +127,12 @@ Compute the single-site entanglement entropy (SEE) spectrum of a trained MPS.
 
 The single-site entanglement entropy (SEE) quantifies the entanglement at each site of the MPS. It is computed as:
 
-[ SEE = -tr(ρ ⋅ log(ρ)) ]
+``\\text{SEE} = -\\text{tr}(\\rho \\log(\\rho))``
 
-where ρ is the single-site reduced density matrix (RDM).
+where ``\\rho`` is the single-site reduced density matrix (RDM).
 
 # Arguments
-- `mps::TrainedMPS`: A trained Matrix Product State (MPS) object, which includes the MPS and associated labels.
+- `mps::TrainedMPS`: A trained Matrix Product State (MPS).
 
 # Returns
 A vector of vectors, where the outer vector corresponds to each label in the expanded MPS, and the inner vectors contain the SEE values for the respective sites.