From b3dc10f2870aafed048cdb5bb83f99a8a3e877d4 Mon Sep 17 00:00:00 2001
From: efischer <ej@eugenefischer.com>
Date: Sat, 26 Feb 2022 08:15:48 -0600
Subject: [PATCH] add graph attributes to graphml writer

---
 src/main/java/GraphMLFileWriter.java | 30 ++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/src/main/java/GraphMLFileWriter.java b/src/main/java/GraphMLFileWriter.java
index b767a8f..f357828 100644
--- a/src/main/java/GraphMLFileWriter.java
+++ b/src/main/java/GraphMLFileWriter.java
@@ -18,7 +18,18 @@ public class GraphMLFileWriter {
 
     String filename;
     SimpleWeightedGraph graph;
+    GraphWithMapData data;
+    Map<String, Attribute> graphAttributes = new HashMap<>();
 
+    public GraphMLFileWriter(String filename, GraphWithMapData data) {
+        if(!filename.matches(".*\\.graphml")){
+            filename = filename + ".graphml";
+        }
+        this.filename = filename;
+        this.data = data;
+        this.graph = data.getGraph();
+        createGraphAttributes();
+    }
 
     public GraphMLFileWriter(String filename, SimpleWeightedGraph<Vertex, DefaultWeightedEdge> graph) {
         if(!filename.matches(".*\\.graphml")){
@@ -28,6 +39,23 @@ public class GraphMLFileWriter {
         this.graph = graph;
     }
 
+    private void createGraphAttributes(){
+        //Sample plate filename
+        graphAttributes.put("sample plate filename", DefaultAttribute.createAttribute(data.getSourceFilename()));
+        // Number of wells
+        graphAttributes.put("well count", DefaultAttribute.createAttribute(data.getNumWells().toString()));
+        //Well populations
+        Integer[] wellPopulations = data.getWellPopulations();
+        StringBuilder populationsStringBuilder = new StringBuilder();
+        populationsStringBuilder.append(wellPopulations[0].toString());
+        for(int i = 1; i < wellPopulations.length; i++){
+            populationsStringBuilder.append(", ");
+            populationsStringBuilder.append(wellPopulations[i].toString());
+        }
+        String wellPopulationsString = populationsStringBuilder.toString();
+        graphAttributes.put("well populations", DefaultAttribute.createAttribute(wellPopulationsString));
+    }
+
     public void writeGraphToFile() {
         try(BufferedWriter writer = Files.newBufferedWriter(Path.of(filename), StandardOpenOption.CREATE_NEW);
         ){
@@ -35,6 +63,8 @@ public class GraphMLFileWriter {
             GraphMLExporter<Vertex, SimpleWeightedGraph<Vertex, DefaultWeightedEdge>> exporter = new GraphMLExporter<>(Vertex::getVertexLabel);
             //set to export weights
             exporter.setExportEdgeWeights(true);
+            //Set graph attributes
+            exporter.setGraphAttributeProvider( () -> graphAttributes);
             //set type, sequence, and occupancy attributes for each vertex
             exporter.setVertexAttributeProvider( v -> {
                 Map<String, Attribute> attributes = new HashMap<>();