113 lines
3 KiB
Python
113 lines
3 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
|
|
import sys
|
|
import os
|
|
import argparse
|
|
import numpy as np
|
|
|
|
|
|
program_description = '''\
|
|
Compute conditional mean of v given condition c by data binning
|
|
'''
|
|
|
|
# Commandline argument parser
|
|
parser = argparse.ArgumentParser(description=program_description)
|
|
parser.add_argument("-x", "--xindex", help="sampling x index", type=int, required=True)
|
|
parser.add_argument("-v", "--variable", help="file containing variable to calculate conditional mean", required=True)
|
|
parser.add_argument("-n", "--nbins", help="number of bins per condtion c interval", default=45, )
|
|
parser.add_argument("-b", "--boundary", help="width of boundary region for fine mesh", default=0.06, )
|
|
args = parser.parse_args()
|
|
params = vars(args)
|
|
|
|
|
|
# Parameters
|
|
xidx = int(params["xindex"])
|
|
vfilename = params["variable"]
|
|
vname = os.path.splitext(os.path.basename(vfilename))[0]
|
|
|
|
nbins = int(params["nbins"])
|
|
boundary = float(params["boundary"])
|
|
|
|
cfile = "c.dat"
|
|
ufile = vfilename # "u.dat"
|
|
|
|
|
|
|
|
print ("Computing conditional mean at {}".format(xidx))
|
|
|
|
|
|
|
|
def cmean_binning (c, v, nbins=100, boundary=0.1, quiet=False):
|
|
'''
|
|
Compute conditional mean of v given condition c by data binning
|
|
'''
|
|
c_bins = np.hstack(
|
|
(np.linspace(0, boundary, nbins),
|
|
np.linspace(boundary, 1-boundary, nbins)[1:],
|
|
np.linspace(1-boundary, 1, nbins)[1:],
|
|
)
|
|
)
|
|
c_bins[0] = -np.finfo(np.float).eps
|
|
|
|
cstar = np.zeros(c_bins.size + 1)
|
|
cstar[1:-1] = (c_bins[1:]+c_bins[:-1])/2.
|
|
cstar[-1] = 1.0
|
|
|
|
v_sums = np.zeros(c_bins.size)
|
|
v_counts = np.zeros(c_bins.size, dtype=int)
|
|
|
|
ntimes = c.shape[0]
|
|
|
|
progress_view = 0
|
|
indicator_length = 50
|
|
|
|
if not quiet: print (indicator_length * "|")
|
|
|
|
for tidx, cplane in enumerate(c):
|
|
|
|
progress = indicator_length * (tidx / ntimes)
|
|
if np.floor(progress) > progress_view:
|
|
for j in range(int(np.floor(progress) - progress_view)):
|
|
if not quiet: print("=", end='')
|
|
progress_view = np.floor(progress)
|
|
|
|
v_plane = v[tidx]
|
|
|
|
bin_indices = np.digitize(cplane, c_bins, right=True)
|
|
|
|
for i in range(1, c_bins.size):
|
|
binned = v_plane[bin_indices == i]
|
|
v_sums[i] += binned.sum()
|
|
v_counts[i] += binned.size
|
|
|
|
v_avg = np.zeros(v_sums.size+1)
|
|
v_avg[:-1] = v_sums / v_counts
|
|
|
|
return cstar, v_avg, c_bins, v_counts
|
|
|
|
|
|
# Load data
|
|
|
|
sc = np.memmap(cfile, mode="r", dtype=np.double).reshape((512,-1,256,256))
|
|
sv = np.memmap(ufile, mode="r", dtype=np.double).reshape((512,-1,256,256))
|
|
|
|
c_at_x = sc[xidx]
|
|
v_at_x = sv[xidx]
|
|
|
|
print ("Completed loading data".format(xidx))
|
|
|
|
|
|
# Compute conditional mean
|
|
|
|
cstar, cmean, bin_edges, count = cmean_binning(c_at_x, v_at_x, nbins=nbins, boundary=boundary, quiet=True)
|
|
|
|
|
|
# Save result
|
|
arr_dict = {}
|
|
arr_dict["cstar"] = np.asarray(cstar)
|
|
arr_dict["cmean"] = np.asarray(cmean)
|
|
arr_dict["bins"] = np.asarray(bin_edges)
|
|
arr_dict["count"] = np.asarray(count)
|
|
|
|
np.savez("cmean_{}_given_c_{:03d}".format(vname, xidx), **arr_dict)
|