Repository: dizak/prwlr

Issue:
prwlr.save_network
The highest-level convenience method prwlr.save_network should be implemented in prwlr.core so that whole network can be saved without losing prwlr.profiles.Profile objects.

Discussion:

Patch:
diff --git a/prwlr/core.py b/prwlr/core.py
index f512fc7..eee71c2 100644
--- a/prwlr/core.py
+++ b/prwlr/core.py
@@ -1,3 +1,4 @@
+import gc
 import pandas as _pd
 import numpy as _np
 from . import databases as _databases
@@ -221,6 +222,64 @@ def save_profiles(
         },
     ).to_csv(filename, **kwargs)
 
+def save_network(
+    dataframe,
+    filename,
+    profile_query_column=Columns.PROF_Q,
+    profile_array_column=Columns.PROF_A,
+    profile_reference_column_suffix=Columns.REF,
+    profile_query_column_suffix=Columns.QRY,
+    str_sep='|',
+    **kwargs
+):
+    """
+    Saves pandas.DataFrame with prwlr.profiles.Profile to CSV.
+
+    Parameters
+    -------
+    filename: str, path
+        Filename of the CSV.
+    profile_query_column: str, default prwlr.core.Columns.PROF_Q
+        Column name of the query profile. The default value is taken from the
+        module-wide settings.
+    profile_array_column: str, default prwlr.core.Columns.PROF_A
+        Column name of the array profile. The default value is taken from the
+        module-wide settings.
+    profile_reference_column_suffix: str, default prwlr.core.Columns.REF
+        Column name suffix  of the profile's reference
+        (see prwlr.profiles.Profile).
+    profile_query_column_suffix: str, default prwlr.core.Columns.QRY
+        Column name suffix  of the profile's query
+        (see prwlr.profiles.Profile).
+    str_sep: str, default '|'
+        The internal separator for the profile's reference/query.
+
+    Returns
+    -------
+        pandas.DataFrame
+    """
+    qry_ref_col = '{}_{}'.format(profile_query_column, profile_reference_column_suffix)
+    qry_qry_col = '{}_{}'.format(profile_query_column, profile_query_column_suffix)
+    arr_ref_col = '{}_{}'.format(profile_array_column, profile_reference_column_suffix)
+    arr_qry_col = '{}_{}'.format(profile_array_column, profile_query_column_suffix)
+    df = dataframe.copy()
+    df[qry_ref_col] = df[profile_query_column].apply(
+        lambda x: x.reference,
+    ).str.join(str_sep)
+    df[qry_qry_col] = df[profile_query_column].apply(
+        lambda x: x.query,
+    ).str.join(str_sep)
+    df[arr_ref_col] = df[profile_array_column].apply(
+        lambda x: x.reference,
+    ).str.join(str_sep)
+    df[arr_qry_col] = df[profile_array_column].apply(
+        lambda x: x.query,
+    ).str.join(str_sep)
+    df.drop(
+        columns=[profile_query_column, profile_array_column],
+    ).to_csv(filename, **kwargs)
+    del df; gc.collect()
+
 def merge_sga_profiles(
     sga,
     profiles
@@ -294,8 +353,3 @@ def calculate_pss(
             axis=1,
         )
     return network
-
-def pss(ar1, ar2):
-    return sum(a == b for a, b in zip(ar1, ar2))
-
-pss_vect = _np.vectorize(pss)


Test patch:
diff --git a/test_data/CoreTests/test_network.csv b/test_data/CoreTests/ref_network.csv
similarity index 100%
rename from test_data/CoreTests/test_network.csv
rename to test_data/CoreTests/ref_network.csv
diff --git a/test_data/CoreTests/test_profiles.csv b/test_data/CoreTests/ref_profiles.csv
similarity index 100%
rename from test_data/CoreTests/test_profiles.csv
rename to test_data/CoreTests/ref_profiles.csv
diff --git a/tests/tests.py b/tests/tests.py
index 4c91dd5..88cc936 100644
--- a/tests/tests.py
+++ b/tests/tests.py
@@ -648,9 +648,10 @@ class CoreTests(unittest.TestCase):
         self.ref_query_1 = list('acdfhiklostuz')
         self.ref_reference_2 = list('acefghijklmnprstuwxy')
         self.ref_query_2 = list('abdfhiklostuz')
-        self.test_profiles_filename = 'test_data/CoreTests/test_profiles.csv'
-        self.test_network_filename = 'test_data/CoreTests/test_network.csv'
+        self.ref_profiles_filename = 'test_data/CoreTests/ref_profiles.csv'
+        self.ref_network_filename = 'test_data/CoreTests/ref_network.csv'
         self.test_saved_profiles_filename = 'test_data/CoreTests/test_save_profiles.csv'
+        self.test_saved_network_filename = 'test_data/CoreTests/test_save_network.csv'
 
         self.ref_profile_1, self.ref_profile_2 = (
             prwlr.profiles.Profile(
@@ -683,6 +684,8 @@ class CoreTests(unittest.TestCase):
         """
         if os.path.exists(self.test_saved_profiles_filename):
             os.remove(self.test_saved_profiles_filename)
+        if os.path.exists(self.test_saved_network_filename):
+            os.remove(self.test_saved_network_filename)
 
     def test_read_profiles(self):
         """
@@ -694,7 +697,7 @@ class CoreTests(unittest.TestCase):
         pd.testing.assert_series_equal(
             self.ref_profiles_srs,
             prwlr.core.read_profiles(
-                self.test_profiles_filename,
+                self.ref_profiles_filename,
                 index_col=[0],
             ),
         )
@@ -709,12 +712,32 @@ class CoreTests(unittest.TestCase):
         pd.testing.assert_frame_equal(
             self.ref_network_df,
             prwlr.core.read_network(
-                self.test_network_filename,
+                self.ref_network_filename,
                 index_col=[0],
             ),
             check_like=True,
         )
 
+    def test_save_network(self):
+        """
+        Tests if prwlr.core.save_network saves network suitable for
+        prwlr.core.read_network.
+        """
+        import prwlr.core
+
+        prwlr.core.save_network(
+            self.ref_network_df,
+            self.test_saved_network_filename,
+        )
+        pd.testing.assert_frame_equal(
+            self.ref_network_df,
+            prwlr.core.read_network(
+                self.test_saved_network_filename,
+                index_col=[0],
+                ),
+            check_like=True,
+        )
+
     def test_save_profiles(self):
         """
         Tests if prwlr.core.save_profiles writes CSV file with correct values.