RNet.py

import importlib.util
import os
import ast
import networkx as nx
import matplotlib.pyplot as plt
import argparse
from collections import Counter
import shutil

"""

A script that can construct a reaction network from 
a molecular library and a reaction library containing up to C6 products.

Usage:
python RN_generator {noc} {noo} --{layout}
{noc}:Max number of C atoms.
{noo}:Max number of O atoms.
{--layout}: layout of reaction network images generated by 'networkx'. 
Choices=["spring(Default)", "circular", "shell", "kamada_kawai", "spectral", "random"]

"""

def copy_files_from_utilis():
    source_folder = 'utilis'
    target_folder = '.'

    files_to_copy = ['molecule_db.py', 'material_db.py', 'reaction_db.txt', 'id2chem.txt']

    for file_name in files_to_copy:
        source_file = os.path.join(source_folder, file_name)
        target_file = os.path.join(target_folder, file_name)

        try:
            shutil.copy(source_file, target_file)
        except Exception as e:
            print(f'复制文件时发生错误: {e}')


def filter_molecule_file(noc, noo, file_name='molecule_db.py'):
    try:
        with open(file_name, 'r') as infile:
            lines = infile.readlines()

        with open(file_name, 'w') as outfile:
            for line in lines:
                count_C = line.count("C")
                count_O = line.count("O")

                if count_C <= noc and count_O <= noo:
                    outfile.write(line)
    except FileNotFoundError:
        print(f"Error: The file {file_name} was not found.")
    except IOError as e:
        print(f"Error: An IOError occurred - {e}")


def filter_reactions(file_name='reaction_db.txt', module_path='molecule_db.py'):
    if not os.path.exists(module_path):
        raise FileNotFoundError(f"{module_path} not found.")

    spec = importlib.util.spec_from_file_location("mol_mody", module_path)
    mol_mody = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mol_mody)

    molecule_database = mol_mody.molecule_database
    valid_ids = set(molecule_database.keys())

    temp_file_name = file_name + '.temp'

    try:
        with open(file_name, 'r') as infile, open(temp_file_name, 'w') as outfile:
            for line in infile:
                ids = line.strip().split(',')
                if all(id.strip() in valid_ids for id in ids):
                    outfile.write(line)

        os.replace(temp_file_name, file_name)
    except Exception as e:
        print(f"An error occurred: {e}")
        if os.path.exists(temp_file_name):
            os.remove(temp_file_name)


def extract_nodes_from_database(database_file):
    with open(database_file, 'r') as file:
        content = file.read()
        molecule_database = ast.literal_eval(content.split('=')[1].strip())

    nodes = list(molecule_database.keys())
    return nodes


def id2chem(nodes, file_path):
    id2chem = {}
    with open(file_path, 'r') as file:
        for line in file:
            id_str, chem_symbol = line.strip().split(': ')
            id2chem[id_str] = chem_symbol

    replaced_nodes = [id2chem[node] for node in nodes]

    return replaced_nodes


def reaction2chem(reaction_file_path, id2chem_file_path):
    id2chem = {}
    with open(id2chem_file_path, 'r') as file:
        for line in file:
            id_str, chem_symbol = line.strip().split(': ')
            id2chem[id_str] = chem_symbol

    with open(reaction_file_path, 'r') as file:
        lines = file.readlines()

    replaced_lines = []
    for line in lines:
        numbers = line.strip().split(', ')
        replaced_numbers = [id2chem[number] for number in numbers]
        replaced_line = ', '.join(replaced_numbers)
        replaced_lines.append(replaced_line)

    with open(reaction_file_path, 'w') as file:
        for replaced_line in replaced_lines:
            file.write(replaced_line + '\n')


def replace_mol_db_with_nodes(file_path):
    try:
        # 读取文件内容
        with open(file_path, 'r') as file:
            content = file.readlines()

        # 替换所有'mol.db'为nodes列表的字符串表示形式
        new_content = []
        for line in content:
            if 'mol.db' in line:
                # 将nodes列表转换为字符串，并替换掉'mol.db'
                nodes_str = str(nodes).replace("'", '"')  # 假设你希望使用双引号
                new_line = line.replace('mol.db', nodes_str)
                new_content.append(new_line)
            else:
                new_content.append(line)

        # 写回文件
        with open(file_path, 'w') as file:
            file.writelines(new_content)

        #print(f"File '{file_path}' has been updated successfully.")

    except Exception as e:
        print(f"An error occurred: {e}")



def draw(filename, nodes, layout='spring'):
    def parse_reaction_data(filename):
        reactions = []
        with open(filename, 'r') as file:
            for line in file:
                parts = line.strip().split(', ')
                reactions.append(parts)
        return reactions

    def build_reaction_network(reactions, nodes):
        G = nx.DiGraph()
        G.add_nodes_from(nodes)

        for reaction in reactions:
            reactant = reaction[0]
            products = reaction[1:]
            for product in products:
                G.add_edge(reactant, product)

        return G

    reactions = parse_reaction_data(filename)
    G = build_reaction_network(reactions, nodes)

    def draw_network(G, layout):
        if layout == 'spring':
            pos = nx.spring_layout(G, k=10, iterations=500)
        elif layout == 'circular':
            pos = nx.circular_layout(G)
        elif layout == 'random':
            pos = nx.random_layout(G)
        elif layout == 'shell':
            pos = nx.shell_layout(G)
        else:
            raise ValueError(f"Unknown layout: {layout}")
        plt.figure(figsize=(15, 15))
        nx.draw(G, pos, with_labels=True, node_color='lightblue', width=1.5, node_size=2800, font_size=16, edge_color='gray', arrows=True, arrowstyle='-|>', arrowsize=20)
        plt.savefig("Reaction_Net.png", dpi=300)
        plt.show()

    draw_network(G, layout)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process some integers.")
    parser.add_argument("noc", type=int, help="The maximum number of carbon atoms.")
    parser.add_argument("noo", type=int, help="The maximum number of oxygen atoms.")
    parser.add_argument("--layout", type=str, default="spring", choices=["spring", "circular", "shell", "kamada_kawai", "spectral", "random"],
                        help="The layout for the network graph.")

    args = parser.parse_args()

    copy_files_from_utilis()
    filter_molecule_file(args.noc, args.noo)
    filter_reactions()
    nodes = extract_nodes_from_database('molecule_db.py')
    replaced_nodes = id2chem(nodes, 'id2chem.txt')
    reaction2chem('reaction_db.txt', 'id2chem.txt')
    replace_mol_db_with_nodes('material_db.py')
    draw('reaction_db.txt', replaced_nodes, layout=args.layout)