From 22639905b054de61ee2c43920dda8bb7d6c61358 Mon Sep 17 00:00:00 2001
From: Frederic Perez <frederic.perez@polytechnique.edu>
Date: Sun, 27 Oct 2024 15:21:29 +0100
Subject: [PATCH] slightly faster collisional ionization

---
 benchmarks/collisions/ionization_multiple.py |  2 +-
 doc/Sphinx/Understand/collisions.rst         | 25 ++++++++++++++++---
 src/Collisions/CollisionalIonization.cpp     | 26 ++++++++++----------
 src/Collisions/CollisionalIonization.h       |  2 +-
 4 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/benchmarks/collisions/ionization_multiple.py b/benchmarks/collisions/ionization_multiple.py
index 66b2f75bf..6313a72dd 100755
--- a/benchmarks/collisions/ionization_multiple.py
+++ b/benchmarks/collisions/ionization_multiple.py
@@ -22,7 +22,7 @@
 		
 		timestep2 = int(np.double(S2.namelist.Main.timestep))
 		D += [
-			S2.ParticleBinning(0,sum={"ekin":[0,1]}, linestyle="", marker=".", color=color )
+			S2.ParticleBinning(0,sum={"ekin":[0,1]}, linestyle="", marker=".", color=color, label=elm )
 		]
 
 
diff --git a/doc/Sphinx/Understand/collisions.rst b/doc/Sphinx/Understand/collisions.rst
index 281ff0764..dc89d3e8c 100755
--- a/doc/Sphinx/Understand/collisions.rst
+++ b/doc/Sphinx/Understand/collisions.rst
@@ -309,21 +309,38 @@ the ion:
   \sum\limits_{p=0}^{k-1} R^{i+k}_{i+p} \left(\bar{P}^{i+k} - \bar{P}^{i+p}\right)
   \prod\limits_{j=0,j\ne p}^{k-1} R^{i+p}_{i+j}
   &
-  \quad\mathrm{if}\quad 0<k<k_\mathrm{max}
+  \quad\mathrm{if}\quad 0<k<Z-Z^\star
+  \end{array}
+  \right.
+
+..
   \\
   \sum\limits_{p=0}^{k-1} \left[ 1+R^{i+k}_{i+p}\left(\frac{W_{i+k}}{W_{i+p}}\bar{P}^{i+p} - \bar{P}^{i+k}\right) \right]
   \prod\limits_{j=0,j\ne p}^{k-1} R^{i+p}_{i+j}
   &
   \quad\mathrm{if}\quad k=k_\mathrm{max}
-  \end{array}
-  \right.
 
-where :math:`k_\mathrm{max} = Z-Z^\star`.
+
+To simplify the calculation of :math:`P^i_k` (in particular the second case in the
+equation above) we use the following equivalent expression:
+
+.. math::
+  
+  P^i_k = A_{k-1} \sum\limits_{p=0}^{k-1}  \left(\bar{P}^{i+k} - \bar{P}^{i+p}\right)
+  / B_k^p
+  \quad\mathrm{if}\quad 0<k<Z-Z^\star
+
+where :math:`A_k = \prod\limits_{j=0}^{k} W_{i+j}` and
+:math:`B_k^p = \prod\limits_{j=0,j\ne p}^{k} (W_{i+j}-W_{i+p})`.
+These two quantities can be computed recursively for each :math:`k`.
+
 
 The cumulative probability :math:`F^i_k = \sum_{j=0}^{k} P^i_j` provides an efficient
 way to pick when the ionization stops: we pick a random number :math:`U\in [0,1]` and
 loop from :math:`k=0` to :math:`k_\mathrm{max}`. We stop ionizing when :math:`F^i_k>U`.
 
+
+
 ----
 
 Test cases for ionization
diff --git a/src/Collisions/CollisionalIonization.cpp b/src/Collisions/CollisionalIonization.cpp
index 96d0489b1..3f32af64d 100755
--- a/src/Collisions/CollisionalIonization.cpp
+++ b/src/Collisions/CollisionalIonization.cpp
@@ -23,7 +23,7 @@ CollisionalIonization::CollisionalIonization( int Z, Params *params, int ionizat
 {
     atomic_number = Z;
     rate .resize( Z );
-    irate.resize( Z );
+    rate_product.resize( Z );
     prob .resize( Z );
     ionization_electrons_ = ionization_electrons;
     if( params ) {
@@ -39,7 +39,7 @@ CollisionalIonization::CollisionalIonization( CollisionalIonization *CI )
 {
     atomic_number = CI->atomic_number;
     rate .resize( atomic_number );
-    irate.resize( atomic_number );
+    rate_product.resize( atomic_number );
     prob .resize( atomic_number );
     ionization_electrons_ = CI->ionization_electrons_;
     new_electrons.initialize( 0, CI->new_electrons );
@@ -177,7 +177,7 @@ void CollisionalIonization::calculate( double gamma_s, double gammae, double gam
     // Loop for multiple ionization
     // k+1 is the number of ionizations
     const int kmax = atomic_number-Zstar-1;
-    double cs, w, e, cum_prob = 0;
+    double cs, w, e, cum_prob = 0, A = 1.;
     for( int k = 0; k <= kmax;  k++ ) {
         // Calculate the location x (~log of energy) in the databases
         const double x = a2*log( a1*( gamma_s-1. ) );
@@ -203,24 +203,24 @@ void CollisionalIonization::calculate( double gamma_s, double gammae, double gam
         }
         
         rate[k] = K*cs/gammae  ; // k-th ionization rate
-        irate[k] = 1./rate[k]  ; // k-th ionization inverse rate
         prob[k] = exp( -rate[k] ); // k-th ionization probability
+        rate_product[k] = 1.;
         
         // Calculate the cumulative probability for k-th ionization (Nuter et al, 2011)
         if( k==0 ) {
-            cum_prob = prob[k];
+            cum_prob = prob[k]; // cumulative probability
         } else {
+            double sum = 0.;
             for( int p=0; p<k; p++ ) {
-                double cp = 1. - rate[k]*irate[p];
-                for( int j=0  ; j<p; j++ ) {
-                    cp *= 1.-rate[p]*irate[j];
-                }
-                for( int j=p+1; j<k; j++ ) {
-                    cp *= 1.-rate[p]*irate[j];
-                }
-                cum_prob += ( prob[k]-prob[p] )/cp;
+                const double d = rate[k] - rate[p];
+                rate_product[p] *= d;
+                rate_product[k] *= -d;
+                sum += ( prob[p]-prob[k] ) / rate_product[p];
             }
+            sum *= A;
+            cum_prob += sum;
         }
+        A *= rate[k];
         
         // If no more ionization, leave
         if( U1 < cum_prob ) {
diff --git a/src/Collisions/CollisionalIonization.h b/src/Collisions/CollisionalIonization.h
index 06a9d72b0..90bb90fd0 100755
--- a/src/Collisions/CollisionalIonization.h
+++ b/src/Collisions/CollisionalIonization.h
@@ -73,7 +73,7 @@ class CollisionalIonization : public BinaryProcess
     
     //! Current ionization rate array (one cell per number of ionization events)
     std::vector<double> rate;
-    std::vector<double> irate;
+    std::vector<double> rate_product;
     //! Current ionization probability array (one cell per number of ionization events)
     std::vector<double> prob;