Add calc_potential.py: split out from "calc_mass_potential.py"

1 files changed, 234 insertions, 0 deletions

diff --git a/calc_potential.py b/calc_potential.py
new file mode 100755
index 0000000..f2e3b66
--- /dev/null
+++ b/calc_potential.py
@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+#
+# Aaron LI
+# Created: 2016-07-13
+# Updated: 2016-07-13
+#
+# Change logs:
+#
+
+"""
+Calculate the gravitational potential profile through the
+total mass density profile.
+
+Newton's theorems (ref.[1], Sec. 2.2.1):
+1. A body that is inside a spherical shell of matter experiences no net
+   gravitational force from that shell.
+2. The gravitational force on a body that lies outside a spherical shell
+   of matter is the same as it would be if all the shell's matter were
+   concentrated into a point at its center.
+
+Therefore, the gravitational potential produced by a spherical shell of
+mass 'M' is:
+    phi = (1) - G * M / R;  (r <= R, i.e., inside the shell)
+          (2) - G * M / r;  (r > R, i.e., outside the shell)
+
+The total gravitational potential may be considered to be the sum of the
+potentials of spherical shells of mass
+    dM(r) = 4 * pi * rho(r) * r^2 dr,
+so we may calculate the gravitational potential at 'R' generated by an
+arbitrary spherically symmetric density distribution 'rho(r)' by adding
+the contributions to the potential produced by shells
+    (1) with r < R,
+and
+    (2) with r > R.
+In this way, we obtain
+    phi(R) = - (G/R) * \int_0^R dM(r) - G * \int_R^{\inf} dM(r)/r
+           = - 4*pi*G * ((1/R) * \int_0^R r^2 * rho(r) dr +
+             \int_R^{\inf} r * rho(r) dr)
+
+References:
+[1] Tremaine & Binney, Galactic Dynamics, 2nd edition, 2008
+"""
+
+
+import argparse
+
+import numpy as np
+import astropy.units as au
+import astropy.constants as ac
+import scipy.integrate as integrate
+from configobj import ConfigObj
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+from matplotlib.figure import Figure
+
+from spline import SmoothSpline
+
+plt.style.use("ggplot")
+
+
+config_default = """
+## Configuration for `calc_potential.py`
+
+# total mass density profile
+rho_total_profile = rho_total_profile.txt
+
+# output gravitational potential profile
+potential_profile = potential_profile.txt
+potential_profile_image = potential_profile.png
+"""
+
+
+class DensityProfile:
+    """
+    Cluster's gas/total mass density profile.
+    """
+    # supported types of density profile
+    DENSITY_TYPES = ["total", "gas"]
+    # available splines
+    SPLINES = ["rho_total"]
+    # input density data: [r, r_err, m]
+    r = None
+    r_err = None
+    d = None
+    # redshift of the object
+    redshift = None
+    # fitted smoothing spline to total density profile
+    rho_total_spline = None
+    # calculated potential profile
+    potential = None
+
+    def __init__(self, data, density_type="total"):
+        self.load_data(data=data, density_type=density_type)
+
+    def load_data(self, data, density_type="total"):
+        if density_type not in self.DENSITY_TYPES:
+            raise ValueError("invalid density_types: %s" % density_type)
+        # 3-column density profile: r[kpc], r_err[kpc], rho[g/cm^3]
+        self.r = data[:, 0].copy()
+        self.r_err = data[:, 1].copy()
+        self.d = data[:, 2].copy()
+        self.density_type = density_type
+
+    def calc_potential(self, verbose=True):
+        """
+        Calculate the gravitational potential profile from the
+        total mass density profile.
+
+        NOTE:
+        * The integral in the potential formula can only be integrated
+          to the largest radius of availability.
+        * This limitation will affect the absolute values of the calculated
+          potentials, but not the shape of the potential profile.
+        """
+        def _int_inner(x):
+            return x**2 * self.eval_spline(spline="rho_total", x=x)
+
+        def _int_outer(x):
+            return x * self.eval_spline(spline="rho_total", x=x)
+
+        if self.density_type != "total":
+            raise ValueError("total mass density profile is required")
+
+        if self.rho_total_spline is None:
+            self.fit_spline(spline="rho_total", log10=["x", "y"])
+        potential = np.zeros(self.r.shape)
+        if verbose:
+            print("Calculating the potential profile (#%d): ..." %
+                  len(self.r), end="", flush=True)
+        r_max = max(self.r)
+        c = - 4 * np.pi * au.kpc.to(au.cm)**2
+        for i, r in enumerate(self.r):
+            if verbose and (i+1) % 10 == 0:
+                print("%d..." % (i+1), end="", flush=True)
+            # total gravitational potential [ cm^2/s^2 ]
+            potential[i] = c * ac.G.cgs.value * (
+                (1/r) * integrate.quad(_int_inner, a=0.0, b=r,
+                                       epsrel=1e-2)[0] +
+                integrate.quad(_int_outer, a=r, b=r_max,
+                               epsrel=1e-2)[0])
+        if verbose:
+            print("DONE!", flush=True)
+        self.potential = potential
+        return potential
+
+    def plot(self, ax=None, fig=None):
+        x = self.r
+        xlabel = "3D Radius"
+        xunit = "kpc"
+        xscale = "log"
+        yscale = "log"
+        y = self.potential
+        ylabel = "Gravitational potential"
+        yunit = "cm$^2$/s$^2$"
+        yscale = "linear"
+        #
+        if ax is None:
+            fig, ax = plt.subplots(1, 1)
+        ax.plot(x, y, linewidth=2)
+        ax.set_xscale(xscale)
+        ax.set_yscale(yscale)
+        ax.set_xlim(min(x), max(x))
+        y_min, y_max = min(y), max(y)
+        y_min = y_min/1.2 if y_min > 0 else y_min*1.2
+        y_max = y_max*1.2 if y_max > 0 else y_max/1.2
+        ax.set_ylim(y_min, y_max)
+        ax.set_xlabel(r"%s (%s)" % (xlabel, xunit))
+        ax.set_ylabel(r"%s (%s)" % (ylabel, yunit))
+        fig.tight_layout()
+        return (fig, ax)
+
+    def save(self, outfile):
+        """
+        Save calculated overdensity profile.
+        """
+        data = np.column_stack([self.r,
+                                self.r_err,
+                                self.potential])
+        header = "radius[kpc]  radius_err[kpc]  potential[cm^2/s^2]"
+        np.savetxt(outfile, data, header=header)
+
+    def fit_spline(self, spline, log10=[]):
+        if spline not in self.SPLINES:
+            raise ValueError("invalid spline: %s" % spline)
+        #
+        if spline == "rho_total":
+            # input total mass density profile
+            x = self.r
+            y = self.d
+            spl = "rho_total_spline"
+        else:
+            raise ValueError("invalid spline: %s" % spline)
+        setattr(self, spl, SmoothSpline(x=x, y=y))
+        getattr(self, spl).fit(log10=log10)
+
+    def eval_spline(self, spline, x):
+        """
+        Evaluate the specified spline at the supplied positions.
+        Also check whether the spline was fitted in the log-scale space,
+        and transform the evaluated values if necessary.
+        """
+        if spline == "rho_total":
+            spl = self.rho_total_spline
+        else:
+            raise ValueError("invalid spline: %s" % spline)
+        #
+        return spl.eval(x)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+            description="Calculate the gravitational potential profile")
+    parser.add_argument("config", nargs="?",
+                        help="additional config")
+    args = parser.parse_args()
+
+    config = ConfigObj(config_default.splitlines())
+    if args.config is not None:
+        config_user = ConfigObj(open(args.config))
+        config.merge(config_user)
+
+    density_data = np.loadtxt(config["rho_total_profile"])
+    density_profile = DensityProfile(data=density_data, density_type="total")
+    density_profile.calc_potential(verbose=True)
+    density_profile.save(outfile=config["potential_profile"])
+    fig = Figure(figsize=(10, 8))
+    FigureCanvas(fig)
+    ax = fig.add_subplot(1, 1, 1)
+    density_profile.plot(ax=ax, fig=fig)
+    fig.savefig(config["potential_profile_image"], dpi=150)
+
+
+if __name__ == "__main__":
+    main()