separating i and j cases

getkeops · May 15, 2024 · 8ea200c · 8ea200c
1 parent 107e03d
commit 8ea200c
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 7 deletions.
diff --git a/keopscore/keopscore/config/config.py b/keopscore/keopscore/config/config.py
@@ -10,7 +10,8 @@
 # global parameters can be set here :
 use_cuda = True  # use cuda if possible
 use_OpenMP = True  # use OpenMP if possible (see function set_OpenMP below)
-lim_dim_local_var = 0#10000
+lim_dim_local_var_i = 0#10000
+lim_dim_local_var_j = 10000
 
 # System Path
 base_dir_path = os.path.abspath(join(os.path.dirname(os.path.realpath(__file__)), ".."))

diff --git a/keopscore/keopscore/sandbox/bench_global_var.py b/keopscore/keopscore/sandbox/bench_global_var.py
@@ -18,7 +18,8 @@ def fun_keops(A, I, J):
 
 
 def bench(Ptab, lim_dim_local_var):
-    keopscore.config.config.lim_dim_local_var = lim_dim_local_var
+    keopscore.config.config.lim_dim_local_var_i = lim_dim_local_var
+    keopscore.config.config.lim_dim_local_var_j = lim_dim_local_var
     pykeops.clean_pykeops()
     M, N = 100000, 100000
     times = []

diff --git a/keopscore/keopscore/utils/code_gen_utils.py b/keopscore/keopscore/utils/code_gen_utils.py
@@ -96,13 +96,17 @@ def __init__(self, red_formula, force_all_local):
             dims_sorted, inds_sorted, cats_sorted = zip(
                 *sorted(zip(self.dims, self.inds, self.cats))
             )
-            dimcur = 0
+            dimcur_i = 0
+            dimcur_j = 0
             for k in range(len(dims_sorted)):
-                dimcur += dims_sorted[k]
-                if dimcur < keopscore.config.config.lim_dim_local_var:
-                    self.is_local_var[inds_sorted[k]] = True
+                if cats_sorted[k] in [0,2]:
+                    dimcur_i += dims_sorted[k]
+                    if dimcur_i < keopscore.config.config.lim_dim_local_var_i:
+                        self.is_local_var[inds_sorted[k]] = True
                 else:
-                    break
+                    dimcur_j += dims_sorted[k]
+                    if dimcur_j < keopscore.config.config.lim_dim_local_var_j:
+                        self.is_local_var[inds_sorted[k]] = True
 
             cnt = [0, 0, 0]
             for k in range(len(dims_sorted)):