|
import os |
|
import sys |
|
import csv |
|
import json |
|
import networkx as nx |
|
from collections import defaultdict |
|
|
|
import matplotlib |
|
import matplotlib.pyplot as plt |
|
import warnings |
|
|
|
warnings.filterwarnings('ignore', category=matplotlib.cbook.deprecation.MatplotlibDeprecationWarning) |
|
|
|
|
|
def load_alerts(_conf_json): |
|
_g = nx.DiGraph() |
|
_bank_accts = defaultdict(list) |
|
|
|
with open(_conf_json, "r") as rf: |
|
conf = json.load(rf) |
|
|
|
data_dir = os.path.join(conf["output"]["directory"], conf["general"]["simulation_name"]) |
|
acct_csv = os.path.join(data_dir, conf["output"]["alert_members"]) |
|
tx_csv = os.path.join(data_dir, conf["output"]["alert_transactions"]) |
|
|
|
input_dir = conf["input"]["directory"] |
|
schema_json = os.path.join(input_dir, conf["input"]["schema"]) |
|
with open(schema_json, "r") as rf: |
|
schema = json.load(rf) |
|
|
|
acct_idx = None |
|
bank_idx = None |
|
orig_idx = None |
|
bene_idx = None |
|
amt_idx = None |
|
date_idx = None |
|
for i, col in enumerate(schema["alert_member"]): |
|
if col.get("dataType") == "account_id": |
|
acct_idx = i |
|
elif col.get("dataType") == "bank_id": |
|
bank_idx = i |
|
for i, col in enumerate(schema["alert_tx"]): |
|
if col.get("dataType") == "orig_id": |
|
orig_idx = i |
|
elif col.get("dataType") == "dest_id": |
|
bene_idx = i |
|
elif col.get("dataType") == "amount": |
|
amt_idx = i |
|
elif col.get("dataType") == "timestamp": |
|
date_idx = i |
|
|
|
with open(acct_csv, "r") as rf: |
|
reader = csv.reader(rf) |
|
next(reader) |
|
for row in reader: |
|
acct_id = row[acct_idx] |
|
bank_id = row[bank_idx] |
|
_g.add_node(acct_id, bank_id=bank_id) |
|
_bank_accts[bank_id].append(acct_id) |
|
|
|
with open(tx_csv, "r") as rf: |
|
reader = csv.reader(rf) |
|
next(reader) |
|
for row in reader: |
|
orig_id = row[orig_idx] |
|
bene_id = row[bene_idx] |
|
amount = row[amt_idx] |
|
date = row[date_idx].split("T")[0] |
|
label = amount + "\n" + date |
|
_g.add_edge(orig_id, bene_id, amount=amount, date=date, label=label) |
|
|
|
return _g, _bank_accts |
|
|
|
|
|
def plot_alerts(_g, _bank_accts, _output_png): |
|
bank_ids = _bank_accts.keys() |
|
cmap = plt.get_cmap("tab10") |
|
pos = nx.nx_agraph.graphviz_layout(_g) |
|
|
|
plt.figure(figsize=(12.0, 8.0)) |
|
plt.axis('off') |
|
|
|
for i, bank_id in enumerate(bank_ids): |
|
color = cmap(i) |
|
members = _bank_accts[bank_id] |
|
nx.draw_networkx_nodes(_g, pos, members, node_size=300, node_color=color, label=bank_id) |
|
nx.draw_networkx_labels(_g, pos, {n: n for n in members}, font_size=10) |
|
|
|
edge_labels = nx.get_edge_attributes(_g, "label") |
|
nx.draw_networkx_edges(_g, pos) |
|
nx.draw_networkx_edge_labels(_g, pos, edge_labels, font_size=6) |
|
|
|
plt.legend(numpoints=1) |
|
plt.subplots_adjust(left=0, right=1, bottom=0, top=1) |
|
plt.savefig(_output_png, dpi=120) |
|
|
|
|
|
if __name__ == "__main__": |
|
argv = sys.argv |
|
|
|
if len(argv) < 3: |
|
print("Usage: python3 %s [ConfJSON] [OutputPNG]" % argv[0]) |
|
exit(1) |
|
|
|
conf_json = argv[1] |
|
output_png = argv[2] |
|
g, bank_accts = load_alerts(conf_json) |
|
plot_alerts(g, bank_accts, output_png) |
|
|