用图机器学习探索 A 股个股相关性变化

import org.jgrapht.*;
import org.jgrapht.graph.*;
import org.jgrapht.nio.*;
import org.jgrapht.nio.dot.*;
import org.jgrapht.traverse.*;

import java.io.*;
import java.net.*;
import java.util.*;

Graph<URI, DefaultEdge> g = new DefaultDirectedGraph<>(DefaultEdge.class);

URI google = new URI("http://www.google.com");
URI wikipedia = new URI("http://www.wikipedia.org");
URI jgrapht = new URI("http://www.jgrapht.org");

// add the vertices
g.addVertex(google);
g.addVertex(wikipedia);
g.addVertex(jgrapht);

// add edges to create linking structure
g.addEdge(jgrapht, wikipedia);
g.addEdge(google, jgrapht);
g.addEdge(google, wikipedia);
g.addEdge(wikipedia, google);

// space 为待扫描的图空间名称，returnCols 为需要读取的点/边及其属性列，
// returnCols 参数格式：{tag1Name: prop1, prop2, tag2Name: prop3, prop4, prop5}
Iterator<ScanVertexResponse> scanVertex(
            String space, Map<String, List<String>> returnCols);
Iterator<ScanEdgeResponse> scanEdge(
            String space, Map<String, List<String>> returnCols);

MetaClientImpl metaClientImpl = new MetaClientImpl(metaHost, metaPort);
metaClientImpl.connect();
StorageClient storageClient = new StorageClientImpl(metaClientImpl);
Processor processor = new ScanVertexProcessor(metaClientImpl);

Iterator<ScanVertexResponse> iterator =
                storageClient.scanVertex(spaceName, returnCols);

while (iterator.hasNext()) {
  ScanVertexResponse response = iterator.next();
  if (response == null) {
    log.error("Error occurs while scan vertex");
    break;
  }
  
  Result result =  processor.process(spaceName, response);
  results.addAll(result.getRows(TAGNAME));
}

Graph<String, MyEdge> graph = GraphTypeBuilder
                .undirected()
    .weighted(true)
    .allowingMultipleEdges(true)
    .allowingSelfLoops(false)
    .vertexSupplier(SupplierUtil.createStringSupplier())
    .edgeSupplier(SupplierUtil.createSupplier(MyEdge.class))
    .buildGraph();

for (VertexDomain vertex : vertexDomainList){
    graph.addVertex(vertex.getVid().toString());
    stockIdToName.put(vertex.getVid().toString(), vertex);
}

for (EdgeDomain edgeDomain : edgeDomainList){
    graph.addEdge(edgeDomain.getSrcid().toString(), edgeDomain.getDstid().toString());
    MyEdge newEdge = graph.getEdge(edgeDomain.getSrcid().toString(), edgeDomain.getDstid().toString());
    graph.setEdgeWeight(newEdge, edgeDomain.getWeight());
}

SpanningTreeAlgorithm.SpanningTree pMST = new PrimMinimumSpanningTree(graph).getSpanningTree();

Legend.drawGraph(pMST.getEdges(), filename, stockIdToName);

public class Legend {
  
...
  
  public static void drawGraph(Set<MyEdge> edges, String filename, Map<String, VertexDomain> idVertexMap) throws IOException {
     // Creates graph with model
     mxGraph graph = new mxGraph();
     Object parent = graph.getDefaultParent();

     // set style
     graph.getModel().beginUpdate();
     mxStylesheet myStylesheet =  graph.getStylesheet();
     graph.setStylesheet(setMsStylesheet(myStylesheet));

     Map<String, Object> idMap = new HashMap<>();
     Map<String, String> industryColor = new HashMap<>();

     int colorIndex = 0;

     for (MyEdge edge : edges) {
       Object src, dst;
       if (!idMap.containsKey(edge.getSrc())) {
         VertexDomain srcNode = idVertexMap.get(edge.getSrc());
         String nodeColor;
         if (industryColor.containsKey(srcNode.getIndustry())){
           nodeColor = industryColor.get(srcNode.getIndustry());
         }else {
           nodeColor = COLOR_LIST[colorIndex++];
           industryColor.put(srcNode.getIndustry(), nodeColor);
         }
         src = graph.insertVertex(parent, null, srcNode.getName(), 0, 0, 105, 50, "fillColor=" + nodeColor);
         idMap.put(edge.getSrc(), src);
       } else {
         src = idMap.get(edge.getSrc());
       }

       if (!idMap.containsKey(edge.getDst())) {
         VertexDomain dstNode = idVertexMap.get(edge.getDst());

         String nodeColor;
         if (industryColor.containsKey(dstNode.getIndustry())){
           nodeColor = industryColor.get(dstNode.getIndustry());
         }else {
           nodeColor = COLOR_LIST[colorIndex++];
           industryColor.put(dstNode.getIndustry(), nodeColor);
         }

         dst = graph.insertVertex(parent, null, dstNode.getName(), 0, 0, 105, 50, "fillColor=" + nodeColor);
         idMap.put(edge.getDst(), dst);
       } else {
         dst = idMap.get(edge.getDst());
       }
       graph.insertEdge(parent, null, "", src, dst);
     }


     log.info("vertice " + idMap.size());
     log.info("colorsize " + industryColor.size());

     mxFastOrganicLayout layout = new mxFastOrganicLayout(graph);
     layout.setMaxIterations(2000);
     //layout.setMinDistanceLimit(10D);
     layout.execute(parent);

     graph.getModel().endUpdate();

     // Creates an image than can be saved using ImageIO
     BufferedImage image = createBufferedImage(graph, null, 1, Color.WHITE,
                                               true, null);

     // For the sake of this example we display the image in a window
     // Save as JPEG
     File file = new File(filename);
     ImageIO.write(image, "JPEG", file);

   }
  
  ...
    
}