From 03c74e8dc274b232227d1d8dd786036945ee32f6 Mon Sep 17 00:00:00 2001
From: anpolol <polinspb@ya.ru>
Date: Mon, 4 Dec 2023 15:20:43 +0300
Subject: [PATCH] improve pep style

---
 stable_gnn/explain.py                    |  2 +-
 stable_gnn/fairness.py                   | 23 +++++++---------
 stable_gnn/model_gc.py                   |  2 +-
 tests/test_general/test_fairness.py      |  3 ++-
 tests/test_general/test_link_predict.py  |  0
 tutorials/tutorial_link_prediction.ipynb | 34 ++++++++++++++++++++++++
 6 files changed, 47 insertions(+), 17 deletions(-)
 create mode 100644 tests/test_general/test_link_predict.py
 create mode 100644 tutorials/tutorial_link_prediction.ipynb

diff --git a/stable_gnn/explain.py b/stable_gnn/explain.py
index a65bb36..50f3550 100644
--- a/stable_gnn/explain.py
+++ b/stable_gnn/explain.py
@@ -317,4 +317,4 @@ def _search_m_k(data: pd.DataFrame, target: str, nodes: List[str]) -> List[str]:
                         if count == len(m_b):
                             return m_b
         else:
-            return m_b
\ No newline at end of file
+            return m_b
diff --git a/stable_gnn/fairness.py b/stable_gnn/fairness.py
index 9088c80..0d67fba 100644
--- a/stable_gnn/fairness.py
+++ b/stable_gnn/fairness.py
@@ -1,15 +1,14 @@
-from sklearn.linear_model import LogisticRegression
-from sklearn.tree import DecisionTreeClassifier
-from sklearn.svm import SVC
-from sklearn.neighbors import KNeighborsClassifier
-
 import numpy as np
 import pandas as pd
 from scipy.optimize import linprog
-from sklearn.model_selection import train_test_split
+from sklearn.base import BaseEstimator
 from sklearn.ensemble import RandomForestClassifier
+from sklearn.linear_model import LogisticRegression
 from sklearn.metrics import accuracy_score
-from sklearn.base import BaseEstimator
+from sklearn.model_selection import train_test_split
+from sklearn.neighbors import KNeighborsClassifier
+from sklearn.svm import SVC
+from sklearn.tree import DecisionTreeClassifier
 
 
 class Fair:
@@ -34,7 +33,7 @@ def run(
         random_state: int = None,
     ):
         """
-        Main function running fairness correction and calculating accuracy and fairness values
+        Correct fairness and calculate accuracy and fairness values
 
         :param number_iterations: (int) Number of iterations of interior algorithm
         :param prefit: (bool) Flag if passed estimator fitted already or not (default: False)
@@ -91,7 +90,7 @@ def _ml_model(self, df, random_state, estimator, prefit):
         x = df["target"]
 
         y_train, y_test, x_train, x_test = train_test_split(y, x, random_state=random_state)
-        if prefit == False:
+        if prefit is False:
             estimator.fit(y_train, x_train)
         estimator_pred = estimator.predict(y_test)
         accuracy_estimator = accuracy_score(estimator_pred, x_test)
@@ -161,13 +160,9 @@ def _lp_solver(
         multiplier=1,
         random_state=None,
     ):
-        group = multiplier * d["group"]
         one_group = multiplier * d["one_group"]
         zero_group = multiplier * d["zero_group"]
 
-        one_train_probs = d["one_train_probs"]
-        zero_train_probs = d["zero_train_probs"]
-
         bounds = []
         for i in range(3 * one_group + 3 * zero_group):
             bounds.append((0, 1))
@@ -193,7 +188,7 @@ def _lp_solver(
         if verbose:
             print("Start fitting")
         for k in range(number_iterations):
-            if random_state == None:
+            if random_state is None:
                 one_sample = d["one_train_probs"].sample(one_group)
                 zero_sample = d["zero_train_probs"].sample(zero_group)
             else:
diff --git a/stable_gnn/model_gc.py b/stable_gnn/model_gc.py
index d7c1915..6b6b990 100644
--- a/stable_gnn/model_gc.py
+++ b/stable_gnn/model_gc.py
@@ -293,4 +293,4 @@ def self_supervised_loss(deg_pred: Tensor, batch: Batch) -> Tensor:
             batch_ptr.numpy()[1:] - batch_ptr.numpy()[: len(batch_ptr) - 1]
         )  # после экстраполяции мы получили очень плотный граф, хотим степень снизить в N раз для более хороших предсказаний
         true = degree(batch.edge_index[0], batch.x.shape[0])[indices] / ratio
-        return F.mse_loss(deg_pred, true)
\ No newline at end of file
+        return F.mse_loss(deg_pred, true)
diff --git a/tests/test_general/test_fairness.py b/tests/test_general/test_fairness.py
index 7a81143..f94b652 100644
--- a/tests/test_general/test_fairness.py
+++ b/tests/test_general/test_fairness.py
@@ -1,9 +1,10 @@
 import numpy as np
 import pandas as pd
-from stable_gnn.fairness import Fair
 from sklearn.linear_model import LogisticRegression
 from sklearn.model_selection import train_test_split
 
+from stable_gnn.fairness import Fair
+
 
 def synthetic_dataset(size=1000, influence=True):
     """
diff --git a/tests/test_general/test_link_predict.py b/tests/test_general/test_link_predict.py
new file mode 100644
index 0000000..e69de29
diff --git a/tutorials/tutorial_link_prediction.ipynb b/tutorials/tutorial_link_prediction.ipynb
new file mode 100644
index 0000000..9e60cbb
--- /dev/null
+++ b/tutorials/tutorial_link_prediction.ipynb
@@ -0,0 +1,34 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}