summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorVolker Hoffmann <volker@cheleb.net>2013-04-04 12:04:08 +0200
committerVolker Hoffmann <volker@cheleb.net>2013-04-04 12:04:08 +0200
commit121e5925004b8ead9fb0ebc45eacd23f180a2444 (patch)
treedd9a64e37d93f8342a47e042ba89e3f877054579
initial commit
-rw-r--r--.gitignore3
-rw-r--r--Disk.py296
-rw-r--r--demo_plot.py5
-rw-r--r--demo_reduce.py5
-rw-r--r--helpers.py167
5 files changed, 476 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..ca621a2
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+*.npz
+*.pyc
+
diff --git a/Disk.py b/Disk.py
new file mode 100644
index 0000000..67f3641
--- /dev/null
+++ b/Disk.py
@@ -0,0 +1,296 @@
+"""
+Disk Post-Processing Classes.
+
+Loads Raw Ramses Data. Processes F(X,Y,Z) -> F(X,Y) -> F(R).
+Loads/Saves Reduced Data to Numpy File.
+Plots Reduced F(X,Y), F(R) Data.
+
+@todo - Plot w/ global figure handle.
+@todo - Mass flow.
+@todo - Status messages.
+@todo - RZ-Integrate (rho)
+@todo - Save plots.
+
+Volker Hoffmann <volker@cheleb.net>
+04 April 2013
+"""
+
+from pymses import RamsesOutput
+from pymses.analysis import sample_points
+from helpers import mkpoints_xyz, prof1d
+from scipy.integrate import simps
+import numpy as np
+from math import pi
+import matplotlib.pyplot as plt
+
+class DiskBase():
+ """Declarations and Initialization Methods."""
+
+ def __init__(self, iout=1):
+ # Declarations
+ self.pxy = None # Coords for Sampling, (0,1)
+ self.pxyz = None
+ self.xy = {"x": None, "y": None} # Wrapped Coords
+ self.xyz = {"x": None, "y": None, "z": None} # (0,boxlen)
+ self.xy0 = {"x": None, "y": None} # (-boxlen/2, boxlen/2)
+ self.xyz0 = {"x": None, "y": None}
+ self.rtz = {"r": None, "theta": None, "z": None}
+ self.rt = {"r": None, "theta": None}
+ self.info = None
+ self.Q_xy = None
+ self.Q_r = None
+ self.P_xy = None
+ self.P_r = None
+ self.rho_xy = None
+ self.rho_r = None
+ self.vcart_xyz = {"vx": None, "vy": None, "vz": None}
+ self.vcart_xy = {"vx": None, "vy": None, "vz": None}
+ self.vcyl_xyz = {"vr": None, "vtheta": None, "vz": None}
+ self.vcyl_xy = {"vr": None, "vtheta": None, "vz": None}
+ self.dset = None
+ self.dl = {"dx": None, "dy": None, "dz": None}
+ self.rbins_edges = None
+ self.rbins_centers = None
+ # Output
+ self.iout = iout
+ # Sampling Volume, Options
+ self.nx = 64; self.ny = 64; self.nz = 128; self.nrbins = 64
+ self.center = 0.5; self.radius = 0.4; self.thickness = 0.2
+ # self.rbin_min = 0; self.rbin_max = 5
+
+ def init_cartesian_coords(self):
+ self.xyz["x"] = self.pxyz[:,0] * self.info["boxlen"]
+ self.xyz["y"] = self.pxyz[:,1] * self.info["boxlen"]
+ self.xyz["z"] = self.pxyz[:,2] * self.info["boxlen"]
+ self.xy["x"] = self.pxy[:,0] * self.info["boxlen"]
+ self.xy["y"] = self.pxy[:,1] * self.info["boxlen"]
+ self.xyz0["x"] = (self.pxyz[:,0] - 0.5) * self.info["boxlen"]
+ self.xyz0["y"] = (self.pxyz[:,1] - 0.5) * self.info["boxlen"]
+ self.xyz0["z"] = (self.pxyz[:,2] - 0.5) * self.info["boxlen"]
+ self.xy0["x"] = (self.pxy[:,0] - 0.5) * self.info["boxlen"]
+ self.xy0["y"] = (self.pxy[:,1] - 0.5) * self.info["boxlen"]
+
+ def init_cylindrical_coords(self):
+ self.rt["r"] = np.sqrt(self.xy0["x"]**2. + self.xy0["y"]**2.)
+ self.rt["theta"] = np.arctan2(self.xy0["y"], self.xy0["x"])
+ self.rtz["r"] = np.sqrt(self.xyz0["x"]**2. + self.xyz0["y"]**2.)
+ self.rtz["theta"] = np.arctan2(self.xyz0["y"], self.xyz0["x"])
+ self.rtz["z"] = self.xyz0["z"]
+
+ def init_velocities(self):
+ self.vcart_xyz["vx"] = self.dset["vel"][:,0]
+ self.vcart_xyz["vy"] = self.dset["vel"][:,1]
+ self.vcart_xyz["vz"] = self.dset["vel"][:,2]
+
+ def init_rbins(self):
+ self.rbin_min = 0.
+ self.rbin_max = np.max(self.xy0["x"])
+ self.rbins_edges = np.linspace(self.rbin_min, self.rbin_max,\
+ self.nrbins + 1)
+ self.rbins_centers = (self.rbins_edges[1:] + self.rbins_edges[:-1])/2.
+
+class DiskIo(DiskBase):
+ """Load/Save Methods."""
+
+ def load_ramses(self):
+ """Load **Raw** Data."""
+ # Sampling Points
+ self.pxyz, self.pxy, _, _, _, self.dl = \
+ mkpoints_xyz(self.center,
+ self.radius,
+ self.thickness,
+ self.nx, self.ny, self.nz)
+ # Read Data
+ output = RamsesOutput(".", self.iout)
+ source = output.amr_source(["rho", "vel", "P"])
+ self.dset = sample_points(source, self.pxyz)
+ # Populate Object Fields
+ self.info = output.info
+ for key in self.dl: self.dl[key] = self.dl[key] * self.info["boxlen"]
+ self.init_cartesian_coords()
+ self.init_cylindrical_coords()
+ self.init_velocities()
+ self.init_rbins()
+ self.convert_velocities()
+
+ def load_npz(self):
+ """Load **Reduced** Data.."""
+ npz = np.load("Disk_%05d.npz" % self.iout)
+ self.nx = npz["nx"]; self.ny = npz["ny"]; self.nz = npz["nz"]
+ self.xy = npz["xy"][()]; self.xy0 = npz["xy0"][()];
+ self.rt = npz["rt"][()]
+ self.rbins_edges = npz["rbins_edges"]
+ self.rbins_centers = npz["rbins_centers"]
+ self.rbin_min = npz["rbin_min"]; self.rbin_max = npz["rbin_max"]
+ self.vcart_xy = npz["vcart_xy"][()]; self.vcyl_xy = npz["vcyl_xy"][()]
+ self.rho_xy = npz["rho_xy"]; self.rho_r = npz["rho_r"]
+ self.Q_xy = npz["Q_xy"]; self.Q_r = npz["Q_r"]
+ self.info = npz["info"]
+
+ def save_npz(self):
+ """Save **Reduced** Data."""
+ np.savez("Disk_%05d.npz" % self.iout, \
+ nx = self.nx, ny = self.ny, nz = self.nz, \
+ xy = self.xy, xy0 = self.xy0, rt = self.rt, \
+ rbins_edges = self.rbins_edges, \
+ rbins_centers = self.rbins_centers, \
+ rbin_min = self.rbin_min, rbin_max = self.rbin_max, \
+ vcart_xy = self.vcart_xy, vcyl_xy = self.vcyl_xy, \
+ rho_xy = self.rho_xy, rho_r = self.rho_r, \
+ Q_xy = self.Q_xy, Q_r = self.Q_r, \
+ info = self.info )
+
+class DiskReduceBase(DiskIo):
+ """General Reduction Methods."""
+
+ def integrate(self, data_xyz):
+ """Numerical Z-Integrator."""
+ """@todo - Weight function support."""
+ data_xy = np.zeros(self.nx * self.ny)
+ idx_lo = 0
+ for ii in range(self.pxy.shape[0]):
+ idx_hi = idx_lo + self.nz
+ data_xy[ii] = simps(data_xyz[idx_lo:idx_hi], dx=self.dl["dz"])
+ idx_lo = idx_hi
+ return data_xy
+
+ def average(self, data_xyz, weights_xyz=None):
+ """Z-Averaging. Supports Weight Function."""
+ if weights_xyz == None:
+ weights_xyz = self.weights(np.ones(data_xyz.shape))
+ data_xy = np.zeros(self.nx * self.ny)
+ idx_lo = 0
+ for ii in range(self.pxy.shape[0]):
+ idx_hi = idx_lo + self.nz
+ data_xy[ii] = np.sum(data_xyz[idx_lo:idx_hi] * \
+ weights_xyz[idx_lo:idx_hi])
+ idx_lo = idx_hi
+ return data_xy
+
+ def weights(self, weights_xyz):
+ """Create Weight Function."""
+ weights_xyz = weights_xyz.reshape(self.nx, self.ny, self.nz)
+ weights_sum = np.sum(weights_xyz, axis=2)
+ weights_sum = weights_sum[:,:,None]
+ weights_sum = weights_sum * np.ones([self.nx, self.ny, self.nz])
+ weights_xyz = weights_xyz / weights_sum
+ weights_xyz = weights_xyz.reshape(self.nz * self.ny * self.nx)
+ return weights_xyz
+
+class DiskReduce(DiskReduceBase):
+ """Specific Reduction Methods."""
+
+ def integrate_to_rho_xy(self):
+ """Z-Integrate Volume Density. Gives Surface Density."""
+ self.rho_xy = self.integrate(self.dset["rho"])
+
+ def average_to_rho_r(self):
+ """Theta-Averaged Surface Density."""
+ self.rho_r = prof1d(self.rt["r"], self.rho_xy, self.rbins_edges)
+
+ def integrate_to_P_xy(self):
+ """Z-Integrate 3D Pressure Density. Gives 2D Pressure."""
+ self.P_xy = self.integrate(self.dset["P"])
+
+ def convert_velocities(self):
+ """Cylindrical Coordinate Components of Velocities."""
+ phi = np.arctan2(self.vcart_xyz["vy"], self.vcart_xyz["vx"])
+ alpha = self.rtz["theta"] - phi
+ vxy = np.sqrt(self.vcart_xyz["vx"]**2. + self.vcart_xyz["vy"]**2.)
+ self.vcyl_xyz["vr"] = vxy * np.cos(alpha)
+ self.vcyl_xyz["vtheta"] = vxy * np.sin(alpha)
+ self.vcyl_xyz["vz"] = self.vcart_xyz["vz"]
+
+ def average_velocities(self):
+ """Z-Average Velocity Components. Density Weighted."""
+ weights_xyz = self.weights(self.dset["rho"])
+ self.vcart_xy["vx"] = self.average(self.vcart_xyz["vx"], weights_xyz)
+ self.vcart_xy["vy"] = self.average(self.vcart_xyz["vy"], weights_xyz)
+ self.vcart_xy["vz"] = self.average(self.vcart_xyz["vz"], weights_xyz)
+ self.vcyl_xy["vr"] = self.average(self.vcyl_xyz["vr"], weights_xyz)
+ self.vcyl_xy["vtheta"] = self.average(self.vcyl_xyz["vtheta"], weights_xyz)
+ self.vcyl_xy["vz"] = self.average(self.vcyl_xyz["vz"], weights_xyz)
+
+ def compute_MdotR_xy(self):
+ """@todo Integrated mass flow (rtz->rt->r). See notes."""
+ pass
+
+ def compute_Q_xy(self):
+ """Toomre-Q."""
+ G = 1.0
+ Omega_xy = self.vcyl_xy["vtheta"] / self.rt["r"]
+ cs_xy = np.sqrt(self.P_xy / self.rho_xy)
+ self.Q_xy = cs_xy * Omega_xy / pi / G / self.rho_xy
+
+ def average_to_Q_r(self):
+ """Theta-Averaged Toomre-Q."""
+ self.Q_r = prof1d(self.rt["r"], self.Q_xy, self.rbins_edges)
+
+ def reduce_all(self):
+ self.average_velocities()
+ self.integrate_to_P_xy()
+ self.integrate_to_rho_xy()
+ self.compute_Q_xy()
+ self.average_to_rho_r()
+ self.average_to_Q_r()
+
+class DiskPlots(DiskReduce):
+ """Plotting Routines."""
+ def plot_rho_xy(self):
+ ext = [np.min(self.xy0["x"]), np.max(self.xy0["x"]), \
+ np.min(self.xy0["y"]), np.max(self.xy0["y"])]
+ fig = plt.figure()
+ ax = fig.add_subplot(1,1,1)
+ im = ax.imshow(np.log10(self.rho_xy.reshape(self.nx, self.ny)), \
+ extent=ext)
+ ax.set_xlabel('X [AU]')
+ ax.set_ylabel('Y [AU]')
+ ax.set_title('Surface Density [Mstar/AU^2]')
+ plt.colorbar(im)
+ plt.show()
+
+ def plot_rho_r(self):
+ # Cut Out Empty Bins
+ rho_r = self.rho_r
+ rbins_centers = self.rbins_centers
+ rho_r[rho_r==0.] = np.nan
+ rbins_centers = rbins_centers[~np.isnan(rho_r)]
+ rho_r = rho_r[~np.isnan(rho_r)]
+ # Plot
+ fig = plt.figure()
+ ax = fig.add_subplot(1,1,1)
+ ax.plot(rbins_centers, rho_r, 'bs-')
+ ax.set_yscale('log')
+ ax.grid()
+ plt.show()
+
+ def plot_Q_xy(self):
+ Q_xy = self.Q_xy
+ Q_xy[Q_xy>10] = np.nan
+ Q_xy[Q_xy<=0] = np.nan
+ fig = plt.figure()
+ ax = fig.add_subplot(1,1,1)
+ im = ax.imshow(Q_xy.reshape(self.nx, self.ny))
+ plt.colorbar(im)
+ plt.show()
+
+ def plot_Q_r(self):
+ # Cut Out Empty Bins and Useless Values
+ Q_r = self.Q_r
+ rbins_centers = self.rbins_centers
+ Q_r[Q_r>10] = np.nan
+ Q_r[Q_r<=0] = np.nan
+ rbins_centers = rbins_centers[~np.isnan(Q_r)]
+ Q_r = Q_r[~np.isnan(Q_r)]
+ # Plot
+ fig = plt.figure()
+ ax = fig.add_subplot(1,1,1)
+ ax.plot(rbins_centers, Q_r, 'bs-')
+ ax.grid()
+ ax.set_xlim([self.rbin_min, self.rbin_max])
+ ax.set_ylim([0, 10])
+ plt.show()
+
+class Disk(DiskPlots):
+ """Wrapper."""
+ pass
diff --git a/demo_plot.py b/demo_plot.py
new file mode 100644
index 0000000..4026b02
--- /dev/null
+++ b/demo_plot.py
@@ -0,0 +1,5 @@
+from Disk import Disk
+
+disk = Disk(1)
+disk.load_npz()
+disk.plot_rho_xy() \ No newline at end of file
diff --git a/demo_reduce.py b/demo_reduce.py
new file mode 100644
index 0000000..0c5f2c6
--- /dev/null
+++ b/demo_reduce.py
@@ -0,0 +1,5 @@
+from Disk import Disk
+
+disk = Disk(1)
+disk.load_ramses()
+disk.reduce_all()
diff --git a/helpers.py b/helpers.py
new file mode 100644
index 0000000..4a6e208
--- /dev/null
+++ b/helpers.py
@@ -0,0 +1,167 @@
+"""
+Various Helper Functions.
+
+Copied from Viz2. Fat trimmed.
+
+Volker Hoffmann <volker@cheleb.net>
+04 April 2013
+"""
+
+from math import pi
+# from time import gmtime, strftime
+import numpy as np
+import os
+
+def prof1d(x, data, xbins):
+ """Compute 1D Profile.
+
+ Initially, xbins contains bin edges.
+ After binning, we compute bin centres and return those in xbins.
+ """
+
+ # Flat Arrays?
+ if x.ndim > 1 or data.ndim > 1 or xbins.ndim > 1:
+ raise Exception("Non-Flat Input Arrays for Binning. Terminating.")
+
+ # Allocate Memory
+ bval = np.zeros(xbins.shape[0])
+ bcount = np.zeros_like(bval)
+
+ # Binning
+ # print "(%s UTC) Binning 1D Profile.." % strftime("%H:%M:%S", gmtime())
+ for idata in range(data.shape[0]):
+ for ibin in range(xbins.shape[0] - 1):
+ if x[idata] > xbins[ibin] and x[idata] <= xbins[ibin + 1]:
+ bval[ibin] += data[idata]
+ bcount[ibin] += 1
+ # print "(%s UTC) Done Binning." % strftime("%H:%M:%S", gmtime())
+
+ # Drop Last Entry
+ bval = bval[:-1]
+ bcount = bcount[:-1]
+
+ # Transpose Array
+ bval = bval.T
+ bcount = bcount.T
+
+ # Normalize
+ bcount[bcount == 0] = 1
+ bval /= bcount
+
+ return bval
+
+def mkvec(x, y, z):
+ """Takes Three Vectors.
+ Builds Three Vectors Combining All Possible Combinations.
+
+ Example In:
+ - x = [ 0 1 ]
+ - y = [ 2 3 ]
+ - z = [ 4 5 ]
+
+ Example Out:
+ - xx = [ 0 1 0 1 0 1 0 1 ]
+ - yy = [ 2 2 3 3 2 2 3 3 ]
+ - zz = [ 4 4 4 4 5 5 5 5 ]
+ """
+
+ # Length Counters
+ xlen = len(x)
+ ylen = len(y)
+ zlen = len(z)
+
+ xylen = xlen * ylen
+ xyzlen = xylen * zlen
+
+ # Allocate Output Arrays
+ xx = np.zeros(xyzlen)
+ yy = np.zeros_like(xx)
+ zz = np.zeros_like(xx)
+
+ # Run Through Arrays
+ xcount = 0
+ ycount = 0
+ zcount = 0
+ for ii in range(xyzlen):
+
+ # Assign Values
+ xx[ii] = x[xcount]
+ yy[ii] = y[ycount]
+ zz[ii] = z[zcount]
+
+ # Increment or Reset Counters
+ if xcount == xlen - 1:
+ xcount = 0
+ if ycount == ylen - 1:
+ ycount = 0
+ zcount += 1
+ else:
+ ycount += 1
+ else:
+ xcount += 1
+
+ return xx, yy, zz
+
+def mkpoints_xyz(center=0.5, radius=0.4, thickness=0.2, nx=128, ny=128, nz=64):
+ """Generate Sampling Points. Sampling Defined in (x,y,z) Coordinates."""
+
+ # Compute Characters
+ cchar = str(center).translate(None, ".")
+ rchar = str(radius).translate(None, ".")
+ tchar = str(thickness).translate(None, ".")
+ nxchar = str(nx)
+ nychar = str(ny)
+ nzchar = str(nz)
+
+ # Build File Name, File Location
+ fname = "points_c%s_r%s_t%s_nx%s_ny%s_nz%s.npz" % \
+ ( cchar, rchar, tchar, nxchar, nychar, nzchar )
+ floc = "%s/%s" % \
+ ( os.path.dirname(os.path.abspath(__file__)), fname )
+
+ # File Exist?
+ if os.path.isfile(floc):
+
+ # Load Points
+ npz = np.load(floc)
+ points_xyz = npz["points_xyz"]
+ points_xy = npz["points_xy"]
+ x = npz["x"]
+ y = npz["y"]
+ z = npz["z"]
+ dl = npz["dl"]
+ dl = dl[()] # Recover dict from 0d-array
+ # http://stackoverflow.com/questions/8361561/recover-dict-from-0-d-numpy-array
+ npz.close()
+
+ else:
+
+ # Generate Points
+ x, dx = np.linspace(center - radius, center + radius, nx, retstep=True)
+ y, dy = np.linspace(center - radius, center + radius, ny, retstep=True)
+ z, dz = np.linspace(center - thickness/2., center + thickness/2., nz, retstep=True)
+
+ # Point Spacing
+ dl = {"dx": dx, "dy": dy, "dz": dz}
+
+ # Reshape
+ points_xyz = mkpoints_pymses(x, y, z)
+ points_xy = mkpoints_pymses(x, y, np.array([center]))
+
+ # Save
+ np.savez(floc, \
+ points_xyz=points_xyz, points_xy=points_xy, \
+ x=x, y=y, z=z, dl=dl)
+
+ return points_xyz, points_xy, x, y, z, dl
+
+def mkpoints_pymses(x, y, z):
+ """Generate Points Array for Pymses."""
+
+ zz, yy, xx = mkvec(z, y, x)
+ points = np.zeros((xx.shape[0], 3))
+
+ for ii in range(xx.shape[0]):
+ points[ii,:] = np.array([xx[ii], yy[ii], zz[ii]])
+
+ return points