Commits

dan mackinlay  committed 79a412d

better binning options for the wicks case

  • Participants
  • Parent commits 5eb6f36

Comments (0)

Files changed (1)

File trader_stats.py

     use a binning based on known characteristics of the isotropic case.
     see _vel_loc_self_mi_adaptive for more info"""
     if n_bins == "wicks":
-        # bin using either double Wick's criterion (radius) or, if that has
+        # bin using either Wick's criterion (radius) or, if that has
         # too many bins, Cochrane's.
         n_bins = int(min(
             2.0/traderset.radius+0.5,
         loc_binner=bin_naive,
         vel_binner=bin_angle_naive,
         estimator="ince",
+        binning='disc',
         squish_bins=True,
         **mi_kwargs):
     """does loc determine angle? We work this out in the special 2D case using
     Wick's actual method, which involves the use of trignometric 2d angles,
     and idiosyncratic bin numbers and placement.
     """
-    #import pdb; pdb.set_trace()
     #from IPython.core.debugger import Tracer
     #Tracer()()
+    n_bins = int(min(
+        2.0/traderset.radius+0.5,
+        math.sqrt(choose_n_bins(vel_slice.shape[0]*vel_slice.shape[1], test=False))
+    ))
+    
     n_bins = int(np.floor(2./(traderset.radius)))
     n_loc_bins = n_bins
     n_vel_bins = n_bins
+
+    if binning =='disc':
+        locs = loc_binner(loc_slice[:,:,0:1], n_bins)
+        vels = vel_binner(vel_angle_slice, n_bins)
+        if squish_bins:
+            locs, n_loc_bins = squish(locs)
+            vels, n_vel_bins = squish(vels)
+        if estimator=="ince":
+            estimator = lambda x,y,**kwargs: ince_mi_dist_disc(x,y,n_loc_bins, n_vel_bins,**kwargs)
+        else:
+            estimator = lambda x,y,**kwargs: plugin_mi_dist_disc(x,y,n_loc_bins, n_vel_bins,**kwargs)
+    else:
+        locs = loc_slice[:,:,0:1]
+        vels = vel_angle_slice
+        if estimator=="ince":
+            estimator = lambda x,y,**kwargs: ince_mi_dist_cont(x,y, n_bins=n_bins, **kwargs)
+        else:
+            estimator = lambda x,y,**kwargs: plugin_mi_dist_cont(x,y, n_bins=n_bins, **kwargs)
     
-    locs = loc_binner(loc_slice[:,:,0:1], n_bins)
-    vels = vel_binner(vel_angle_slice, n_bins)
-    
-    if squish_bins:
-        locs, n_loc_bins = squish(locs)
-        vels, n_vel_bins = squish(vels)
-    
-    if estimator=="ince":
-        estimator = ince_mi_dist_disc
-    else:
-        estimator = plugin_mi_dist_disc
-            
     est = estimator(
         locs, vels,
-        n_loc_bins, n_vel_bins, **mi_kwargs
+        **mi_kwargs
     )
-    return d(est=est)
+    return d(est=est, n_bins=(n_loc_bins, n_vel_bins))
     
 def mi_wicks_hi_d(*args, **kwargs):
     return _do_slicewise_stat(_mi_wicks_hi_d,
     particle_len = loc_slice.shape[1]
     dim_len = loc_slice.shape[2]
     if n_bins == "wicks":
-        # bin using either double Wick's criterion (radius) or, if that has
-        # too many bins, Cochrane's.
+        # bin using either Wick's criterion (radius) 
         n_bins = int(np.floor(2./(traderset.radius)))
     elif n_bins == "cochrane":
         # bin using the cochrane criterion, which has as many bins as possible
-        n_bins = int(np.floor(np.sqrt(choose_n_bins(time_len*oversample*time_len*oversample))))
+        n_bins = int(np.floor(math.sqrt(choose_n_bins(time_len*oversample*time_len*oversample))))
     if binning =='disc':
         if estimator=="ince":
             estimator = ince_mi_dist_disc
     
     est = estimator(
         locs.ravel(), vels.ravel(),
-        n_loc_bins, n_vel_bins, **mi_kwargs
+        x_bins=n_loc_bins, y_bins=n_vel_bins, **mi_kwargs
     )
     return d(est=est)