Add cache for score function

Author: chenweiDelight

causallearn/score/LocalScoreFunctionClass.py

@@ -0,0 +1,27 @@
+import math
+from functools import lru_cache
+from typing import List, Dict, Any, Callable
+
+import pandas as pd
+from numpy import ndarray
+from causallearn.score.LocalScoreFunction import local_score_BDeu, local_score_BIC, local_score_cv_multi, local_score_marginal_multi, local_score_marginal_general, local_score_cv_general
+
+from causallearn.utils.ScoreUtils import *
+
+
+class LocalScoreClass(object):
+
+    def __init__(self, data: ndarray, local_score_fun: Callable[[ndarray, int, List[int], Any], float], parameters=None):
+        self.data = data
+        self.local_score_fun = local_score_fun
+        self.parameters = parameters
+        self.score_cache = {}
+
+    def score(self, i: int, PAi: List[int]) -> float:
+        hash_key = f'i_{str(i)}_PAi_{str(PAi)}'
+        if self.score_cache.__contains__(hash_key):
+            return self.score_cache[hash_key]
+        else:
+            res = self.local_score_fun(self.data, i, PAi, self.parameters)
+            self.score_cache[hash_key] = res
+            return res

causallearn/search/ScoreBased/GES.py

@@ -1,5 +1,5 @@
 from typing import Optional
-
+from causallearn.score.LocalScoreFunctionClass import LocalScoreClass
 from causallearn.graph.GeneralGraph import GeneralGraph
 from causallearn.graph.GraphNode import GraphNode
 from causallearn.utils.DAG2CPDAG import dag2cpdag
@@ -46,12 +46,14 @@ def ges(X: ndarray, score_func: str = 'local_score_BIC', maxP: Optional[float] =
         if maxP is None:
             maxP = X.shape[1] / 2  # maximum number of parents
         N = X.shape[1]  # number of variables
+        localScoreClass = LocalScoreClass(data=X, local_score_fun=local_score_cv_general, parameters=parameters)
 
     elif score_func == 'local_score_marginal_general':  # negative marginal likelihood based on regression in RKHS
         parameters = {}
         if maxP is None:
             maxP = X.shape[1] / 2  # maximum number of parents
         N = X.shape[1]  # number of variables
+        localScoreClass = LocalScoreClass(data=X, local_score_fun=local_score_marginal_general, parameters=parameters)
 
     elif score_func == 'local_score_CV_multi':  # k-fold negative cross validated likelihood based on regression in RKHS
         # for data with multi-variate dimensions
@@ -62,6 +64,7 @@ def ges(X: ndarray, score_func: str = 'local_score_BIC', maxP: Optional[float] =
         if maxP is None:
             maxP = len(parameters['dlabel']) / 2
         N = len(parameters['dlabel'])
+        localScoreClass = LocalScoreClass(data=X, local_score_fun=local_score_cv_multi, parameters=parameters)
 
     elif score_func == 'local_score_marginal_multi':  # negative marginal likelihood based on regression in RKHS
         # for data with multi-variate dimensions
@@ -72,19 +75,23 @@ def ges(X: ndarray, score_func: str = 'local_score_BIC', maxP: Optional[float] =
         if maxP is None:
             maxP = len(parameters['dlabel']) / 2
         N = len(parameters['dlabel'])
+        localScoreClass = LocalScoreClass(data=X, local_score_fun=local_score_marginal_multi, parameters=parameters)
 
     elif score_func == 'local_score_BIC':  # Greedy equivalence search with BIC score
         if maxP is None:
             maxP = X.shape[1] / 2
         N = X.shape[1]  # number of variables
+        localScoreClass = LocalScoreClass(data=X, local_score_fun=local_score_BIC, parameters=None)
 
     elif score_func == 'local_score_BDeu':  # Greedy equivalence search with BDeu score
         if maxP is None:
             maxP = X.shape[1] / 2
         N = X.shape[1]  # number of variables
+        localScoreClass = LocalScoreClass(data=X, local_score_fun=local_score_BDeu, parameters=None)
 
     else:
         raise Exception('Unknown function!')
+    score_func = localScoreClass
 
     node_names = [("x%d" % i) for i in range(N)]
     nodes = []

causallearn/utils/GESUtils.py

@@ -10,20 +10,7 @@
 
 
 def feval(parameters: list):
-    if parameters[0] == 'local_score_CV_general':
-        return local_score_cv_general(parameters[1], parameters[2], parameters[3], parameters[4])
-    elif parameters[0] == 'local_score_marginal_general':
-        return local_score_marginal_general(parameters[1], parameters[2], parameters[3], parameters[4])
-    elif parameters[0] == 'local_score_CV_multi':
-        return local_score_cv_multi(parameters[1], parameters[2], parameters[3], parameters[4])
-    elif parameters[0] == 'local_score_marginal_multi':
-        return local_score_marginal_multi(parameters[1], parameters[2], parameters[3], parameters[4])
-    elif parameters[0] == 'local_score_BIC':
-        return local_score_BIC(parameters[1], parameters[2], parameters[3], parameters[4])
-    elif parameters[0] == 'local_score_BDeu':
-        return local_score_BDeu(parameters[1], parameters[2], parameters[3], parameters[4])
-    else:
-        raise Exception('Undefined function')
+    return parameters[0].score(parameters[2], parameters[3])
 
 
 def kernel(x, xKern, theta):