modified tree.py and pde.py

codes/PDE_find.py

@@ -28,10 +28,11 @@ def FiniteDiff(u, dx):
     ux = np.zeros(n)
 
     for i in range(1, n - 1):
-        ux[i] = (u[i + 1] - u[i - 1]) / (2 * dx)
+        ux[i] = (u[i + 1] - u[i - 1]) / (2 * dx) # centered difference
 
-    ux[0] = (-3.0 / 2 * u[0] + 2 * u[1] - u[2] / 2) / dx
-    ux[n - 1] = (3.0 / 2 * u[n - 1] - 2 * u[n - 2] + u[n - 3] / 2) / dx
+    ux[0] = (-3.0 / 2 * u[0] + 2 * u[1] - u[2] / 2) / dx # forward difference
+    ux[n - 1] = (3.0 / 2 * u[n - 1] - 2 * u[n - 2] + u[n - 3] / 2) / dx # backward difference
+
     return ux
 
 
@@ -41,16 +42,17 @@ def FiniteDiff2(u, dx):
     ux = np.zeros(n)
 
     for i in range(1, n - 1):
-        ux[i] = (u[i + 1] - 2 * u[i] + u[i - 1]) / dx ** 2
+        ux[i] = (u[i + 1] - 2 * u[i] + u[i - 1]) / dx ** 2 # centered difference
 
-    ux[0] = (2 * u[0] - 5 * u[1] + 4 * u[2] - u[3]) / dx ** 2
-    ux[n - 1] = (2 * u[n - 1] - 5 * u[n - 2] + 4 * u[n - 3] - u[n - 4]) / dx ** 2
+    ux[0] = (2 * u[0] - 5 * u[1] + 4 * u[2] - u[3]) / dx ** 2 # forward difference
+    ux[n - 1] = (2 * u[n - 1] - 5 * u[n - 2] + 4 * u[n - 3] - u[n - 4]) / dx ** 2 # backward difference
+
     return ux
 
 
 def Diff(u, dxt, name):
     """
-    Here dx is a scalar, name is a str indicating what it is
+    Here dxt is a scalar, name is a str indicating what it is
     """
 
     n, m = u.shape
@@ -72,7 +74,7 @@ def Diff(u, dxt, name):
 
 def Diff2(u, dxt, name):
     """
-    Here dx is a scalar, name is a str indicating what it is
+    Here dxt is a scalar, name is a str indicating what it is
     """
 
     n, m = u.shape

codes/pde.py

@@ -8,62 +8,58 @@
 
 class PDE:
 
-    def __init__(self, depth, max_width, p_var):
+    def __init__(self, depth, max_num_trees, p_var):
         self.depth = depth
         self.p_var = p_var
-        self.W = np.random.randint(max_width)+1  # 1 -- width
-        self.elements = []
-        for i in range(0, self.W):
-            # 产生W个tree，也就是W个项
-            one_tree = Tree(depth, p_var)
-            # while 'd u t' in tree.preorder:# 没用，挡不住如(sin x + u) d t；不如直接看mse，太小就扔掉
-            #     tree = Tree(depth, p_var)
-            self.elements.append(one_tree)
+        self.num_trees = np.random.randint(max_num_trees+1)
+        self.tree_list = []
+        for i in range(self.num_trees):
+            self.tree_list.append(Tree(depth, p_var))
 
     def mutate(self, p_mute):
-        for i in range(0, self.W):  # 0 -- W-1
-            self.elements[i].mutate(p_mute)
-
-    def replace(self): # 直接产生一个新的tree，替换pde中的一项
-        # print('replace!')
-        one_tree = Tree(self.depth, self.p_var)
-        ix = np.random.randint(self.W)  # 0 -- W-1
-        if len(self.elements) == 0:
+        for i in range(self.num_trees):  
+            self.tree_list[i].mutate(p_mute)
+
+    def replace(self):
+        if len(self.tree_list) == 0:
             NotImplementedError('replace error')
-        self.elements[ix] = one_tree
+
+        new_tree = Tree(self.depth, self.p_var)
+        idx = np.random.randint(self.num_trees)  
+        self.tree_list[idx] = new_tree
 
-    def visualize(self): # 写出SGA产生的项的形式，包含产生的所有项，未去除系数小的项。
+    def visualize(self): 
         name = ''
-        for i in range(len(self.elements)):
+        for i in range(self.num_trees):
             if i != 0:
                 name += '+'
-            name += self.elements[i].inorder
+            name += self.tree_list[i].inorder
         return name
 
-    def concise_visualize(self): # 写出所有项的形式，包含固定候选集和SGA，且包含系数。会区分是来自于固定候选集的还是来自于SGA生成的候选集的。如果是来自于SGA生成的候选集，需要用inorder来写出可理解的项。
+    def concise_visualize(self): 
         name = ''
-        elements = copy.deepcopy(self.elements)
-        elements, coefficients = evaluate_mse(elements, True)
+        tree_list = copy.deepcopy(self.tree_list)
+        tree_list, coefficients = evaluate_mse(tree_list, True)
         coefficients = coefficients[:, 0]
-        # print(len(elements), len(coefficients))
+        # print(len(tree_list), len(coefficients))
         for i in range(len(coefficients)):
-            if np.abs(coefficients[i]) < 1e-4: # 忽略过于小的系数
+            if np.abs(coefficients[i]) < 1e-4: 
                 continue
             if i != 0 and name != '':
                 name += ' + '
             name += str(round(np.real(coefficients[i]), 4))
-            if i < num_default: # num_default中为一定包含的候选集
+            if i < num_default: 
                 name += default_names[i]
             else:
-                name += elements[i-num_default].inorder # element是SGA生成的候选集
+                name += tree_list[i-num_default].inorder 
         return name
 
 
-def evaluate_mse(a_pde, is_term=False):
-    if is_term:
-        terms = a_pde
+def evaluate_mse(pde, is_list=False):
+    if is_list:
+        terms = pde
     else:
-        terms = a_pde.elements
+        terms = pde.tree_list
     terms_values = np.zeros((u.shape[0] * u.shape[1], len(terms)))
     delete_ix = []
     for ix, term in enumerate(terms):
@@ -132,12 +128,11 @@ def evaluate_mse(a_pde, is_term=False):
         if is_term:
             terms.pop(ixx - move)
         else:
-            a_pde.elements.pop(ixx-move)
-            a_pde.W -= 1  # 实际宽度减一
+            pde.tree_list.pop(ixx-move)
+            pde.num_trees -= 1 
         terms_values = np.delete(terms_values, ixx-move, axis=1)
-        move += 1  # pop以后index左移
+        move += 1  
 
-    # 检查是否存在inf或者nan，或者terms_values是否被削没了
     if False in np.isfinite(terms_values) or terms_values.shape[1] == 0:
         # pdb.set_trace()
         error = np.inf

codes/tree.py

@@ -34,7 +34,6 @@ class Node:
     """
 
     def __init__(self, depth, idx, parent_idx, name, child_num, var, full):
-
         self.depth = depth
         self.idx = idx
         self.parent_idx = parent_idx
@@ -64,24 +63,25 @@ def __init__(self, max_depth, p_var):
 
         self.max_depth = max_depth
         self.p_var = p_var
-        self.tree = [[] for i in range(max_depth)]
+        self.node_list = [[] for i in range(max_depth)]
         self.inorder = None
 
         self.add_root_node()
 
         for depth in range(1,self.max_depth):
-            for parent_idx in range(len(self.tree[depth - 1])): 
-                parent = self.tree[depth - 1][parent_idx] 
+            for parent_idx in range(len(self.node_list[depth - 1])):
+
+                parent = self.node_list[depth - 1][parent_idx] 
                 self.add_nodes(depth, parent)
 
-        model_tree = copy.deepcopy(self.tree)
+        model_tree = copy.deepcopy(self.node_list)
         self.inorder = tree2str_merge(model_tree)
 
     def add_root_node(self):
         root = ROOT[np.random.randint(0, len(ROOT))] 
         node = Node(depth = 0, idx = 0, parent_idx = None, name = root[0], 
                     child_num = int(root[1]), var = root[2], full = root)
-        self.tree[0].append(node)
+        self.node_list[0].append(node)
 
     def add_nodes(self, depth, parent):
         if parent.child_num == 0: 
@@ -90,52 +90,52 @@ def add_nodes(self, depth, parent):
             # rule 1: right child of a derivative operator must be an independent variable
             if parent.name in ['d', 'd^2'] and j == 1: 
                 node = DENOM[np.random.randint(0, len(DENOM))] 
-                node = Node(depth = depth, idx = len(self.tree[depth]), parent_idx = parent.idx, name = node[0],
+                node = Node(depth = depth, idx = len(self.node_list[depth]), parent_idx = parent.idx, name = node[0],
                             child_num = int(node[1]), var = node[2], full = node)
-                self.tree[depth].append(node)
+                self.node_list[depth].append(node)
             # rule 2: leaf nodes must be variables/operands (rather than operators)
             elif depth == self.max_depth - 1:
                 node = VARS[np.random.randint(0, len(VARS))]
-                node = Node(depth = depth, idx = len(self.tree[depth]), parent_idx = parent.idx, name = node[0],
+                node = Node(depth = depth, idx = len(self.node_list[depth]), parent_idx = parent.idx, name = node[0],
                             child_num = int(node[1]), var = node[2], full = node)
-                self.tree[depth].append(node)
+                self.node_list[depth].append(node)
             else:
             # rule 3: if rules 1 and 2 do not apply, make the next node a variable/operand with probability p_var
             # (and an operator with probability 1-p_var)
                 if np.random.random() <= self.p_var:
                     node = VARS[np.random.randint(0, len(VARS))]
-                    node = Node(depth = depth, idx = len(self.tree[depth]), parent_idx = parent.idx, name = node[0],
+                    node = Node(depth = depth, idx = len(self.node_list[depth]), parent_idx = parent.idx, name = node[0],
                                 child_num = int(node[1]), var = node[2], full = node)
-                    self.tree[depth].append(node)
+                    self.node_list[depth].append(node)
                 else:
                     node = OPS[np.random.randint(0, len(OPS))]
-                    node = Node(depth = depth, idx = len(self.tree[depth]), parent_idx = parent.idx, name = node[0],
+                    node = Node(depth = depth, idx = len(self.node_list[depth]), parent_idx = parent.idx, name = node[0],
                                 child_num = int(node[1]), var = node[2], full = node)
-                    self.tree[depth].append(node)
+                    self.node_list[depth].append(node)
 
     def get_child_idx(self, node):
         child_idx = 0
         for i in range(node.idx):
-            child_idx = child_idx + self.tree[node.depth][i].child_num
+            child_idx = child_idx + self.node_list[node.depth][i].child_num
         return child_idx
 
     def get_right_child_idx(self, node):
         return self.get_child_idx(node) + 1
 
     def mutate(self, p_mute): 
         global see_tree
-        see_tree = copy.deepcopy(self.tree)
+        see_tree = copy.deepcopy(self.node_list)
         for depth in range(1,self.max_depth):
-            for parent_idx in range(len(self.tree[depth - 1])):
-                parent = self.tree[depth - 1][parent_idx]
+            for parent_idx in range(len(self.node_list[depth - 1])):
+                parent = self.node_list[depth - 1][parent_idx]
                 if parent.child_num == 0:
                     continue
                 for j in range(parent.child_num):  
                     mute = np.random.choice([True, False], p=([p_mute, 1-p_mute]))
                     # rule 1: 
                     if mute == False:
                         continue
-                    current = self.tree[depth][self.get_child_idx(parent) + j]
+                    current = self.node_list[depth][self.get_child_idx(parent) + j]
                     temp = current.name
                     num_child = current.child_num  
                     # print('mutate!')
@@ -147,25 +147,25 @@ def mutate(self, p_mute):
                             node = VARS[np.random.randint(0, len(VARS))] 
                         new_node = Node(depth=depth, idx=self.get_child_idx(parent) + j, parent_idx=parent_idx, name=node[0],
                                     child_num=int(node[1]), var=node[2], full=node)
-                        self.tree[depth][self.get_child_idx(parent) + j] = new_node 
+                        self.node_list[depth][self.get_child_idx(parent) + j] = new_node 
                     else: 
                         if num_child == 1:
                             node = OP1[np.random.randint(0, len(OP1))]
                             while node[0] == temp:  
                                 node = OP1[np.random.randint(0, len(OP1))]
                         elif num_child == 2:
                             node = OP2[np.random.randint(0, len(OP2))]
-                            right = self.tree[depth + 1][self.get_right_child_idx(current)].name
+                            right = self.node_list[depth + 1][self.get_right_child_idx(current)].name
                             while node[0] == temp or (node[0] in ['d', 'd^2'] and right not in DENOM[:, 0]): # rule 4
                                 node = OP2[np.random.randint(0, len(OP2))]
                         else:
                             raise NotImplementedError("Error occurs!")
 
                         new_node = Node(depth=depth, idx=self.get_child_idx(parent) + j, parent_idx=parent_idx, name=node[0],
                                     child_num=int(node[1]), var=node[2], full=node)
-                        self.tree[depth][self.get_child_idx(parent) + j] = new_node
+                        self.node_list[depth][self.get_child_idx(parent) + j] = new_node
 
-        model_tree = copy.deepcopy(self.tree)
+        model_tree = copy.deepcopy(self.node_list)
         self.inorder = tree2str_merge(model_tree)