Merge pull request #46 from antvis/feat/add-one-hot

feat: add one-hot data preprocessing

Author: Yanyan-Wang

CHANGELOG.md

@@ -1,4 +1,9 @@
 # ChangeLog
+
+#### 0.1.18
+
+- feat: add one-hot data preprocessing
+
 #### 0.1.17
 
 - feat: add consine-similarity algorithm and nodes-consine-similarity algorithm;

packages/graph/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@antv/algorithm",
-  "version": "0.1.17",
+  "version": "0.1.18",
   "description": "graph algorithm",
   "keywords": [
     "graph",

packages/graph/src/louvain.ts

@@ -2,7 +2,8 @@ import { clone } from '@antv/util';
 import getAdjMatrix from './adjacent-matrix';
 import { NodeConfig, ClusterData, GraphData, ClusterMap } from './types';
 import Vector from './utils/vector';
-import { getPropertyWeight } from './utils/node-properties';
+import { getAllProperties } from './utils/node-properties';
+import { oneHot } from './utils/data-preprocessing';
 
 const getModularity = (
   nodes: NodeConfig[],
@@ -118,8 +119,10 @@ const louvain = (
         node.properties.nodeType = nodeTypeInfo.findIndex(nodeType => nodeType === node.nodeType);
       })
     }
-    // 所有节点属性特征向量集合
-    allPropertiesWeight = getPropertyWeight(nodes);
+    // 所有节点属性集合
+    const properties = getAllProperties(nodes);
+    // 所有节点属性one-hot特征向量集合
+    allPropertiesWeight = oneHot(properties);
   }
 
   let uniqueId = 1;

packages/graph/src/nodes-cosine-similarity.ts

@@ -1,23 +1,30 @@
 import { clone } from '@antv/util';
 import { NodeConfig } from './types';
-import { getPropertyWeight } from './utils/node-properties';
+import { getAllProperties } from './utils/node-properties';
+import { oneHot } from './utils/data-preprocessing';
 import cosineSimilarity from './cosine-similarity';
 /**
  *  nodes-cosine-similarity算法 基于节点属性计算余弦相似度(基于种子节点寻找相似节点)
  * @param nodes 图节点数据
  * @param seedNode 种子节点
+ * @param involvedKeys 参与计算的key集合
+ * @param uninvolvedKeys 不参与计算的key集合
  */
 const nodesCosineSimilarity = (
   nodes: NodeConfig[] = [],
   seedNode: NodeConfig,
+  involvedKeys: string[] = [],
+  uninvolvedKeys: string[] = [],
 ): {
   allCosineSimilarity: number[],
   similarNodes: NodeConfig[],
 } => {
   const similarNodes = clone(nodes.filter(node => node.id !== seedNode.id));
   const seedNodeIndex = nodes.findIndex(node => node.id === seedNode.id);
-  // 所有节点属性特征向量集合
-  const allPropertiesWeight = getPropertyWeight(nodes);
+  // 所有节点属性集合
+  const properties = getAllProperties(nodes);
+  // 所有节点属性one-hot特征向量集合s
+  const allPropertiesWeight = oneHot(properties, involvedKeys, uninvolvedKeys);
   // 种子节点属性
   const seedNodeProperties = allPropertiesWeight[seedNodeIndex];
 

packages/graph/src/types.ts

@@ -49,6 +49,9 @@ export interface DegreeType {
   }
 }
 
+export interface PlainObject {
+  [key: string]: any;
+}
 
 export interface IAlgorithm {
   getAdjMatrix: (graphData: GraphData, directed?: boolean) => Matrix[],

packages/graph/src/utils/data-preprocessing.ts

@@ -0,0 +1,74 @@
+import { isEmpty, uniq } from '@antv/util';
+import { PlainObject } from '../types';
+
+/**
+ * 获取数据中所有的属性及其对应的值
+ * @param dataList 数据集
+ * @param involvedKeys 参与计算的key集合
+ * @param uninvolvedKeys 不参与计算的key集合
+ */
+export const getAllKeyValueMap = (dataList: PlainObject[], involvedKeys?: string[], uninvolvedKeys?: string[]) => {
+  let keys = [];
+  // 指定了参与计算的keys时，使用指定的keys
+  if (involvedKeys?.length) {
+    keys = involvedKeys;
+  } else {
+    // 未指定抽取的keys时，提取数据中所有的key
+    dataList.forEach(data => {
+      keys = keys.concat(Object.keys(data));
+    })
+    keys = uniq(keys);
+  }
+  // 获取所有值非空的key的value数组
+  const allKeyValueMap = {};
+  keys.forEach(key => {
+    let value = [];
+    dataList.forEach(data => {
+      if (data[key] !== undefined && data[key] !== '') {
+        value.push(data[key]);
+      }
+    })
+    if (value.length && !uninvolvedKeys?.includes(key)) {
+      allKeyValueMap[key] = uniq(value);
+    }
+  })
+
+  return allKeyValueMap;
+}
+
+/**
+ * one-hot编码：数据特征提取
+ * @param dataList 数据集
+ * @param involvedKeys 参与计算的的key集合
+ * @param uninvolvedKeys 不参与计算的key集合
+ */
+export const oneHot = (dataList: PlainObject[], involvedKeys?: string[], uninvolvedKeys?: string[]) => {
+  // 获取数据中所有的属性及其对应的值
+  const allKeyValueMap = getAllKeyValueMap(dataList, involvedKeys, uninvolvedKeys);
+  const oneHotCode = [];
+  // 对数据进行one-hot编码
+  dataList.forEach((data, index) => {
+    let code = [];
+    Object.keys(allKeyValueMap).forEach(key => {
+      const keyValue = data[key];
+      const allKeyValue = allKeyValueMap[key];
+      const valueIndex = allKeyValue.findIndex(value => keyValue === value);
+      let subCode = [];
+      for(let i = 0; i < allKeyValue.length; i++) {
+        if (i === valueIndex) {
+          subCode.push(1);
+        } else {
+          subCode.push(0);
+        }
+      }
+      code = code.concat(subCode);
+    })
+    oneHotCode[index] = code;
+  })
+  return oneHotCode;
+}
+
+export default {
+  getAllKeyValueMap,
+  oneHot,
+}

packages/graph/src/utils/node-properties.ts

@@ -56,7 +56,20 @@ export const getPropertyWeight = (nodes: NodeConfig[]) => {
   return allPropertiesWeight;
 }
 
+// 获取所有节点的属性集合
+export const getAllProperties = (nodes, key='properties') => {
+  const allProperties = [];
+  nodes.forEach(node => {
+    if (!node.properties) {
+      return;
+    }
+    allProperties.push(node[key]);
+  })
+  return allProperties;
+}
+
 export default {
   getAllSortProperties,
   getPropertyWeight,
-}
+  getAllProperties
+}

packages/graph/tests/unit/nodesCosineSimilarity-spec.ts

@@ -74,4 +74,33 @@ describe('nodesCosineSimilarity normal demo', () => {
       expect(data).toBeLessThanOrEqual(1);
     })
   });
+
+
+  it('demo use involvedKeys: ', () => {
+    const involvedKeys = ['amount', 'city'];
+    const { nodes } = propertiesGraphData;
+    const { allCosineSimilarity, similarNodes } = nodesCosineSimilarity(nodes as NodeConfig[], nodes[16], involvedKeys);
+    expect(allCosineSimilarity.length).toBe(16);
+    expect(similarNodes.length).toBe(16);
+    allCosineSimilarity.forEach(data => {
+      expect(data).toBeGreaterThanOrEqual(0);
+      expect(data).toBeLessThanOrEqual(1);
+    })
+    expect(Number(Math.max.apply(null, allCosineSimilarity).toString().match(/^\d+(?:\.\d{0,2})?/))).toBe(0.99);
+    expect(similarNodes[0].id).toBe('node-11');
+  });
+
+  it('demo use uninvolvedKeys: ', () => {
+    const uninvolvedKeys = ['amount'];
+    const { nodes } = propertiesGraphData;
+    const { allCosineSimilarity, similarNodes } = nodesCosineSimilarity(nodes as NodeConfig[], nodes[16], [], uninvolvedKeys);
+    expect(allCosineSimilarity.length).toBe(16);
+    expect(similarNodes.length).toBe(16);
+    allCosineSimilarity.forEach(data => {
+      expect(data).toBeGreaterThanOrEqual(0);
+      expect(data).toBeLessThanOrEqual(1);
+    })
+    expect(Number(Math.max.apply(null, allCosineSimilarity).toString().match(/^\d+(?:\.\d{0,2})?/))).toBe(0.66);
+    expect(similarNodes[0].id).toBe('node-11');
+  });
 });