clusters: Accept 1D numpy array and calculate values for all gamma's

Signed-off-by: Aaron LI <aly@aaronly.me>

2 files changed, 27 insertions, 21 deletions

diff --git a/fg21sim/extragalactic/clusters/halo.py b/fg21sim/extragalactic/clusters/halo.py
index 75789b6..ce4bcaf 100644
--- a/fg21sim/extragalactic/clusters/halo.py
+++ b/fg21sim/extragalactic/clusters/halo.py
@@ -182,7 +182,7 @@ class RadioHalo:
 
         Returns
         -------
-        electron_spec : `~numpy.ndarray`
+        electron_spec : float 1D `~numpy.ndarray`
             The solved electron spectrum at ``zend``.
             Unit: [cm^-3]
         """
@@ -196,8 +196,7 @@ class RadioHalo:
             tstop = COSMO.age(zend)
         if n0_e is None:
             # Accumulated constantly injected electrons until ``tstart``.
-            n_inj = np.array([self.fp_injection(gm)
-                              for gm in self.gamma])
+            n_inj = self.fp_injection(self.gamma)
             n0_e = n_inj * tstart
 
         self.electron_spec = self.fpsolver.solve(u0=n0_e, tstart=tstart,
@@ -218,8 +217,8 @@ class RadioHalo:
 
         Parameters
         ----------
-        gamma : float
-            Lorentz factor of electrons
+        gamma : float, or float 1D `~numpy.ndarray`
+            Lorentz factors of electrons
         t : None
             Currently a constant injection rate is assumed, therefore
             this parameter is not used.  Keep it for the consistency
@@ -227,8 +226,8 @@ class RadioHalo:
 
         Returns
         -------
-        Qe : float
-            Current electron injection rate at specified energy (gamma).
+        Qe : float, or float 1D `~numpy.ndarray`
+            Current electron injection rate at specified energies (gamma).
             Unit: [cm^-3 Gyr^-1]
 
         References
@@ -251,16 +250,16 @@ class RadioHalo:
 
         Parameters
         ----------
-        gamma : float
-            The Lorentz factor of electrons
+        gamma : float, or float 1D `~numpy.ndarray`
+            The Lorentz factors of electrons
         t : float
             Current (cosmic) time when solving the equation
             Unit: [Gyr]
 
         Returns
         -------
-        diffusion : float
-            Diffusion coefficient
+        diffusion : float, or float 1D `~numpy.ndarray`
+            Diffusion coefficients
             Unit: [Gyr^-1]
 
         References
@@ -283,8 +282,8 @@ class RadioHalo:
 
         Returns
         -------
-        advection : float
-            Advection coefficient, describing the energy loss/gain rate.
+        advection : float, or float 1D `~numpy.ndarray`
+            Advection coefficients, describing the energy loss/gain rates.
             Unit: [Gyr^-1]
         """
         advection = (abs(self._loss_ion(gamma, t)) +
@@ -394,16 +393,16 @@ class RadioHalo:
 
         Parameters
         ----------
-        gamma : float
-            The Lorentz factor of electrons
+        gamma : float, or float 1D `~numpy.ndarray`
+            The Lorentz factors of electrons
         t : float
             The cosmic time/age
             Unit: [Gyr]
 
         Returns
         -------
-        loss : float
-            The energy loss rate
+        loss : float, or float 1D `~numpy.ndarray`
+            The energy loss rates
             Unit: [Gyr^-1]
 
         References
diff --git a/fg21sim/extragalactic/clusters/solver.py b/fg21sim/extragalactic/clusters/solver.py
index dea7f3f..d7660b0 100644
--- a/fg21sim/extragalactic/clusters/solver.py
+++ b/fg21sim/extragalactic/clusters/solver.py
@@ -104,6 +104,13 @@ class FokkerPlanckSolver:
 
     NOTE
     ----
+    All above functions should accept two parameters: ``(x, t)``,
+    where ``x`` is an 1D float `~numpy.ndarray` representing the adopted
+    logarithmic grid points along the spatial/energy axis, ``t`` is the
+    time of each solving step.
+
+    NOTE
+    ----
     The diffusion coefficients (i.e., calculated by ``f_diffusion()``)
     should be *positive* (i.e., C(x) > 0), otherwise unstable or wrong
     results may occur, due to the current numerical scheme/algorithm
@@ -281,9 +288,9 @@ class FokkerPlanckSolver:
         dx_phalf = self.dx_phalf
         dx_mhalf = self.dx_mhalf
         dt = self.tstep
-        B = np.array([self.f_advection(x_, tc) for x_ in x])
-        C = np.array([self.f_diffusion(x_, tc) for x_ in x])
-        Q = np.array([self.f_injection(x_, tc) for x_ in x])
+        B = self.f_advection(x, tc)
+        C = self.f_diffusion(x, tc)
+        Q = self.f_injection(x, tc)
         #
         B_phalf = self.X_phalf(B)
         B_mhalf = self.X_mhalf(B)
@@ -307,7 +314,7 @@ class FokkerPlanckSolver:
         b[-1] = 1 + (dt/dx[-1]) * (C_mhalf[-1]/dx_mhalf[-1])*Wplus_mhalf[-1]
         # Escape from the system
         if self.f_escape is not None:
-            T = np.array([self.f_escape(x_, tc) for x_ in x])
+            T = self.f_escape(x, tc)
             b += dt / T
         # Right-hand side
         r = dt * Q + uc