added saving option for alaina (#91)

Co-authored-by: Niklas Abraham - INFlux <[email protected]>

pyeed/network/network.py

@@ -1,11 +1,12 @@
 import copy
-from typing import List, Optional
+from typing import List, Optional, Any
 
 import matplotlib.pyplot as plt
 import networkx as nx
+from networkx.readwrite import json_graph
 import py4cytoscape as p4c
 from py4cytoscape import gen_node_size_map, scheme_c_number_continuous
-from pydantic import BaseModel, Field, PrivateAttr
+from pydantic import BaseModel, Field, PrivateAttr, field_serializer, field_validator
 from requests import RequestException
 
 from pyeed.align.pairwise import PairwiseAligner
@@ -56,7 +57,7 @@ class Config:
     )
 
     network: Optional[nx.Graph] = Field(
-        default=nx.Graph(),
+        default=None,
         description="Network graph with networkx",
     )
 
@@ -72,13 +73,25 @@ class Config:
         default="http://cytoscape:1234/v1",
     )
 
-    def post_init(self):
-        self._create_graph()
+    def model_post_init(self, __context):
+        if self.network is None:
+            self._create_graph()
 
     def add_to_targets(self, target: SequenceRecord):
         if target.id not in self.targets:
             self.targets.append(target.id)
 
+    # custom method for serialization to save the network as a json
+    @field_serializer("network")
+    def serialize_network(self, G: nx.Graph):
+        return json_graph.node_link_data(G)
+
+    @field_validator("network", mode="before")
+    @classmethod
+    def parse_network(cls, G: Any):
+        assert isinstance(G, dict), "Network has to be a dictionary"
+        return json_graph.node_link_graph(G)
+
     def _create_graph(self):
         """
         Initializes the nx.Graph object and adds nodes and edges based on the sequences.

tests/unit/network_tests/test_newtork_graph_build.py

@@ -1,3 +1,5 @@
+import json
+import pytest
 import pandas as pd
 import networkx as nx
 
@@ -27,6 +29,19 @@ def test_general_build_networkx(self):
         )
 
     def test_graph_build(self):
+        mat_accessions = [
+            "MBP1912539.1",
+            "SEV92896.1",
+            "MBO8174569.1",
+            "WP_042680787.1",
+        ]
+        mats = ProteinRecord.get_ids(mat_accessions)
+        # Create network
+        network = SequenceNetwork(sequences=mats)
+
+        assert len(list(network.network.nodes)) == len(mats)
+
+    def test_graph_save(self):
         mat_accessions = [
             "MBP1912539.1",
             "SEV92896.1",
@@ -40,4 +55,39 @@ def test_graph_build(self):
             weight="identity",
         )
 
-        assert len(list(network.network.nodes)) == len(mats)
+        # this is the data as a string, here one could write this to a file
+        data_string = network.model_dump_json()
+        # to read in the data again
+        data_json = json.loads(data_string)
+        # to load it into a Sequence Network object
+        network2 = SequenceNetwork(**data_json)
+
+        assert len(data_json['network']['nodes']) == len(mats)
+        assert network2.network.number_of_nodes() == network.network.number_of_nodes()
+        assert network2.network.number_of_edges() == network.network.number_of_edges()
+
+    def test_graph_network_input_interger(self):
+        mat_accessions = [
+            "MBP1912539.1",
+            "SEV92896.1",
+            "MBO8174569.1",
+            "WP_042680787.1",
+        ]
+        mats = ProteinRecord.get_ids(mat_accessions)
+        # Create network
+        network = SequenceNetwork(
+            sequences=mats,
+            weight="identity",
+        )
+
+        # this is the data as a string, here one could write this to a file
+        data_string = network.model_dump_json()
+        # to read in the data again
+        data_json = json.loads(data_string)
+
+        data_json['network'] = 1
+
+        # expecte an error
+        with pytest.raises(ValueError):
+            network2 = SequenceNetwork(**data_json)
+