conditional mean code using binning

binning_v_given_c.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import sys
+import argparse
+import numpy as np
+
+
+# Commandline argument parser
+parser = argparse.ArgumentParser()
+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"]
+
+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 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_u_given_c_{:03d}".format(xidx), **arr_dict)