init

INTRODUCTION.md

@@ -0,0 +1,9 @@
+# Foundation Models for Materials - FM4M
+
+FM4M adopts a modular architecture designed for flexible extensibility. As illustrated in the figure below, it comprises both uni-modal and fused models. Each uni-modal model is pre-trained independently for its respective modality (e.g., SMILES), and users can access individual functionalities directly from the corresponding model directory (e.g., smi-ted/). Some of these uni-modal models can be "late-fused" using fusion algorithms, creating a more powerful multi-modal feature representation for downstream predictions. To simplify usage, we provide fm4m-kit, a wrapper that enables users to easily access the capabilities of all models through straightforward methods. These models are also available on Hugging Face, where they can be accessed via an intuitive and user-friendly GUI.
+
+<p align="center">
+  <img src="gradio_api/file=img/introduction.png" alt="FM4M Overview" width="50%"/>
+</p>
+
+

app.py

@@ -0,0 +1,523 @@
+import gradio as gr
+import numpy as np
+import pandas as pd
+from tempfile import NamedTemporaryFile
+from PIL import Image
+from rdkit import RDLogger
+from sklearn.model_selection import train_test_split
+
+from molecule_generation_helpers import *
+from property_prediction_helpers import *
+
+DEBUG_VISIBLE = False
+RDLogger.logger().setLevel(RDLogger.ERROR)
+
+# Predefined dataset paths (these should be adjusted to your file paths)
+predefined_datasets = {
+    " ": " ",
+    "BACE": "./data/bace/train.csv, ./data/bace/test.csv, smiles, Class",
+    "ESOL": "./data/esol/train.csv, ./data/esol/test.csv, smiles, prop",
+}
+
+# Models
+models_enabled = [
+    "MorganFingerprint",
+    "SMI-TED",
+    "SELFIES-TED",
+    "MHG-GED",
+]
+
+blank_df = pd.DataFrame({"id": [], "Model": [], "Score": []})
+
+
+# Function to load a predefined dataset from the local path
+def load_predefined_dataset(dataset_name):
+    val = predefined_datasets.get(dataset_name)
+    if val:
+        try:
+            df = pd.read_csv(val.split(",")[0])
+            return (
+                df.head(),
+                gr.update(choices=list(df.columns), value=None),
+                gr.update(choices=list(df.columns), value=None),
+                dataset_name.lower(),
+            )
+        except:
+            pass
+    else:
+        dataset_name = "Custom"
+    return (
+        pd.DataFrame(),
+        gr.update(choices=[], value=None),
+        gr.update(choices=[], value=None),
+        dataset_name.lower(),
+    )
+
+
+# Function to handle dataset selection (predefined or custom)
+def handle_dataset_selection(selected_dataset, state):
+    state["dataset_name"] = (
+        selected_dataset if selected_dataset in predefined_datasets else "CUSTOM"
+    )
+    # Show file upload fields for train and test datasets if "Custom Dataset" is selected
+    task_type = (
+        "Classification"
+        if selected_dataset == "BACE"
+        else "Regression" if selected_dataset == "ESOL" else None
+    )
+    return (
+        gr.update(visible=selected_dataset not in predefined_datasets or DEBUG_VISIBLE),
+        task_type,
+    )
+
+
+# Function to select input and output columns and display a message
+def select_columns(input_column, output_column, train_data, test_data, state):
+    if train_data and test_data and input_column and output_column:
+        return f"{train_data.name},{test_data.name},{input_column},{output_column},{state['dataset_name']}"
+    return gr.update()
+
+
+# Function to display the head of the uploaded CSV file
+def display_csv_head(file):
+    if file is not None:
+        # Load the CSV file into a DataFrame
+        df = pd.read_csv(file.name)
+        return (
+            df.head(),
+            gr.update(choices=list(df.columns)),
+            gr.update(choices=list(df.columns)),
+        )
+    return pd.DataFrame(), gr.update(choices=[]), gr.update(choices=[])
+
+
+def process_custom_file(file, selected_dataset):
+    if file and os.path.getsize(file.name) < 50 * 1024:
+        df = pd.read_csv(file.name)
+        if "input" in df.columns and "output" in df.columns:
+            train, test = train_test_split(df, test_size=0.2)
+            with NamedTemporaryFile(
+                prefix="fm4m-train-", suffix=".csv", delete=False
+            ) as train_file:
+                train.to_csv(train_file.name, index=False)
+            with NamedTemporaryFile(
+                prefix="fm4m-test-", suffix=".csv", delete=False
+            ) as test_file:
+                test.to_csv(test_file.name, index=False)
+            task_type = (
+                "Classification" if df["output"].dtype == np.int64 else "Regression"
+            )
+            return train_file.name, test_file.name, "input", "output", task_type
+
+    return (
+        None,
+        None,
+        None,
+        None,
+        gr.update() if selected_dataset in predefined_datasets else None,
+    )
+
+
+def update_plot_choices(current, state):
+    choices = []
+    if state.get("roc_auc") is not None:
+        choices.append("ROC-AUC")
+    if state.get("RMSE") is not None:
+        choices.append("Parity Plot")
+    if state.get("x_batch") is not None:
+        choices.append("Latent Space")
+    if current in choices:
+        return gr.update(choices=choices)
+    return gr.update(choices=choices, value=None if len(choices) == 0 else choices[0])
+
+
+def log_selected(df: pd.DataFrame, evt: gr.SelectData, state):
+    state.update(state["results"].get(df.at[evt.index[0], 'id'], {}))
+
+
+# Dictionary for SMILES strings and corresponding images (you can replace with your actual image paths)
+smiles_image_mapping = {
+    # Example SMILES for ethanol
+    "Mol 1": {
+        "smiles": "C=C(C)CC(=O)NC[C@H](CO)NC(=O)C=Cc1ccc(C)c(Cl)c1",
+        "image": "img/img1.png",
+    },
+    # Example SMILES for butane
+    "Mol 2": {
+        "smiles": "C=CC1(CC(=O)NC[C@@H](CCCC)NC(=O)c2cc(Cl)cc(Br)c2)CC1",
+        "image": "img/img2.png",
+    },
+    # Example SMILES for ethylamine
+    "Mol 3": {
+        "smiles": "C=C(C)C[C@H](NC(C)=O)C(=O)N1CC[C@H](NC(=O)[C@H]2C[C@@]2(C)Br)C(C)(C)C1",
+        "image": "img/img3.png",
+    },
+    # Example SMILES for diethyl ether
+    "Mol 4": {
+        "smiles": "C=C1CC(CC(=O)N[C@H]2CCN(C(=O)c3ncccc3SC)C23CC3)C1",
+        "image": "img/img4.png",
+    },
+    # Example SMILES for chloroethane
+    "Mol 5": {
+        "smiles": "C=CCS[C@@H](C)CC(=O)OCC",
+        "image": "img/img5.png",
+    },
+}
+
+
+# Load images for selection
+def load_image(path):
+    try:
+        return Image.open(smiles_image_mapping[path]["image"])
+    except:
+        pass
+
+
+# Function to handle image selection
+def handle_image_selection(image_key):
+    if not image_key:
+        return None, None
+    smiles = smiles_image_mapping[image_key]["smiles"]
+    mol_image = smiles_to_image(smiles)
+    return smiles, mol_image
+
+
+# Introduction
+with gr.Blocks() as introduction:
+    with open("INTRODUCTION.md") as f:
+        gr.Markdown(f.read(), sanitize_html=False)
+
+# Property Prediction
+with gr.Blocks() as property_prediction:
+    state = gr.State({"model_name": "Default - Auto", "results": {}})
+    gr.HTML(
+        '''
+    <p style="text-align: center">
+        Task : Property Prediction
+        <br>
+        Models are finetuned with different combination of modalities on the uploaded or selected built data set.
+    </p>
+    '''
+    )
+    with gr.Row():
+        with gr.Column():
+            # Dropdown menu for predefined datasets including "Custom Dataset" option
+            dataset_selector = gr.Dropdown(
+                label="Select Dataset",
+                choices=list(predefined_datasets.keys()) + ["Custom Dataset"],
+            )
+            # Display the message for selected columns
+            selected_columns_message = gr.Textbox(
+                label="Selected Columns Info", visible=DEBUG_VISIBLE
+            )
+
+            with gr.Accordion(
+                "Custom Dataset Settings", open=True, visible=DEBUG_VISIBLE
+            ) as settings:
+                # File upload options for custom dataset (train and test)
+                custom_file = gr.File(
+                    label="Upload Custom Dataset",
+                    file_types=[".csv"],
+                )
+                train_file = gr.File(
+                    label="Upload Custom Train Dataset",
+                    file_types=[".csv"],
+                    visible=False,
+                )
+                train_display = gr.Dataframe(
+                    label="Train Dataset Preview (First 5 Rows)",
+                    interactive=False,
+                    visible=DEBUG_VISIBLE,
+                )
+
+                test_file = gr.File(
+                    label="Upload Custom Test Dataset",
+                    file_types=[".csv"],
+                    visible=False,
+                )
+                test_display = gr.Dataframe(
+                    label="Test Dataset Preview (First 5 Rows)",
+                    interactive=False,
+                    visible=DEBUG_VISIBLE,
+                )
+
+                # Predefined dataset displays
+                predefined_display = gr.Dataframe(
+                    label="Predefined Dataset Preview (First 5 Rows)",
+                    interactive=False,
+                    visible=DEBUG_VISIBLE,
+                )
+
+                # Dropdowns for selecting input and output columns for the custom dataset
+                input_column_selector = gr.Dropdown(
+                    label="Select Input Column",
+                    choices=[],
+                    allow_custom_value=True,
+                    visible=DEBUG_VISIBLE,
+                )
+                output_column_selector = gr.Dropdown(
+                    label="Select Output Column",
+                    choices=[],
+                    allow_custom_value=True,
+                    visible=DEBUG_VISIBLE,
+                )
+
+                # When a custom train file is uploaded, display its head and update column selectors
+                train_file.change(
+                    display_csv_head,
+                    inputs=train_file,
+                    outputs=[
+                        train_display,
+                        input_column_selector,
+                        output_column_selector,
+                    ],
+                )
+
+                # When a custom test file is uploaded, display its head
+                test_file.change(
+                    display_csv_head,
+                    inputs=test_file,
+                    outputs=[
+                        test_display,
+                        input_column_selector,
+                        output_column_selector,
+                    ],
+                )
+
+            model_checkbox = gr.CheckboxGroup(
+                choices=models_enabled, label="Select Model", visible=DEBUG_VISIBLE
+            )
+
+            task_radiobutton = gr.Radio(
+                choices=["Classification", "Regression"],
+                label="Task Type",
+                visible=DEBUG_VISIBLE,
+            )
+
+            # When a dataset is selected, show either file upload fields (for custom) or load predefined datasets
+            # When a predefined dataset is selected, load its head and update column selectors
+            dataset_selector.change(lambda: None, outputs=custom_file).then(
+                handle_dataset_selection,
+                inputs=[dataset_selector, state],
+                outputs=[settings, task_radiobutton],
+            ).then(
+                load_predefined_dataset,
+                inputs=dataset_selector,
+                outputs=[
+                    predefined_display,
+                    input_column_selector,
+                    output_column_selector,
+                    selected_columns_message,
+                ],
+            )
+
+            custom_file.change(
+                process_custom_file,
+                inputs=[custom_file, dataset_selector],
+                outputs=[
+                    train_file,
+                    test_file,
+                    input_column_selector,
+                    output_column_selector,
+                    task_radiobutton,
+                ],
+            )
+            eval_clear_button = gr.Button("Clear")
+            eval_button = gr.Button("Submit", variant="primary")
+            step_slider = gr.Slider(
+                minimum=0,
+                maximum=8,
+                value=0,
+                label="Progress",
+                show_label=True,
+                interactive=False,
+                visible=False,
+            )
+
+        # Right Column
+        with gr.Column():
+            log_table = gr.Dataframe(value=blank_df, interactive=False)
+
+            plot_radio = gr.Radio(choices=[], label="Select Plot Type")
+
+            plot_output = gr.Plot(label="Visualization")
+
+            log_table.select(log_selected, [log_table, state]).success(
+                update_plot_choices, inputs=[plot_radio, state], outputs=plot_radio
+            ).then(display_plot, inputs=[plot_radio, state], outputs=plot_output)
+
+            def clear_eval(state):
+                state["results"] = {}
+                return None, gr.update(choices=[], value=None), blank_df
+
+            def eval_part(part, step, selector, show_progress=False):
+                return (
+                    part.then(
+                        lambda: [models_enabled[x] for x in selector],
+                        outputs=model_checkbox,
+                    )
+                    .then(
+                        evaluate_and_log,
+                        inputs=[
+                            model_checkbox,
+                            selected_columns_message,
+                            task_radiobutton,
+                            log_table,
+                            state,
+                        ],
+                        outputs=log_table,
+                        show_progress=show_progress,
+                    )
+                    .then(lambda: step, outputs=step_slider, show_progress=False)
+                )
+
+            part = (
+                eval_button.click(
+                    lambda: (
+                        gr.update(interactive=False),
+                        gr.update(interactive=False),
+                    ),
+                    outputs=[eval_clear_button, eval_button],
+                )
+                .then(
+                    select_columns,
+                    inputs=[
+                        input_column_selector,
+                        output_column_selector,
+                        train_file,
+                        test_file,
+                        state,
+                    ],
+                    outputs=selected_columns_message,
+                )
+                .then(
+                    clear_eval,
+                    inputs=state,
+                    outputs=[
+                        plot_output,
+                        plot_radio,
+                        log_table,
+                    ],
+                )
+            )
+            part = part.then(
+                lambda: gr.update(value=0, visible=True),
+                outputs=step_slider,
+                show_progress=False,
+            )
+            part = eval_part(part, 1, [0], True)
+            part = eval_part(part, 2, [1])
+            part = eval_part(part, 3, [2])
+            part = eval_part(part, 4, [3])
+            part = eval_part(part, 5, [1, 2])
+            part = eval_part(part, 6, [2, 3])
+            part = eval_part(part, 7, [1, 3])
+            part = eval_part(part, 8, [1, 2, 3])
+            part = part.then(
+                lambda: gr.update(visible=False),
+                outputs=step_slider,
+                show_progress=False,
+            )
+            part.then(
+                lambda: (
+                    gr.update(interactive=True),
+                    gr.update(interactive=True),
+                ),
+                outputs=[eval_clear_button, eval_button],
+            )
+
+            plot_radio.change(
+                display_plot, inputs=[plot_radio, state], outputs=plot_output
+            )
+
+        eval_clear_button.click(
+            clear_eval,
+            inputs=state,
+            outputs=[
+                plot_output,
+                plot_radio,
+                log_table,
+            ],
+        ).then(lambda: " ", outputs=dataset_selector)
+
+
+# Molecule Generation
+with gr.Blocks() as molecule_generation:
+    gr.HTML(
+        '''
+    <p style="text-align: center">
+        Task : Molecule Generation
+        <br>
+        Generate a new molecule similar to the initial molecule with better drug-likeness and synthetic accessibility.
+    </p>
+    '''
+    )
+    with gr.Row():
+        with gr.Column():
+            smiles_input = gr.Textbox(label="Input SMILES String")
+            image_display = gr.Image(label="Molecule Image", height=250, width=250)
+            # Show images for selection
+            with gr.Accordion("Select from sample molecules", open=False):
+                image_selector = gr.Radio(
+                    choices=list(smiles_image_mapping.keys()),
+                    label="Select from sample molecules",
+                    value=None,
+                )
+                image_selector.change(load_image, image_selector, image_display)
+            clear_button = gr.Button("Clear")
+            generate_button = gr.Button("Submit", variant="primary")
+
+        # Right Column
+        with gr.Column():
+            gen_image_display = gr.Image(
+                label="Generated Molecule Image", height=250, width=250
+            )
+            generated_output = gr.Textbox(label="Generated Output")
+            property_table = gr.Dataframe(label="Molecular Properties Comparison")
+
+            # Handle image selection
+            image_selector.change(
+                handle_image_selection,
+                inputs=image_selector,
+                outputs=[smiles_input, image_display],
+            )
+            smiles_input.change(
+                smiles_to_image, inputs=smiles_input, outputs=image_display
+            )
+
+            # Generate button to display canonical SMILES and molecule image
+            generate_button.click(
+                lambda: (
+                    gr.update(interactive=False),
+                    gr.update(interactive=False),
+                ),
+                outputs=[clear_button, generate_button],
+            ).then(
+                generate_canonical,
+                inputs=smiles_input,
+                outputs=[property_table, generated_output, gen_image_display],
+            ).then(
+                lambda: (
+                    gr.update(interactive=True),
+                    gr.update(interactive=True),
+                ),
+                outputs=[clear_button, generate_button],
+            )
+        clear_button.click(
+            lambda: (None, None, None, None, None, None),
+            outputs=[
+                smiles_input,
+                image_display,
+                image_selector,
+                gen_image_display,
+                generated_output,
+                property_table,
+            ],
+        )
+
+
+# Render with tabs
+gr.TabbedInterface(
+    [introduction, property_prediction, molecule_generation],
+    ["Introduction", "Property Prediction", "Molecule Generation"],
+).launch(server_name="0.0.0.0", allowed_paths=["./"])

data/esol/test.csv

@@ -0,0 +1,109 @@
+,selfies,prop,smiles
+0,[Cl] [C] [=C] [Branch1] [C] [Cl] [C] [Branch1] [C] [Cl] [C] [C] [C] [C] [Branch1] [Branch2] [C] [O] [C] [Ring1] [=Branch1] [Ring1] [Ring1] [C] [Ring1] [Branch2] [C] [Ring1] [=C] [Branch1] [C] [Cl] [C] [Ring1] [=N] [Branch1] [C] [Cl] [Cl],-4.533,ClC4=C(Cl)C5(Cl)C3C1CC(C2OC12)C3C4(Cl)C5(Cl)Cl
+1,[C] [C] [C] [C] [C] [=O],-1.103,CCCCC=O
+2,[O] [C] [C] [C] [C] [=C],-0.7909999999999999,OCCCC=C
+3,[C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [N] [N] [=C] [C] [Branch1] [C] [N] [=C] [Branch1] [C] [Br] [C] [Ring1] [Branch2] [=O],-3.005,c1ccccc1n2ncc(N)c(Br)c2(=O)
+4,[N] [C] [=C] [C] [=C] [Branch1] [C] [O] [C] [=C] [Ring1] [#Branch1],-1.231,Nc1ccc(O)cc1
+5,[C] [C] [Branch1] [C] [C] [C] [C] [O] [C] [=Branch1] [C] [=O] [C],-1.817,CC(C)CCOC(=O)C
+6,[C] [O] [P] [=Branch1] [C] [=S] [Branch1] [Ring1] [O] [C] [S] [C] [C] [=Branch1] [C] [=O] [N] [Branch1] [C] [C] [C] [=O],-2.087,COP(=S)(OC)SCC(=O)N(C)C=O
+7,[Cl] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=Branch1] [Ring2] [=C] [Ring1] [#Branch1] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [Branch1] [C] [Cl] [=C] [Ring1] [Branch2],-6.312,Clc1ccc(Cl)c(c1)c2ccc(Cl)c(Cl)c2
+8,[C] [Branch1] [C] [Cl] [=C] [Branch1] [C] [Cl] [C] [Branch1] [C] [Cl] [=C] [N] [=C] [C] [=N] [C] [Ring1] [=Branch1] [=C] [Ring1] [=N] [Cl],-4.438,c2(Cl)c(Cl)c(Cl)c1nccnc1c2(Cl) 
+9,[C] [C] [C] [C] [C] [=C] [Branch1] [C] [C] [N] [=C] [Branch1] [=Branch1] [N] [=C] [Ring1] [#Branch1] [O] [N] [Branch1] [C] [C] [C],-3.57,CCCCc1c(C)nc(nc1O)N(C)C 
+10,[C] [C] [O] [C] [=Branch1] [C] [=O] [C] [C] [=Branch1] [C] [=O] [O] [C] [C],-1.413,CCOC(=O)CC(=O)OCC
+11,[C] [C] [Branch1] [C] [C] [Branch1] [C] [C] [C] [=C] [C] [=C] [Branch1] [C] [O] [C] [=C] [Ring1] [#Branch1],-3.192,CC(C)(C)c1ccc(O)cc1
+12,[C] [C] [=C] [C] [=C] [C] [Branch1] [C] [C] [=C] [Ring1] [#Branch1],-3.035,Cc1cccc(C)c1
+13,[C] [C] [C] [O] [C] [=Branch1] [C] [=O] [C],-1.125,CCCOC(=O)C
+14,[C] [S] [C] [=N] [N] [=C] [Branch1] [=Branch2] [C] [=Branch1] [C] [=O] [N] [Ring1] [#Branch1] [N] [C] [Branch1] [C] [C] [Branch1] [C] [C] [C],-2.324,CSc1nnc(c(=O)n1N)C(C)(C)C
+15,[Cl] [C] [=C] [C] [=C] [Branch1] [Branch1] [C] [=C] [Ring1] [=Branch1] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [Cl],-5.142,Clc1ccc(cc1)c2ccccc2Cl 
+16,[C] [C] [C] [C] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [C] [Branch1] [Ring2] [C] [Ring1] [Branch2] [C] [Branch1] [C] [O] [C] [C] [C] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [C] [Ring1] [Branch2],-1.5319999999999998,CC1CC(C)C(=O)C(C1)C(O)CC2CC(=O)NC(=O)C2 
+17,[C] [N] [C] [=Branch1] [C] [=O] [O] [C] [=C] [C] [=C] [C] [Branch1] [Branch2] [N] [=C] [N] [Branch1] [C] [C] [C] [=C] [Ring1] [O],-1.846,CNC(=O)Oc1cccc(N=CN(C)C)c1
+18,[C] [C] [=C] [C] [=N] [C] [N] [Branch1] [=Branch1] [C] [C] [C] [Ring1] [Ring1] [C] [=N] [C] [=C] [C] [=C] [Ring1] [=Branch1] [C] [=Branch1] [C] [=O] [N] [C] [Ring2] [Ring1] [Ring1] [=Ring1] [#C],-3.397,Cc3ccnc4N(C1CC1)c2ncccc2C(=O)Nc34 
+19,[C] [C] [N] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1],-2.389,CCNc1ccccc1
+20,[C] [C] [=C] [C] [=C] [C] [=C] [C] [Ring1] [=Branch1] [=C] [Branch1] [C] [C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [Ring1] [=Branch1] [=C] [Ring2] [Ring1] [Ring1] [Ring1] [#Branch2],-6.297000000000001,Cc1c2ccccc2c(C)c3ccc4ccccc4c13
+21,[F] [C] [=C] [C] [=C] [C] [Branch1] [C] [F] [=C] [Ring1] [#Branch1] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [Branch1] [C] [Cl] [=C] [Branch1] [C] [F] [C] [Branch1] [C] [Cl] [=C] [Ring1] [=Branch2] [F],-5.462000000000001,Fc1cccc(F)c1C(=O)NC(=O)Nc2cc(Cl)c(F)c(Cl)c2F 
+22,[C] [O] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=C] [Ring1] [#Branch1],-3.057,COc1ccc(Cl)cc1
+23,[O] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=N] [Ring1] [=Branch1],-4.2010000000000005,o1c2ccccc2c3ccccc13
+24,[C] [=C] [C] [=C] [N] [=C] [C] [=C] [C] [=C] [C] [Ring1] [=Branch1] [=C] [C] [Ring1] [#Branch2] [=C] [Ring1] [=C],-3.846,c3ccc2nc1ccccc1cc2c3
+25,[C] [C] [C] [C] [=Branch1] [C] [=O] [C] [C] [Branch1] [P] [C] [C] [C] [=C] [C] [=Branch1] [C] [=O] [C] [C] [C] [Ring1] [O] [Ring1] [#Branch1] [C] [C] [Ring1] [P] [C] [C] [C] [Ring2] [Ring1] [Ring2] [Branch1] [C] [O] [C] [=Branch1] [C] [=O] [C] [O],-2.893,CC12CC(=O)C3C(CCC4=CC(=O)CCC34C)C2CCC1(O)C(=O)CO
+26,[C] [C] [C] [=C] [C] [=C] [C] [Branch1] [Ring1] [C] [C] [=C] [Ring1] [Branch2] [N] [Branch1] [Ring2] [C] [O] [C] [C] [=Branch1] [C] [=O] [C] [Cl],-3.319,CCc1cccc(CC)c1N(COC)C(=O)CCl 
+27,[C] [C] [C] [C] [N] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [Branch1] [C] [Cl] [=C] [Ring1] [Branch2],-4.157,CCCCN(C)C(=O)Nc1ccc(Cl)c(Cl)c1
+28,[C] [S] [C] [=Branch1] [C] [=S] [N] [C] [Ring1] [=Branch1] [=O],-0.396,C1SC(=S)NC1(=O)
+29,[O] [C] [=C] [C] [=C] [Branch1] [Branch2] [C] [Branch1] [C] [O] [=C] [Ring1] [#Branch1] [C] [O] [C] [=C] [C] [Branch1] [C] [O] [=C] [C] [Branch1] [C] [O] [=C] [Ring1] [Branch2] [C] [=Branch1] [C] [=O] [C] [=Ring1] [=N] [O],-2.7310000000000003,Oc1ccc(c(O)c1)c3oc2cc(O)cc(O)c2c(=O)c3O 
+30,[C] [N] [Branch1] [C] [C] [C] [=N] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=C] [Ring1] [#Branch1] [C],-3.164,CN(C)C=Nc1ccc(Cl)cc1C
+31,[N] [C] [=Branch1] [C] [=O] [N] [C] [N] [C] [=Branch1] [C] [=O] [N] [C] [Ring1] [=Branch1] [=O],0.652,NC(=O)NC1NC(=O)NC1=O 
+32,[Cl] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [Ring1] [#Branch2] [Ring1] [=Branch1],-4.063,Clc1cccc2ccccc12
+33,[O] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [Branch1] [C] [Cl] [=C] [Ring1] [Branch2],-3.352,Oc1ccc(Cl)c(Cl)c1
+34,[C] [C] [Branch1] [C] [C] [C] [Branch1] [#Branch1] [C] [=C] [Branch1] [C] [Cl] [Cl] [C] [Ring1] [Branch2] [C] [=Branch1] [C] [=O] [O] [C] [Branch1] [Ring1] [C] [#N] [C] [=C] [C] [=C] [C] [Branch1] [#Branch2] [O] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [=C] [Ring1] [=N],-6.775,CC1(C)C(C=C(Cl)Cl)C1C(=O)OC(C#N)c2cccc(Oc3ccccc3)c2
+35,[C] [=C] [C] [=C] [NH1] [N] [=N] [C] [Ring1] [Branch1] [=C] [Ring1] [=Branch2],-2.21,c2ccc1[nH]nnc1c2
+36,[C] [C] [Branch1] [C] [C] [C] [Branch2] [Ring1] [Branch1] [N] [C] [=C] [C] [=C] [Branch1] [=Branch1] [C] [=C] [Ring1] [=Branch1] [Cl] [C] [Branch1] [C] [F] [Branch1] [C] [F] [F] [C] [=Branch1] [C] [=O] [O] [C] [Branch1] [Ring1] [C] [#N] [C] [=C] [C] [=C] [C] [Branch1] [#Branch2] [O] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [=C] [Ring1] [=N],-8.057,CC(C)C(Nc1ccc(cc1Cl)C(F)(F)F)C(=O)OC(C#N)c2cccc(Oc3ccccc3)c2
+37,[C] [C] [C],-1.5530000000000002,CCC
+38,[C] [C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [Ring1] [O] [=C] [Ring1] [#Branch2] [Ring1] [=Branch1],-3.792,C1Cc2cccc3cccc1c23
+39,[C] [C] [C] [#C],-1.092,CCC#C
+40,[Cl] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=C] [Ring1] [#Branch1],-3.5580000000000003,Clc1ccc(Cl)cc1
+41,[C] [C] [=C] [NH1] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch2] [Ring1] [=Branch1],-2.9810000000000003,Cc1c[nH]c2ccccc12 
+42,[C] [C] [#N],0.152,CC#N
+43,[C] [C] [C] [C] [O],-0.688,CCCCO
+44,[C] [C] [=Branch1] [C] [=C] [C] [=Branch1] [C] [=C] [C],-2.052,CC(=C)C(=C)C
+45,[C] [C] [C] [Branch1] [C] [C] [C] [C] [O],-1.308,CCC(C)CCO
+46,[Cl] [C] [=C] [C] [=C] [Branch1] [=Branch2] [C] [Branch1] [C] [Cl] [=C] [Ring1] [#Branch1] [Cl] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [Branch1] [C] [Cl] [=C] [Ring1] [Branch2] [Cl],-7.192,Clc1ccc(c(Cl)c1Cl)c2ccc(Cl)c(Cl)c2Cl 
+47,[C] [C] [=C] [C] [=Branch2] [Ring1] [=Branch1] [=C] [C] [=C] [Ring1] [=Branch1] [N] [S] [=Branch1] [C] [=O] [=Branch1] [C] [=O] [C] [Branch1] [C] [F] [Branch1] [C] [F] [F] [S] [=Branch1] [C] [=O] [=Branch1] [C] [=O] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1],-4.945,Cc1cc(ccc1NS(=O)(=O)C(F)(F)F)S(=O)(=O)c2ccccc2
+48,[O] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=C] [Ring1] [#Branch1] [Cl],-3.22,Oc1ccc(Cl)cc1Cl
+49,[C] [N] [C] [=Branch2] [Ring1] [Ring2] [=C] [Branch1] [C] [O] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [S] [Ring1] [O] [=Branch1] [C] [=O] [=O] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [C] [=N] [Ring1] [=Branch1],-3.4730000000000003,CN2C(=C(O)c1ccccc1S2(=O)=O)C(=O)Nc3ccccn3 
+50,[C] [C] [C] [C] [C] [C] [Branch1] [S] [C] [C] [C] [=C] [C] [Branch1] [C] [O] [=C] [C] [=C] [Ring1] [O] [Ring1] [#Branch1] [C] [Ring1] [#C] [C] [C] [C] [Ring2] [Ring1] [C] [=O],-3.872,CC12CCC3C(CCc4cc(O)ccc34)C2CCC1=O
+51,[C] [C] [=C] [C] [=C] [C] [=C] [Branch1] [C] [C] [C] [=C] [C] [=C] [Ring1] [O] [Ring1] [#Branch1],-4.147,Cc1cccc2c(C)cccc12
+52,[N] [S] [=Branch1] [C] [=O] [=Branch1] [C] [=O] [C] [=C] [C] [=C] [Branch1] [O] [N] [C] [N] [S] [Ring1] [=Branch1] [=Branch1] [C] [=O] [=O] [C] [=C] [Ring1] [N] [Cl],-1.72,NS(=O)(=O)c2cc1c(NCNS1(=O)=O)cc2Cl 
+53,[O] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [N] [=C] [Ring1] [#Branch2] [Ring1] [=Branch1],-2.725,Oc1cccc2cccnc12 
+54,[C] [C] [C] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [C] [Ring1] [#Branch2],-3.447,C1CCc2ccccc2C1
+55,[C] [C] [O] [C] [Branch1] [C] [C] [O] [C] [C],-0.899,CCOC(C)OCC
+56,[C] [C] [C] [C] [Ring1] [Ring1] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [N] [Branch1] [Branch1] [C] [Ring1] [Branch2] [=O] [C] [=C] [C] [Branch1] [C] [Cl] [=C] [C] [Branch1] [C] [Cl] [=C] [Ring1] [Branch2],-3.464,CC12CC2(C)C(=O)N(C1=O)c3cc(Cl)cc(Cl)c3
+57,[C] [C] [=C] [C] [=C] [C] [=C] [C] [Ring1] [=Branch1] [=C] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=C] [Ring1] [=Branch1],-4.87,Cc1c2ccccc2cc3ccccc13
+58,[C] [C] [C] [C] [O] [C],-1.072,CCCCOC
+59,[C] [C] [C] [C] [C] [=Branch1] [C] [=O] [C] [=C] [Ring1] [#Branch1] [C] [C] [C] [C] [C] [C] [C] [Branch1] [#Branch1] [C] [=Branch1] [C] [=O] [C] [O] [C] [Ring1] [=Branch2] [Branch1] [N] [C] [C] [Branch1] [C] [O] [C] [Ring2] [Ring1] [#Branch1] [Ring1] [=C] [C] [=O],-3.0660000000000003,CC13CCC(=O)C=C1CCC4C2CCC(C(=O)CO)C2(CC(O)C34)C=O 
+60,[C] [C] [C] [Branch1] [=Branch1] [C] [Branch1] [C] [C] [C] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [N] [C] [Ring1] [O] [=O],-1.6030000000000002,CCC1(C(C)C)C(=O)NC(=O)NC1=O 
+61,[C] [C] [O] [C] [=Branch1] [C] [=O] [C] [=C] [C] [=C] [Branch1] [C] [O] [C] [=C] [Ring1] [#Branch1],-2.761,CCOC(=O)c1ccc(O)cc1
+62,[C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [C] [=C] [Ring2] [Ring1] [C] [C] [=C] [Ring2] [Ring1] [C] [C] [Ring1] [S] [=C] [Ring1] [=C] [C] [Ring1] [N] [=C] [Ring1] [#Branch2] [Ring1] [=Branch1],-6.885,c1cc2ccc3ccc4ccc5ccc6ccc1c7c2c3c4c5c67
+63,[C] [C] [N] [C] [=C] [C] [Branch1] [=Branch1] [N] [Branch1] [C] [C] [C] [=C] [C] [Branch1] [C] [C] [=C] [Ring1] [#Branch2] [N] [C] [=Branch1] [C] [=O] [C] [=C] [C] [=C] [N] [=C] [Ring2] [Ring1] [Ring2] [Ring1] [=Branch1],-4.408,CCN2c1cc(N(C)C)cc(C)c1NC(=O)c3cccnc23 
+64,[C] [N] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [Branch1] [C] [Cl] [=C] [Ring1] [Branch2],-3.301,CN(C)C(=O)Nc1ccc(Cl)c(Cl)c1
+65,[C] [C] [C] [C] [C] [C] [Branch1] [C] [C] [C],-3.3080000000000003,CCCCCC(C)C
+66,[C] [O] [C] [=C] [C] [=C] [Branch1] [C] [N] [N] [=C] [Branch1] [#C] [N] [=C] [Ring1] [#Branch1] [C] [Branch1] [Ring1] [O] [C] [=C] [Ring1] [=N] [O] [C] [N] [C] [C] [N] [Branch1] [Branch1] [C] [C] [Ring1] [=Branch1] [C] [=Branch1] [C] [=O] [O] [C] [C] [Branch1] [C] [C] [Branch1] [C] [C] [O],-3.958,COc2cc1c(N)nc(nc1c(OC)c2OC)N3CCN(CC3)C(=O)OCC(C)(C)O 
+67,[C] [=C] [C] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [C] [Ring1] [#Branch1] [C] [=C] [Ring1] [O],-0.636,c1cC2C(=O)NC(=O)C2cc1
+68,[C] [C] [C] [=O],-0.3939999999999999,CCC=O
+69,[Cl] [C] [=C] [C] [=C] [Branch2] [Ring1] [=Branch2] [C] [N] [Branch1] [Branch2] [C] [C] [C] [C] [C] [Ring1] [Branch1] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [C] [=C] [Ring2] [Ring1] [=Branch1],-5.126,Clc1ccc(CN(C2CCCC2)C(=O)Nc3ccccc3)cc1
+70,[C] [C] [C] [C] [C] [Branch1] [Ring1] [C] [C] [C] [=O],-2.232,CCCCC(CC)C=O
+71,[O] [=C] [N] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [C] [Ring1] [Branch2] [Branch1] [Ring1] [C] [C] [C] [C] [C] [Branch1] [C] [C] [C],-2.312,O=C1NC(=O)NC(=O)C1(CC)CCC(C)C
+72,[C] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1],-1.857,CC(=O)Nc1ccccc1
+73,[C] [=N] [C] [=C] [C] [Branch1] [#Branch1] [C] [=Branch1] [C] [=O] [N] [N] [=C] [Ring1] [#Branch2],-0.7170000000000001,c1nccc(C(=O)NN)c1
+74,[C] [C] [Branch1] [C] [C] [C] [C] [C] [C] [Branch1] [C] [C] [Branch1] [Ring2] [C] [Ring1] [=Branch1] [C] [Ring1] [=Branch2] [=O],-2.158,CC2(C)C1CCC(C)(C1)C2=O
+75,[C] [O] [C] [=C] [N] [=C] [C] [=N] [C] [=N] [C] [Ring1] [=Branch1] [=N] [Ring1] [#Branch2],-1.589,COc2cnc1cncnc1n2
+76,[C] [N] [C] [=Branch1] [C] [=O] [C] [=C] [Branch1] [C] [C] [O] [P] [=Branch1] [C] [=O] [Branch1] [Ring1] [O] [C] [O] [C],-0.949,CNC(=O)C=C(C)OP(=O)(OC)OC
+77,[O] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [N] [Branch1] [Ring1] [C] [C] [C] [=Branch1] [C] [=O] [C] [=C] [C] [=C] [C] [=C] [Ring2] [Ring1] [C] [Ring1] [=Branch1],-3.784,O2c1ccccc1N(CC)C(=O)c3ccccc23 
+78,[C] [=C] [C] [=C] [C] [=C] [Branch1] [Ring1] [O] [C] [C] [Branch1] [Branch2] [C] [C] [=C] [Branch1] [C] [C] [C] [=C] [Ring1] [=N] [O] [C] [Ring1] [P] [=O],-4.0760000000000005,c1cc2ccc(OC)c(CC=C(C)(C))c2oc1=O 
+79,[C] [C] [C] [S] [C] [C] [C],-2.307,CCCSCCC
+80,[C] [O] [N] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=C] [Ring1] [#Branch1],-2.948,CON(C)C(=O)Nc1ccc(Cl)cc1
+81,[C] [C] [O] [C] [C],-0.718,CCOCC
+82,[C] [C] [C] [C] [C] [C] [Branch1] [S] [C] [C] [C] [=C] [C] [Branch1] [C] [O] [=C] [C] [=C] [Ring1] [O] [Ring1] [#Branch1] [C] [Ring1] [#C] [C] [C] [Branch1] [C] [O] [C] [Ring2] [Ring1] [Ring1] [O],-3.858,CC34CCC1C(CCc2cc(O)ccc12)C3CC(O)C4O 
+83,[C] [C] [N] [C] [=N] [C] [Branch1] [C] [Cl] [=N] [C] [Branch1] [O] [N] [C] [Branch1] [C] [C] [Branch1] [C] [C] [C] [#N] [=N] [Ring1] [=N],-2.49,CCNc1nc(Cl)nc(NC(C)(C)C#N)n1
+84,[C] [C] [Branch1] [C] [C] [C] [C] [Branch1] [C] [C] [Branch1] [C] [C] [O],-1.6469999999999998,CC(C)CC(C)(C)O
+85,[Cl] [C] [=C] [C] [=C] [C] [Branch1] [C] [Br] [=C] [Ring1] [#Branch1],-3.928,Clc1cccc(Br)c1
+86,[C] [C] [C] [C] [C] [C] [Branch1] [C] [O] [C] [C],-2.033,CCCCCC(O)CC
+87,[O] [=C] [N] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [C] [Ring1] [Branch2] [Branch1] [Ring1] [C] [C] [C] [C] [=C] [Branch1] [C] [C] [C],-2.126,O=C1NC(=O)NC(=O)C1(CC)CC=C(C)C
+88,[C] [C] [C] [Branch1] [C] [C] [C] [Branch1] [#Branch1] [C] [C] [Branch1] [C] [Br] [=C] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [N] [C] [Ring1] [N] [=O],-2.766,CCC(C)C1(CC(Br)=C)C(=O)NC(=O)NC1=O 
+89,[C] [O] [C] [=Branch1] [C] [=O] [C],-0.416,COC(=O)C
+90,[C] [C] [Branch1] [C] [C] [C] [=C] [C] [=C] [Branch1] [C] [C] [C] [=C] [Ring1] [#Branch1] [O],-3.129,CC(C)c1ccc(C)cc1O
+91,[C],-0.636,C
+92,[N] [C] [=N] [C] [Branch1] [C] [O] [=N] [C] [N] [=C] [NH1] [C] [Ring1] [#Branch2] [=Ring1] [Branch1],-1.74,Nc1nc(O)nc2nc[nH]c12 
+93,[F] [C] [=C] [C] [=C] [C] [Branch1] [C] [F] [=C] [Ring1] [#Branch1] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=C] [Ring1] [#Branch1],-4.692,Fc1cccc(F)c1C(=O)NC(=O)Nc2ccc(Cl)cc2
+94,[C] [C] [C] [C] [C] [Branch1] [Branch1] [C] [C] [Ring1] [=Branch1] [C] [Branch1] [C] [C] [Branch1] [C] [C] [O] [Ring1] [#Branch2],-2.579,CC12CCC(CC1)C(C)(C)O2
+95,[C] [C] [O],0.02,CCO
+96,[C] [=C] [Branch2] [Ring1] [C] [N] [C] [=Branch1] [C] [=O] [O] [C] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [N] [C] [C] [C] [=C] [C] [=C] [Ring1] [P],-2.29,c1c(NC(=O)OC(C)C(=O)NCC)cccc1
+97,[C] [C] [Branch1] [C] [C] [=C] [C] [C] [Branch2] [Ring1] [#Branch2] [C] [=Branch1] [C] [=O] [O] [C] [C] [=C] [C] [=C] [C] [Branch1] [#Branch2] [O] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1] [=C] [Ring1] [=N] [C] [Ring2] [Ring1] [Ring2] [Branch1] [C] [C] [C],-6.763,CC(C)=CC3C(C(=O)OCc2cccc(Oc1ccccc1)c2)C3(C)C 
+98,[C] [C] [C] [C] [N] [C] [=Branch1] [C] [=O] [N] [C] [Branch1] [Branch2] [N] [C] [=Branch1] [C] [=O] [O] [C] [=N] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=C] [Ring1] [=Branch1],-2.902,CCCCNC(=O)n1c(NC(=O)OC)nc2ccccc12
+99,[C] [N] [Branch1] [C] [C] [C] [=C] [C] [=C] [C] [=C] [Ring1] [=Branch1],-2.542,CN(C)c1ccccc1
+100,[C] [O] [C] [=Branch1] [C] [=O] [C] [=C],-0.878,COC(=O)C=C
+101,[C] [N] [Branch1] [C] [C] [C] [=Branch1] [C] [=O] [N] [C] [=C] [C] [=C] [Branch1] [=N] [O] [C] [=C] [C] [=C] [Branch1] [C] [Cl] [C] [=C] [Ring1] [#Branch1] [C] [=C] [Ring1] [=C],-4.477,CN(C)C(=O)Nc2ccc(Oc1ccc(Cl)cc1)cc2
+102,[O] [=C] [N] [C] [=Branch1] [C] [=O] [N] [C] [=Branch1] [C] [=O] [C] [Ring1] [Branch2] [Branch1] [=Branch1] [C] [Branch1] [C] [C] [C] [C] [C] [=C] [Branch1] [C] [C] [C],-2.465,O=C1NC(=O)NC(=O)C1(C(C)C)CC=C(C)C
+103,[C] [C] [=C] [C] [=C] [Branch1] [C] [O] [C] [=C] [Ring1] [#Branch1] [C],-2.6210000000000004,Cc1ccc(O)cc1C
+104,[Cl] [C] [=C] [C] [=C] [C] [=Branch1] [Ring2] [=N] [Ring1] [=Branch1] [C] [Branch1] [C] [Cl] [Branch1] [C] [Cl] [Cl],-3.833,Clc1cccc(n1)C(Cl)(Cl)Cl
+105,[C] [C] [=Branch1] [C] [=O] [O] [C] [Branch2] [Ring1] [=C] [C] [C] [C] [C] [C] [C] [C] [=C] [C] [=Branch1] [C] [=O] [C] [C] [C] [Ring1] [#Branch1] [C] [Ring1] [O] [C] [C] [C] [Ring2] [Ring1] [C] [Ring1] [#C] [C] [C] [#C],-4.2410000000000005,CC(=O)OC3(CCC4C2CCC1=CC(=O)CCC1C2CCC34C)C#C
+106,[C] [N] [C] [=Branch1] [C] [=O] [O] [N] [=C] [Branch1] [Ring2] [C] [S] [C] [C] [Branch1] [C] [C] [Branch1] [C] [C] [C],-2.7,CNC(=O)ON=C(CSC)C(C)(C)C 
+107,[C] [C] [C] [C] [C] [C] [C] [Branch1] [C] [C] [O],-2.033,CCCCCCC(C)O

data/lce/test.csv

@@ -0,0 +1,31 @@
+smi1,conc1,smi2,conc2,smi3,conc3,smi4,conc4,smi5,conc5,smi6,conc6,LCE
+C1C(OC(=O)O1)F,0.733,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.267,O,0.0,O,0.0,O,0.0,O,0.0,1.629
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.431,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,O,0.0,1.085
+COC(=O)OC,0.299,C(C(F)(F)F)OCC(F)(F)F,0.598,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.103,O,0.0,O,0.0,O,0.0,2.056
+COCCOC,0.358,O1CCOC1,0.532,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.074,[Li+].[N+](=O)([O-])[O-],,O,0.0,O,0.0,1.658
+C1COC(=O)O1,0.197,COC(=O)OC,0.156,COCCOCCOCCOCCOC,0.59,[Li+].F[P-](F)(F)(F)(F)F,0.026,[Li+].[N+](=O)([O-])[O-],0.031,O,0.0,1.638
+C1COC(=O)O1,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.002,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.276
+O1CCOC1,0.368,COCCOC,0.547,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.076,CSi(C)(C)([N+]).C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)C(F)(F)C(F)(F)C(F)(F)F,0.008,O,0.0,O,0.0,1.569
+COCCOC,0.507,COC(C(F)(F)F)C(F)(F)F,0.399,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.095,O,0.0,O,0.0,O,0.0,2.268
+C1COC(=O)O1,0.425,O=C(OCC)OCC(F)(F)F,0.481,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.094,O,0.0,O,0.0,O,0.0,1.602
+C1C(OC(=O)O1)F,0.318,CCOC(=O)OC,0.504,COC(=O)OC,0.094,B(O[Si](C)(C)C)(O[Si](C)(C)C)O[Si](C)(C),0.083,[Li+].F[P-](F)(F)(F)(F)F,0.001,O,0.0,1.678
+O=S1(=O)CCCC1,0.359,C(C(F)(F)F)OC(C(F)F)(F)F,0.504,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.133,[Li+].[N+](=O)([O-])[O-],0.004,O,0.0,O,0.0,2.0
+C1COC(=O)O1,0.594,O=C(OCC)OCC,0.327,[Li+].F[P-](F)(F)(F)(F)F,0.079,O,0.0,O,0.0,O,0.0,0.921
+C1COC(=O)O1,0.331,O=C(OCC)OCC,0.577,[Li+].[B-]1(OC(=O)C(=O)O1)(F)F,0.092,O,0.0,O,0.0,O,0.0,1.301
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].C(C(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(C(F)(F)F)(F)F)(F)(F)F,0.069,O,0.0,O,0.0,0.854
+C1C(OC(=O)O1)F,0.107,C1COC(=O)O1,0.526,O=C(OCC)OCC,0.289,[Li+].F[P-](F)(F)(F)(F)F,0.078,O,0.0,O,0.0,1.108
+O1CCOC1,0.322,COCCOC,0.478,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)C(F)(F)C(F)(F)C(F)(F)F,0.2,O,0.0,O,0.0,O,0.0,1.523
+CC1COC(=O)O1,0.595,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.405,O,0.0,O,0.0,O,0.0,O,0.0,1.921
+CC1COC(=O)O1,0.702,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.298,O,0.0,O,0.0,O,0.0,O,0.0,1.602
+O1CCOC1,0.375,COCCOC,0.557,[Li+][S-]SSS[S-][Li+],,[Li+].[N+](=O)([O-])[O-],0.008,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.061,O,0.0,1.523
+COC(=O)OC,0.161,FC(F)C(F)(F)COC(F)(F)C(F)F,0.355,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.484,O,0.0,O,0.0,O,0.0,2.155
+C1COC(=O)O1,0.338,COC(=O)OC,0.625,[Li+].[O-]P(=O)(F)F,0.008,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.03,O,0.0,O,0.0,1.26
+CN(C)C(=O)C(F)(F)F,0.362,C1C(OC(=O)O1)F,0.556,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.081,O,0.0,O,0.0,O,0.0,2.155
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.0,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.225
+COCCOC,0.231,FC1CCCCC1,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0.0,O,0.0,O,0.0,2.155
+COCCOC,0.277,FC(F)C(F)(F)COC(F)(F)C(F)F,0.555,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.168,O,0.0,O,0.0,O,0.0,2.155
+O1C(C)CCC1,0.331,FC(F)C(F)(F)COC(F)(F)C(F)F,0.498,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.171,O,0.0,O,0.0,O,0.0,2.301
+COCC(F)(F)C(F)(F)COC,0.864,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.136,O,0.0,O,0.0,O,0.0,O,0.0,1.991
+COC(=O)OC,0.29,C(C(F)(F)F)OCC(F)(F)F,0.589,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.121,O,0.0,O,0.0,O,0.0,2.301
+C1COC(=O)O1,0.425,O=C(OCC)OCC,0.234,[Li+].F[P-](F)(F)(F)(F)F,0.34,O,0.0,O,0.0,O,0.0,1.398
+COCCOC,0.707,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.147,[Li+].[B-]1(OC(=O)C(=O)O1)(F)F,0.147,O,0.0,O,0.0,O,0.0,1.268

data/lce/test_data.csv

@@ -0,0 +1,14 @@
+smiles1,conc1,mol1,smiles2,conc2,mol2,smiles3,conc3,mol3,smiles4,conc4,mol4,smiles5,conc5,mol5,smiles6,conc6,LCE_Predicted,LCE
+C1COC(=O)O1,0.519,51.92400559,COC(=O)OC,0.411,41.14791596,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.069,6.928078454,O,0,0,O,0,0,O,0,1.187,1.094
+COCCOC,0.596,59.5609428,COCCOCCOCCOCCOC,0.281,28.07124115,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.124,12.36781605,O,0,0,O,0,0,O,0,1.691,1.384
+C1COC(=O)O1,0.285,28.50894036,C1C(OC(=O)O1)F,0.261,26.07552384,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.228,22.82322096,COC(=O)OC,0.226,22.59231484,O,0,0,O,0,1.508,1.468
+COCCOC,0.434,43.4423376,COCCOCCOCCOCCOC,0.205,20.47449683,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.361,36.08316557,O,0,0,O,0,0,O,0,1.882,1.71
+C1C(OC(=O)O1)F,0.187,18.72872664,COC(=O)OC,0.162,16.22691423,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.109,10.92850826,FC(F)C(F)(F)COC(F)(F)C(F)F,0.541,54.11585087,O,0,0,O,0,2.103,1.832
+C1COC(=O)O1,0.134,13.35070843,C1C(OC(=O)O1)F,0.122,12.2111419,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.107,10.72028474,COC(=O)OC,0.106,10.57995858,FC(F)C(F)(F)COC(F)(F)C(F)F,0.531,53.13790635,O,0,2.077,2.104
+COCCOC,0.096,9.614613177,COCCOCCOCCOCCOC,0.045,4.53139444,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.12,12.01491409,C1COCO1,0.143,14.28400162,FC(F)C(F)(F)COC(F)(F)C(F)F,0.596,59.55507668,O,0,2.211,2.274
+C1COC(=O)O1,0.519,51.92400559,COC(=O)OC,0.411,41.14791596,[Li+].F[P-](F)(F)(F)(F)F,0.069,6.928078454,O,0,0,O,0,0,O,0,1.17,1.071
+C1COC(=O)O1,0.519,51.92400559,COC(=O)OC,0.411,41.14791596,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.069,6.928078454,O,0,0,O,0,0,O,0,1.077,1.166
+C1COC(=O)O1,0.519,51.85215842,COC(=O)OC,0.411,41.09097965,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.069,6.918492083,[Li+].[N+](=O)([O-])[O-],0.001,0.138369842,O,0,0,O,0,1.19,1.335
+C1COC(=O)O1,0.513,51.33049845,COC(=O)OC,0.407,40.6775828,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.069,6.9173773,C1=COC(=O)O1,0.011,1.07454145,O,0,0,O,0,1.114,1.129
+COCCOC,0.53,53.00533987,COCCOCCOCCOCCOC,0.25,24.98156691,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.22,22.01309322,O,0,0,O,0,0,O,0,1.758,1.501
+COCCOC,0.477,47.74974224,COCCOCCOCCOCCOC,0.225,22.50458884,[Li+].[N-](S(=O)(=O)F)S(=O)(=O)F,0.297,29.74566892,O,0,0,O,0,0,O,0,1.821,1.663

data/lce/train.csv

@@ -0,0 +1,121 @@
+smi1,conc1,smi2,conc2,smi3,conc3,smi4,conc4,smi5,conc5,smi6,conc6,LCE
+C1COC(=O)O1,0.327,O=C(OCC)OCC,0.594,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.079,O,0.0,O,0.0,O,0.0,1.155
+C1COC(=O)O1,0.356,COC(=O)OC,0.566,FC(F)(F)COB(OCC(F)(F)F)OCC(F)(F)F,0.007,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.046
+O=S1(=O)CCCC1,0.25,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.75,O,0.0,O,0.0,O,0.0,O,0.0,1.569
+C1COC(=O)O1,0.331,O=C(OCC)OCC,0.577,[Li+].F[P-](F)(F)(F)(F)F,0.092,O,0.0,O,0.0,O,0.0,0.886
+COCCOC,0.763,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.237,O,0.0,O,0.0,O,0.0,O,0.0,1.367
+COCCOC,0.2,FC(F)C(F)(F)COC(F)(F)C(F)F,0.6,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.2,O,0.0,O,0.0,O,0.0,2.301
+C1C(OC(=O)O1)F,0.873,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.127,O,0.0,O,0.0,O,0.0,O,0.0,1.489
+COCCOC,0.706,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.008,[Li+].[O-]P(=O)(F)F,0.286,O,0.0,O,0.0,O,0.0,1.244
+C1COC(=O)O1,0.3,CCOC(=O)OC,0.593,C1=COC(=O)O1,0.026,[Li+].F[P-](F)(F)(F)(F)F,0.081,O,0.0,O,0.0,0.745
+COCCOC,0.763,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.174,[Li+].[O-]P(=O)(F)F,0.063,O,0.0,O,0.0,O,0.0,1.292
+CCOCC,0.313,C(C(F)(F)F)OCC(F)(F)F,0.51,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.177,O,0.0,O,0.0,O,0.0,2.301
+O=S1(=O)CCCC1,0.75,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.25,O,0.0,O,0.0,O,0.0,O,0.0,1.745
+COC(=O)OC,0.29,C(C(F)(F)F)OCC(F)(F)F,0.589,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.121,O,0.0,O,0.0,O,0.0,1.745
+C1COC(=O)O1,0.682,CCOC(=O)OC,0.247,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.043,[Li+].O=C1O[B-]2(OC1=O)OC(=O)C(=O)O2,0.028,O,0.0,O,0.0,1.076
+C1COC(=O)O1,0.359,COC(=O)OC,0.569,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,O,0.0,0.854
+C1COC(=O)O1,0.305,COC(=O)OC,0.242,COCCOCCOCCOCCOC,0.392,[Li+].F[P-](F)(F)(F)(F)F,0.041,[Li+].[N+](=O)([O-])[O-],0.02,O,0.0,1.678
+FC(F)(F)COCCOCC,0.838,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.162,O,0.0,O,0.0,O,0.0,O,0.0,2.155
+CC#N,0.882,FC,0.065,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,,O,0.0,O,0.0,O,0.0,2.222
+COC(C)C(C)OC,0.879,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.121,O,0.0,O,0.0,O,0.0,O,0.0,1.638
+CCOP(=O)(OCC)OCC,0.728,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.272,O,0.0,O,0.0,O,0.0,O,0.0,2.0
+COC(=O)OC,0.375,FC(F)C(F)(F)COC(F)(F)C(F)F,0.375,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.25,O,0.0,O,0.0,O,0.0,1.854
+O1CCOC1,0.371,COCCOC,0.552,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.077,O,0.0,O,0.0,O,0.0,1.959
+C1C(OC(=O)O1)F,0.774,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.226,O,0.0,O,0.0,O,0.0,O,0.0,1.587
+CC1COC(=O)O1,0.875,C1C(OC(=O)O1)F,0.051,[Li+].[O-]Cl(=O)(=O)=O,0.074,O,0.0,O,0.0,O,0.0,0.699
+C1C(OC(=O)O1)F,0.264,COC(=O)OCCF,0.479,C(C(F)(F)F)OC(C(F)F)(F)F,0.155,[Li+].F[P-](F)(F)(F)(F)F,0.103,O,0.0,O,0.0,2.097
+C1C(OC(=O)O1)F,0.413,O=C(OCC)OCC,0.497,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.09,O,0.0,O,0.0,O,0.0,1.59
+C1C(OC(=O)O1)F,0.106,C1COC(=O)O1,0.522,O=C(OCC)OCC,0.287,[Li+].F[P-](F)(F)(F)(F)F,0.077,[Rb+].[O-][N+]([O-])=O,0.004,O1CCOCCOCCOCCOCCOCC1,0.004,1.252
+COCCOC,0.259,B(OCC(F)(F)F)(OCC(F)(F)F)OCC(F)(F)F,0.556,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.185,O,0.0,O,0.0,O,0.0,1.337
+C1CCOC1,0.925,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.075,O,0.0,O,0.0,O,0.0,O,0.0,1.377
+C1C(OC(=O)O1)F,0.82,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.18,O,0.0,O,0.0,O,0.0,O,0.0,1.544
+CCOP(=O)(OCC)OCC,0.5,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.5,O,0.0,O,0.0,O,0.0,O,0.0,2.097
+COCCOC,0.731,[Li+].[O-]P(=O)(F)F,0.064,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.205,O,0.0,O,0.0,O,0.0,1.215
+COCCOCCOCCOCCOC,0.819,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.181,O,0.0,O,0.0,O,0.0,O,0.0,1.222
+C1COC(=O)O1,0.338,COC(=O)OC,0.625,[Li+].[O-]P(=O)(F)F,0.008,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.03,O,0.0,O,0.0,1.194
+O1CCOC1,0.463,COCCOC,0.312,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.194,[Li+].[N+](=O)([O-])[O-],0.03,O,0.0,O,0.0,1.824
+C1C(OC(=O)O1)F,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.002,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.333
+O1CCOC1,0.539,COCCOC,0.363,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.075,[Li+].[N+](=O)([O-])[O-],0.023,O,0.0,O,0.0,1.824
+COCCOC,0.257,C(C(F)(F)F)OCC(F)(F)F,0.508,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.235,O,0.0,O,0.0,O,0.0,2.051
+COCCOC,0.906,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.047,[Li+].FP(F)(=O)([O-]),0.047,O,0.0,O,0.0,O,0.0,1.444
+O1CCOC1,0.478,COCCOC,0.322,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.134,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.067,O,0.0,O,0.0,1.854
+CCOCC,0.707,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.293,O,0.0,O,0.0,O,0.0,O,0.0,2.046
+C1COC(=O)O1,0.563,O=C(OCC)OCC,0.31,C1C(OC(=O)O1)F,0.052,[Li+].F[P-](F)(F)(F)(F)F,0.075,O,0.0,O,0.0,1.301
+C1CCOC1,0.942,FC,0.029,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,,O,0.0,O,0.0,O,0.0,2.222
+O1CCOC1,0.478,COCCOC,0.322,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.2,O,0.0,O,0.0,O,0.0,1.903
+COCCOC,0.906,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.094,O,0.0,O,0.0,O,0.0,O,0.0,1.561
+C1C(OC(=O)O1)F,0.149,COC(=O)OCCF,0.178,C(C(F)(F)F)OC(C(F)F)(F)F,0.564,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.108,O,0.0,O,0.0,1.735
+FC(F)COCCOCC(F)(F),0.845,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.155,O,0.0,O,0.0,O,0.0,O,0.0,2.301
+C1C(OC(=O)O1)F,0.495,COC(=O)OC,0.429,O1CCOCCOCCOCC1,0.003,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.498
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.069,O,0.0,O,0.0,0.745
+O=S1(=O)CCCC1,0.758,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.235,[Li+].[N+](=O)([O-])[O-],0.007,O,0.0,O,0.0,O,0.0,1.824
+CCOCC,0.856,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.144,O,0.0,O,0.0,O,0.0,O,0.0,2.0
+O=C(OCC)C,0.105,ClCCl,0.64,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.255,O,0.0,O,0.0,O,0.0,1.456
+COCCOCCOCC(F)(F)OC(F)(F)OC(F)(F)COCCOCCOC,0.708,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.292,O,0.0,O,0.0,O,0.0,O,0.0,1.301
+COCCOC,0.583,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.278,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.139,O,0.0,O,0.0,O,0.0,1.678
+C1C(OC(=O)O1)F,0.662,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.338,O,0.0,O,0.0,O,0.0,O,0.0,1.646
+O1CCOC1,0.397,COCCOC,0.589,[Li+][S-]SSS[S-][Li+],,[Li+].[N+](=O)([O-])[O-],0.012,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.002,O,0.0,1.301
+C1COC(=O)O1,0.308,O=C(OCC)OCC(F)(F)F,0.349,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.343,O,0.0,O,0.0,O,0.0,2.046
+C1COC(=O)O1,0.362,O=C(OCC)OCC,0.548,[Li+].F[P-](F)(F)(F)(F)F,0.09,O,0.0,O,0.0,O,0.0,0.788
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.001,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.373
+O1CCOCC1,0.912,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.088,O,0.0,O,0.0,O,0.0,O,0.0,1.602
+CC#N,0.621,C1=COC(=O)O1,0.056,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.323,O,0.0,O,0.0,O,0.0,1.854
+COC(=O)OC,0.684,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.316,O,0.0,O,0.0,O,0.0,O,0.0,2.097
+O=S1(=O)CCCC1,0.714,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.286,O,0.0,O,0.0,O,0.0,O,0.0,1.699
+FC(F)(F)COCCOCC(F)(F)(F),0.838,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.162,O,0.0,O,0.0,O,0.0,O,0.0,2.155
+CCOCC,0.64,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.36,O,0.0,O,0.0,O,0.0,O,0.0,2.208
+COC(=O)OC,0.6,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.4,O,0.0,O,0.0,O,0.0,O,0.0,1.77
+CC1COC(=O)O1,0.887,[Li+].F[As-](F)(F)(F)(F)F,0.113,O,0.0,O,0.0,O,0.0,O,0.0,0.824
+C1COC(=O)O1,0.5,CCOC(=O)OC,0.423,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.046,[Li+].O=C1O[B-]2(OC1=O)OC(=O)C(=O)O2,0.031,O,0.0,O,0.0,0.924
+CCOP(=O)(OCC)OCC,0.214,C(C(F)(F)F)OCC(F)(F)F,0.642,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.144,O,0.0,O,0.0,O,0.0,2.097
+COCCOC,0.682,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.318,O,0.0,O,0.0,O,0.0,O,0.0,2.108
+CC1COC(=O)O1,0.922,[LI+].F[B-](F)(F)OC(C(F)(F)(F))(C(F)(F)(F))C(F)(F)(F),0.078,O,0.0,O,0.0,O,0.0,O,0.0,0.712
+C1COC(=O)O1,0.854,CCOC(=O)OC,0.08,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.039,[Li+].O=C1O[B-]2(OC1=O)OC(=O)C(=O)O2,0.026,O,0.0,O,0.0,1.081
+C1COC(=O)O1,0.519,O=C(OCC)OCC,0.387,[Li+].F[P-](F)(F)(F)(F)F,0.082,[Li+].[O-]P(=O)(F)F,0.012,O,0.0,O,0.0,1.319
+COC(=O)CC(F)(F)F,0.768,C1C(OC(=O)O1)F,0.134,[Li+].F[P-](F)(F)(F)(F)F,0.098,O,0.0,O,0.0,O,0.0,1.62
+C1C(OC(=O)O1)F,0.144,COC(=O)OCCF,0.173,C(C(F)(F)F)OC(C(F)F)(F)F,0.548,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.135,O,0.0,O,0.0,2.222
+C1COC(=O)O1,0.326,COC(=O)OC,0.602,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,O,0.0,0.777
+CCOCC,0.877,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.123,O,0.0,O,0.0,O,0.0,O,0.0,2.018
+COC(=O)OC,0.664,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.336,O,0.0,O,0.0,O,0.0,O,0.0,1.886
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].F[B-](F)(F)F,0.069,O,0.0,O,0.0,0.699
+CCOP(=O)(OCC)OCC,0.648,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.352,O,0.0,O,0.0,O,0.0,O,0.0,1.569
+C1C(OC(=O)O1)F,0.481,O=C(OCC)OCC,0.432,[Li+].F[P-](F)(F)(F)(F)F,0.087,O,0.0,O,0.0,O,0.0,1.523
+COCCOC,0.231,FC(F)C(F)(F)COC(F)(F)C(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0.0,O,0.0,O,0.0,2.155
+C1C(OC(=O)O1)F,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.001,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.488
+O1CCOC1,0.453,COCCOC,0.305,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.127,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.063,[Li+].[N+](=O)([O-])[O-],0.051,O,0.0,2.046
+C1C(OC(=O)O1)F,0.932,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.068,O,0.0,O,0.0,O,0.0,O,0.0,1.41
+COCCOC,0.139,COCC(F)(F)C(F)(F)C(F)(F)C(F)(F)COC,0.692,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.169,O,0.0,O,0.0,O,0.0,2.222
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.431,O1CCOCCOCCOCC1,0.0,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.559
+COCCOC,0.231,FC(COC(OCC(F)(F)F)OCC(F)(F)F)(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0.0,O,0.0,O,0.0,2.301
+CN(C)S(=O)(=O)F,0.921,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.079,O,0.0,O,0.0,O,0.0,O,0.0,1.672
+C1C(OC(=O)O1)F,0.105,C1COC(=O)O1,0.518,O=C(OCC)OCC,0.285,[Li+].F[P-](F)(F)(F)(F)F,0.077,[Rb+].[O-][N+]([O-])=O,0.008,O1CCOCCOCCOCCOCCOCC1,0.008,1.538
+CC1CCC(C)O1,0.893,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.107,O,0.0,O,0.0,O,0.0,O,0.0,1.796
+C1C(OC(=O)O1)F,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.002,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0.0,O,0.0,1.355
+C1COC(=O)O1,0.444,C1COS(=O)O1,0.497,[Li+].[O-]Cl(=O)(=O)=O,0.059,O,0.0,O,0.0,O,0.0,1.523
+COCCOC,0.371,O1CCOC1,0.552,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.031,[Li+].[N+](=O)([O-])[O-],0.046,O,0.0,O,0.0,1.78
+O=S1(=O)CCCC1,0.764,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.236,O,0.0,O,0.0,O,0.0,O,0.0,1.456
+O1C(C)CCC1,0.908,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.092,O,0.0,O,0.0,O,0.0,O,0.0,1.745
+O1CCOC1,0.362,C(C(F)(F)F)OCC(F)(F)F,0.59,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.048,O,0.0,O,0.0,O,0.0,1.967
+COC(=O)OC,0.543,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.457,O,0.0,O,0.0,O,0.0,O,0.0,2.097
+COCCOC,0.73,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.27,O,0.0,O,0.0,O,0.0,O,0.0,1.143
+O1CCOC1,0.552,COCCOC,0.371,[Li+].[N+](=O)([O-])[O-],0.039,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.039,O,0.0,O,0.0,1.523
+COCCOC,0.242,FC(COC(OCC(F)(F)F)OCC(F)(F)F)(F)F,0.604,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.154,O,0.0,O,0.0,O,0.0,2.301
+CCOP(=O)(OCC)OCC,0.6,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.4,O,0.0,O,0.0,O,0.0,O,0.0,2.155
+C1C(OC(=O)O1)F,0.318,CCOC(=O)OC,0.504,COC(=O)OC,0.094,[Li+].F[P-](F)(F)(F)(F)F,0.083,O,0.0,O,0.0,1.301
+COCCOC,0.231,C(C(F)(F)F)OCC(F)(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0.0,O,0.0,O,0.0,2.222
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].F[P-](F)(F)(F)(F)F,0.069,O,0.0,O,0.0,0.699
+COCCOC,0.231,C(C(F)(F)F)OC(=O)OCC(F)(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0.0,O,0.0,O,0.0,1.495
+C1COC(=O)O1,0.32,COC(=O)OC,0.253,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.427,O,0.0,O,0.0,O,0.0,2.155
+C1C(OC(=O)O1)F,0.312,O=C1OCCC1,0.599,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.068,[Li+].[N+](=O)([O-])[O-],0.021,O,0.0,O,0.0,1.921
+COC(=O)OC,0.478,FC(F)C(F)(F)COC(F)(F)C(F)F,0.322,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.067,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.134,O,0.0,O,0.0,1.886
+CCOP(=O)(OCC)OCC,0.259,FC(F)C(F)(F)COC(F)(F)C(F)F,0.556,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.185,O,0.0,O,0.0,O,0.0,2.046
+COCCOC,0.677,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.323,O,0.0,O,0.0,O,0.0,O,0.0,1.745
+C1C(OC(=O)O1)F,0.696,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.304,O,0.0,O,0.0,O,0.0,O,0.0,1.633
+C1CCOC1,0.47,O1C(C)CCC1,0.378,[Li+].F[P-](F)(F)(F)(F)F,0.152,O,0.0,O,0.0,O,0.0,2.097
+FC(F)COCCOCC(F)(F)(F),0.838,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.162,O,0.0,O,0.0,O,0.0,O,0.0,2.301
+C1COC(=O)O1,0.496,COC(=O)OC,0.393,C1C(OC(=O)O1)F,0.045,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.066,O,0.0,O,0.0,1.108
+C1C(OC(=O)O1)F,0.62,C(C(F)(F)F)OC(=O)OCC(F)(F)F,0.291,[Li+].F[P-](F)(F)(F)(F)F,0.089,O,0.0,O,0.0,O,0.0,1.62
+CCOCC,0.906,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.094,O,0.0,O,0.0,O,0.0,O,0.0,1.959
+C1COC(=O)O1,0.526,O=C(OCC)OCC,0.392,[Li+].F[P-](F)(F)(F)(F)F,0.083,O,0.0,O,0.0,O,0.0,1.013
+C1COC(=O)O1,0.05,CCOC(=O)OC,0.237,C(C(F)(F)F)OCC(F)(F)F,0.575,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.123,[Li+].[B-]1(OC(=O)C(=O)O1)(F)F,0.015,O,0.0,1.824
+O=S1(=O)CCCC1,0.429,FC(F)C(F)(F)COC(F)(F)C(F)F,0.429,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.143,O,0.0,O,0.0,O,0.0,1.921

data/lce/train_data.csv

@@ -0,0 +1,148 @@
+smiles1,conc1,smiles2,conc2,smiles3,conc3,smiles4,conc4,smiles5,conc5,smiles6,conc6,LCE
+CC1COC(=O)O1,0.875,C1C(OC(=O)O1)F,0.051,[Li+].[O-]Cl(=O)(=O)=O,0.074,O,0,O,0,O,0,0.699
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].F[P-](F)(F)(F)(F)F,0.069,O,0,O,0,0.699
+FC(F)COCCOCC(F)(F),0.845,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.155,O,0,O,0,O,0,O,0,2.301
+FC(F)COCCOCC(F)(F)(F),0.838,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.162,O,0,O,0,O,0,O,0,2.301
+CN(C)C(=O)C(F)(F)F,0.362,C1C(OC(=O)O1)F,0.556,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.081,O,0,O,0,O,0,2.155
+COCCOC,0.231,FC1CCCCC1,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0,O,0,O,0,2.155
+CCOP(=O)(OCC)OCC,0.6,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.4,O,0,O,0,O,0,O,0,2.155
+O1CCOC1,0.362,C(C(F)(F)F)OCC(F)(F)F,0.59,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.048,O,0,O,0,O,0,1.967
+COCC(F)(F)C(F)(F)COC,0.864,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.136,O,0,O,0,O,0,O,0,1.991
+C1C(OC(=O)O1)F,0.662,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.338,O,0,O,0,O,0,O,0,1.646
+COCCOC,0.358,O1CCOC1,0.532,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.074,[Li+].[N+](=O)([O-])[O-],0.035,O,0,O,0,1.658
+CN(C)S(=O)(=O)F,0.921,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.079,O,0,O,0,O,0,O,0,1.672
+C1C(OC(=O)O1)F,0.106,C1COC(=O)O1,0.522,O=C(OCC)OCC,0.287,[Li+].F[P-](F)(F)(F)(F)F,0.077,[Rb+].[O-][N+]([O-])=O,0.004,O1CCOCCOCCOCCOCCOCC1,0.004,1.252
+C1COC(=O)O1,0.32,COC(=O)OC,0.253,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.427,O,0,O,0,O,0,2.155
+COCCOC,0.277,FC(F)C(F)(F)COC(F)(F)C(F)F,0.555,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.168,O,0,O,0,O,0,2.155
+COC(=O)OC,0.161,FC(F)C(F)(F)COC(F)(F)C(F)F,0.355,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.484,O,0,O,0,O,0,2.155
+FC(F)(F)COCCOCC,0.838,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.162,O,0,O,0,O,0,O,0,2.155
+FC(F)(F)COCCOCC(F)(F)(F),0.838,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.162,O,0,O,0,O,0,O,0,2.155
+CCOCC,0.64,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.36,O,0,O,0,O,0,O,0,2.208
+C1C(OC(=O)O1)F,0.144,COC(=O)OCCF,0.173,C(C(F)(F)F)OC(C(F)F)(F)F,0.548,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.135,O,0,O,0,2.222
+CC#N,0.882,FC,0.065,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.054,O,0,O,0,O,0,2.222
+C1CCOC1,0.942,FC,0.029,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.029,O,0,O,0,O,0,2.222
+COCCOC,0.139,COCC(F)(F)C(F)(F)C(F)(F)C(F)(F)COC,0.692,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.169,O,0,O,0,O,0,2.222
+COCCOC,0.231,C(C(F)(F)F)OCC(F)(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0,O,0,O,0,2.222
+COCCOC,0.507,COC(C(F)(F)F)C(F)(F)F,0.399,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.095,O,0,O,0,O,0,2.268
+CCOCC,0.313,C(C(F)(F)F)OCC(F)(F)F,0.51,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.177,O,0,O,0,O,0,2.301
+COC(=O)OC,0.29,C(C(F)(F)F)OCC(F)(F)F,0.589,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.121,O,0,O,0,O,0,2.301
+COCCOC,0.242,FC(COC(OCC(F)(F)F)OCC(F)(F)F)(F)F,0.604,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.154,O,0,O,0,O,0,2.301
+O1C(C)CCC1,0.331,FC(F)C(F)(F)COC(F)(F)C(F)F,0.498,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.171,O,0,O,0,O,0,2.301
+COCCOC,0.2,FC(F)C(F)(F)COC(F)(F)C(F)F,0.6,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.2,O,0,O,0,O,0,2.301
+COCCOC,0.231,FC(COC(OCC(F)(F)F)OCC(F)(F)F)(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0,O,0,O,0,2.301
+O=S1(=O)CCCC1,0.359,C(C(F)(F)F)OC(C(F)F)(F)F,0.504,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.133,[Li+].[N+](=O)([O-])[O-],0.004,O,0,O,0,2
+CCOCC,0.856,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.144,O,0,O,0,O,0,O,0,2
+CCOCC,0.877,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.123,O,0,O,0,O,0,O,0,2.018
+CCOCC,0.707,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.293,O,0,O,0,O,0,O,0,2.046
+C1COC(=O)O1,0.308,O=C(OCC)OCC(F)(F)F,0.349,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.343,O,0,O,0,O,0,2.046
+O1CCOC1,0.453,COCCOC,0.305,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.127,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.063,[Li+].[N+](=O)([O-])[O-],0.051,O,0,2.046
+CCOP(=O)(OCC)OCC,0.259,FC(F)C(F)(F)COC(F)(F)C(F)F,0.556,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.185,O,0,O,0,O,0,2.046
+COCCOC,0.257,C(C(F)(F)F)OCC(F)(F)F,0.508,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.235,O,0,O,0,O,0,2.051
+COC(=O)OC,0.299,C(C(F)(F)F)OCC(F)(F)F,0.598,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.103,O,0,O,0,O,0,2.056
+CCOP(=O)(OCC)OCC,0.214,C(C(F)(F)F)OCC(F)(F)F,0.642,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.144,O,0,O,0,O,0,2.097
+COC(=O)OC,0.684,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.316,O,0,O,0,O,0,O,0,2.097
+C1CCOC1,0.47,O1C(C)CCC1,0.378,[Li+].F[P-](F)(F)(F)(F)F,0.152,O,0,O,0,O,0,2.097
+C1C(OC(=O)O1)F,0.264,COC(=O)OCCF,0.479,C(C(F)(F)F)OC(C(F)F)(F)F,0.155,[Li+].F[P-](F)(F)(F)(F)F,0.103,O,0,O,0,2.097
+CCOP(=O)(OCC)OCC,0.5,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.5,O,0,O,0,O,0,O,0,2.097
+COC(=O)OC,0.543,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.457,O,0,O,0,O,0,O,0,2.097
+COCCOC,0.682,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.318,O,0,O,0,O,0,O,0,2.108
+COCCOC,0.231,FC(F)C(F)(F)COC(F)(F)C(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0,O,0,O,0,2.155
+CCOP(=O)(OCC)OCC,0.728,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.272,O,0,O,0,O,0,O,0,2
+COCCOC,0.583,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.278,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.139,O,0,O,0,O,0,1.678
+C1COC(=O)O1,0.305,COC(=O)OC,0.242,COCCOCCOCCOCCOC,0.392,[Li+].F[P-](F)(F)(F)(F)F,0.041,[Li+].[N+](=O)([O-])[O-],0.02,O,0,1.678
+C1C(OC(=O)O1)F,0.318,CCOC(=O)OC,0.504,COC(=O)OC,0.094,B(O[Si](C)(C)C)(O[Si](C)(C)C)O[Si](C)(C),0.083,[Li+].F[P-](F)(F)(F)(F)F,0.001,O,0,1.678
+O=S1(=O)CCCC1,0.714,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.286,O,0,O,0,O,0,O,0,1.699
+C1C(OC(=O)O1)F,0.149,COC(=O)OCCF,0.178,C(C(F)(F)F)OC(C(F)F)(F)F,0.564,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.108,O,0,O,0,1.735
+O=S1(=O)CCCC1,0.75,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.25,O,0,O,0,O,0,O,0,1.745
+COC(=O)OC,0.29,C(C(F)(F)F)OCC(F)(F)F,0.589,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.121,O,0,O,0,O,0,1.745
+COCCOC,0.677,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.323,O,0,O,0,O,0,O,0,1.745
+O1C(C)CCC1,0.908,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.092,O,0,O,0,O,0,O,0,1.745
+COC(=O)OC,0.6,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.4,O,0,O,0,O,0,O,0,1.77
+COCCOC,0.371,O1CCOC1,0.552,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.031,[Li+].[N+](=O)([O-])[O-],0.046,O,0,O,0,1.78
+CC1CCC(C)O1,0.893,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.107,O,0,O,0,O,0,O,0,1.796
+C1COC(=O)O1,0.05,CCOC(=O)OC,0.237,C(C(F)(F)F)OCC(F)(F)F,0.575,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.123,[Li+].[B-]1(OC(=O)C(=O)O1)(F)F,0.015,O,0,1.824
+O=S1(=O)CCCC1,0.758,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.235,[Li+].[N+](=O)([O-])[O-],0.007,O,0,O,0,O,0,1.824
+O1CCOC1,0.463,COCCOC,0.312,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.194,[Li+].[N+](=O)([O-])[O-],0.03,O,0,O,0,1.824
+O1CCOC1,0.539,COCCOC,0.363,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.075,[Li+].[N+](=O)([O-])[O-],0.023,O,0,O,0,1.824
+COC(=O)OC,0.375,FC(F)C(F)(F)COC(F)(F)C(F)F,0.375,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.25,O,0,O,0,O,0,1.854
+O1CCOC1,0.478,COCCOC,0.322,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.134,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.067,O,0,O,0,1.854
+CC#N,0.621,C1=COC(=O)O1,0.056,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.323,O,0,O,0,O,0,1.854
+COC(=O)OC,0.478,FC(F)C(F)(F)COC(F)(F)C(F)F,0.322,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.067,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.134,O,0,O,0,1.886
+COC(=O)OC,0.664,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.336,O,0,O,0,O,0,O,0,1.886
+O1CCOC1,0.478,COCCOC,0.322,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.2,O,0,O,0,O,0,1.903
+O=S1(=O)CCCC1,0.429,FC(F)C(F)(F)COC(F)(F)C(F)F,0.429,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.143,O,0,O,0,O,0,1.921
+C1C(OC(=O)O1)F,0.312,O=C1OCCC1,0.599,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.068,[Li+].[N+](=O)([O-])[O-],0.021,O,0,O,0,1.921
+CC1COC(=O)O1,0.595,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.405,O,0,O,0,O,0,O,0,1.921
+O1CCOC1,0.371,COCCOC,0.552,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.077,O,0,O,0,O,0,1.959
+CCOCC,0.906,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.094,O,0,O,0,O,0,O,0,1.959
+C1CCOC1,0.925,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.075,O,0,O,0,O,0,O,0,1.377
+C1COC(=O)O1,0.425,O=C(OCC)OCC,0.234,[Li+].F[P-](F)(F)(F)(F)F,0.34,O,0,O,0,O,0,1.398
+C1C(OC(=O)O1)F,0.932,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.068,O,0,O,0,O,0,O,0,1.41
+COCCOC,0.906,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.047,[Li+].FP(F)(=O)([O-]),0.047,O,0,O,0,O,0,1.444
+O=S1(=O)CCCC1,0.764,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.236,O,0,O,0,O,0,O,0,1.456
+O=C(OCC)C,0.105,ClCCl,0.64,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.255,O,0,O,0,O,0,1.456
+C1C(OC(=O)O1)F,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.001,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.488
+C1C(OC(=O)O1)F,0.873,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.127,O,0,O,0,O,0,O,0,1.489
+COCCOC,0.231,C(C(F)(F)F)OC(=O)OCC(F)(F)F,0.577,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.192,O,0,O,0,O,0,1.495
+C1C(OC(=O)O1)F,0.495,COC(=O)OC,0.429,O1CCOCCOCCOCC1,0.003,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.498
+C1C(OC(=O)O1)F,0.481,O=C(OCC)OCC,0.432,[Li+].F[P-](F)(F)(F)(F)F,0.087,O,0,O,0,O,0,1.523
+O1CCOC1,0.322,COCCOC,0.478,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)C(F)(F)C(F)(F)C(F)(F)F,0.2,O,0,O,0,O,0,1.523
+O1CCOC1,0.552,COCCOC,0.371,[Li+].[N+](=O)([O-])[O-],0.039,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.039,O,0,O,0,1.523
+C1COC(=O)O1,0.444,C1COS(=O)O1,0.497,[Li+].[O-]Cl(=O)(=O)=O,0.059,O,0,O,0,O,0,1.523
+C1C(OC(=O)O1)F,0.105,C1COC(=O)O1,0.518,O=C(OCC)OCC,0.285,[Li+].F[P-](F)(F)(F)(F)F,0.077,[Rb+].[O-][N+]([O-])=O,0.008,O1CCOCCOCCOCCOCCOCC1,0.008,1.538
+C1C(OC(=O)O1)F,0.82,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.18,O,0,O,0,O,0,O,0,1.544
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.431,O1CCOCCOCCOCC1,0,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.559
+COCCOC,0.906,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.094,O,0,O,0,O,0,O,0,1.561
+CCOP(=O)(OCC)OCC,0.648,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.352,O,0,O,0,O,0,O,0,1.569
+O=S1(=O)CCCC1,0.25,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.75,O,0,O,0,O,0,O,0,1.569
+C1C(OC(=O)O1)F,0.774,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.226,O,0,O,0,O,0,O,0,1.587
+C1C(OC(=O)O1)F,0.413,O=C(OCC)OCC,0.497,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.09,O,0,O,0,O,0,1.59
+C1COC(=O)O1,0.425,O=C(OCC)OCC(F)(F)F,0.481,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.094,O,0,O,0,O,0,1.602
+CC1COC(=O)O1,0.702,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.298,O,0,O,0,O,0,O,0,1.602
+O1CCOCC1,0.912,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.088,O,0,O,0,O,0,O,0,1.602
+C1C(OC(=O)O1)F,0.62,C(C(F)(F)F)OC(=O)OCC(F)(F)F,0.291,[Li+].F[P-](F)(F)(F)(F)F,0.089,O,0,O,0,O,0,1.62
+COC(=O)CC(F)(F)F,0.768,C1C(OC(=O)O1)F,0.134,[Li+].F[P-](F)(F)(F)(F)F,0.098,O,0,O,0,O,0,1.62
+C1C(OC(=O)O1)F,0.733,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.267,O,0,O,0,O,0,O,0,1.629
+C1C(OC(=O)O1)F,0.696,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.304,O,0,O,0,O,0,O,0,1.633
+COC(C)C(C)OC,0.879,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.121,O,0,O,0,O,0,O,0,1.638
+C1COC(=O)O1,0.197,COC(=O)OC,0.156,COCCOCCOCCOCCOC,0.59,[Li+].F[P-](F)(F)(F)(F)F,0.026,[Li+].[N+](=O)([O-])[O-],0.031,O,0,1.638
+C1COC(=O)O1,0.338,COC(=O)OC,0.625,[Li+].[O-]P(=O)(F)F,0.008,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.03,O,0,O,0,1.26
+COCCOC,0.707,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.147,[Li+].[B-]1(OC(=O)C(=O)O1)(F)F,0.147,O,0,O,0,O,0,1.268
+C1COC(=O)O1,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.002,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.276
+COCCOC,0.763,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.174,[Li+].[O-]P(=O)(F)F,0.063,O,0,O,0,O,0,1.292
+C1COC(=O)O1,0.563,O=C(OCC)OCC,0.31,C1C(OC(=O)O1)F,0.052,[Li+].F[P-](F)(F)(F)(F)F,0.075,O,0,O,0,1.301
+COCCOCCOCC(F)(F)OC(F)(F)OC(F)(F)COCCOCCOC,0.708,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.292,O,0,O,0,O,0,O,0,1.301
+C1COC(=O)O1,0.331,O=C(OCC)OCC,0.577,[Li+].[B-]1(OC(=O)C(=O)O1)(F)F,0.092,O,0,O,0,O,0,1.301
+C1C(OC(=O)O1)F,0.318,CCOC(=O)OC,0.504,COC(=O)OC,0.094,[Li+].F[P-](F)(F)(F)(F)F,0.083,O,0,O,0,1.301
+C1COC(=O)O1,0.519,O=C(OCC)OCC,0.387,[Li+].F[P-](F)(F)(F)(F)F,0.082,[Li+].[O-]P(=O)(F)F,0.012,O,0,O,0,1.319
+C1C(OC(=O)O1)F,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.002,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.333
+COCCOC,0.259,B(OCC(F)(F)F)(OCC(F)(F)F)OCC(F)(F)F,0.556,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.185,O,0,O,0,O,0,1.337
+C1C(OC(=O)O1)F,0.496,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.002,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.355
+COCCOC,0.763,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.237,O,0,O,0,O,0,O,0,1.367
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0.001,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.373
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].F[B-](F)(F)F,0.069,O,0,O,0,0.699
+CC1COC(=O)O1,0.922,[Li+].F[B-](F)(F)OC(C(F)(F)(F))(C(F)(F)(F))C(F)(F)(F),0.078,O,0,O,0,O,0,O,0,0.712
+C1COC(=O)O1,0.3,CCOC(=O)OC,0.593,C1=COC(=O)O1,0.026,[Li+].F[P-](F)(F)(F)(F)F,0.081,O,0,O,0,0.745
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.069,O,0,O,0,0.745
+C1COC(=O)O1,0.326,COC(=O)OC,0.602,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,O,0,0.777
+C1COC(=O)O1,0.362,O=C(OCC)OCC,0.548,[Li+].F[P-](F)(F)(F)(F)F,0.09,O,0,O,0,O,0,0.788
+CC1COC(=O)O1,0.887,[Li+].F[As-](F)(F)(F)(F)F,0.113,O,0,O,0,O,0,O,0,0.824
+C1COC(=O)O1,0.507,COC(=O)OC,0.402,C1=COC(=O)O1,0.022,[Li+].C(C(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(C(F)(F)F)(F)F)(F)(F)F,0.069,O,0,O,0,0.854
+C1COC(=O)O1,0.359,COC(=O)OC,0.569,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,O,0,0.854
+C1COC(=O)O1,0.331,O=C(OCC)OCC,0.577,[Li+].F[P-](F)(F)(F)(F)F,0.092,O,0,O,0,O,0,0.886
+C1COC(=O)O1,0.594,O=C(OCC)OCC,0.327,[Li+].F[P-](F)(F)(F)(F)F,0.079,O,0,O,0,O,0,0.921
+C1COC(=O)O1,0.5,CCOC(=O)OC,0.423,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.046,[Li+].O=C1O[B-]2(OC1=O)OC(=O)C(=O)O2,0.031,O,0,O,0,0.924
+C1COC(=O)O1,0.526,O=C(OCC)OCC,0.392,[Li+].F[P-](F)(F)(F)(F)F,0.083,O,0,O,0,O,0,1.013
+C1COC(=O)O1,0.356,COC(=O)OC,0.566,FC(F)(F)COB(OCC(F)(F)F)OCC(F)(F)F,0.007,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.046
+C1COC(=O)O1,0.682,CCOC(=O)OC,0.247,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.043,[Li+].O=C1O[B-]2(OC1=O)OC(=O)C(=O)O2,0.028,O,0,O,0,1.076
+C1COC(=O)O1,0.854,CCOC(=O)OC,0.08,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.039,[Li+].O=C1O[B-]2(OC1=O)OC(=O)C(=O)O2,0.026,O,0,O,0,1.081
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.431,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,O,0,1.085
+C1C(OC(=O)O1)F,0.107,C1COC(=O)O1,0.526,O=C(OCC)OCC,0.289,[Li+].F[P-](F)(F)(F)(F)F,0.078,O,0,O,0,1.108
+C1COC(=O)O1,0.496,COC(=O)OC,0.393,C1C(OC(=O)O1)F,0.045,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.066,O,0,O,0,1.108
+COCCOC,0.73,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.27,O,0,O,0,O,0,O,0,1.143
+C1COC(=O)O1,0.327,O=C(OCC)OCC,0.594,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.079,O,0,O,0,O,0,1.155
+C1COC(=O)O1,0.338,COC(=O)OC,0.625,[Li+].[O-]P(=O)(F)F,0.008,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.03,O,0,O,0,1.194
+COCCOC,0.731,[Li+].[O-]P(=O)(F)F,0.064,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.205,O,0,O,0,O,0,1.215
+COCCOCCOCCOCCOC,0.819,FS([N-]S(F)(=O)=O)(=O)=O.[Li+],0.181,O,0,O,0,O,0,O,0,1.222
+C1C(OC(=O)O1)F,0.497,COC(=O)OC,0.43,O1CCOCCOCCOCC1,0,[Li+].F[P-](F)(F)(F)(F)F,0.072,O,0,O,0,1.225
+COCCOC,0.706,[Li+].C(F)(F)(F)S(=O)(=O)[N-]S(=O)(=O)C(F)(F)F,0.008,[Li+].[O-]P(=O)(F)F,0.286,O,0,O,0,O,0,1.244

models/.gitattributes

@@ -0,0 +1,3 @@
+*.csv filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.pdf filter=lfs diff=lfs merge=lfs -text

models/fm4m.py

@@ -0,0 +1,964 @@
+from sklearn.metrics import roc_auc_score, roc_curve
+
+import datetime
+import os
+import umap
+import numpy as np
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import pickle
+import json
+
+from xgboost import XGBClassifier, XGBRegressor
+import xgboost as xgb
+from sklearn.metrics import roc_auc_score, mean_squared_error
+import xgboost as xgb
+from sklearn.svm import SVR
+from sklearn.linear_model import LinearRegression
+from sklearn.kernel_ridge import KernelRidge
+import json
+from sklearn.compose import TransformedTargetRegressor
+from sklearn.preprocessing import MinMaxScaler
+
+
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+import sys
+sys.path.append("models/")
+
+from models.selfies_ted.load import SELFIES as bart
+from models.mhg_model import load as mhg
+from models.smi_ted.smi_ted_light.load import load_smi_ted
+
+import mordred
+from mordred import Calculator, descriptors
+from rdkit import Chem
+from rdkit.Chem import AllChem
+
+datasets = {}
+models = {}
+downstream_models ={}
+
+
+def avail_models_data():
+    global datasets
+    global models
+
+    datasets = [{"Dataset": "hiv", "Input": "smiles", "Output": "HIV_active", "Path": "data/hiv", "Timestamp": "2024-06-26 11:27:37"},
+  {"Dataset": "esol", "Input": "smiles", "Output": "ESOL predicted log solubility in mols per litre", "Path": "data/esol", "Timestamp": "2024-06-26 11:31:46"},
+  {"Dataset": "freesolv", "Input": "smiles", "Output": "expt", "Path": "data/freesolv", "Timestamp": "2024-06-26 11:33:47"},
+  {"Dataset": "lipo", "Input": "smiles", "Output": "y", "Path": "data/lipo", "Timestamp": "2024-06-26 11:34:37"},
+  {"Dataset": "bace", "Input": "smiles", "Output": "Class", "Path": "data/bace", "Timestamp": "2024-06-26 11:36:40"},
+  {"Dataset": "bbbp", "Input": "smiles", "Output": "p_np", "Path": "data/bbbp", "Timestamp": "2024-06-26 11:39:23"},
+  {"Dataset": "clintox", "Input": "smiles", "Output": "CT_TOX", "Path": "data/clintox", "Timestamp": "2024-06-26 11:42:43"}]
+
+
+    models = [{"Name": "bart","Model Name": "SELFIES-TED","Description": "BART model for string based SELFIES modality", "Timestamp": "2024-06-21 12:32:20"},
+  {"Name": "mol-xl","Model Name": "MolFormer", "Description": "MolFormer model for string based SMILES modality", "Timestamp": "2024-06-21 12:35:56"},
+  {"Name": "mhg", "Model Name": "MHG-GED","Description": "Molecular hypergraph model", "Timestamp": "2024-07-10 00:09:42"},
+  {"Name": "smi-ted", "Model Name": "SMI-TED","Description": "SMILES based encoder decoder model", "Timestamp": "2024-07-10 00:09:42"}]
+
+
+def avail_models(raw=False):
+    global models
+
+    models = [{"Name": "smi-ted", "Model Name": "SMI-TED","Description": "SMILES based encoder decoder model"},
+              {"Name": "bart","Model Name": "SELFIES-TED","Description": "BART model for string based SELFIES modality"},
+              {"Name": "mol-xl","Model Name": "MolFormer", "Description": "MolFormer model for string based SMILES modality"},
+              {"Name": "mhg", "Model Name": "MHG-GED","Description": "Molecular hypergraph model"},
+              {"Name": "Mordred", "Model Name": "Mordred","Description": "Baseline: A descriptor-calculation software application that can calculate more than 1800 two- and three-dimensional descriptors"},
+              {"Name": "MorganFingerprint", "Model Name": "MorganFingerprint","Description": "Baseline: Circular atom environments based descriptor"}              
+  ]
+
+
+
+    if raw: return models
+    else:
+        return pd.DataFrame(models).drop('Name', axis=1)
+
+    return models
+
+def avail_downstream_models(raw=False):
+    global downstream_models
+
+    downstream_models = [{"Name": "XGBClassifier", "Task Type": "Classfication"},
+                         {"Name": "DefaultClassifier", "Task Type": "Classfication"},
+                        {"Name": "SVR", "Task Type": "Regression"},
+                        {"Name": "Kernel Ridge", "Task Type": "Regression"},
+                        {"Name": "Linear Regression", "Task Type": "Regression"},
+                        {"Name": "DefaultRegressor", "Task Type": "Regression"},
+                         ]
+
+    if raw: return downstream_models
+    else:
+        return pd.DataFrame(downstream_models)
+
+
+
+def avail_datasets():
+    global datasets
+
+    datasets = [{"Dataset": "hiv", "Input": "smiles", "Output": "HIV_active", "Path": "data/hiv",
+                 "Timestamp": "2024-06-26 11:27:37"},
+                {"Dataset": "esol", "Input": "smiles", "Output": "ESOL predicted log solubility in mols per litre",
+                 "Path": "data/esol", "Timestamp": "2024-06-26 11:31:46"},
+                {"Dataset": "freesolv", "Input": "smiles", "Output": "expt", "Path": "data/freesolv",
+                 "Timestamp": "2024-06-26 11:33:47"},
+                {"Dataset": "lipo", "Input": "smiles", "Output": "y", "Path": "data/lipo",
+                 "Timestamp": "2024-06-26 11:34:37"},
+                {"Dataset": "bace", "Input": "smiles", "Output": "Class", "Path": "data/bace",
+                 "Timestamp": "2024-06-26 11:36:40"},
+                {"Dataset": "bbbp", "Input": "smiles", "Output": "p_np", "Path": "data/bbbp",
+                 "Timestamp": "2024-06-26 11:39:23"},
+                {"Dataset": "clintox", "Input": "smiles", "Output": "CT_TOX", "Path": "data/clintox",
+                 "Timestamp": "2024-06-26 11:42:43"}]
+
+    return datasets
+
+def reset():
+
+    """datasets = {"esol": ["smiles", "ESOL predicted log solubility in mols per litre", "data/esol", "2024-06-26 11:36:46.509324"],
+           "freesolv": ["smiles", "expt", "data/freesolv", "2024-06-26 11:37:37.393273"],
+           "lipo": ["smiles", "y", "data/lipo", "2024-06-26 11:37:37.393273"],
+           "hiv": ["smiles", "HIV_active", "data/hiv",  "2024-06-26 11:37:37.393273"],
+           "bace": ["smiles", "Class", "data/bace", "2024-06-26 11:38:40.058354"],
+           "bbbp": ["smiles", "p_np", "data/bbbp","2024-06-26 11:38:40.058354"],
+           "clintox": ["smiles", "CT_TOX", "data/clintox","2024-06-26 11:38:40.058354"],
+           "sider": ["smiles","1:", "data/sider","2024-06-26 11:38:40.058354"],
+           "tox21": ["smiles",":-2", "data/tox21","2024-06-26 11:38:40.058354"]
+           }"""
+
+    datasets = [
+      {"Dataset": "hiv", "Input": "smiles", "Output": "HIV_active", "Path": "data/hiv", "Timestamp": "2024-06-26 11:27:37"},
+      {"Dataset": "esol", "Input": "smiles", "Output": "ESOL predicted log solubility in mols per litre", "Path": "data/esol", "Timestamp": "2024-06-26 11:31:46"},
+      {"Dataset": "freesolv", "Input": "smiles", "Output": "expt", "Path": "data/freesolv", "Timestamp": "2024-06-26 11:33:47"},
+      {"Dataset": "lipo", "Input": "smiles", "Output": "y", "Path": "data/lipo", "Timestamp": "2024-06-26 11:34:37"},
+      {"Dataset": "bace", "Input": "smiles", "Output": "Class", "Path": "data/bace", "Timestamp": "2024-06-26 11:36:40"},
+      {"Dataset": "bbbp", "Input": "smiles", "Output": "p_np", "Path": "data/bbbp", "Timestamp": "2024-06-26 11:39:23"},
+      {"Dataset": "clintox", "Input": "smiles", "Output": "CT_TOX", "Path": "data/clintox", "Timestamp": "2024-06-26 11:42:43"},
+      #{"Dataset": "sider", "Input": "smiles", "Output": "1:", "path": "data/sider", "Timestamp": "2024-06-26 11:38:40.058354"},
+      #{"Dataset": "tox21", "Input": "smiles", "Output": ":-2", "path": "data/tox21", "Timestamp": "2024-06-26 11:38:40.058354"}
+    ]
+
+    models = [{"Name": "bart", "Description": "BART model for string based SELFIES modality",
+      "Timestamp": "2024-06-21 12:32:20"},
+     {"Name": "mol-xl", "Description": "MolFormer model for string based SMILES modality",
+      "Timestamp": "2024-06-21 12:35:56"},
+     {"Name": "mhg", "Description": "MHG", "Timestamp": "2024-07-10 00:09:42"},
+     {"Name": "spec-gru", "Description": "Spectrum modality with GRU", "Timestamp": "2024-07-10 00:09:42"},
+     {"Name": "spec-lstm", "Description": "Spectrum modality with LSTM", "Timestamp": "2024-07-10 00:09:54"},
+     {"Name": "3d-vae", "Description": "VAE model for 3D atom positions", "Timestamp": "2024-07-10 00:10:08"}]
+
+
+    downstream_models = [
+        {"Name": "XGBClassifier", "Description": "XG Boost Classifier",
+         "Timestamp": "2024-06-21 12:31:20"},
+        {"Name": "XGBRegressor", "Description": "XG Boost Regressor",
+         "Timestamp": "2024-06-21 12:32:56"},
+        {"Name": "2-FNN", "Description": "A two layer feedforward network",
+         "Timestamp": "2024-06-24 14:34:16"},
+        {"Name": "3-FNN", "Description": "A three layer feedforward network",
+         "Timestamp": "2024-06-24 14:38:37"},
+    ]
+
+    with open("datasets.json", "w") as outfile:
+        json.dump(datasets, outfile)
+
+    with open("models.json", "w") as outfile:
+        json.dump(models, outfile)
+
+    with open("downstream_models.json", "w") as outfile:
+        json.dump(downstream_models, outfile)
+
+def update_data_list(list_data):
+    #datasets[list_data[0]] = list_data[1:]
+
+    with open("datasets.json", "w") as outfile:
+        json.dump(datasets, outfile)
+
+    avail_models_data()
+
+def update_model_list(list_model):
+    #models[list_model[0]] = list_model[1]
+
+    with open("models.json", "w") as outfile:
+        json.dump(list_model, outfile)
+
+    avail_models_data()
+
+def update_downstream_model_list(list_model):
+    #models[list_model[0]] = list_model[1]
+
+    with open("downstream_models.json", "w") as outfile:
+        json.dump(list_model, outfile)
+
+    avail_models_data()
+
+avail_models_data()
+
+
+
+def get_representation(train_data,test_data,model_type, return_tensor=True):
+    alias = {"MHG-GED": "mhg", "SELFIES-TED": "bart", "MolFormer": "mol-xl", "Molformer": "mol-xl", "SMI-TED": "smi-ted"}
+    if model_type in alias.keys():
+        model_type = alias[model_type]
+
+    if model_type == "mhg":
+        model = mhg.load("../models/mhg_model/pickles/mhggnn_pretrained_model_0724_2023.pickle")
+        with torch.no_grad():
+            train_emb = model.encode(train_data)
+            x_batch = torch.stack(train_emb)
+
+            test_emb = model.encode(test_data)
+            x_batch_test = torch.stack(test_emb)
+        if not return_tensor:
+            x_batch = pd.DataFrame(x_batch)
+            x_batch_test = pd.DataFrame(x_batch_test)
+
+
+    elif model_type == "bart":
+        model = bart()
+        model.load()
+        x_batch = model.encode(train_data, return_tensor=return_tensor)
+        x_batch_test = model.encode(test_data, return_tensor=return_tensor)
+
+    elif model_type == "smi-ted":
+        # model = load_smi_ted(folder='../models/smi_ted/smi_ted_light', ckpt_filename='smi-ted-Light_40.pt')
+        folder = os.path.join(os.path.dirname(__file__), '../models/smi_ted/smi_ted_light')
+        model = load_smi_ted(folder=folder, ckpt_filename='smi-ted-Light_40.pt')
+        with torch.no_grad():
+            x_batch = model.encode(train_data, return_torch=return_tensor)
+            x_batch_test = model.encode(test_data, return_torch=return_tensor)
+
+    elif model_type == "mol-xl":
+        model = AutoModel.from_pretrained("ibm/MoLFormer-XL-both-10pct", deterministic_eval=True,
+                                          trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained("ibm/MoLFormer-XL-both-10pct", trust_remote_code=True)
+
+        if type(train_data) == list:
+            inputs = tokenizer(train_data, padding=True, return_tensors="pt")
+        else:
+            inputs = tokenizer(list(train_data.values), padding=True, return_tensors="pt")
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        x_batch = outputs.pooler_output
+
+        if type(test_data) == list:
+            inputs = tokenizer(test_data, padding=True, return_tensors="pt")
+        else:
+            inputs = tokenizer(list(test_data.values), padding=True, return_tensors="pt")
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        x_batch_test = outputs.pooler_output
+
+        if not return_tensor:
+            x_batch = pd.DataFrame(x_batch)
+            x_batch_test = pd.DataFrame(x_batch_test)
+    
+    elif model_type == 'Mordred':
+        all_data = train_data + test_data
+        calc = Calculator(descriptors, ignore_3D=True)
+        mol_list = [Chem.MolFromSmiles(sm) for sm in all_data]
+        x_all = calc.pandas(mol_list)
+        print (f'original mordred fv dim: {x_all.shape}')
+        
+        for j in x_all.columns:
+            for k in range(len(x_all[j])):
+                i = x_all.loc[k, j]
+                if type(i) is mordred.error.Missing or type(i) is mordred.error.Error:
+                    x_all.loc[k, j] = np.nan
+                    
+        x_all.dropna(how="any", axis = 1, inplace=True)
+        print (f'Nan excluded mordred fv dim: {x_all.shape}')
+        
+        x_batch = x_all.iloc[:len(train_data)]
+        x_batch_test = x_all.iloc[len(train_data):]
+        # print(f'x_batch: {len(x_batch)}, x_batch_test: {len(x_batch_test)}')
+        
+    elif model_type == 'MorganFingerprint':
+        params = {'radius':2, 'nBits':1024}
+        
+        mol_train = [Chem.MolFromSmiles(sm) for sm in train_data]
+        mol_test = [Chem.MolFromSmiles(sm) for sm in test_data]
+        
+        x_batch = []
+        for mol in mol_train:
+            info = {}
+            fp = AllChem.GetMorganFingerprintAsBitVect(mol, **params, bitInfo=info)
+            vector = list(fp)
+            x_batch.append(vector)
+        x_batch = pd.DataFrame(x_batch)
+        
+        x_batch_test = []
+        for mol in mol_test:
+            info = {}
+            fp = AllChem.GetMorganFingerprintAsBitVect(mol, **params, bitInfo=info)
+            vector = list(fp)
+            x_batch_test.append(vector)
+        x_batch_test = pd.DataFrame(x_batch_test)
+
+    return x_batch, x_batch_test
+
+def single_modal(model,dataset=None, downstream_model=None, params=None, x_train=None, x_test=None, y_train=None, y_test=None):
+    print(model)
+    alias = {"MHG-GED":"mhg", "SELFIES-TED": "bart", "MolFormer":"mol-xl", "Molformer": "mol-xl", "SMI-TED": "smi-ted"}
+    data = avail_models(raw=True)
+    df = pd.DataFrame(data)
+    #print(list(df["Name"].values))
+    
+    if model in list(df["Name"].values):
+        model_type = model
+    elif alias[model] in list(df["Name"].values):
+            model_type = alias[model]
+    else:
+        print("Model not available")
+        return
+    
+
+    data = avail_datasets()
+    df = pd.DataFrame(data)
+    #print(list(df["Dataset"].values))
+
+    if dataset in list(df["Dataset"].values):
+        task = dataset
+        with open(f"representation/{task}_{model_type}.pkl", "rb") as f1:
+            x_batch, y_batch, x_batch_test, y_batch_test = pickle.load(f1)
+        print(f" Representation loaded successfully")
+
+    elif x_train==None:
+
+        print("Custom Dataset")
+        #return
+        components = dataset.split(",")
+        train_data = pd.read_csv(components[0])[components[2]]
+        test_data = pd.read_csv(components[1])[components[2]]
+
+        y_batch = pd.read_csv(components[0])[components[3]]
+        y_batch_test = pd.read_csv(components[1])[components[3]]
+
+
+        x_batch,  x_batch_test = get_representation(train_data,test_data,model_type)
+
+
+
+        print(f" Representation loaded successfully")
+
+    else:
+
+        y_batch = y_train
+        y_batch_test = y_test
+        x_batch, x_batch_test = get_representation(x_train, x_test, model_type)
+    
+    # exclude row containing Nan value
+    if isinstance(x_batch, torch.Tensor):
+        x_batch = pd.DataFrame(x_batch)    
+    nan_indices = x_batch.index[x_batch.isna().any(axis=1)]
+    if len(nan_indices) > 0:
+        x_batch.dropna(inplace = True)
+        for index in sorted(nan_indices, reverse=True):
+            del y_batch[index]
+        print(f'x_batch Nan index: {nan_indices}')
+        print(f'x_batch shape: {x_batch.shape}, y_batch len: {len(y_batch)}')
+            
+    if isinstance(x_batch_test, torch.Tensor):
+        x_batch_test = pd.DataFrame(x_batch_test)
+    nan_indices = x_batch_test.index[x_batch_test.isna().any(axis=1)]
+    if len(nan_indices) > 0:
+        x_batch_test.dropna(inplace = True)
+        for index in sorted(nan_indices, reverse=True):
+            del y_batch_test[index]
+        print(f'x_batch_test Nan index: {nan_indices}')
+        print(f'x_batch_test shape: {x_batch_test.shape}, y_batch_test len: {len(y_batch_test)}')
+
+    print(f" Calculating ROC AUC Score ...")
+
+    if downstream_model == "XGBClassifier":
+        if params == None:
+            xgb_predict_concat = XGBClassifier()
+        else:
+            xgb_predict_concat = XGBClassifier(**params) # n_estimators=5000, learning_rate=0.01, max_depth=10
+        xgb_predict_concat.fit(x_batch, y_batch)
+
+        y_prob = xgb_predict_concat.predict_proba(x_batch_test)[:, 1]
+
+        roc_auc = roc_auc_score(y_batch_test, y_prob)
+        fpr, tpr, _ = roc_curve(y_batch_test, y_prob)
+        print(f"ROC-AUC Score: {roc_auc:.4f}")
+
+        try:
+            with open(f"plot_emb/{task}_{model_type}.pkl", "rb") as f1:
+                class_0,class_1 = pickle.load(f1)
+        except:
+            print("Generating latent plots")
+            reducer = umap.UMAP(metric='euclidean', n_neighbors=10, n_components=2, low_memory=True, min_dist=0.1,
+                                verbose=False)
+            n_samples = np.minimum(1000, len(x_batch))
+
+            try:x = y_batch.values[:n_samples]
+            except: x = y_batch[:n_samples]
+            index_0 = [index for index in range(len(x)) if x[index] == 0]
+            index_1 = [index for index in range(len(x)) if x[index] == 1]
+
+            try:
+                features_umap = reducer.fit_transform(x_batch[:n_samples])
+                class_0 = features_umap[index_0]
+                class_1 = features_umap[index_1]
+            except:
+                class_0 = []
+                class_1 = []
+            print("Generating latent plots : Done")
+
+        #vizualize(roc_auc,fpr, tpr, x_batch, y_batch )
+
+        result = f"ROC-AUC Score: {roc_auc:.4f}"
+
+        return result, roc_auc,fpr, tpr, class_0, class_1
+
+    elif downstream_model == "DefaultClassifier":
+        xgb_predict_concat = XGBClassifier() # n_estimators=5000, learning_rate=0.01, max_depth=10
+        xgb_predict_concat.fit(x_batch, y_batch)
+
+        y_prob = xgb_predict_concat.predict_proba(x_batch_test)[:, 1]
+
+        roc_auc = roc_auc_score(y_batch_test, y_prob)
+        fpr, tpr, _ = roc_curve(y_batch_test, y_prob)
+        print(f"ROC-AUC Score: {roc_auc:.4f}")
+
+        try:
+            with open(f"plot_emb/{task}_{model_type}.pkl", "rb") as f1:
+                class_0,class_1 = pickle.load(f1)
+        except:
+            print("Generating latent plots")
+            reducer = umap.UMAP(metric='euclidean', n_neighbors=  10, n_components=2, low_memory=True, min_dist=0.1, verbose=False)
+            n_samples = np.minimum(1000,len(x_batch))
+
+            try:
+                x = y_batch.values[:n_samples]
+            except:
+                x = y_batch[:n_samples]
+
+            try:
+                features_umap = reducer.fit_transform(x_batch[:n_samples])
+                index_0 = [index for index in range(len(x)) if x[index] == 0]
+                index_1 = [index for index in range(len(x)) if x[index] == 1]
+
+                class_0 = features_umap[index_0]
+                class_1 = features_umap[index_1]
+            except:
+                class_0 = []
+                class_1 = []
+
+            print("Generating latent plots : Done")
+
+        #vizualize(roc_auc,fpr, tpr, x_batch, y_batch )
+
+        result = f"ROC-AUC Score: {roc_auc:.4f}"
+
+        return result, roc_auc,fpr, tpr, class_0, class_1
+    
+    elif downstream_model == "SVR":
+        if params == None:
+            regressor = SVR()
+        else:            
+            regressor = SVR(**params)
+        model = TransformedTargetRegressor(regressor= regressor,
+                                                transformer = MinMaxScaler(feature_range=(-1, 1))
+                                                ).fit(x_batch,y_batch)
+        
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+        
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+
+        print("Generating latent plots")
+        reducer = umap.UMAP(metric='euclidean', n_neighbors=10, n_components=2, low_memory=True, min_dist=0.1,
+                            verbose=False)
+        n_samples = np.minimum(1000, len(x_batch))
+
+        try: x = y_batch.values[:n_samples]
+        except: x = y_batch[:n_samples]
+        #index_0 = [index for index in range(len(x)) if x[index] == 0]
+        #index_1 = [index for index in range(len(x)) if x[index] == 1]
+
+        try:
+            features_umap = reducer.fit_transform(x_batch[:n_samples])
+            class_0 = features_umap#[index_0]
+            class_1 = features_umap#[index_1]
+        except:
+            class_0 = []
+            class_1 = []
+        print("Generating latent plots : Done")
+        
+        return result, RMSE_score,y_batch_test, y_prob, class_0, class_1
+
+    elif downstream_model == "Kernel Ridge":
+        if params == None:
+            regressor = KernelRidge()
+        else:
+            regressor = KernelRidge(**params)
+        model = TransformedTargetRegressor(regressor=regressor,
+                                           transformer=MinMaxScaler(feature_range=(-1, 1))
+                                           ).fit(x_batch, y_batch)
+
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+
+        print("Generating latent plots")
+        reducer = umap.UMAP(metric='euclidean', n_neighbors=10, n_components=2, low_memory=True, min_dist=0.1,
+                            verbose=False)
+        n_samples = np.minimum(1000, len(x_batch))
+        features_umap = reducer.fit_transform(x_batch[:n_samples])
+        try: x = y_batch.values[:n_samples]
+        except: x = y_batch[:n_samples]
+        # index_0 = [index for index in range(len(x)) if x[index] == 0]
+        # index_1 = [index for index in range(len(x)) if x[index] == 1]
+
+        class_0 = features_umap#[index_0]
+        class_1 = features_umap#[index_1]
+        print("Generating latent plots : Done")
+
+        return result, RMSE_score, y_batch_test, y_prob, class_0, class_1
+
+
+    elif downstream_model == "Linear Regression":
+        if params == None:
+            regressor = LinearRegression()
+        else:
+            regressor = LinearRegression(**params)
+        model = TransformedTargetRegressor(regressor=regressor,
+                                           transformer=MinMaxScaler(feature_range=(-1, 1))
+                                           ).fit(x_batch, y_batch)
+
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+
+        print("Generating latent plots")
+        reducer = umap.UMAP(metric='euclidean', n_neighbors=10, n_components=2, low_memory=True, min_dist=0.1,
+                            verbose=False)
+        n_samples = np.minimum(1000, len(x_batch))
+        features_umap = reducer.fit_transform(x_batch[:n_samples])
+        try:x = y_batch.values[:n_samples]
+        except: x = y_batch[:n_samples]
+        # index_0 = [index for index in range(len(x)) if x[index] == 0]
+        # index_1 = [index for index in range(len(x)) if x[index] == 1]
+
+        class_0 = features_umap#[index_0]
+        class_1 = features_umap#[index_1]
+        print("Generating latent plots : Done")
+
+        return result, RMSE_score, y_batch_test, y_prob, class_0, class_1
+
+
+    elif downstream_model == "DefaultRegressor":
+        regressor = SVR(kernel="rbf", degree=3, C=5, gamma="scale", epsilon=0.01)
+        model = TransformedTargetRegressor(regressor=regressor,
+                                           transformer=MinMaxScaler(feature_range=(-1, 1))
+                                           ).fit(x_batch, y_batch)
+
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+
+        print("Generating latent plots")
+        reducer = umap.UMAP(metric='euclidean', n_neighbors=10, n_components=2, low_memory=True, min_dist=0.1,
+                            verbose=False)
+        n_samples = np.minimum(1000, len(x_batch))
+        features_umap = reducer.fit_transform(x_batch[:n_samples])
+        try:x = y_batch.values[:n_samples]
+        except: x = y_batch[:n_samples]
+        # index_0 = [index for index in range(len(x)) if x[index] == 0]
+        # index_1 = [index for index in range(len(x)) if x[index] == 1]
+
+        class_0 = features_umap#[index_0]
+        class_1 = features_umap#[index_1]
+        print("Generating latent plots : Done")
+
+        return result, RMSE_score, y_batch_test, y_prob, class_0, class_1
+        
+
+def multi_modal(model_list,dataset=None, downstream_model=None,params=None, x_train=None, x_test=None, y_train=None, y_test=None):
+    #print(model_list)
+    data = avail_datasets()
+    df = pd.DataFrame(data)
+    list(df["Dataset"].values)
+
+    if dataset in list(df["Dataset"].values):
+        task = dataset
+        predefined = True
+    elif x_train==None:
+        predefined = False
+        components = dataset.split(",")
+        train_data = pd.read_csv(components[0])[components[2]]
+        test_data = pd.read_csv(components[1])[components[2]]
+
+        y_batch = pd.read_csv(components[0])[components[3]]
+        y_batch_test = pd.read_csv(components[1])[components[3]]
+
+        print("Custom Dataset loaded")
+    else:
+        predefined = False
+        y_batch = y_train
+        y_batch_test = y_test
+        train_data = x_train
+        test_data = x_test
+
+    data = avail_models(raw=True)
+    df = pd.DataFrame(data)
+    list(df["Name"].values)
+
+    alias = {"MHG-GED":"mhg", "SELFIES-TED": "bart", "MolFormer":"mol-xl",  "Molformer": "mol-xl","SMI-TED":"smi-ted", "Mordred": "Mordred", "MorganFingerprint": "MorganFingerprint"}
+    #if set(model_list).issubset(list(df["Name"].values)):
+    if set(model_list).issubset(list(alias.keys())):
+        for i, model in enumerate(model_list):
+            if model in alias.keys():
+                model_type = alias[model]
+            else:
+                model_type = model
+
+            if i == 0:
+                if predefined:
+                    with open(f"representation/{task}_{model_type}.pkl", "rb") as f1:
+                        x_batch, y_batch, x_batch_test, y_batch_test = pickle.load(f1)
+                    print(f" Loaded representation/{task}_{model_type}.pkl")
+                else:
+                    x_batch, x_batch_test = get_representation(train_data, test_data, model_type)
+                    x_batch = pd.DataFrame(x_batch)
+                    x_batch_test = pd.DataFrame(x_batch_test)
+
+            else:
+                if predefined:
+                    with open(f"representation/{task}_{model_type}.pkl", "rb") as f1:
+                        x_batch_1, y_batch_1, x_batch_test_1, y_batch_test_1 = pickle.load(f1)
+                        print(f" Loaded representation/{task}_{model_type}.pkl")
+                else:
+                    x_batch_1, x_batch_test_1 = get_representation(train_data, test_data, model_type)
+                    x_batch_1 = pd.DataFrame(x_batch_1)
+                    x_batch_test_1 = pd.DataFrame(x_batch_test_1)
+
+                x_batch = pd.concat([x_batch, x_batch_1], axis=1)
+                x_batch_test = pd.concat([x_batch_test, x_batch_test_1], axis=1)
+
+    else:
+        print("Model not available")
+        return
+
+    num_columns = x_batch_test.shape[1]
+    x_batch_test.columns = [f'{i + 1}' for i in range(num_columns)]
+
+    num_columns = x_batch.shape[1]
+    x_batch.columns = [f'{i + 1}' for i in range(num_columns)]
+    
+    # exclude row containing Nan value
+    if isinstance(x_batch, torch.Tensor):
+        x_batch = pd.DataFrame(x_batch)    
+    nan_indices = x_batch.index[x_batch.isna().any(axis=1)]
+    if len(nan_indices) > 0:
+        x_batch.dropna(inplace = True)
+        for index in sorted(nan_indices, reverse=True):
+            del y_batch[index]
+        print(f'x_batch Nan index: {nan_indices}')
+        print(f'x_batch shape: {x_batch.shape}, y_batch len: {len(y_batch)}')
+            
+    if isinstance(x_batch_test, torch.Tensor):
+        x_batch_test = pd.DataFrame(x_batch_test)
+    nan_indices = x_batch_test.index[x_batch_test.isna().any(axis=1)]
+    if len(nan_indices) > 0:
+        x_batch_test.dropna(inplace = True)
+        for index in sorted(nan_indices, reverse=True):
+            del y_batch_test[index]
+        print(f'x_batch_test Nan index: {nan_indices}')
+        print(f'x_batch_test shape: {x_batch_test.shape}, y_batch_test len: {len(y_batch_test)}')
+
+    print(f"Representations loaded successfully")
+    try:
+        with open(f"plot_emb/{task}_multi.pkl", "rb") as f1:
+            class_0, class_1 = pickle.load(f1)
+    except:
+        print("Generating latent plots")
+        reducer = umap.UMAP(metric='euclidean', n_neighbors=10, n_components=2, low_memory=True, min_dist=0.1,
+                            verbose=False)
+        n_samples = np.minimum(1000, len(x_batch))
+        features_umap = reducer.fit_transform(x_batch[:n_samples])
+
+        if "Classifier" in downstream_model:
+            try: x = y_batch.values[:n_samples]
+            except: x = y_batch[:n_samples]
+            index_0 = [index for index in range(len(x)) if x[index] == 0]
+            index_1 = [index for index in range(len(x)) if x[index] == 1]
+
+            class_0 = features_umap[index_0]
+            class_1 = features_umap[index_1]
+
+        else:
+            class_0 = features_umap
+            class_1 = features_umap
+
+        print("Generating latent plots : Done")
+
+    print(f" Calculating ROC AUC Score ...")
+
+
+    if downstream_model == "XGBClassifier":
+        if params == None:
+            xgb_predict_concat = XGBClassifier()
+        else:            
+            xgb_predict_concat = XGBClassifier(**params)#n_estimators=5000, learning_rate=0.01, max_depth=10)
+        xgb_predict_concat.fit(x_batch, y_batch)
+
+        y_prob = xgb_predict_concat.predict_proba(x_batch_test)[:, 1]
+
+
+        roc_auc = roc_auc_score(y_batch_test, y_prob)
+        fpr, tpr, _ = roc_curve(y_batch_test, y_prob)
+        print(f"ROC-AUC Score: {roc_auc:.4f}")
+
+        #vizualize(roc_auc,fpr, tpr, x_batch, y_batch )
+
+        #vizualize(x_batch_test, y_batch_test)
+        print(f"ROC-AUC Score: {roc_auc:.4f}")
+        result = f"ROC-AUC Score: {roc_auc:.4f}"
+
+        return result, roc_auc,fpr, tpr, class_0, class_1
+
+    elif downstream_model == "DefaultClassifier":
+        xgb_predict_concat = XGBClassifier()#n_estimators=5000, learning_rate=0.01, max_depth=10)
+        xgb_predict_concat.fit(x_batch, y_batch)
+
+        y_prob = xgb_predict_concat.predict_proba(x_batch_test)[:, 1]
+
+
+        roc_auc = roc_auc_score(y_batch_test, y_prob)
+        fpr, tpr, _ = roc_curve(y_batch_test, y_prob)
+        print(f"ROC-AUC Score: {roc_auc:.4f}")
+
+        #vizualize(roc_auc,fpr, tpr, x_batch, y_batch )
+
+        #vizualize(x_batch_test, y_batch_test)
+        print(f"ROC-AUC Score: {roc_auc:.4f}")
+        result = f"ROC-AUC Score: {roc_auc:.4f}"
+
+        return result, roc_auc,fpr, tpr, class_0, class_1
+
+    elif downstream_model == "SVR":
+        if params == None:
+            regressor = SVR()
+        else:
+            regressor = SVR(**params)
+        model = TransformedTargetRegressor(regressor= regressor,
+                                                transformer = MinMaxScaler(feature_range=(-1, 1))
+                                                ).fit(x_batch,y_batch)
+        
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+        
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+        
+        return result, RMSE_score,y_batch_test, y_prob, class_0, class_1
+
+    elif downstream_model == "Linear Regression":
+        if params == None:
+            regressor = LinearRegression()
+        else:
+            regressor = LinearRegression(**params)
+        model = TransformedTargetRegressor(regressor=regressor,
+                                           transformer=MinMaxScaler(feature_range=(-1, 1))
+                                           ).fit(x_batch, y_batch)
+
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+
+        return result, RMSE_score, y_batch_test, y_prob, class_0, class_1
+
+    elif downstream_model == "Kernel Ridge":
+        if params == None:
+            regressor = KernelRidge()
+        else:
+            regressor = KernelRidge(**params)
+        model = TransformedTargetRegressor(regressor=regressor,
+                                           transformer=MinMaxScaler(feature_range=(-1, 1))
+                                           ).fit(x_batch, y_batch)
+
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+
+        return result, RMSE_score, y_batch_test, y_prob, class_0, class_1
+
+    elif downstream_model == "DefaultRegressor":
+        regressor = SVR(kernel="rbf", degree=3, C=5, gamma="scale", epsilon=0.01)
+        model = TransformedTargetRegressor(regressor=regressor,
+                                           transformer=MinMaxScaler(feature_range=(-1, 1))
+                                           ).fit(x_batch, y_batch)
+
+        y_prob = model.predict(x_batch_test)
+        RMSE_score = np.sqrt(mean_squared_error(y_batch_test, y_prob))
+
+        print(f"RMSE Score: {RMSE_score:.4f}")
+        result = f"RMSE Score: {RMSE_score:.4f}"
+
+        return result, RMSE_score, y_batch_test, y_prob, class_0, class_1
+
+
+
+def finetune_optuna(x_batch,y_batch, x_batch_test, y_test ):
+    print(f" Finetuning with Optuna and calculating ROC AUC Score ...")
+    X_train = x_batch.values
+    y_train = y_batch.values
+    X_test = x_batch_test.values
+    y_test = y_test.values
+    def objective(trial):
+        # Define parameters to be optimized
+        params = {
+            # 'objective': 'binary:logistic',
+            'eval_metric': 'auc',
+            'verbosity': 0,
+            'n_estimators': trial.suggest_int('n_estimators', 1000, 10000),
+            # 'booster': trial.suggest_categorical('booster', ['gbtree', 'gblinear', 'dart']),
+            # 'lambda': trial.suggest_loguniform('lambda', 1e-8, 1.0),
+            'alpha': trial.suggest_loguniform('alpha', 1e-8, 1.0),
+            'max_depth': trial.suggest_int('max_depth', 1, 12),
+            # 'eta': trial.suggest_loguniform('eta', 1e-8, 1.0),
+            # 'gamma': trial.suggest_loguniform('gamma', 1e-8, 1.0),
+            # 'grow_policy': trial.suggest_categorical('grow_policy', ['depthwise', 'lossguide']),
+            # "subsample": trial.suggest_float("subsample", 0.05, 1.0),
+            # "colsample_bytree": trial.suggest_float("colsample_bytree", 0.05, 1.0),
+        }
+
+        # Train XGBoost model
+        dtrain = xgb.DMatrix(X_train, label=y_train)
+        dtest = xgb.DMatrix(X_test, label=y_test)
+
+        model = xgb.train(params, dtrain)
+
+        # Predict probabilities
+        y_pred = model.predict(dtest)
+
+        # Calculate ROC AUC score
+        roc_auc = roc_auc_score(y_test, y_pred)
+        print("ROC_AUC : ", roc_auc)
+
+        return roc_auc
+
+def add_new_model():
+    models = avail_models(raw=True)
+
+    # Function to display models
+    def display_models():
+        for model in models:
+            model_display = f"Name: {model['Name']}, Description: {model['Description']}, Timestamp: {model['Timestamp']}"
+            print(model_display)
+
+    # Function to update models
+    def update_models(new_name, new_description, new_path):
+        new_model = {
+            "Name": new_name,
+            "Description": new_description,
+            "Timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            #"path": new_path
+        }
+        models.append(new_model)
+        with open("models.json", "w") as outfile:
+            json.dump(models, outfile)
+
+        print("Model uploaded and updated successfully!")
+        list_models()
+        #display_models()
+
+    # Widgets
+    name_text = widgets.Text(description="Name:", layout=Layout(width='50%'))
+    description_text = widgets.Text(description="Description:", layout=Layout(width='50%'))
+    path_text = widgets.Text(description="Path:", layout=Layout(width='50%'))
+
+    def browse_callback(b):
+        root = tk.Tk()
+        root.withdraw()  # Hide the main window
+        file_path = filedialog.askopenfilename(title="Select a Model File")
+        if file_path:
+            path_text.value = file_path
+
+    browse_button = widgets.Button(description="Browse")
+    browse_button.on_click(browse_callback)
+
+    def submit_callback(b):
+        update_models(name_text.value, description_text.value, path_text.value)
+
+    submit_button = widgets.Button(description="Submit")
+    submit_button.on_click(submit_callback)
+
+    # Display widgets
+    display(VBox([name_text, description_text, path_text, browse_button, submit_button]))
+
+
+def add_new_dataset():
+    # Sample data
+    datasets = avail_datasets()
+
+    # Function to display models
+    def display_datasets():
+        for dataset in datasets:
+            dataset_display = f"Name: {dataset['Dataset']}, Input: {dataset['Input']},Output: {dataset['Output']},Path: {dataset['Path']}, Timestamp: {dataset['Timestamp']}"
+
+    # Function to update models
+    def update_datasets(new_dataset, new_input, new_output, new_path):
+        new_model = {
+            "Dataset": new_dataset,
+            "Input": new_input,
+            "Output": new_output,
+            "Timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            "Path": os.path.basename(new_path)
+        }
+        datasets.append(new_model)
+        with open("datasets.json", "w") as outfile:
+            json.dump(datasets, outfile)
+
+        print("Dataset uploaded and updated successfully!")
+        list_data()
+
+
+    # Widgets
+    dataset_text = widgets.Text(description="Dataset:", layout=Layout(width='50%'))
+    input_text = widgets.Text(description="Input:", layout=Layout(width='50%'))
+    output_text = widgets.Text(description="Output:", layout=Layout(width='50%'))
+    path_text = widgets.Text(description="Path:", layout=Layout(width='50%'))
+
+    def browse_callback(b):
+        root = tk.Tk()
+        root.withdraw()  # Hide the main window
+        file_path = filedialog.askopenfilename(title="Select a Dataset File")
+        if file_path:
+            path_text.value = file_path
+
+    browse_button = widgets.Button(description="Browse")
+    browse_button.on_click(browse_callback)
+
+    def submit_callback(b):
+        update_datasets(dataset_text.value, input_text.value, output_text.value, path_text.value)
+
+    submit_button = widgets.Button(description="Submit")
+    submit_button.on_click(submit_callback)
+
+    display(VBox([dataset_text, input_text, output_text, path_text, browse_button, submit_button]))
+
+
+

models/mhg_model/README.md

@@ -0,0 +1,75 @@
+# mhg-gnn
+
+This repository provides PyTorch source code assosiated with our publication, "MHG-GNN: Combination of Molecular Hypergraph Grammar with Graph Neural Network"
+
+**Paper:** [Arxiv Link](https://arxiv.org/pdf/2309.16374)
+
+![mhg-gnn](images/mhg_example1.png)
+
+## Introduction
+
+We present MHG-GNN, an autoencoder architecture
+that has an encoder based on GNN and a decoder based on a sequential model with MHG.
+Since the encoder is a GNN variant, MHG-GNN can accept any molecule as input, and  
+demonstrate high predictive performance on molecular graph data.
+In addition, the decoder inherits the theoretical guarantee of MHG on always generating a structurally valid molecule as output.
+
+## Table of Contents
+
+1. [Getting Started](#getting-started)
+    1. [Pretrained Models and Training Logs](#pretrained-models-and-training-logs)
+    2. [Installation](#installation)
+2. [Feature Extraction](#feature-extraction)
+
+## Getting Started
+
+**This code and environment have been tested on Intel E5-2667 CPUs at 3.30GHz and NVIDIA A100 Tensor Core GPUs.**
+
+### Pretrained Models and Training Logs
+
+We provide checkpoints of the MHG-GNN model pre-trained on a dataset of ~1.34M molecules curated from PubChem. (later) For model weights: [HuggingFace Link](https://huggingface.co/ibm/materials.mhg-ged/blob/main/mhggnn_pretrained_model_0724_2023.pickle)
+
+Add the MHG-GNN `pre-trained weights.pt` to the `models/` directory according to your needs. 
+
+### Installation
+
+We recommend to create a virtual environment. For example:
+
+```
+python3 -m venv .venv
+. .venv/bin/activate
+```
+
+Type the following command once the virtual environment is activated:
+
+```
+git clone [email protected]:CMD-TRL/mhg-gnn.git
+cd ./mhg-gnn
+pip install .
+```
+
+## Feature Extraction
+
+The example notebook [mhg-gnn_encoder_decoder_example.ipynb](notebooks/mhg-gnn_encoder_decoder_example.ipynb) contains code to load checkpoint files and use the pre-trained model for encoder and decoder tasks.
+
+To load mhg-gnn, you can simply use:
+
+```python
+import torch
+import load
+
+model = load.load()
+```
+
+To encode SMILES into embeddings, you can use:
+
+```python
+with torch.no_grad():
+    repr = model.encode(["CCO", "O=C=O", "OC(=O)c1ccccc1C(=O)O"])
+```
+
+For decoder, you can use the function, so you can return from embeddings to SMILES strings:
+
+```python
+orig = model.decode(repr)
+```

models/mhg_model/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding:utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#

models/mhg_model/graph_grammar/__init__.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 1 2018"
+

models/mhg_model/graph_grammar/algo/__init__.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+# -*- coding:utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 1 2018"
+

models/mhg_model/graph_grammar/algo/tree_decomposition.py

@@ -0,0 +1,821 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2017"
+__version__ = "0.1"
+__date__ = "Dec 11 2017"
+
+from copy import deepcopy
+from itertools import combinations
+from ..hypergraph import Hypergraph
+import networkx as nx
+import numpy as np
+
+
+class CliqueTree(nx.Graph):
+    ''' clique tree object
+
+    Attributes
+    ----------
+    hg : Hypergraph
+        This hypergraph will be decomposed.
+    root_hg : Hypergraph
+        Hypergraph on the root node.
+    ident_node_dict : dict
+        ident_node_dict[key_node] gives a list of nodes that are identical (i.e., the adjacent hyperedges are common)
+    '''
+    def __init__(self, hg=None, **kwargs):
+        self.hg = deepcopy(hg)
+        if self.hg is not None:
+            self.ident_node_dict = self.hg.get_identical_node_dict()
+        else:
+            self.ident_node_dict = {}
+        super().__init__(**kwargs)
+
+    @property
+    def root_hg(self):
+        ''' return the hypergraph on the root node
+        '''
+        return self.nodes[0]['subhg']
+
+    @root_hg.setter
+    def root_hg(self, hypergraph):
+        ''' set the hypergraph on the root node
+        '''
+        self.nodes[0]['subhg'] = hypergraph
+
+    def insert_subhg(self, subhypergraph: Hypergraph) -> None:
+        ''' insert a subhypergraph, which is extracted from a root hypergraph, into the tree.
+
+        Parameters
+        ----------
+        subhg : Hypergraph
+        '''
+        num_nodes = self.number_of_nodes()
+        self.add_node(num_nodes, subhg=subhypergraph)
+        self.add_edge(num_nodes, 0)
+        adj_nodes = deepcopy(list(self.adj[0].keys()))
+        for each_node in adj_nodes:
+            if len(self.nodes[each_node]["subhg"].nodes.intersection(
+                    self.nodes[num_nodes]["subhg"].nodes)\
+                   - self.root_hg.nodes) != 0 and each_node != num_nodes:
+                self.remove_edge(0, each_node)
+                self.add_edge(each_node, num_nodes)
+
+    def to_irredundant(self) -> None:
+        ''' convert the clique tree to be irredundant
+        '''
+        for each_node in self.hg.nodes:
+            subtree = self.subgraph([
+                each_tree_node for each_tree_node in self.nodes()\
+                if each_node in self.nodes[each_tree_node]["subhg"].nodes]).copy()
+            leaf_node_list = [x for x in subtree.nodes() if subtree.degree(x)==1]
+            redundant_leaf_node_list = []
+            for each_leaf_node in leaf_node_list:
+                if len(self.nodes[each_leaf_node]["subhg"].adj_edges(each_node)) == 0:
+                    redundant_leaf_node_list.append(each_leaf_node)
+            for each_red_leaf_node in redundant_leaf_node_list:
+                current_node = each_red_leaf_node
+                while subtree.degree(current_node) == 1 \
+                      and len(subtree.nodes[current_node]["subhg"].adj_edges(each_node)) == 0:
+                    self.nodes[current_node]["subhg"].remove_node(each_node)
+                    remove_node = current_node
+                    current_node = list(dict(subtree[remove_node]).keys())[0]
+                    subtree.remove_node(remove_node)
+
+        fixed_node_set = deepcopy(self.nodes)
+        for each_node in fixed_node_set:
+            if self.nodes[each_node]["subhg"].num_edges == 0:
+                if len(self[each_node]) == 1:
+                    self.remove_node(each_node)
+                elif len(self[each_node]) == 2:
+                    self.add_edge(*self[each_node])
+                    self.remove_node(each_node)
+                else:
+                    pass
+            else:
+                pass
+
+        redundant = True
+        while redundant:
+            redundant = False
+            fixed_edge_set = deepcopy(self.edges)
+            remove_node_set = set()
+            for node_1, node_2 in fixed_edge_set:
+                if node_1 in remove_node_set or node_2 in remove_node_set:
+                    pass
+                else:
+                    if self.nodes[node_1]['subhg'].is_subhg(self.nodes[node_2]['subhg']):
+                        redundant = True
+                        adj_node_list = set(self.adj[node_1]) - {node_2}
+                        self.remove_node(node_1)
+                        remove_node_set.add(node_1)
+                        for each_node in adj_node_list:
+                            self.add_edge(node_2, each_node)
+
+                    elif self.nodes[node_2]['subhg'].is_subhg(self.nodes[node_1]['subhg']):
+                        redundant = True
+                        adj_node_list = set(self.adj[node_2]) - {node_1}
+                        self.remove_node(node_2)
+                        remove_node_set.add(node_2)
+                        for each_node in adj_node_list:
+                            self.add_edge(node_1, each_node)
+
+    def node_update(self, key_node: str, subhg) -> None:
+        """ given a pair of a hypergraph, H, and its subhypergraph, sH, return a hypergraph H\sH.
+
+        Parameters
+        ----------
+        key_node : str
+            key node that must be removed.
+        subhg : Hypegraph
+        """
+        for each_edge in subhg.edges:
+            self.root_hg.remove_edge(each_edge)
+        self.root_hg.remove_nodes(self.ident_node_dict[key_node])
+
+        adj_node_list = list(subhg.nodes)
+        for each_node in subhg.nodes:
+            if each_node not in self.ident_node_dict[key_node]:
+                if set(self.root_hg.adj_edges(each_node)).issubset(subhg.edges):
+                    self.root_hg.remove_node(each_node)
+                    adj_node_list.remove(each_node)
+            else:
+                adj_node_list.remove(each_node)
+
+        for each_node_1, each_node_2 in combinations(adj_node_list, 2):
+            if not self.root_hg.is_adj(each_node_1, each_node_2):
+                self.root_hg.add_edge(set([each_node_1, each_node_2]), attr_dict=dict(tmp=True))
+        
+        subhg.remove_edges_with_attr({'tmp' : True})
+        self.insert_subhg(subhg)
+
+    def update(self, subhg, remove_nodes=False):
+        """ given a pair of a hypergraph, H, and its subhypergraph, sH, return a hypergraph H\sH.
+
+        Parameters
+        ----------
+        subhg : Hypegraph
+        """
+        for each_edge in subhg.edges:
+            self.root_hg.remove_edge(each_edge)
+        if remove_nodes:
+            remove_edge_list = []
+            for each_edge in self.root_hg.edges:
+                if set(self.root_hg.nodes_in_edge(each_edge)).issubset(subhg.nodes)\
+                   and self.root_hg.edge_attr(each_edge).get('tmp', False):
+                    remove_edge_list.append(each_edge)
+            self.root_hg.remove_edges(remove_edge_list)
+
+        adj_node_list = list(subhg.nodes)
+        for each_node in subhg.nodes:
+            if self.root_hg.degree(each_node) == 0:
+                self.root_hg.remove_node(each_node)
+                adj_node_list.remove(each_node)
+
+        if len(adj_node_list) != 1 and not remove_nodes:
+            self.root_hg.add_edge(set(adj_node_list), attr_dict=dict(tmp=True))
+        '''
+        else:
+            for each_node_1, each_node_2 in combinations(adj_node_list, 2):
+                if not self.root_hg.is_adj(each_node_1, each_node_2):
+                    self.root_hg.add_edge(
+                        [each_node_1, each_node_2], attr_dict=dict(tmp=True))
+        '''
+        subhg.remove_edges_with_attr({'tmp':True})
+        self.insert_subhg(subhg)
+
+
+def _get_min_deg_node(hg, ident_node_dict: dict, mode='mol'):
+    if mode == 'standard':
+        degree_dict = hg.degrees()
+        min_deg_node = min(degree_dict, key=degree_dict.get)
+        min_deg_subhg = hg.adj_subhg(min_deg_node, ident_node_dict)
+        return min_deg_node, min_deg_subhg
+    elif mode == 'mol':
+        degree_dict = hg.degrees()
+        min_deg = min(degree_dict.values())
+        min_deg_node_list = [each_node for each_node in hg.nodes if degree_dict[each_node]==min_deg]
+        min_deg_subhg_list = [hg.adj_subhg(each_min_deg_node, ident_node_dict)
+                              for each_min_deg_node in min_deg_node_list]
+        best_score = np.inf
+        best_idx = -1
+        for each_idx in range(len(min_deg_subhg_list)):
+            if min_deg_subhg_list[each_idx].num_nodes < best_score:
+                best_idx = each_idx
+        return min_deg_node_list[each_idx], min_deg_subhg_list[each_idx]
+    else:
+        raise ValueError
+
+
+def tree_decomposition(hg, irredundant=True):
+    """ compute a tree decomposition of the input hypergraph
+
+    Parameters
+    ----------
+    hg : Hypergraph
+        hypergraph to be decomposed
+    irredundant : bool
+        if True, irredundant tree decomposition will be computed.
+
+    Returns
+    -------
+    clique_tree : nx.Graph
+        each node contains a subhypergraph of `hg`
+    """
+    org_hg = hg.copy()
+    ident_node_dict = hg.get_identical_node_dict()
+    clique_tree = CliqueTree(org_hg)
+    clique_tree.add_node(0, subhg=org_hg)
+    while True:
+        degree_dict = org_hg.degrees()
+        min_deg_node = min(degree_dict, key=degree_dict.get)
+        min_deg_subhg = org_hg.adj_subhg(min_deg_node, ident_node_dict)
+        if org_hg.nodes == min_deg_subhg.nodes:
+            break
+
+        # org_hg and min_deg_subhg are divided
+        clique_tree.node_update(min_deg_node, min_deg_subhg)
+
+    clique_tree.root_hg.remove_edges_with_attr({'tmp' : True})
+
+    if irredundant:
+        clique_tree.to_irredundant()
+    return clique_tree
+
+
+def tree_decomposition_with_hrg(hg, hrg, irredundant=True, return_root=False):
+    ''' compute a tree decomposition given a hyperedge replacement grammar.
+    the resultant clique tree should induce a less compact HRG.
+    
+    Parameters
+    ----------
+    hg : Hypergraph
+        hypergraph to be decomposed
+    hrg : HyperedgeReplacementGrammar
+        current HRG
+    irredundant : bool
+        if True, irredundant tree decomposition will be computed.
+
+    Returns
+    -------
+    clique_tree : nx.Graph
+        each node contains a subhypergraph of `hg`
+    '''
+    org_hg = hg.copy()
+    ident_node_dict = hg.get_identical_node_dict()
+    clique_tree = CliqueTree(org_hg)
+    clique_tree.add_node(0, subhg=org_hg)
+    root_node = 0
+    
+    # construct a clique tree using HRG
+    success_any = True
+    while success_any:
+        success_any = False
+        for each_prod_rule in hrg.prod_rule_list:
+            org_hg, success, subhg = each_prod_rule.revert(org_hg, True)
+            if success:
+                if each_prod_rule.is_start_rule: root_node = clique_tree.number_of_nodes()
+                success_any = True
+                subhg.remove_edges_with_attr({'terminal' : False})
+                clique_tree.root_hg = org_hg
+                clique_tree.insert_subhg(subhg)
+    
+    clique_tree.root_hg = org_hg
+    
+    for each_edge in deepcopy(org_hg.edges):
+        if not org_hg.edge_attr(each_edge)['terminal']:
+            node_list = org_hg.nodes_in_edge(each_edge)
+            org_hg.remove_edge(each_edge)
+            
+            for each_node_1, each_node_2 in combinations(node_list, 2):
+                if not org_hg.is_adj(each_node_1, each_node_2):
+                    org_hg.add_edge([each_node_1, each_node_2], attr_dict=dict(tmp=True))
+
+    # construct a clique tree using the existing algorithm
+    degree_dict = org_hg.degrees()
+    if degree_dict:
+        while True:
+            min_deg_node, min_deg_subhg = _get_min_deg_node(org_hg, ident_node_dict)
+            if org_hg.nodes == min_deg_subhg.nodes: break
+
+            # org_hg and min_deg_subhg are divided
+            clique_tree.node_update(min_deg_node, min_deg_subhg)
+
+    clique_tree.root_hg.remove_edges_with_attr({'tmp' : True})
+    if irredundant:
+        clique_tree.to_irredundant()
+
+    if return_root:
+        if root_node == 0 and 0 not in clique_tree.nodes:
+            root_node = clique_tree.number_of_nodes()
+            while root_node not in clique_tree.nodes:
+                root_node -= 1
+        elif root_node not in clique_tree.nodes:
+            while root_node not in clique_tree.nodes:
+                root_node -= 1
+        else:
+            pass
+        return clique_tree, root_node
+    else:
+        return clique_tree
+
+
+def tree_decomposition_from_leaf(hg, irredundant=True):
+    """ compute a tree decomposition of the input hypergraph
+
+    Parameters
+    ----------
+    hg : Hypergraph
+        hypergraph to be decomposed
+    irredundant : bool
+        if True, irredundant tree decomposition will be computed.
+
+    Returns
+    -------
+    clique_tree : nx.Graph
+        each node contains a subhypergraph of `hg`
+    """
+    def apply_normal_decomposition(clique_tree):
+        degree_dict = clique_tree.root_hg.degrees()
+        min_deg_node = min(degree_dict, key=degree_dict.get)
+        min_deg_subhg = clique_tree.root_hg.adj_subhg(min_deg_node, clique_tree.ident_node_dict)
+        if clique_tree.root_hg.nodes == min_deg_subhg.nodes:
+            return clique_tree, False
+        clique_tree.node_update(min_deg_node, min_deg_subhg)
+        return clique_tree, True
+
+    def apply_min_edge_deg_decomposition(clique_tree):
+        edge_degree_dict = clique_tree.root_hg.edge_degrees()
+        non_tmp_edge_list = [each_edge for each_edge in clique_tree.root_hg.edges \
+                             if not clique_tree.root_hg.edge_attr(each_edge).get('tmp')]
+        if not non_tmp_edge_list:
+            return clique_tree, False
+        min_deg_edge = None
+        min_deg = np.inf
+        for each_edge in non_tmp_edge_list:
+            if min_deg > edge_degree_dict[each_edge]:
+                min_deg_edge = each_edge
+                min_deg = edge_degree_dict[each_edge]
+        node_list = clique_tree.root_hg.nodes_in_edge(min_deg_edge)
+        min_deg_subhg = clique_tree.root_hg.get_subhg(
+            node_list, [min_deg_edge], clique_tree.ident_node_dict)
+        if clique_tree.root_hg.nodes == min_deg_subhg.nodes:
+            return clique_tree, False
+        clique_tree.update(min_deg_subhg)
+        return clique_tree, True
+
+    org_hg = hg.copy()
+    clique_tree = CliqueTree(org_hg)
+    clique_tree.add_node(0, subhg=org_hg)
+
+    success = True
+    while success:
+        clique_tree, success = apply_min_edge_deg_decomposition(clique_tree)
+        if not success:
+            clique_tree, success = apply_normal_decomposition(clique_tree)
+
+    clique_tree.root_hg.remove_edges_with_attr({'tmp' : True})
+    if irredundant:
+        clique_tree.to_irredundant()
+    return clique_tree
+
+def topological_tree_decomposition(
+        hg, irredundant=True, rip_labels=True, shrink_cycle=False, contract_cycles=False):
+    ''' compute a tree decomposition of the input hypergraph
+
+    Parameters
+    ----------
+    hg : Hypergraph
+        hypergraph to be decomposed
+    irredundant : bool
+        if True, irredundant tree decomposition will be computed.
+
+    Returns
+    -------
+    clique_tree : CliqueTree
+        each node contains a subhypergraph of `hg`
+    '''
+    def _contract_tree(clique_tree):
+        ''' contract a single leaf
+
+        Parameters
+        ----------
+        clique_tree : CliqueTree
+
+        Returns
+        -------
+        CliqueTree, bool
+            bool represents whether this operation succeeds or not.
+        '''
+        edge_degree_dict = clique_tree.root_hg.edge_degrees()
+        leaf_edge_list = [each_edge for each_edge in clique_tree.root_hg.edges \
+                          if (not clique_tree.root_hg.edge_attr(each_edge).get('tmp'))\
+                          and edge_degree_dict[each_edge] == 1]
+        if not leaf_edge_list:
+            return clique_tree, False
+        min_deg_edge = leaf_edge_list[0]
+        node_list = clique_tree.root_hg.nodes_in_edge(min_deg_edge)
+        min_deg_subhg = clique_tree.root_hg.get_subhg(
+            node_list, [min_deg_edge], clique_tree.ident_node_dict)
+        if clique_tree.root_hg.nodes == min_deg_subhg.nodes:
+            return clique_tree, False
+        clique_tree.update(min_deg_subhg)
+        return clique_tree, True
+
+    def _rip_labels_from_cycles(clique_tree, org_hg):
+        ''' rip hyperedge-labels off
+
+        Parameters
+        ----------
+        clique_tree : CliqueTree
+        org_hg : Hypergraph
+
+        Returns
+        -------
+        CliqueTree, bool
+            bool represents whether this operation succeeds or not.
+        '''
+        ident_node_dict = clique_tree.ident_node_dict #hg.get_identical_node_dict()
+        for each_edge in clique_tree.root_hg.edges:
+            if each_edge in org_hg.edges:
+                if org_hg.in_cycle(each_edge):
+                    node_list = clique_tree.root_hg.nodes_in_edge(each_edge)
+                    subhg = clique_tree.root_hg.get_subhg(
+                        node_list, [each_edge], ident_node_dict)
+                    if clique_tree.root_hg.nodes == subhg.nodes:
+                        return clique_tree, False
+                    clique_tree.update(subhg)
+                    '''
+                    in_cycle_dict = {each_node: org_hg.node_attr(each_node)['is_in_ring'] for each_node in node_list}
+                    if not all(in_cycle_dict.values()):
+                        node_not_in_cycle = [each_node for each_node in in_cycle_dict.keys() if not in_cycle_dict[each_node]][0]
+                        node_list = [node_not_in_cycle]
+                        node_list.extend(clique_tree.root_hg.adj_nodes(node_not_in_cycle))
+                        edge_list = clique_tree.root_hg.adj_edges(node_not_in_cycle)
+                        import pdb; pdb.set_trace()
+                        subhg = clique_tree.root_hg.get_subhg(
+                            node_list, edge_list, ident_node_dict)
+                        
+                        clique_tree.update(subhg)
+                    '''
+                    return clique_tree, True
+        return clique_tree, False
+
+    def _shrink_cycle(clique_tree):
+        ''' shrink a cycle
+
+        Parameters
+        ----------
+        clique_tree : CliqueTree
+
+        Returns
+        -------
+        CliqueTree, bool
+            bool represents whether this operation succeeds or not.
+        '''
+        def filter_subhg(subhg, hg, key_node):
+            num_nodes_cycle = 0
+            nodes_in_cycle_list = []
+            for each_node in subhg.nodes:
+                if hg.in_cycle(each_node):
+                    num_nodes_cycle += 1
+                    if each_node != key_node:
+                        nodes_in_cycle_list.append(each_node)
+                if num_nodes_cycle > 3:
+                    break
+            if num_nodes_cycle != 3:
+                return False
+            else:
+                for each_edge in hg.edges:
+                    if set(nodes_in_cycle_list).issubset(hg.nodes_in_edge(each_edge)):
+                        return False
+                return True
+
+        #ident_node_dict = hg.get_identical_node_dict()
+        ident_node_dict = clique_tree.ident_node_dict
+        for each_node in clique_tree.root_hg.nodes:
+            if clique_tree.root_hg.in_cycle(each_node)\
+               and filter_subhg(clique_tree.root_hg.adj_subhg(each_node, ident_node_dict),
+                                clique_tree.root_hg,
+                                each_node):
+                target_node = each_node
+                target_subhg = clique_tree.root_hg.adj_subhg(target_node, ident_node_dict)
+                if clique_tree.root_hg.nodes == target_subhg.nodes:
+                    return clique_tree, False
+                clique_tree.update(target_subhg)
+                return clique_tree, True
+        return clique_tree, False
+
+    def _contract_cycles(clique_tree):
+        '''
+        remove a subhypergraph that looks like a cycle on a leaf.
+
+        Parameters
+        ----------
+        clique_tree : CliqueTree
+
+        Returns
+        -------
+        CliqueTree, bool
+            bool represents whether this operation succeeds or not.
+        '''
+        def _divide_hg(hg):
+            ''' divide a hypergraph into subhypergraphs such that
+            each subhypergraph is connected to each other in a tree-like way.
+
+            Parameters
+            ----------
+            hg : Hypergraph
+
+            Returns
+            -------
+            list of Hypergraphs
+                each element corresponds to a subhypergraph of `hg`
+            '''
+            for each_node in hg.nodes:
+                if hg.is_dividable(each_node):
+                    adj_edges_dict = {each_edge: hg.in_cycle(each_edge) for each_edge in hg.adj_edges(each_node)}
+                    '''
+                    if any(adj_edges_dict.values()):
+                        import pdb; pdb.set_trace()
+                        edge_in_cycle = [each_key for each_key, each_val in adj_edges_dict.items() if each_val][0]
+                        subhg1, subhg2, subhg3 = hg.divide(each_node, edge_in_cycle)
+                        return _divide_hg(subhg1) + _divide_hg(subhg2) + _divide_hg(subhg3)
+                    else:
+                    '''
+                    subhg1, subhg2 = hg.divide(each_node)
+                    return _divide_hg(subhg1) + _divide_hg(subhg2)
+            return [hg]
+
+        def _is_leaf(hg, divided_subhg) -> bool:
+            ''' judge whether subhg is a leaf-like in the original hypergraph
+
+            Parameters
+            ----------
+            hg : Hypergraph
+            divided_subhg : Hypergraph
+                `divided_subhg` is a subhypergraph of `hg`
+
+            Returns
+            -------
+            bool
+            '''
+            '''
+            adj_edges_set = set([])
+            for each_node in divided_subhg.nodes:
+                adj_edges_set.update(set(hg.adj_edges(each_node)))
+
+
+            _hg = deepcopy(hg)
+            _hg.remove_subhg(divided_subhg)
+            if nx.is_connected(_hg.hg) != (len(adj_edges_set - divided_subhg.edges) == 1):
+                import pdb; pdb.set_trace()
+            return len(adj_edges_set - divided_subhg.edges) == 1
+            '''
+            _hg = deepcopy(hg)
+            _hg.remove_subhg(divided_subhg)
+            return nx.is_connected(_hg.hg)
+        
+        subhg_list = _divide_hg(clique_tree.root_hg)
+        if len(subhg_list) == 1:
+            return clique_tree, False
+        else:
+            while len(subhg_list) > 1:
+                max_leaf_subhg = None
+                for each_subhg in subhg_list:
+                    if _is_leaf(clique_tree.root_hg, each_subhg):
+                        if max_leaf_subhg is None:
+                            max_leaf_subhg = each_subhg
+                        elif max_leaf_subhg.num_nodes < each_subhg.num_nodes:
+                            max_leaf_subhg = each_subhg
+                clique_tree.update(max_leaf_subhg)
+                subhg_list.remove(max_leaf_subhg)
+            return clique_tree, True
+
+    org_hg = hg.copy()
+    clique_tree = CliqueTree(org_hg)
+    clique_tree.add_node(0, subhg=org_hg)
+
+    success = True
+    while success:
+        '''
+        clique_tree, success = _rip_labels_from_cycles(clique_tree, hg)
+        if not success:
+            clique_tree, success = _contract_cycles(clique_tree)
+        '''
+        clique_tree, success = _contract_tree(clique_tree)
+        if not success:
+            if rip_labels:
+                clique_tree, success = _rip_labels_from_cycles(clique_tree, hg)
+            if not success:
+                if shrink_cycle:
+                    clique_tree, success = _shrink_cycle(clique_tree)
+                if not success:
+                    if contract_cycles:
+                        clique_tree, success = _contract_cycles(clique_tree)
+    clique_tree.root_hg.remove_edges_with_attr({'tmp' : True})
+    if irredundant:
+        clique_tree.to_irredundant()
+    return clique_tree
+
+def molecular_tree_decomposition(hg, irredundant=True):
+    """ compute a tree decomposition of the input molecular hypergraph
+
+    Parameters
+    ----------
+    hg : Hypergraph
+        molecular hypergraph to be decomposed
+    irredundant : bool
+        if True, irredundant tree decomposition will be computed.
+
+    Returns
+    -------
+    clique_tree : CliqueTree
+        each node contains a subhypergraph of `hg`
+    """
+    def _divide_hg(hg):
+        ''' divide a hypergraph into subhypergraphs such that
+        each subhypergraph is connected to each other in a tree-like way.
+
+        Parameters
+        ----------
+        hg : Hypergraph
+
+        Returns
+        -------
+        list of Hypergraphs
+            each element corresponds to a subhypergraph of `hg`
+        '''
+        is_ring = False
+        for each_node in hg.nodes:
+            if hg.node_attr(each_node)['is_in_ring']:
+                is_ring = True
+            if not hg.node_attr(each_node)['is_in_ring'] \
+               and hg.degree(each_node) == 2:
+                subhg1, subhg2 = hg.divide(each_node)
+                return _divide_hg(subhg1) + _divide_hg(subhg2)
+
+        if is_ring:
+            subhg_list = []
+            remove_edge_list = []
+            remove_node_list = []
+            for each_edge in hg.edges:
+                node_list = hg.nodes_in_edge(each_edge)
+                subhg = hg.get_subhg(node_list, [each_edge], hg.get_identical_node_dict())
+                subhg_list.append(subhg)
+                remove_edge_list.append(each_edge)
+                for each_node in node_list:
+                    if not hg.node_attr(each_node)['is_in_ring']:
+                        remove_node_list.append(each_node)
+            hg.remove_edges(remove_edge_list)
+            hg.remove_nodes(remove_node_list, False)
+            return subhg_list + [hg]
+        else:
+            return [hg]
+
+    org_hg = hg.copy()
+    clique_tree = CliqueTree(org_hg)
+    clique_tree.add_node(0, subhg=org_hg)
+
+    subhg_list = _divide_hg(deepcopy(clique_tree.root_hg))
+    #_subhg_list = deepcopy(subhg_list)
+    if len(subhg_list) == 1:
+        pass
+    else:
+        while len(subhg_list) > 1:
+            max_leaf_subhg = None
+            for each_subhg in subhg_list:
+                if _is_leaf(clique_tree.root_hg, each_subhg) and not _is_ring(each_subhg):
+                    if max_leaf_subhg is None:
+                        max_leaf_subhg = each_subhg
+                    elif max_leaf_subhg.num_nodes < each_subhg.num_nodes:
+                        max_leaf_subhg = each_subhg
+
+            if max_leaf_subhg is None:
+                for each_subhg in subhg_list:
+                    if _is_ring_label(clique_tree.root_hg, each_subhg):
+                        if max_leaf_subhg is None:
+                            max_leaf_subhg = each_subhg
+                        elif max_leaf_subhg.num_nodes < each_subhg.num_nodes:
+                            max_leaf_subhg = each_subhg
+            if max_leaf_subhg is not None:
+                clique_tree.update(max_leaf_subhg)
+                subhg_list.remove(max_leaf_subhg)
+            else:
+                for each_subhg in subhg_list:
+                    if _is_leaf(clique_tree.root_hg, each_subhg):
+                        if max_leaf_subhg is None:
+                            max_leaf_subhg = each_subhg
+                        elif max_leaf_subhg.num_nodes < each_subhg.num_nodes:
+                            max_leaf_subhg = each_subhg
+                if max_leaf_subhg is not None:
+                    clique_tree.update(max_leaf_subhg, True)
+                    subhg_list.remove(max_leaf_subhg)
+                else:
+                    break
+    if len(subhg_list) > 1:
+        '''
+        for each_idx, each_subhg in enumerate(subhg_list):
+            each_subhg.draw(f'{each_idx}', True)
+        clique_tree.root_hg.draw('root', True)
+        import pickle
+        with open('buggy_hg.pkl', 'wb') as f:
+            pickle.dump(hg, f)
+        return clique_tree, subhg_list, _subhg_list
+        '''
+        raise RuntimeError('bug in tree decomposition algorithm')
+    clique_tree.root_hg.remove_edges_with_attr({'tmp' : True})
+
+    '''
+    for each_tree_node in clique_tree.adj[0]:
+        subhg = clique_tree.nodes[each_tree_node]['subhg']
+        for each_edge in subhg.edges:
+            if set(subhg.nodes_in_edge(each_edge)).issubset(clique_tree.root_hg.nodes):
+                clique_tree.root_hg.add_edge(set(subhg.nodes_in_edge(each_edge)), attr_dict=dict(tmp=True))
+    '''
+    if irredundant:
+        clique_tree.to_irredundant()
+    return clique_tree #, _subhg_list
+
+def _is_leaf(hg, subhg) -> bool:
+    ''' judge whether subhg is a leaf-like in the original hypergraph
+
+    Parameters
+    ----------
+    hg : Hypergraph
+    subhg : Hypergraph
+        `subhg` is a subhypergraph of `hg`
+
+    Returns
+    -------
+    bool
+    '''
+    if len(subhg.edges) == 0:
+        adj_edge_set = set([])
+        subhg_edge_set = set([])
+        for each_edge in hg.edges:
+            if set(hg.nodes_in_edge(each_edge)).issubset(subhg.nodes) and hg.edge_attr(each_edge).get('tmp', False):
+                subhg_edge_set.add(each_edge)
+        for each_node in subhg.nodes:
+            adj_edge_set.update(set(hg.adj_edges(each_node)))
+        if subhg_edge_set.issubset(adj_edge_set) and len(adj_edge_set.difference(subhg_edge_set)) == 1:
+            return True
+        else:
+            return False
+    elif len(subhg.edges) == 1:
+        adj_edge_set = set([])
+        subhg_edge_set = subhg.edges
+        for each_node in subhg.nodes:
+            for each_adj_edge in hg.adj_edges(each_node):
+                adj_edge_set.add(each_adj_edge)
+        if subhg_edge_set.issubset(adj_edge_set) and len(adj_edge_set.difference(subhg_edge_set)) == 1:
+            return True
+        else:
+            return False
+    else:
+        raise ValueError('subhg should be nodes only or one-edge hypergraph.')
+
+def _is_ring_label(hg, subhg):
+    if len(subhg.edges) != 1:
+        return False
+    edge_name = list(subhg.edges)[0]
+    #assert edge_name in hg.edges, f'{edge_name}'
+    is_in_ring = False
+    for each_node in subhg.nodes:
+        if subhg.node_attr(each_node)['is_in_ring']:
+            is_in_ring = True
+        else:
+            adj_edge_list = list(hg.adj_edges(each_node))
+            adj_edge_list.remove(edge_name)
+            if len(adj_edge_list) == 1:
+                if not hg.edge_attr(adj_edge_list[0]).get('tmp', False):
+                    return False
+            elif len(adj_edge_list) == 0:
+                pass
+            else:
+                raise ValueError
+    if is_in_ring:
+        return True
+    else:
+        return False
+
+def _is_ring(hg):
+    for each_node in hg.nodes:
+        if not hg.node_attr(each_node)['is_in_ring']:
+            return False
+    return True        
+

models/mhg_model/graph_grammar/graph_grammar/__init__.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 1 2018"
+

models/mhg_model/graph_grammar/graph_grammar/base.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2017"
+__version__ = "0.1"
+__date__ = "Dec 11 2017"
+
+from abc import ABCMeta, abstractmethod
+
+class GraphGrammarBase(metaclass=ABCMeta):
+    @abstractmethod
+    def learn(self):
+        pass
+    
+    @abstractmethod
+    def sample(self):
+        pass

models/mhg_model/graph_grammar/graph_grammar/corpus.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jun 4 2018"
+
+from collections import Counter
+from functools import partial
+from .utils import _easy_node_match, _edge_match, _node_match, common_node_list, _node_match_prod_rule
+from networkx.algorithms.isomorphism import GraphMatcher
+import os
+
+
+class CliqueTreeCorpus(object):
+
+    ''' clique tree corpus
+
+    Attributes
+    ----------
+    clique_tree_list : list of CliqueTree
+    subhg_list : list of Hypergraph
+    '''
+
+    def __init__(self):
+        self.clique_tree_list = []
+        self.subhg_list = []
+
+    @property
+    def size(self):
+        return len(self.subhg_list)
+
+    def add_clique_tree(self, clique_tree):
+        for each_node in clique_tree.nodes:
+            subhg = clique_tree.nodes[each_node]['subhg']
+            subhg_idx = self.add_subhg(subhg)
+            clique_tree.nodes[each_node]['subhg_idx'] = subhg_idx
+        self.clique_tree_list.append(clique_tree)
+
+    def add_to_subhg_list(self, clique_tree, root_node):
+        parent_node_dict = {}
+        current_node = None
+        parent_node_dict[root_node] = None
+        stack = [root_node]
+        while stack:
+            current_node = stack.pop()
+            current_subhg = clique_tree.nodes[current_node]['subhg']
+            for each_child in clique_tree.adj[current_node]:
+                if each_child != parent_node_dict[current_node]:
+                    stack.append(each_child)
+                    parent_node_dict[each_child] = current_node
+            if parent_node_dict[current_node] is not None:
+                parent_subhg = clique_tree.nodes[parent_node_dict[current_node]]['subhg']
+                common, _ = common_node_list(parent_subhg, current_subhg)
+                parent_subhg.add_edge(set(common), attr_dict={'tmp': True})
+
+        parent_node_dict = {}
+        current_node = None
+        parent_node_dict[root_node] = None
+        stack = [root_node]
+        while stack:
+            current_node = stack.pop()
+            current_subhg = clique_tree.nodes[current_node]['subhg']
+            for each_child in clique_tree.adj[current_node]:
+                if each_child != parent_node_dict[current_node]:
+                    stack.append(each_child)
+                    parent_node_dict[each_child] = current_node
+            if parent_node_dict[current_node] is not None:
+                parent_subhg = clique_tree.nodes[parent_node_dict[current_node]]['subhg']
+                common, _ = common_node_list(parent_subhg, current_subhg)
+                for each_idx, each_node in enumerate(common):
+                    current_subhg.set_node_attr(each_node, {'ext_id': each_idx})
+
+            subhg_idx, is_new = self.add_subhg(current_subhg)
+            clique_tree.nodes[current_node]['subhg_idx'] = subhg_idx
+        return clique_tree
+
+    def add_subhg(self, subhg):
+        if len(self.subhg_list) == 0:
+            node_dict = {}
+            for each_node in subhg.nodes:
+                node_dict[each_node] = subhg.node_attr(each_node)['symbol'].__hash__()
+            node_list = []
+            for each_key, _ in sorted(node_dict.items(), key=lambda x:x[1]):
+                node_list.append(each_key)
+            for each_idx, each_node in enumerate(node_list):
+                subhg.node_attr(each_node)['order4hrg'] = each_idx
+            self.subhg_list.append(subhg)
+            return 0, True
+        else:
+            match = False
+            subhg_bond_symbol_counter \
+                = Counter([subhg.node_attr(each_node)['symbol'] \
+                           for each_node in subhg.nodes])
+            subhg_atom_symbol_counter \
+                = Counter([subhg.edge_attr(each_edge).get('symbol', None) \
+                           for each_edge in subhg.edges])
+            for each_idx, each_subhg in enumerate(self.subhg_list):
+                each_bond_symbol_counter \
+                    = Counter([each_subhg.node_attr(each_node)['symbol'] \
+                               for each_node in each_subhg.nodes])
+                each_atom_symbol_counter \
+                    = Counter([each_subhg.edge_attr(each_edge).get('symbol', None) \
+                               for each_edge in each_subhg.edges])
+                if not match \
+                   and (subhg.num_nodes == each_subhg.num_nodes
+                        and subhg.num_edges == each_subhg.num_edges
+                        and subhg_bond_symbol_counter == each_bond_symbol_counter
+                        and subhg_atom_symbol_counter == each_atom_symbol_counter):
+                    gm = GraphMatcher(each_subhg.hg,
+                                      subhg.hg,
+                                      node_match=_easy_node_match,
+                                      edge_match=_edge_match)
+                    try:
+                        isomap = next(gm.isomorphisms_iter())
+                        match = True
+                        for each_node in each_subhg.nodes:
+                            subhg.node_attr(isomap[each_node])['order4hrg'] \
+                                = each_subhg.node_attr(each_node)['order4hrg']
+                            if 'ext_id' in each_subhg.node_attr(each_node):
+                                subhg.node_attr(isomap[each_node])['ext_id'] \
+                                    = each_subhg.node_attr(each_node)['ext_id']
+                        return each_idx, False
+                    except StopIteration:
+                        match = False    
+            if not match:
+                node_dict = {}
+                for each_node in subhg.nodes:
+                    node_dict[each_node] = subhg.node_attr(each_node)['symbol'].__hash__()
+                node_list = []
+                for each_key, _ in sorted(node_dict.items(), key=lambda x:x[1]):
+                    node_list.append(each_key)
+                for each_idx, each_node in enumerate(node_list):
+                    subhg.node_attr(each_node)['order4hrg'] = each_idx
+
+                #for each_idx, each_node in enumerate(subhg.nodes):
+                #    subhg.node_attr(each_node)['order4hrg'] = each_idx
+                self.subhg_list.append(subhg)
+                return len(self.subhg_list) - 1, True

models/mhg_model/graph_grammar/graph_grammar/hrg.py

@@ -0,0 +1,1065 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2017"
+__version__ = "0.1"
+__date__ = "Dec 11 2017"
+
+from .corpus import CliqueTreeCorpus
+from .base import GraphGrammarBase
+from .symbols import TSymbol, NTSymbol, BondSymbol
+from .utils import _node_match, _node_match_prod_rule, _edge_match, masked_softmax, common_node_list
+from ..hypergraph import Hypergraph
+from collections import Counter
+from copy import deepcopy
+from ..algo.tree_decomposition import (
+    tree_decomposition,
+    tree_decomposition_with_hrg,
+    tree_decomposition_from_leaf,
+    topological_tree_decomposition,
+    molecular_tree_decomposition)
+from functools import partial
+from networkx.algorithms.isomorphism import GraphMatcher
+from typing import List, Dict, Tuple
+import networkx as nx
+import numpy as np
+import torch
+import os
+import random
+
+DEBUG = False
+
+
+class ProductionRule(object):
+    """ A class of a production rule
+
+    Attributes
+    ----------
+    lhs : Hypergraph or None
+        the left hand side of the production rule.
+        if None, the rule is a starting rule.
+    rhs : Hypergraph
+        the right hand side of the production rule.
+    """
+    def __init__(self, lhs, rhs):
+        self.lhs = lhs
+        self.rhs = rhs
+
+    @property
+    def is_start_rule(self) -> bool:
+        return self.lhs.num_nodes == 0
+
+    @property
+    def ext_node(self) -> Dict[int, str]:
+        """ return a dict of external nodes
+        """
+        if self.is_start_rule:
+            return {}
+        else:
+            ext_node_dict = {}
+            for each_node in self.lhs.nodes:
+                ext_node_dict[self.lhs.node_attr(each_node)["ext_id"]] = each_node
+            return ext_node_dict
+
+    @property
+    def lhs_nt_symbol(self) -> NTSymbol:
+        if self.is_start_rule:
+            return NTSymbol(degree=0, is_aromatic=False, bond_symbol_list=[])
+        else:
+            return self.lhs.edge_attr(list(self.lhs.edges)[0])['symbol']
+
+    def rhs_adj_mat(self, node_edge_list):
+        ''' return the adjacency matrix of rhs of the production rule
+        '''
+        return nx.adjacency_matrix(self.rhs.hg, node_edge_list)
+        
+    def draw(self, file_path=None):
+        return self.rhs.draw(file_path)
+
+    def is_same(self, prod_rule, ignore_order=False):
+        """ judge whether this production rule is
+        the same as the input one, `prod_rule`
+
+        Parameters
+        ----------
+        prod_rule : ProductionRule
+            production rule to be compared
+
+        Returns
+        -------
+        is_same : bool
+        isomap : dict
+            isomorphism of nodes and hyperedges.
+            ex) {'bond_42': 'bond_37', 'bond_2': 'bond_1',
+                 'e36': 'e11', 'e16': 'e12', 'e25': 'e18',
+                 'bond_40': 'bond_38', 'e26': 'e21', 'bond_41': 'bond_39'}.
+            key comes from `prod_rule`, value comes from `self`.
+        """
+        if self.is_start_rule:
+            if not prod_rule.is_start_rule:
+                return False, {}
+        else:
+            if prod_rule.is_start_rule:
+                return False, {}
+            else:
+                if prod_rule.lhs.num_nodes != self.lhs.num_nodes:
+                    return False, {}
+
+        if prod_rule.rhs.num_nodes != self.rhs.num_nodes:
+            return False, {}
+        if prod_rule.rhs.num_edges != self.rhs.num_edges:
+            return False, {}
+
+        subhg_bond_symbol_counter \
+            = Counter([prod_rule.rhs.node_attr(each_node)['symbol'] \
+                       for each_node in prod_rule.rhs.nodes])
+        each_bond_symbol_counter \
+            = Counter([self.rhs.node_attr(each_node)['symbol'] \
+                       for each_node in self.rhs.nodes])
+        if subhg_bond_symbol_counter != each_bond_symbol_counter:
+            return False, {}
+
+        subhg_atom_symbol_counter \
+            = Counter([prod_rule.rhs.edge_attr(each_edge)['symbol'] \
+                       for each_edge in prod_rule.rhs.edges])
+        each_atom_symbol_counter \
+            = Counter([self.rhs.edge_attr(each_edge)['symbol'] \
+                       for each_edge in self.rhs.edges])
+        if subhg_atom_symbol_counter != each_atom_symbol_counter:
+            return False, {}
+
+        gm = GraphMatcher(prod_rule.rhs.hg,
+                          self.rhs.hg,
+                          partial(_node_match_prod_rule,
+                                  ignore_order=ignore_order),
+                          partial(_edge_match,
+                                  ignore_order=ignore_order))
+        try:
+            return True, next(gm.isomorphisms_iter())
+        except StopIteration:
+            return False, {}
+
+    def applied_to(self,
+                   hg: Hypergraph,
+                   edge: str) -> Tuple[Hypergraph, List[str]]:
+        """ augment `hg` by replacing `edge` with `self.rhs`.
+
+        Parameters
+        ----------
+        hg : Hypergraph
+        edge : str
+            `edge` must belong to `hg`
+
+        Returns
+        -------
+        hg : Hypergraph
+            resultant hypergraph
+        nt_edge_list : list
+            list of non-terminal edges
+        """
+        nt_edge_dict = {}
+        if self.is_start_rule:
+            if (edge is not None) or (hg is not None):
+                ValueError("edge and hg must be None for this prod rule.")
+            hg = Hypergraph()
+            node_map_rhs = {} # node id in rhs -> node id in hg, where rhs is augmented.
+            for num_idx, each_node in enumerate(self.rhs.nodes):
+                hg.add_node(f"bond_{num_idx}",
+                            #attr_dict=deepcopy(self.rhs.node_attr(each_node)))
+                            attr_dict=self.rhs.node_attr(each_node))
+                node_map_rhs[each_node] = f"bond_{num_idx}"
+            for each_edge in self.rhs.edges:
+                node_list = []
+                for each_node in self.rhs.nodes_in_edge(each_edge):
+                    node_list.append(node_map_rhs[each_node])
+                if isinstance(self.rhs.nodes_in_edge(each_edge), set):
+                    node_list = set(node_list)
+                edge_id = hg.add_edge(
+                    node_list,
+                    #attr_dict=deepcopy(self.rhs.edge_attr(each_edge)))
+                    attr_dict=self.rhs.edge_attr(each_edge))
+                if "nt_idx" in hg.edge_attr(edge_id):
+                    nt_edge_dict[hg.edge_attr(edge_id)["nt_idx"]] = edge_id
+            nt_edge_list = [nt_edge_dict[key] for key in range(len(nt_edge_dict))]
+            return hg, nt_edge_list
+        else:
+            if edge not in hg.edges:
+                raise ValueError("the input hyperedge does not exist.")
+            if hg.edge_attr(edge)["terminal"]:
+                raise ValueError("the input hyperedge is terminal.")
+            if hg.edge_attr(edge)['symbol'] != self.lhs_nt_symbol:
+                print(hg.edge_attr(edge)['symbol'], self.lhs_nt_symbol)
+                raise ValueError("the input hyperedge and lhs have inconsistent number of nodes.")
+            if DEBUG:
+                for node_idx, each_node in enumerate(hg.nodes_in_edge(edge)):
+                    other_node = self.lhs.nodes_in_edge(list(self.lhs.edges)[0])[node_idx]
+                    attr = deepcopy(self.lhs.node_attr(other_node))
+                    attr.pop('ext_id')
+                    if hg.node_attr(each_node) != attr:
+                        raise ValueError('node attributes are inconsistent.')
+
+            # order of nodes that belong to the non-terminal edge in hg
+            nt_order_dict = {}  # hg_node -> order ("bond_17" : 1)
+            nt_order_dict_inv = {} # order -> hg_node
+            for each_idx, each_node in enumerate(hg.nodes_in_edge(edge)):
+                nt_order_dict[each_node] = each_idx
+                nt_order_dict_inv[each_idx] = each_node
+
+            # construct a node_map_rhs: rhs -> new hg
+            node_map_rhs = {} # node id in rhs -> node id in hg, where rhs is augmented.
+            node_idx = hg.num_nodes
+            for each_node in self.rhs.nodes:
+                if "ext_id" in self.rhs.node_attr(each_node):
+                    node_map_rhs[each_node] \
+                        = nt_order_dict_inv[
+                            self.rhs.node_attr(each_node)["ext_id"]]
+                else:
+                    node_map_rhs[each_node] = f"bond_{node_idx}"
+                    node_idx += 1
+
+            # delete non-terminal
+            hg.remove_edge(edge)
+
+            # add nodes to hg
+            for each_node in self.rhs.nodes:
+                hg.add_node(node_map_rhs[each_node],
+                            attr_dict=self.rhs.node_attr(each_node))
+
+            # add hyperedges to hg
+            for each_edge in self.rhs.edges:
+                node_list_hg = []
+                for each_node in self.rhs.nodes_in_edge(each_edge):
+                    node_list_hg.append(node_map_rhs[each_node])
+                edge_id = hg.add_edge(
+                    node_list_hg,
+                    attr_dict=self.rhs.edge_attr(each_edge))#deepcopy(self.rhs.edge_attr(each_edge)))
+                if "nt_idx" in hg.edge_attr(edge_id):
+                    nt_edge_dict[hg.edge_attr(edge_id)["nt_idx"]] = edge_id
+            nt_edge_list = [nt_edge_dict[key] for key in range(len(nt_edge_dict))]
+            return hg, nt_edge_list
+
+    def revert(self, hg: Hypergraph, return_subhg=False):
+        ''' revert applying this production rule.
+        i.e., if there exists a subhypergraph that matches the r.h.s. of this production rule,
+        this method replaces the subhypergraph with a non-terminal hyperedge.
+
+        Parameters
+        ----------
+        hg : Hypergraph
+            hypergraph to be reverted
+        return_subhg : bool
+            if True, the removed subhypergraph will be returned.
+
+        Returns
+        -------
+        hg : Hypergraph
+            the resultant hypergraph. if it cannot be reverted, the original one is returned without any replacement.
+        success : bool
+            this indicates whether reverting is successed or not.
+        '''
+        gm = GraphMatcher(hg.hg, self.rhs.hg, node_match=_node_match_prod_rule,
+                          edge_match=_edge_match)
+        try:
+            # in case when the matched subhg is connected to the other part via external nodes and more.
+            not_iso = True
+            while not_iso:
+                isomap = next(gm.subgraph_isomorphisms_iter())
+                adj_node_set = set([]) # reachable nodes from the internal nodes
+                subhg_node_set = set(isomap.keys()) # nodes in subhg
+                for each_node in subhg_node_set:
+                    adj_node_set.add(each_node)
+                    if isomap[each_node] not in self.ext_node.values():
+                        adj_node_set.update(hg.hg.adj[each_node])
+                if adj_node_set == subhg_node_set:
+                    not_iso = False
+                else:
+                    if return_subhg:
+                        return hg, False, Hypergraph()
+                    else:
+                        return hg, False
+            inv_isomap = {v: k for k, v in isomap.items()}
+            '''
+            isomap = {'e35': 'e8', 'bond_13': 'bond_18', 'bond_14': 'bond_19',
+                      'bond_15': 'bond_17', 'e29': 'e23', 'bond_12': 'bond_20'}
+            where keys come from `hg` and values come from `self.rhs`
+            '''
+        except StopIteration:
+            if return_subhg:
+                return hg, False, Hypergraph()
+            else:
+                return hg, False
+
+        if return_subhg:
+            subhg = Hypergraph()
+            for each_node in hg.nodes:
+                if each_node in isomap:
+                    subhg.add_node(each_node, attr_dict=hg.node_attr(each_node))
+            for each_edge in hg.edges:
+                if each_edge in isomap:
+                    subhg.add_edge(hg.nodes_in_edge(each_edge),
+                                   attr_dict=hg.edge_attr(each_edge),
+                                   edge_name=each_edge)
+            subhg.edge_idx = hg.edge_idx
+
+        # remove subhg except for the externael nodes
+        for each_key, each_val in isomap.items():
+            if each_key.startswith('e'):
+                hg.remove_edge(each_key)
+        for each_key, each_val in isomap.items():
+            if each_key.startswith('bond_'):
+                if each_val not in self.ext_node.values():
+                    hg.remove_node(each_key)
+
+        # add non-terminal hyperedge
+        nt_node_list = []
+        for each_ext_id in self.ext_node.keys():
+            nt_node_list.append(inv_isomap[self.ext_node[each_ext_id]])
+
+        hg.add_edge(nt_node_list,
+                    attr_dict=dict(
+                        terminal=False,
+                        symbol=self.lhs_nt_symbol))
+        if return_subhg:
+            return hg, True, subhg
+        else:
+            return hg, True
+
+
+class ProductionRuleCorpus(object):
+
+    '''
+    A corpus of production rules.
+    This class maintains 
+        (i) list of unique production rules,
+        (ii) list of unique edge symbols (both terminal and non-terminal), and
+        (iii) list of unique node symbols.
+
+    Attributes
+    ----------
+    prod_rule_list : list
+        list of unique production rules
+    edge_symbol_list : list
+        list of unique symbols (including both terminal and non-terminal)
+    node_symbol_list : list
+        list of node symbols
+    nt_symbol_list : list
+        list of unique lhs symbols
+    ext_id_list : list
+        list of ext_ids
+    lhs_in_prod_rule : array
+        a matrix of lhs vs prod_rule (= lhs_in_prod_rule)
+    '''
+
+    def __init__(self):
+        self.prod_rule_list = []
+        self.edge_symbol_list = []
+        self.edge_symbol_dict = {}
+        self.node_symbol_list = []
+        self.node_symbol_dict = {}
+        self.nt_symbol_list = []
+        self.ext_id_list = []
+        self._lhs_in_prod_rule = None
+        self.lhs_in_prod_rule_row_list = []
+        self.lhs_in_prod_rule_col_list = []
+
+    @property
+    def lhs_in_prod_rule(self):
+        if self._lhs_in_prod_rule is None:
+            self._lhs_in_prod_rule = torch.sparse.FloatTensor(
+                torch.LongTensor(list(zip(self.lhs_in_prod_rule_row_list, self.lhs_in_prod_rule_col_list))).t(),
+                torch.FloatTensor([1.0]*len(self.lhs_in_prod_rule_col_list)),
+                torch.Size([len(self.nt_symbol_list), len(self.prod_rule_list)])
+            ).to_dense()
+        return self._lhs_in_prod_rule
+        
+    @property
+    def num_prod_rule(self):
+        ''' return the number of production rules
+
+        Returns
+        -------
+        int : the number of unique production rules
+        '''
+        return len(self.prod_rule_list)
+
+    @property
+    def start_rule_list(self):
+        ''' return a list of start rules
+
+        Returns
+        -------
+        list : list of start rules
+        '''
+        start_rule_list = []
+        for each_prod_rule in self.prod_rule_list:
+            if each_prod_rule.is_start_rule:
+                start_rule_list.append(each_prod_rule)
+        return start_rule_list
+
+    @property
+    def num_edge_symbol(self):
+        return len(self.edge_symbol_list)
+
+    @property
+    def num_node_symbol(self):
+        return len(self.node_symbol_list)
+
+    @property
+    def num_ext_id(self):
+        return len(self.ext_id_list)
+
+    def construct_feature_vectors(self):
+        ''' this method constructs feature vectors for the production rules collected so far.
+        currently, NTSymbol and TSymbol are treated in the same manner.
+        '''
+        feature_id_dict = {}
+        feature_id_dict['TSymbol'] = 0
+        feature_id_dict['NTSymbol'] = 1
+        feature_id_dict['BondSymbol'] = 2
+        for each_edge_symbol in self.edge_symbol_list:
+            for each_attr in each_edge_symbol.__dict__.keys():
+                each_val = each_edge_symbol.__dict__[each_attr]
+                if isinstance(each_val, list):
+                    each_val = tuple(each_val)
+                if (each_attr, each_val) not in feature_id_dict:
+                    feature_id_dict[(each_attr, each_val)] = len(feature_id_dict)
+
+        for each_node_symbol in self.node_symbol_list:
+            for each_attr in each_node_symbol.__dict__.keys():
+                each_val = each_node_symbol.__dict__[each_attr]
+                if isinstance(each_val, list):
+                    each_val = tuple(each_val)
+                if (each_attr, each_val) not in feature_id_dict:
+                    feature_id_dict[(each_attr, each_val)] = len(feature_id_dict)
+        for each_ext_id in self.ext_id_list:
+            feature_id_dict[('ext_id', each_ext_id)] = len(feature_id_dict)
+        dim = len(feature_id_dict)
+
+        feature_dict = {}
+        for each_edge_symbol in self.edge_symbol_list:
+            idx_list = []
+            idx_list.append(feature_id_dict[each_edge_symbol.__class__.__name__])
+            for each_attr in each_edge_symbol.__dict__.keys():
+                each_val = each_edge_symbol.__dict__[each_attr]
+                if isinstance(each_val, list):
+                    each_val = tuple(each_val)
+                idx_list.append(feature_id_dict[(each_attr, each_val)])
+            feature = torch.sparse.LongTensor(
+                torch.LongTensor([idx_list]),
+                torch.ones(len(idx_list)),
+                torch.Size([len(feature_id_dict)])
+            )
+            feature_dict[each_edge_symbol] = feature
+
+        for each_node_symbol in self.node_symbol_list:
+            idx_list = []
+            idx_list.append(feature_id_dict[each_node_symbol.__class__.__name__])
+            for each_attr in each_node_symbol.__dict__.keys():
+                each_val = each_node_symbol.__dict__[each_attr]
+                if isinstance(each_val, list):
+                    each_val = tuple(each_val)
+                idx_list.append(feature_id_dict[(each_attr, each_val)])
+            feature = torch.sparse.LongTensor(
+                torch.LongTensor([idx_list]),
+                torch.ones(len(idx_list)),
+                torch.Size([len(feature_id_dict)])
+            )
+            feature_dict[each_node_symbol] = feature
+        for each_ext_id in self.ext_id_list:
+            idx_list = [feature_id_dict[('ext_id', each_ext_id)]]
+            feature_dict[('ext_id', each_ext_id)] \
+                = torch.sparse.LongTensor(
+                    torch.LongTensor([idx_list]),
+                    torch.ones(len(idx_list)),
+                    torch.Size([len(feature_id_dict)])
+                )
+        return feature_dict, dim
+
+    def edge_symbol_idx(self, symbol):
+        return self.edge_symbol_dict[symbol]
+
+    def node_symbol_idx(self, symbol):
+        return self.node_symbol_dict[symbol]
+
+    def append(self, prod_rule: ProductionRule) -> Tuple[int, ProductionRule]:
+        """ return whether the input production rule is new or not, and its production rule id.
+        Production rules are regarded as the same if 
+            i) there exists a one-to-one mapping of nodes and edges, and
+            ii) all the attributes associated with nodes and hyperedges are the same.
+
+        Parameters
+        ----------
+        prod_rule : ProductionRule
+
+        Returns
+        -------
+        prod_rule_id : int
+            production rule index. if new, a new index will be assigned.
+        prod_rule : ProductionRule
+        """
+        num_lhs = len(self.nt_symbol_list)
+        for each_idx, each_prod_rule in enumerate(self.prod_rule_list):
+            is_same, isomap = prod_rule.is_same(each_prod_rule)
+            if is_same:
+                # we do not care about edge and node names, but care about the order of non-terminal edges.
+                for key, val in isomap.items(): # key : edges & nodes in each_prod_rule.rhs , val : those in prod_rule.rhs
+                    if key.startswith("bond_"):
+                        continue
+
+                    # rewrite `nt_idx` in `prod_rule` for further processing
+                    if "nt_idx" in prod_rule.rhs.edge_attr(val).keys():
+                        if "nt_idx" not in each_prod_rule.rhs.edge_attr(key).keys():
+                            raise ValueError
+                        prod_rule.rhs.set_edge_attr(
+                            val,
+                            {'nt_idx': each_prod_rule.rhs.edge_attr(key)["nt_idx"]})
+                return each_idx, prod_rule
+        self.prod_rule_list.append(prod_rule)
+        self._update_edge_symbol_list(prod_rule)
+        self._update_node_symbol_list(prod_rule)
+        self._update_ext_id_list(prod_rule)
+
+        lhs_idx = self.nt_symbol_list.index(prod_rule.lhs_nt_symbol)
+        self.lhs_in_prod_rule_row_list.append(lhs_idx)
+        self.lhs_in_prod_rule_col_list.append(len(self.prod_rule_list)-1)
+        self._lhs_in_prod_rule = None
+        return len(self.prod_rule_list)-1, prod_rule
+
+    def get_prod_rule(self, prod_rule_idx: int) -> ProductionRule:
+        return self.prod_rule_list[prod_rule_idx]
+
+    def sample(self, unmasked_logit_array, nt_symbol, deterministic=False):
+        ''' sample a production rule whose lhs is `nt_symbol`, followihng `unmasked_logit_array`.
+
+        Parameters
+        ----------
+        unmasked_logit_array : array-like, length `num_prod_rule`
+        nt_symbol : NTSymbol
+        '''
+        if not isinstance(unmasked_logit_array, np.ndarray):
+            unmasked_logit_array = unmasked_logit_array.numpy().astype(np.float64)
+        if deterministic:
+            prob = masked_softmax(unmasked_logit_array,
+                                  self.lhs_in_prod_rule[self.nt_symbol_list.index(nt_symbol)].numpy().astype(np.float64))
+            return self.prod_rule_list[np.argmax(prob)]
+        else:
+            return np.random.choice(
+                self.prod_rule_list, 1,
+                p=masked_softmax(unmasked_logit_array,
+                                 self.lhs_in_prod_rule[self.nt_symbol_list.index(nt_symbol)].numpy().astype(np.float64)))[0]
+
+    def masked_logprob(self, unmasked_logit_array, nt_symbol):
+        if not isinstance(unmasked_logit_array, np.ndarray):
+            unmasked_logit_array = unmasked_logit_array.numpy().astype(np.float64)
+        prob = masked_softmax(unmasked_logit_array,
+                              self.lhs_in_prod_rule[self.nt_symbol_list.index(nt_symbol)].numpy().astype(np.float64))
+        return np.log(prob)
+
+    def _update_edge_symbol_list(self, prod_rule: ProductionRule):
+        ''' update edge symbol list
+
+        Parameters
+        ----------
+        prod_rule : ProductionRule
+        '''
+        if prod_rule.lhs_nt_symbol not in self.nt_symbol_list:
+            self.nt_symbol_list.append(prod_rule.lhs_nt_symbol)
+
+        for each_edge in prod_rule.rhs.edges:
+            if prod_rule.rhs.edge_attr(each_edge)['symbol'] not in self.edge_symbol_dict:
+                edge_symbol_idx = len(self.edge_symbol_list)
+                self.edge_symbol_list.append(prod_rule.rhs.edge_attr(each_edge)['symbol'])
+                self.edge_symbol_dict[prod_rule.rhs.edge_attr(each_edge)['symbol']] = edge_symbol_idx
+            else:
+                edge_symbol_idx = self.edge_symbol_dict[prod_rule.rhs.edge_attr(each_edge)['symbol']]
+            prod_rule.rhs.edge_attr(each_edge)['symbol_idx'] = edge_symbol_idx
+        pass
+
+    def _update_node_symbol_list(self, prod_rule: ProductionRule):
+        ''' update node symbol list
+
+        Parameters
+        ----------
+        prod_rule : ProductionRule
+        '''
+        for each_node in prod_rule.rhs.nodes:
+            if prod_rule.rhs.node_attr(each_node)['symbol'] not in self.node_symbol_dict:
+                node_symbol_idx = len(self.node_symbol_list)
+                self.node_symbol_list.append(prod_rule.rhs.node_attr(each_node)['symbol'])
+                self.node_symbol_dict[prod_rule.rhs.node_attr(each_node)['symbol']] = node_symbol_idx
+            else:
+                node_symbol_idx = self.node_symbol_dict[prod_rule.rhs.node_attr(each_node)['symbol']]
+            prod_rule.rhs.node_attr(each_node)['symbol_idx'] = node_symbol_idx
+
+    def _update_ext_id_list(self, prod_rule: ProductionRule):
+        for each_node in prod_rule.rhs.nodes:
+            if 'ext_id' in prod_rule.rhs.node_attr(each_node):
+                if prod_rule.rhs.node_attr(each_node)['ext_id'] not in self.ext_id_list:
+                    self.ext_id_list.append(prod_rule.rhs.node_attr(each_node)['ext_id'])
+
+
+class HyperedgeReplacementGrammar(GraphGrammarBase):
+    """
+    Learn a hyperedge replacement grammar from a set of hypergraphs.
+
+    Attributes
+    ----------
+    prod_rule_list : list of ProductionRule
+        production rules learned from the input hypergraphs
+    """
+    def __init__(self,
+                 tree_decomposition=molecular_tree_decomposition,
+                 ignore_order=False, **kwargs):
+        from functools import partial
+        self.prod_rule_corpus = ProductionRuleCorpus()
+        self.clique_tree_corpus = CliqueTreeCorpus()
+        self.ignore_order = ignore_order
+        self.tree_decomposition = partial(tree_decomposition, **kwargs)
+
+    @property
+    def num_prod_rule(self):
+        ''' return the number of production rules
+
+        Returns
+        -------
+        int : the number of unique production rules
+        '''
+        return self.prod_rule_corpus.num_prod_rule
+
+    @property
+    def start_rule_list(self):
+        ''' return a list of start rules
+
+        Returns
+        -------
+        list : list of start rules
+        '''
+        return self.prod_rule_corpus.start_rule_list
+
+    @property
+    def prod_rule_list(self):
+        return self.prod_rule_corpus.prod_rule_list
+
+    def learn(self, hg_list, logger=print, max_mol=np.inf, print_freq=500):
+        """ learn from a list of hypergraphs
+
+        Parameters
+        ----------
+        hg_list : list of Hypergraph
+
+        Returns
+        -------
+        prod_rule_seq_list : list of integers
+            each element corresponds to a sequence of production rules to generate each hypergraph.
+        """
+        prod_rule_seq_list = []
+        idx = 0
+        for each_idx, each_hg in enumerate(hg_list):
+            clique_tree = self.tree_decomposition(each_hg)
+
+            # get a pair of myself and children
+            root_node = _find_root(clique_tree)
+            clique_tree = self.clique_tree_corpus.add_to_subhg_list(clique_tree, root_node)
+            prod_rule_seq = []
+            stack = []
+
+            children = sorted(list(clique_tree[root_node].keys()))
+
+            # extract a temporary production rule
+            prod_rule = extract_prod_rule(
+                None,
+                clique_tree.nodes[root_node]["subhg"],
+                [clique_tree.nodes[each_child]["subhg"]
+                 for each_child in children],
+                clique_tree.nodes[root_node].get('subhg_idx', None))
+
+            # update the production rule list
+            prod_rule_id, prod_rule = self.update_prod_rule_list(prod_rule)
+            children = reorder_children(root_node,
+                                        children,
+                                        prod_rule,
+                                        clique_tree)
+            stack.extend([(root_node, each_child) for each_child in children[::-1]])
+            prod_rule_seq.append(prod_rule_id)
+
+            while len(stack) != 0:
+                # get a triple of parent, myself, and children
+                parent, myself = stack.pop()
+                children = sorted(list(dict(clique_tree[myself]).keys()))
+                children.remove(parent)
+
+                # extract a temp prod rule
+                prod_rule = extract_prod_rule(
+                    clique_tree.nodes[parent]["subhg"],
+                    clique_tree.nodes[myself]["subhg"],
+                    [clique_tree.nodes[each_child]["subhg"]
+                     for each_child in children],
+                    clique_tree.nodes[myself].get('subhg_idx', None))
+
+                # update the prod rule list
+                prod_rule_id, prod_rule = self.update_prod_rule_list(prod_rule)
+                children = reorder_children(myself,
+                                            children,
+                                            prod_rule,
+                                            clique_tree)
+                stack.extend([(myself, each_child)
+                              for each_child in children[::-1]])
+                prod_rule_seq.append(prod_rule_id)
+            prod_rule_seq_list.append(prod_rule_seq)
+            if (each_idx+1) % print_freq == 0:
+                msg = f'#(molecules processed)={each_idx+1}\t'\
+                        f'#(production rules)={self.prod_rule_corpus.num_prod_rule}\t#(subhg in corpus)={self.clique_tree_corpus.size}'
+                logger(msg)
+            if each_idx > max_mol:
+                break
+
+        print(f'corpus_size = {self.clique_tree_corpus.size}')
+        return prod_rule_seq_list
+
+    def sample(self, z, deterministic=False):
+        """ sample a new hypergraph from HRG.
+
+        Parameters
+        ----------
+        z : array-like, shape (len, num_prod_rule)
+            logit
+        deterministic : bool
+            if True, deterministic sampling
+
+        Returns
+        -------
+        Hypergraph
+        """
+        seq_idx = 0
+        stack = []
+        z = z[:, :-1]
+        init_prod_rule = self.prod_rule_corpus.sample(z[0], NTSymbol(degree=0,
+                                                                     is_aromatic=False,
+                                                                     bond_symbol_list=[]),
+                                                      deterministic=deterministic)
+        hg, nt_edge_list = init_prod_rule.applied_to(None, None)
+        stack = deepcopy(nt_edge_list[::-1])
+        while len(stack) != 0 and seq_idx < z.shape[0]-1:
+            seq_idx += 1
+            nt_edge = stack.pop()
+            nt_symbol = hg.edge_attr(nt_edge)['symbol']
+            prod_rule = self.prod_rule_corpus.sample(z[seq_idx], nt_symbol, deterministic=deterministic)
+            hg, nt_edge_list = prod_rule.applied_to(hg, nt_edge)
+            stack.extend(nt_edge_list[::-1])
+        if len(stack) != 0:
+            raise RuntimeError(f'{len(stack)} non-terminals are left.')
+        return hg
+
+    def construct(self, prod_rule_seq):
+        """ construct a hypergraph following `prod_rule_seq`
+
+        Parameters
+        ----------
+        prod_rule_seq : list of integers
+            a sequence of production rules.
+
+        Returns
+        -------
+        UndirectedHypergraph
+        """
+        seq_idx = 0
+        init_prod_rule = self.prod_rule_corpus.get_prod_rule(prod_rule_seq[seq_idx])
+        hg, nt_edge_list = init_prod_rule.applied_to(None, None)
+        stack = deepcopy(nt_edge_list[::-1])
+        while len(stack) != 0:
+            seq_idx += 1
+            nt_edge = stack.pop()
+            hg, nt_edge_list = self.prod_rule_corpus.get_prod_rule(prod_rule_seq[seq_idx]).applied_to(hg, nt_edge)
+            stack.extend(nt_edge_list[::-1])            
+        return hg
+
+    def update_prod_rule_list(self, prod_rule):
+        """ return whether the input production rule is new or not, and its production rule id.
+        Production rules are regarded as the same if 
+            i) there exists a one-to-one mapping of nodes and edges, and
+            ii) all the attributes associated with nodes and hyperedges are the same.
+
+        Parameters
+        ----------
+        prod_rule : ProductionRule
+
+        Returns
+        -------
+        is_new : bool
+            if True, this production rule is new
+        prod_rule_id : int
+            production rule index. if new, a new index will be assigned.
+        """
+        return self.prod_rule_corpus.append(prod_rule)
+
+
+class IncrementalHyperedgeReplacementGrammar(HyperedgeReplacementGrammar):
+    '''
+    This class learns HRG incrementally leveraging the previously obtained production rules.
+    '''
+    def __init__(self, tree_decomposition=tree_decomposition_with_hrg, ignore_order=False):
+        self.prod_rule_list = []
+        self.tree_decomposition = tree_decomposition
+        self.ignore_order = ignore_order
+
+    def learn(self, hg_list):
+        """ learn from a list of hypergraphs
+
+        Parameters
+        ----------
+        hg_list : list of UndirectedHypergraph
+
+        Returns
+        -------
+        prod_rule_seq_list : list of integers
+            each element corresponds to a sequence of production rules to generate each hypergraph.
+        """
+        prod_rule_seq_list = []
+        for each_hg in hg_list:
+            clique_tree, root_node = tree_decomposition_with_hrg(each_hg, self, return_root=True)
+
+            prod_rule_seq = []
+            stack = []
+
+            # get a pair of myself and children
+            children = sorted(list(clique_tree[root_node].keys()))
+            
+            # extract a temporary production rule
+            prod_rule = extract_prod_rule(None, clique_tree.nodes[root_node]["subhg"],
+                                          [clique_tree.nodes[each_child]["subhg"] for each_child in children])
+            
+            # update the production rule list
+            prod_rule_id, prod_rule = self.update_prod_rule_list(prod_rule)
+            children = reorder_children(root_node, children, prod_rule, clique_tree)
+            stack.extend([(root_node, each_child) for each_child in children[::-1]])
+            prod_rule_seq.append(prod_rule_id)
+            
+            while len(stack) != 0:
+                # get a triple of parent, myself, and children
+                parent, myself = stack.pop()
+                children = sorted(list(dict(clique_tree[myself]).keys()))
+                children.remove(parent)
+
+                # extract a temp prod rule
+                prod_rule = extract_prod_rule(
+                    clique_tree.nodes[parent]["subhg"], clique_tree.nodes[myself]["subhg"],
+                    [clique_tree.nodes[each_child]["subhg"] for each_child in children])
+
+                # update the prod rule list
+                prod_rule_id, prod_rule = self.update_prod_rule_list(prod_rule)
+                children = reorder_children(myself, children, prod_rule, clique_tree)
+                stack.extend([(myself, each_child) for each_child in children[::-1]])
+                prod_rule_seq.append(prod_rule_id)
+            prod_rule_seq_list.append(prod_rule_seq)
+        self._compute_stats()
+        return prod_rule_seq_list
+
+
+def reorder_children(myself, children, prod_rule, clique_tree):
+    """ reorder children so that they match the order in `prod_rule`.
+
+    Parameters
+    ----------
+    myself : int
+    children : list of int
+    prod_rule : ProductionRule
+    clique_tree : nx.Graph
+
+    Returns
+    -------
+    new_children : list of str
+        reordered children
+    """
+    perm = {} # key : `nt_idx`, val : child
+    for each_edge in prod_rule.rhs.edges:
+        if "nt_idx" in prod_rule.rhs.edge_attr(each_edge).keys():
+            for each_child in children:
+                common_node_set = set(
+                    common_node_list(clique_tree.nodes[myself]["subhg"],
+                                     clique_tree.nodes[each_child]["subhg"])[0])
+                if set(prod_rule.rhs.nodes_in_edge(each_edge)) == common_node_set:
+                    assert prod_rule.rhs.edge_attr(each_edge)["nt_idx"] not in perm
+                    perm[prod_rule.rhs.edge_attr(each_edge)["nt_idx"]] = each_child
+    new_children = []
+    assert len(perm) == len(children)
+    for i in range(len(perm)):
+        new_children.append(perm[i])
+    return new_children
+
+
+def extract_prod_rule(parent_hg, myself_hg, children_hg_list, subhg_idx=None):
+    """ extract a production rule from a triple of `parent_hg`, `myself_hg`, and `children_hg_list`.
+
+    Parameters
+    ----------
+    parent_hg : Hypergraph
+    myself_hg : Hypergraph
+    children_hg_list : list of Hypergraph
+
+    Returns
+    -------
+    ProductionRule, consisting of
+        lhs : Hypergraph or None
+        rhs : Hypergraph
+    """
+    def _add_ext_node(hg, ext_nodes):
+        """ mark nodes to be external (ordered ids are assigned)
+
+        Parameters
+        ----------
+        hg : UndirectedHypergraph
+        ext_nodes : list of str
+            list of external nodes
+
+        Returns
+        -------
+        hg : Hypergraph
+            nodes in `ext_nodes` are marked to be external
+        """
+        ext_id = 0
+        ext_id_exists = []
+        for each_node in ext_nodes:
+            ext_id_exists.append('ext_id' in hg.node_attr(each_node))
+        if ext_id_exists and any(ext_id_exists) != all(ext_id_exists):
+            raise ValueError
+        if not all(ext_id_exists):
+            for each_node in ext_nodes:
+                hg.node_attr(each_node)['ext_id'] = ext_id
+                ext_id += 1
+        return hg
+
+    def _check_aromatic(hg, node_list):
+        is_aromatic = False
+        node_aromatic_list = []
+        for each_node in node_list:
+            if hg.node_attr(each_node)['symbol'].is_aromatic:
+                is_aromatic = True
+                node_aromatic_list.append(True)
+            else:
+                node_aromatic_list.append(False)
+        return is_aromatic, node_aromatic_list
+
+    def _check_ring(hg):
+        for each_edge in hg.edges:
+            if not ('tmp' in hg.edge_attr(each_edge) or (not hg.edge_attr(each_edge)['terminal'])):
+                return False
+        return True
+
+    if parent_hg is None:
+        lhs = Hypergraph()
+        node_list = []
+    else:
+        lhs = Hypergraph()
+        node_list, edge_exists = common_node_list(parent_hg, myself_hg)
+        for each_node in node_list:
+            lhs.add_node(each_node,
+                         deepcopy(myself_hg.node_attr(each_node)))
+        is_aromatic, _ = _check_aromatic(parent_hg, node_list)
+        for_ring = _check_ring(myself_hg)
+        bond_symbol_list = []
+        for each_node in node_list:
+            bond_symbol_list.append(parent_hg.node_attr(each_node)['symbol'])
+        lhs.add_edge(
+            node_list,
+            attr_dict=dict(
+                terminal=False,
+                edge_exists=edge_exists,
+                symbol=NTSymbol(
+                    degree=len(node_list),
+                    is_aromatic=is_aromatic,
+                    bond_symbol_list=bond_symbol_list,
+                    for_ring=for_ring)))
+        try:
+            lhs = _add_ext_node(lhs, node_list)
+        except ValueError:
+            import pdb; pdb.set_trace()
+
+    rhs = remove_tmp_edge(deepcopy(myself_hg))
+    #rhs = remove_ext_node(rhs)
+    #rhs = remove_nt_edge(rhs)
+    try:
+        rhs = _add_ext_node(rhs, node_list)
+    except ValueError:
+        import pdb; pdb.set_trace()
+
+    nt_idx = 0
+    if children_hg_list is not None:
+        for each_child_hg in children_hg_list:
+            node_list, edge_exists = common_node_list(myself_hg, each_child_hg)
+            is_aromatic, _ = _check_aromatic(myself_hg, node_list)
+            for_ring = _check_ring(each_child_hg)
+            bond_symbol_list = []
+            for each_node in node_list:
+                bond_symbol_list.append(myself_hg.node_attr(each_node)['symbol'])
+            rhs.add_edge(
+                node_list,
+                attr_dict=dict(
+                    terminal=False,
+                    nt_idx=nt_idx,
+                    edge_exists=edge_exists,
+                    symbol=NTSymbol(degree=len(node_list),
+                                    is_aromatic=is_aromatic,
+                                    bond_symbol_list=bond_symbol_list,
+                                    for_ring=for_ring)))
+            nt_idx += 1
+    prod_rule = ProductionRule(lhs, rhs)
+    prod_rule.subhg_idx = subhg_idx
+    if DEBUG:
+        if sorted(list(prod_rule.ext_node.keys())) \
+           != list(np.arange(len(prod_rule.ext_node))):
+            raise RuntimeError('ext_id is not continuous')
+    return prod_rule
+
+
+def _find_root(clique_tree):
+    max_node = None
+    num_nodes_max = -np.inf
+    for each_node in clique_tree.nodes:
+        if clique_tree.nodes[each_node]['subhg'].num_nodes > num_nodes_max:
+            max_node = each_node
+            num_nodes_max = clique_tree.nodes[each_node]['subhg'].num_nodes
+        '''
+        children = sorted(list(clique_tree[each_node].keys()))
+        prod_rule = extract_prod_rule(None,
+                                      clique_tree.nodes[each_node]["subhg"],
+                                      [clique_tree.nodes[each_child]["subhg"]
+                                       for each_child in children])
+        for each_start_rule in start_rule_list:
+            if prod_rule.is_same(each_start_rule):
+                return each_node
+        '''
+    return max_node
+
+def remove_ext_node(hg):
+    for each_node in hg.nodes:
+        hg.node_attr(each_node).pop('ext_id', None)
+    return hg
+
+def remove_nt_edge(hg):
+    remove_edge_list = []
+    for each_edge in hg.edges:
+        if not hg.edge_attr(each_edge)['terminal']:
+            remove_edge_list.append(each_edge)
+    hg.remove_edges(remove_edge_list)
+    return hg
+
+def remove_tmp_edge(hg):
+    remove_edge_list = []
+    for each_edge in hg.edges:
+        if hg.edge_attr(each_edge).get('tmp', False):
+            remove_edge_list.append(each_edge)
+    hg.remove_edges(remove_edge_list)
+    return hg

models/mhg_model/graph_grammar/graph_grammar/symbols.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 1 2018"
+
+from typing import List
+
+class TSymbol(object):
+
+    ''' terminal symbol
+
+    Attributes
+    ----------
+    degree : int
+        the number of nodes in a hyperedge
+    is_aromatic : bool
+        whether or not the hyperedge is in an aromatic ring
+    symbol : str
+        atomic symbol
+    num_explicit_Hs : int
+        the number of hydrogens associated to this hyperedge
+    formal_charge : int
+        charge
+    chirality : int
+        chirality
+    '''
+
+    def __init__(self, degree, is_aromatic,
+                 symbol, num_explicit_Hs, formal_charge, chirality):
+        self.degree = degree
+        self.is_aromatic = is_aromatic
+        self.symbol = symbol
+        self.num_explicit_Hs = num_explicit_Hs
+        self.formal_charge = formal_charge
+        self.chirality = chirality
+
+    @property
+    def terminal(self):
+        return True
+
+    def __eq__(self, other):
+        if not isinstance(other, TSymbol):
+            return False
+        if self.degree != other.degree:
+            return False
+        if self.is_aromatic != other.is_aromatic:
+            return False
+        if self.symbol != other.symbol:
+            return False
+        if self.num_explicit_Hs != other.num_explicit_Hs:
+            return False
+        if self.formal_charge != other.formal_charge:
+            return False
+        if self.chirality != other.chirality:
+            return False
+        return True
+
+    def __hash__(self):
+        return self.__str__().__hash__()
+
+    def __str__(self):
+        return f'degree={self.degree}, is_aromatic={self.is_aromatic}, '\
+            f'symbol={self.symbol}, '\
+            f'num_explicit_Hs={self.num_explicit_Hs}, '\
+            f'formal_charge={self.formal_charge}, chirality={self.chirality}'
+
+
+class NTSymbol(object):
+
+    ''' non-terminal symbol
+
+    Attributes
+    ----------
+    degree : int
+        degree of the hyperedge
+    is_aromatic : bool
+        if True, at least one of the associated bonds must be aromatic.
+    node_aromatic_list : list of bool
+        indicate whether each of the nodes is aromatic or not.
+    bond_type_list : list of int
+        bond type of each node"
+    '''
+
+    def __init__(self, degree: int, is_aromatic: bool,
+                 bond_symbol_list: list,
+                 for_ring=False):
+        self.degree = degree
+        self.is_aromatic = is_aromatic
+        self.for_ring = for_ring
+        self.bond_symbol_list = bond_symbol_list
+
+    @property
+    def terminal(self) -> bool:
+        return False
+
+    @property
+    def symbol(self):
+        return f'NT{self.degree}'
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, NTSymbol):
+            return False
+
+        if self.degree != other.degree:
+            return False
+        if self.is_aromatic != other.is_aromatic:
+            return False
+        if self.for_ring != other.for_ring:
+            return False
+        if len(self.bond_symbol_list) != len(other.bond_symbol_list):
+            return False
+        for each_idx in range(len(self.bond_symbol_list)):
+            if self.bond_symbol_list[each_idx] != other.bond_symbol_list[each_idx]:
+                return False
+        return True
+
+    def __hash__(self):
+        return self.__str__().__hash__()
+
+    def __str__(self) -> str:
+        return f'degree={self.degree}, is_aromatic={self.is_aromatic}, '\
+            f'bond_symbol_list={[str(each_symbol) for each_symbol in self.bond_symbol_list]}'\
+            f'for_ring={self.for_ring}'
+
+
+class BondSymbol(object):
+    
+
+    ''' Bond symbol
+
+    Attributes
+    ----------
+    is_aromatic : bool
+        if True, at least one of the associated bonds must be aromatic.
+    bond_type : int
+        bond type of each node"
+    '''
+
+    def __init__(self, is_aromatic: bool,
+                 bond_type: int,
+                 stereo: int):
+        self.is_aromatic = is_aromatic
+        self.bond_type = bond_type
+        self.stereo = stereo
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, BondSymbol):
+            return False
+
+        if self.is_aromatic != other.is_aromatic:
+            return False
+        if self.bond_type != other.bond_type:
+            return False
+        if self.stereo != other.stereo:
+            return False
+        return True
+
+    def __hash__(self):
+        return self.__str__().__hash__()
+
+    def __str__(self) -> str:
+        return f'is_aromatic={self.is_aromatic}, '\
+            f'bond_type={self.bond_type}, '\
+            f'stereo={self.stereo}, '

models/mhg_model/graph_grammar/graph_grammar/utils.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jun 4 2018"
+
+from ..hypergraph import Hypergraph
+from copy import deepcopy
+from typing import List
+import numpy as np
+
+
+def common_node_list(hg1: Hypergraph, hg2: Hypergraph) -> List[str]:
+    """ return a list of common nodes
+
+    Parameters
+    ----------
+    hg1, hg2 : Hypergraph
+
+    Returns
+    -------
+    list of str
+        list of common nodes
+    """
+    if hg1 is None or hg2 is None:
+        return [], False
+    else:
+        node_set = hg1.nodes.intersection(hg2.nodes)
+        node_dict = {}
+        if 'order4hrg' in hg1.node_attr(list(hg1.nodes)[0]):
+            for each_node in node_set:
+                node_dict[each_node] = hg1.node_attr(each_node)['order4hrg']
+        else:
+            for each_node in node_set:
+                node_dict[each_node] = hg1.node_attr(each_node)['symbol'].__hash__()
+        node_list = []
+        for each_key, _ in sorted(node_dict.items(), key=lambda x:x[1]):
+            node_list.append(each_key)
+        edge_name = hg1.has_edge(node_list, ignore_order=True)
+        if edge_name:
+            if not hg1.edge_attr(edge_name).get('terminal', True):
+                node_list = hg1.nodes_in_edge(edge_name)
+            return node_list, True
+        else:
+            return node_list, False
+
+
+def _node_match(node1, node2):
+    # if the nodes are hyperedges, `atom_attr` determines the match
+    if node1['bipartite'] == 'edge' and node2['bipartite'] == 'edge':
+        return node1["attr_dict"]['symbol'] == node2["attr_dict"]['symbol']
+    elif node1['bipartite'] == 'node' and node2['bipartite'] == 'node':
+        # bond_symbol
+        return node1['attr_dict']['symbol'] == node2['attr_dict']['symbol']
+    else:
+        return False
+
+def _easy_node_match(node1, node2):
+    # if the nodes are hyperedges, `atom_attr` determines the match
+    if node1['bipartite'] == 'edge' and node2['bipartite'] == 'edge':
+        return node1["attr_dict"].get('symbol', None) == node2["attr_dict"].get('symbol', None)
+    elif node1['bipartite'] == 'node' and node2['bipartite'] == 'node':
+        # bond_symbol
+        return node1['attr_dict'].get('ext_id', -1) == node2['attr_dict'].get('ext_id', -1)\
+            and node1['attr_dict']['symbol'] == node2['attr_dict']['symbol']
+    else:
+        return False
+
+
+def _node_match_prod_rule(node1, node2, ignore_order=False):
+    # if the nodes are hyperedges, `atom_attr` determines the match
+    if node1['bipartite'] == 'edge' and node2['bipartite'] == 'edge':
+        return node1["attr_dict"]['symbol'] == node2["attr_dict"]['symbol']
+    elif node1['bipartite'] == 'node' and node2['bipartite'] == 'node':
+        # ext_id, order4hrg, bond_symbol
+        if ignore_order:
+            return node1['attr_dict']['symbol'] == node2['attr_dict']['symbol']
+        else:
+            return node1['attr_dict']['symbol'] == node2['attr_dict']['symbol']\
+                and node1['attr_dict'].get('ext_id', -1) == node2['attr_dict'].get('ext_id', -1)
+    else:
+        return False
+
+
+def _edge_match(edge1, edge2, ignore_order=False):
+    #return True
+    if ignore_order:
+        return True
+    else:
+        return edge1["order"] == edge2["order"]
+
+def masked_softmax(logit, mask):
+    ''' compute a probability distribution from logit
+
+    Parameters
+    ----------
+    logit : array-like, length D
+        each element indicates how each dimension is likely to be chosen
+        (the larger, the more likely)
+    mask : array-like, length D
+        each element is either 0 or 1.
+        if 0, the dimension is ignored
+        when computing the probability distribution.
+
+    Returns
+    -------
+    prob_dist : array, length D
+        probability distribution computed from logit.
+        if `mask[d] = 0`, `prob_dist[d] = 0`.
+    '''
+    if logit.shape != mask.shape:
+        raise ValueError('logit and mask must have the same shape')
+    c = np.max(logit)
+    exp_logit = np.exp(logit - c) * mask
+    sum_exp_logit = exp_logit @ mask
+    return exp_logit / sum_exp_logit

models/mhg_model/graph_grammar/hypergraph.py

@@ -0,0 +1,544 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 31 2018"
+
+from copy import deepcopy
+from typing import List, Dict, Tuple
+import networkx as nx
+import numpy as np
+import os
+
+
+class Hypergraph(object):
+    '''
+    A class of a hypergraph.
+    Each hyperedge can be ordered. For the ordered case,
+    edges adjacent to the hyperedge node are labeled by their orders.
+
+    Attributes
+    ----------
+    hg : nx.Graph
+        a bipartite graph representation of a hypergraph
+    edge_idx : int
+        total number of hyperedges that exist so far
+    '''
+    def __init__(self):
+        self.hg = nx.Graph()
+        self.edge_idx = 0
+        self.nodes = set([])
+        self.num_nodes = 0
+        self.edges = set([])
+        self.num_edges = 0
+        self.nodes_in_edge_dict = {}
+
+    def add_node(self, node: str, attr_dict=None):
+        ''' add a node to hypergraph
+
+        Parameters
+        ----------
+        node : str
+            node name
+        attr_dict : dict
+            dictionary of node attributes
+        '''
+        self.hg.add_node(node, bipartite='node', attr_dict=attr_dict)
+        if node not in self.nodes:
+            self.num_nodes += 1
+        self.nodes.add(node)
+
+    def add_edge(self, node_list: List[str], attr_dict=None, edge_name=None):
+        ''' add an edge consisting of nodes `node_list`
+
+        Parameters
+        ----------
+        node_list : list 
+            ordered list of nodes that consist the edge
+        attr_dict : dict
+            dictionary of edge attributes
+        '''
+        if edge_name is None:
+            edge = 'e{}'.format(self.edge_idx)
+        else:
+            assert edge_name not in self.edges
+            edge = edge_name
+        self.hg.add_node(edge, bipartite='edge', attr_dict=attr_dict)
+        if edge not in self.edges:
+            self.num_edges += 1
+        self.edges.add(edge)
+        self.nodes_in_edge_dict[edge] = node_list
+        if type(node_list) == list:
+            for node_idx, each_node in enumerate(node_list):
+                self.hg.add_edge(edge, each_node, order=node_idx)
+                if each_node not in self.nodes:
+                    self.num_nodes += 1
+                self.nodes.add(each_node)
+
+        elif type(node_list) == set:
+            for each_node in node_list:
+                self.hg.add_edge(edge, each_node, order=-1)
+                if each_node not in self.nodes:
+                    self.num_nodes += 1
+                self.nodes.add(each_node)
+        else:
+            raise ValueError
+        self.edge_idx += 1
+        return edge
+
+    def remove_node(self, node: str, remove_connected_edges=True):
+        ''' remove a node
+
+        Parameters
+        ----------
+        node : str
+            node name
+        remove_connected_edges : bool
+            if True, remove edges that are adjacent to the node
+        '''
+        if remove_connected_edges:
+            connected_edges = deepcopy(self.adj_edges(node))
+            for each_edge in connected_edges:
+                self.remove_edge(each_edge)
+        self.hg.remove_node(node)
+        self.num_nodes -= 1
+        self.nodes.remove(node)
+
+    def remove_nodes(self, node_iter, remove_connected_edges=True):
+        ''' remove a set of nodes
+
+        Parameters
+        ----------
+        node_iter : iterator of strings
+            nodes to be removed
+        remove_connected_edges : bool
+            if True, remove edges that are adjacent to the node        
+        '''
+        for each_node in node_iter:
+            self.remove_node(each_node, remove_connected_edges)
+
+    def remove_edge(self, edge: str):
+        ''' remove an edge
+
+        Parameters
+        ----------
+        edge : str
+            edge to be removed
+        '''
+        self.hg.remove_node(edge)
+        self.edges.remove(edge)
+        self.num_edges -= 1
+        self.nodes_in_edge_dict.pop(edge)
+
+    def remove_edges(self, edge_iter):
+        ''' remove a set of edges
+
+        Parameters
+        ----------
+        edge_iter : iterator of strings
+            edges to be removed
+        '''
+        for each_edge in edge_iter:
+            self.remove_edge(each_edge)
+
+    def remove_edges_with_attr(self, edge_attr_dict):
+        remove_edge_list = []
+        for each_edge in self.edges:
+            satisfy = True
+            for each_key, each_val in edge_attr_dict.items():
+                if not satisfy:
+                    break
+                try:
+                    if self.edge_attr(each_edge)[each_key] != each_val:
+                        satisfy = False
+                except KeyError:
+                    satisfy = False
+            if satisfy:
+                remove_edge_list.append(each_edge)
+        self.remove_edges(remove_edge_list)
+
+    def remove_subhg(self, subhg):
+        ''' remove subhypergraph.
+        all of the hyperedges are removed.
+        each node of subhg is removed if its degree becomes 0 after removing hyperedges.
+
+        Parameters
+        ----------
+        subhg : Hypergraph
+        '''
+        for each_edge in subhg.edges:
+            self.remove_edge(each_edge)
+        for each_node in subhg.nodes:
+            if self.degree(each_node) == 0:
+                self.remove_node(each_node)
+
+    def nodes_in_edge(self, edge):
+        ''' return an ordered list of nodes in a given edge.
+
+        Parameters
+        ----------
+        edge : str
+            edge whose nodes are returned
+
+        Returns
+        -------
+        list or set
+            ordered list or set of nodes that belong to the edge
+        '''
+        if edge.startswith('e'):
+            return self.nodes_in_edge_dict[edge]
+        else:
+            adj_node_list = self.hg.adj[edge]
+            adj_node_order_list = []
+            adj_node_name_list = []
+            for each_node in adj_node_list:
+                adj_node_order_list.append(adj_node_list[each_node]['order'])
+                adj_node_name_list.append(each_node)
+            if adj_node_order_list == [-1] * len(adj_node_order_list):
+                return set(adj_node_name_list)
+            else:
+                return [adj_node_name_list[each_idx] for each_idx
+                        in np.argsort(adj_node_order_list)]
+
+    def adj_edges(self, node):
+        ''' return a dict of adjacent hyperedges
+
+        Parameters
+        ----------
+        node : str
+
+        Returns
+        -------
+        set
+            set of edges that are adjacent to `node`
+        '''
+        return self.hg.adj[node]
+
+    def adj_nodes(self, node):
+        ''' return a set of adjacent nodes
+
+        Parameters
+        ----------
+        node : str
+
+        Returns
+        -------
+        set
+            set of nodes that are adjacent to `node`
+        '''
+        node_set = set([])
+        for each_adj_edge in self.adj_edges(node):
+            node_set.update(set(self.nodes_in_edge(each_adj_edge)))
+        node_set.discard(node)
+        return node_set
+
+    def has_edge(self, node_list, ignore_order=False):
+        for each_edge in self.edges:
+            if ignore_order:
+                if set(self.nodes_in_edge(each_edge)) == set(node_list):
+                    return each_edge
+            else:
+                if self.nodes_in_edge(each_edge) == node_list:
+                    return each_edge
+        return False
+
+    def degree(self, node):
+        return len(self.hg.adj[node])
+
+    def degrees(self):
+        return {each_node: self.degree(each_node) for each_node in self.nodes}
+
+    def edge_degree(self, edge):
+        return len(self.nodes_in_edge(edge))
+
+    def edge_degrees(self):
+        return {each_edge: self.edge_degree(each_edge) for each_edge in self.edges}
+
+    def is_adj(self, node1, node2):
+        return node1 in self.adj_nodes(node2)
+
+    def adj_subhg(self, node, ident_node_dict=None):
+        """ return a subhypergraph consisting of a set of nodes and hyperedges adjacent to `node`.
+        if an adjacent node has a self-loop hyperedge, it will be also added to the subhypergraph.
+
+        Parameters
+        ----------
+        node : str
+        ident_node_dict : dict
+            dict containing identical nodes. see `get_identical_node_dict` for more details
+
+        Returns
+        -------
+        subhg : Hypergraph
+        """
+        if ident_node_dict is None:
+            ident_node_dict = self.get_identical_node_dict()
+        adj_node_set = set(ident_node_dict[node])
+        adj_edge_set = set([])
+        for each_node in ident_node_dict[node]:
+            adj_edge_set.update(set(self.adj_edges(each_node)))
+        fixed_adj_edge_set = deepcopy(adj_edge_set)
+        for each_edge in fixed_adj_edge_set:
+            other_nodes = self.nodes_in_edge(each_edge)
+            adj_node_set.update(other_nodes)
+
+            # if the adjacent node has self-loop edge, it will be appended to adj_edge_list.
+            for each_node in other_nodes:
+                for other_edge in set(self.adj_edges(each_node)) - set([each_edge]):
+                    if len(set(self.nodes_in_edge(other_edge)) \
+                           - set(self.nodes_in_edge(each_edge))) == 0:
+                        adj_edge_set.update(set([other_edge]))
+        subhg = Hypergraph()
+        for each_node in adj_node_set:
+            subhg.add_node(each_node, attr_dict=self.node_attr(each_node))
+        for each_edge in adj_edge_set:
+            subhg.add_edge(self.nodes_in_edge(each_edge),
+                           attr_dict=self.edge_attr(each_edge),
+                           edge_name=each_edge)
+        subhg.edge_idx = self.edge_idx
+        return subhg
+
+    def get_subhg(self, node_list, edge_list, ident_node_dict=None):
+        """ return a subhypergraph consisting of a set of nodes and hyperedges adjacent to `node`.
+        if an adjacent node has a self-loop hyperedge, it will be also added to the subhypergraph.
+
+        Parameters
+        ----------
+        node : str
+        ident_node_dict : dict
+            dict containing identical nodes. see `get_identical_node_dict` for more details
+
+        Returns
+        -------
+        subhg : Hypergraph
+        """
+        if ident_node_dict is None:
+            ident_node_dict = self.get_identical_node_dict()
+        adj_node_set = set([])
+        for each_node in node_list:
+            adj_node_set.update(set(ident_node_dict[each_node]))
+        adj_edge_set = set(edge_list)
+
+        subhg = Hypergraph()
+        for each_node in adj_node_set:
+            subhg.add_node(each_node,
+                           attr_dict=deepcopy(self.node_attr(each_node)))
+        for each_edge in adj_edge_set:
+            subhg.add_edge(self.nodes_in_edge(each_edge),
+                           attr_dict=deepcopy(self.edge_attr(each_edge)),
+                           edge_name=each_edge)
+        subhg.edge_idx = self.edge_idx
+        return subhg
+
+    def copy(self):
+        ''' return a copy of the object
+        
+        Returns
+        -------
+        Hypergraph
+        '''
+        return deepcopy(self)
+
+    def node_attr(self, node):
+        return self.hg.nodes[node]['attr_dict']
+
+    def edge_attr(self, edge):
+        return self.hg.nodes[edge]['attr_dict']
+
+    def set_node_attr(self, node, attr_dict):
+        for each_key, each_val in attr_dict.items():
+            self.hg.nodes[node]['attr_dict'][each_key] = each_val
+
+    def set_edge_attr(self, edge, attr_dict):
+        for each_key, each_val in attr_dict.items():
+            self.hg.nodes[edge]['attr_dict'][each_key] = each_val
+
+    def get_identical_node_dict(self):
+        ''' get identical nodes
+        nodes are identical if they share the same set of adjacent edges.
+        
+        Returns
+        -------
+        ident_node_dict : dict
+            ident_node_dict[node] returns a list of nodes that are identical to `node`.
+        '''
+        ident_node_dict = {}
+        for each_node in self.nodes:
+            ident_node_list = []
+            for each_other_node in self.nodes:
+                if each_other_node == each_node:
+                    ident_node_list.append(each_other_node)
+                elif self.adj_edges(each_node) == self.adj_edges(each_other_node) \
+                   and len(self.adj_edges(each_node)) != 0:
+                    ident_node_list.append(each_other_node)
+            ident_node_dict[each_node] = ident_node_list
+        return ident_node_dict
+    '''
+        ident_node_dict = {}
+        for each_node in self.nodes:
+            ident_node_dict[each_node] = [each_node]
+        return ident_node_dict
+    '''
+
+    def get_leaf_edge(self):
+        ''' get an edge that is incident only to one edge
+
+        Returns
+        -------
+        if exists, return a leaf edge. otherwise, return None.
+        '''
+        for each_edge in self.edges:
+            if len(self.adj_nodes(each_edge)) == 1:
+                if 'tmp' not in self.edge_attr(each_edge):
+                    return each_edge
+        return None
+
+    def get_nontmp_edge(self):
+        for each_edge in self.edges:
+            if 'tmp' not in self.edge_attr(each_edge):
+                return each_edge
+        return None
+
+    def is_subhg(self, hg):
+        ''' return whether this hypergraph is a subhypergraph of `hg`
+
+        Returns
+        -------
+        True if self \in hg,
+        False otherwise.
+        '''
+        for each_node in self.nodes:
+            if each_node not in hg.nodes:
+                return False
+        for each_edge in self.edges:
+            if each_edge not in hg.edges:
+                return False
+        return True
+
+    def in_cycle(self, node, visited=None, parent='', root_node='') -> bool:
+        ''' if `node` is in a cycle, then return True. otherwise, False.
+
+        Parameters
+        ----------
+        node : str
+            node in a hypergraph
+        visited : list
+            list of visited nodes, used for recursion
+        parent : str
+            parent node, used to eliminate a cycle consisting of two nodes and one edge.
+
+        Returns
+        -------
+        bool
+        '''
+        if visited is None:
+            visited = []
+        if parent == '':
+            visited = []
+        if root_node == '':
+            root_node = node
+        visited.append(node)
+        for each_adj_node in self.adj_nodes(node):
+            if each_adj_node not in visited:
+                if self.in_cycle(each_adj_node, visited, node, root_node):
+                    return True
+            elif each_adj_node != parent and each_adj_node == root_node:
+                return True
+        return False
+
+    
+    def draw(self, file_path=None, with_node=False, with_edge_name=False):
+        ''' draw hypergraph
+        '''
+        import graphviz
+        G = graphviz.Graph(format='png')
+        for each_node in self.nodes:
+            if 'ext_id' in self.node_attr(each_node):
+                G.node(each_node, label='',
+                       shape='circle', width='0.1', height='0.1', style='filled',
+                       fillcolor='black')
+            else:
+                if with_node:
+                    G.node(each_node, label='',
+                           shape='circle', width='0.1', height='0.1', style='filled',
+                           fillcolor='gray')
+        edge_list = []
+        for each_edge in self.edges:
+            if self.edge_attr(each_edge).get('terminal', False):
+                G.node(each_edge,
+                       label=self.edge_attr(each_edge)['symbol'].symbol if not with_edge_name \
+                       else self.edge_attr(each_edge)['symbol'].symbol + ', ' + each_edge,
+                       fontcolor='black', shape='square')
+            elif self.edge_attr(each_edge).get('tmp', False):
+                G.node(each_edge, label='tmp' if not with_edge_name else 'tmp, ' + each_edge,
+                       fontcolor='black', shape='square')
+            else:
+                G.node(each_edge,
+                       label=self.edge_attr(each_edge)['symbol'].symbol if not with_edge_name \
+                       else self.edge_attr(each_edge)['symbol'].symbol + ', ' + each_edge,
+                       fontcolor='black', shape='square', style='filled')
+            if with_node:
+                for each_node in self.nodes_in_edge(each_edge):
+                    G.edge(each_edge, each_node)
+            else:
+                for each_node in self.nodes_in_edge(each_edge):
+                    if 'ext_id' in self.node_attr(each_node)\
+                       and set([each_node, each_edge]) not in edge_list:
+                        G.edge(each_edge, each_node)
+                        edge_list.append(set([each_node, each_edge]))
+                for each_other_edge in self.adj_nodes(each_edge):
+                    if set([each_edge, each_other_edge]) not in edge_list:
+                        num_bond = 0
+                        common_node_set = set(self.nodes_in_edge(each_edge))\
+                                          .intersection(set(self.nodes_in_edge(each_other_edge)))
+                        for each_node in common_node_set:
+                            if self.node_attr(each_node)['symbol'].bond_type in [1, 2, 3]:
+                                num_bond += self.node_attr(each_node)['symbol'].bond_type
+                            elif self.node_attr(each_node)['symbol'].bond_type in [12]:
+                                num_bond += 1
+                            else:
+                                raise NotImplementedError('unsupported bond type')
+                        for _ in range(num_bond):
+                            G.edge(each_edge, each_other_edge)
+                        edge_list.append(set([each_edge, each_other_edge]))
+        if file_path is not None:
+            G.render(file_path, cleanup=True)
+            #os.remove(file_path)
+        return G
+
+    def is_dividable(self, node):
+        _hg = deepcopy(self.hg)
+        _hg.remove_node(node)
+        return (not nx.is_connected(_hg))
+
+    def divide(self, node):
+        subhg_list = []
+
+        hg_wo_node = deepcopy(self)
+        hg_wo_node.remove_node(node, remove_connected_edges=False)
+        connected_components = nx.connected_components(hg_wo_node.hg)
+        for each_component in connected_components:
+            node_list = [node]
+            edge_list = []
+            node_list.extend([each_node for each_node in each_component
+                              if each_node.startswith('bond_')])
+            edge_list.extend([each_edge for each_edge in each_component
+                              if each_edge.startswith('e')])
+            subhg_list.append(self.get_subhg(node_list, edge_list))
+            #subhg_list[-1].set_node_attr(node, {'divided': True})
+        return subhg_list
+

models/mhg_model/graph_grammar/io/__init__.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 1 2018"
+

models/mhg_model/graph_grammar/io/smi.py

@@ -0,0 +1,559 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 12 2018"
+
+from copy import deepcopy
+from rdkit import Chem
+from rdkit import RDLogger
+import networkx as nx
+import numpy as np
+from ..hypergraph import Hypergraph
+from ..graph_grammar.symbols import TSymbol, BondSymbol
+
+# supress warnings
+lg = RDLogger.logger()
+lg.setLevel(RDLogger.CRITICAL)
+
+
+class HGGen(object):
+    """
+    load .smi file and yield a hypergraph.
+
+    Attributes
+    ----------
+    path_to_file : str
+        path to .smi file
+    kekulize : bool
+        kekulize or not
+    add_Hs : bool
+        add implicit hydrogens to the molecule or not.
+    all_single : bool
+        if True, all multiple bonds are summarized into a single bond with some attributes
+
+    Yields
+    ------
+    Hypergraph
+    """
+    def __init__(self, path_to_file, kekulize=True, add_Hs=False, all_single=True):
+        self.num_line = 1
+        self.mol_gen = Chem.SmilesMolSupplier(path_to_file, titleLine=False)
+        self.kekulize = kekulize
+        self.add_Hs = add_Hs
+        self.all_single = all_single
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        '''
+        each_mol = None
+        while each_mol is None:
+            each_mol = next(self.mol_gen)
+        '''
+        # not ignoring parse errors
+        each_mol = next(self.mol_gen)
+        if each_mol is None:
+            raise ValueError(f'incorrect smiles in line {self.num_line}')
+        else:
+            self.num_line += 1
+        return mol_to_hg(each_mol, self.kekulize, self.add_Hs)
+
+
+def mol_to_bipartite(mol, kekulize):
+    """
+    get a bipartite representation of a molecule.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        molecule object
+
+    Returns
+    -------
+    nx.Graph
+        a bipartite graph representing which bond is connected to which atoms.
+    """
+    try:
+        mol = standardize_stereo(mol)
+    except KeyError:
+        print(Chem.MolToSmiles(mol))
+        raise KeyError
+        
+    if kekulize:
+        Chem.Kekulize(mol)
+
+    bipartite_g = nx.Graph()
+    for each_atom in mol.GetAtoms():
+        bipartite_g.add_node(f"atom_{each_atom.GetIdx()}",
+                             atom_attr=atom_attr(each_atom, kekulize))
+
+    for each_bond in mol.GetBonds():
+        bond_idx = each_bond.GetIdx()
+        bipartite_g.add_node(
+            f"bond_{bond_idx}",
+            bond_attr=bond_attr(each_bond, kekulize))
+        bipartite_g.add_edge(
+            f"atom_{each_bond.GetBeginAtomIdx()}",
+            f"bond_{bond_idx}")
+        bipartite_g.add_edge(
+            f"atom_{each_bond.GetEndAtomIdx()}",
+            f"bond_{bond_idx}")
+    return bipartite_g
+
+
+def mol_to_hg(mol, kekulize, add_Hs):
+    """
+    get a bipartite representation of a molecule.
+
+    Parameters
+    ----------
+    mol : rdkit.Chem.rdchem.Mol
+        molecule object
+    kekulize : bool
+        kekulize or not
+    add_Hs : bool
+        add implicit hydrogens to the molecule or not.
+
+    Returns
+    -------
+    Hypergraph
+    """
+    if add_Hs:
+        mol = Chem.AddHs(mol)
+
+    if kekulize:
+        Chem.Kekulize(mol)
+
+    bipartite_g = mol_to_bipartite(mol, kekulize)
+    hg = Hypergraph()
+    for each_atom in [each_node for each_node in bipartite_g.nodes()
+                      if each_node.startswith('atom_')]:
+        node_set = set([])
+        for each_bond in bipartite_g.adj[each_atom]:
+            hg.add_node(each_bond,
+                        attr_dict=bipartite_g.nodes[each_bond]['bond_attr'])
+            node_set.add(each_bond)
+        hg.add_edge(node_set,
+                    attr_dict=bipartite_g.nodes[each_atom]['atom_attr'])
+    return hg
+
+
+def hg_to_mol(hg, verbose=False):
+    """ convert a hypergraph into Mol object
+
+    Parameters
+    ----------
+    hg : Hypergraph
+
+    Returns
+    -------
+    mol : Chem.RWMol
+    """
+    mol = Chem.RWMol()
+    atom_dict = {}
+    bond_set = set([])
+    for each_edge in hg.edges:
+        atom = Chem.Atom(hg.edge_attr(each_edge)['symbol'].symbol)
+        atom.SetNumExplicitHs(hg.edge_attr(each_edge)['symbol'].num_explicit_Hs)
+        atom.SetFormalCharge(hg.edge_attr(each_edge)['symbol'].formal_charge)
+        atom.SetChiralTag(
+            Chem.rdchem.ChiralType.values[
+                hg.edge_attr(each_edge)['symbol'].chirality])
+        atom_idx = mol.AddAtom(atom)
+        atom_dict[each_edge] = atom_idx
+
+    for each_node in hg.nodes:
+        edge_1, edge_2 = hg.adj_edges(each_node)
+        if edge_1+edge_2 not in bond_set:
+            if hg.node_attr(each_node)['symbol'].bond_type <= 3:
+                num_bond = hg.node_attr(each_node)['symbol'].bond_type
+            elif hg.node_attr(each_node)['symbol'].bond_type == 12:
+                num_bond = 1
+            else:
+                raise ValueError(f'too many bonds; {hg.node_attr(each_node)["bond_symbol"].bond_type}')
+            _ = mol.AddBond(atom_dict[edge_1],
+                            atom_dict[edge_2],
+                            order=Chem.rdchem.BondType.values[num_bond])
+            bond_idx = mol.GetBondBetweenAtoms(atom_dict[edge_1], atom_dict[edge_2]).GetIdx()
+
+            # stereo
+            mol.GetBondWithIdx(bond_idx).SetStereo(
+                Chem.rdchem.BondStereo.values[hg.node_attr(each_node)['symbol'].stereo])
+            bond_set.update([edge_1+edge_2])
+            bond_set.update([edge_2+edge_1])
+    mol.UpdatePropertyCache()
+    mol = mol.GetMol()
+    not_stereo_mol = deepcopy(mol)
+    if Chem.MolFromSmiles(Chem.MolToSmiles(not_stereo_mol)) is None:
+        raise RuntimeError('no valid molecule was obtained.')
+    try:
+        mol = set_stereo(mol)
+        is_stereo = True
+    except:
+        import traceback
+        traceback.print_exc()
+        is_stereo = False
+    mol_tmp = deepcopy(mol)
+    Chem.SetAromaticity(mol_tmp)
+    if Chem.MolFromSmiles(Chem.MolToSmiles(mol_tmp)) is not None:
+        mol = mol_tmp
+    else:
+        if Chem.MolFromSmiles(Chem.MolToSmiles(mol)) is None:
+            mol = not_stereo_mol
+    mol.UpdatePropertyCache()
+    Chem.GetSymmSSSR(mol)
+    mol = Chem.MolFromSmiles(Chem.MolToSmiles(mol))
+    if verbose:
+        return mol, is_stereo
+    else:
+        return mol
+
+def hgs_to_mols(hg_list, ignore_error=False):
+    if ignore_error:
+        mol_list = []
+        for each_hg in hg_list:
+            try:
+                mol = hg_to_mol(each_hg)
+            except:
+                mol = None
+            mol_list.append(mol)
+    else:
+        mol_list = [hg_to_mol(each_hg) for each_hg in hg_list]
+    return mol_list
+
+def hgs_to_smiles(hg_list, ignore_error=False):
+    mol_list = hgs_to_mols(hg_list, ignore_error)
+    smiles_list = []
+    for each_mol in mol_list:
+        try:
+            smiles_list.append(
+                Chem.MolToSmiles(
+                    Chem.MolFromSmiles(
+                        Chem.MolToSmiles(
+                            each_mol))))
+        except:
+            smiles_list.append(None)
+    return smiles_list
+
+def atom_attr(atom, kekulize):
+    """
+    get atom's attributes
+
+    Parameters
+    ----------
+    atom : rdkit.Chem.rdchem.Atom
+    kekulize : bool
+        kekulize or not
+
+    Returns
+    -------
+    atom_attr : dict
+        "is_aromatic" : bool
+            the atom is aromatic or not.
+        "smarts" : str
+            SMARTS representation of the atom.
+    """
+    if kekulize:
+        return {'terminal': True,
+                'is_in_ring': atom.IsInRing(),
+                'symbol': TSymbol(degree=0,
+                                  #degree=atom.GetTotalDegree(),
+                                  is_aromatic=False,
+                                  symbol=atom.GetSymbol(),
+                                  num_explicit_Hs=atom.GetNumExplicitHs(),
+                                  formal_charge=atom.GetFormalCharge(),
+                                  chirality=atom.GetChiralTag().real
+                )}
+    else:
+        return {'terminal': True,
+                'is_in_ring': atom.IsInRing(),
+                'symbol': TSymbol(degree=0,
+                                  #degree=atom.GetTotalDegree(),
+                                  is_aromatic=atom.GetIsAromatic(),
+                                  symbol=atom.GetSymbol(),
+                                  num_explicit_Hs=atom.GetNumExplicitHs(),
+                                  formal_charge=atom.GetFormalCharge(),
+                                  chirality=atom.GetChiralTag().real
+                )}
+
+def bond_attr(bond, kekulize):
+    """
+    get atom's attributes
+
+    Parameters
+    ----------
+    bond : rdkit.Chem.rdchem.Bond
+    kekulize : bool
+        kekulize or not
+
+    Returns
+    -------
+    bond_attr : dict
+        "bond_type" : int
+        {0: rdkit.Chem.rdchem.BondType.UNSPECIFIED,
+         1: rdkit.Chem.rdchem.BondType.SINGLE,
+         2: rdkit.Chem.rdchem.BondType.DOUBLE,
+         3: rdkit.Chem.rdchem.BondType.TRIPLE,
+         4: rdkit.Chem.rdchem.BondType.QUADRUPLE,
+         5: rdkit.Chem.rdchem.BondType.QUINTUPLE,
+         6: rdkit.Chem.rdchem.BondType.HEXTUPLE,
+         7: rdkit.Chem.rdchem.BondType.ONEANDAHALF,
+         8: rdkit.Chem.rdchem.BondType.TWOANDAHALF,
+         9: rdkit.Chem.rdchem.BondType.THREEANDAHALF,
+         10: rdkit.Chem.rdchem.BondType.FOURANDAHALF,
+         11: rdkit.Chem.rdchem.BondType.FIVEANDAHALF,
+         12: rdkit.Chem.rdchem.BondType.AROMATIC,
+         13: rdkit.Chem.rdchem.BondType.IONIC,
+         14: rdkit.Chem.rdchem.BondType.HYDROGEN,
+         15: rdkit.Chem.rdchem.BondType.THREECENTER,
+         16: rdkit.Chem.rdchem.BondType.DATIVEONE,
+         17: rdkit.Chem.rdchem.BondType.DATIVE,
+         18: rdkit.Chem.rdchem.BondType.DATIVEL,
+         19: rdkit.Chem.rdchem.BondType.DATIVER,
+         20: rdkit.Chem.rdchem.BondType.OTHER,
+         21: rdkit.Chem.rdchem.BondType.ZERO}
+    """
+    if kekulize:
+        is_aromatic = False
+        if bond.GetBondType().real == 12:
+            bond_type = 1
+        else:
+            bond_type = bond.GetBondType().real
+    else:
+        is_aromatic = bond.GetIsAromatic()
+        bond_type = bond.GetBondType().real
+    return {'symbol': BondSymbol(is_aromatic=is_aromatic,
+                                 bond_type=bond_type,
+                                 stereo=int(bond.GetStereo())),
+            'is_in_ring': bond.IsInRing()}
+
+
+def standardize_stereo(mol):
+    '''
+ 0: rdkit.Chem.rdchem.BondDir.NONE,
+ 1: rdkit.Chem.rdchem.BondDir.BEGINWEDGE,
+ 2: rdkit.Chem.rdchem.BondDir.BEGINDASH,
+ 3: rdkit.Chem.rdchem.BondDir.ENDDOWNRIGHT,
+ 4: rdkit.Chem.rdchem.BondDir.ENDUPRIGHT,
+
+    '''
+    # mol = Chem.AddHs(mol) # this removes CIPRank !!!
+    for each_bond in mol.GetBonds():
+        if int(each_bond.GetStereo()) in [2, 3]: #2=Z (same side), 3=E
+            begin_stereo_atom_idx = each_bond.GetBeginAtomIdx()
+            end_stereo_atom_idx = each_bond.GetEndAtomIdx()
+            atom_idx_1 = each_bond.GetStereoAtoms()[0]
+            atom_idx_2 = each_bond.GetStereoAtoms()[1]
+            if mol.GetBondBetweenAtoms(atom_idx_1, begin_stereo_atom_idx):
+                begin_atom_idx = atom_idx_1
+                end_atom_idx = atom_idx_2
+            else:
+                begin_atom_idx = atom_idx_2
+                end_atom_idx = atom_idx_1
+
+            begin_another_atom_idx = None
+            assert len(mol.GetAtomWithIdx(begin_stereo_atom_idx).GetNeighbors()) <= 3
+            for each_neighbor in mol.GetAtomWithIdx(begin_stereo_atom_idx).GetNeighbors():
+                each_neighbor_idx = each_neighbor.GetIdx()
+                if each_neighbor_idx not in [end_stereo_atom_idx, begin_atom_idx]:
+                    begin_another_atom_idx = each_neighbor_idx
+
+            end_another_atom_idx = None
+            assert len(mol.GetAtomWithIdx(end_stereo_atom_idx).GetNeighbors()) <= 3
+            for each_neighbor in mol.GetAtomWithIdx(end_stereo_atom_idx).GetNeighbors():
+                each_neighbor_idx = each_neighbor.GetIdx()
+                if each_neighbor_idx not in [begin_stereo_atom_idx, end_atom_idx]:
+                    end_another_atom_idx = each_neighbor_idx
+
+            ''' 
+            relationship between begin_atom_idx and end_atom_idx is encoded in GetStereo
+            '''
+            begin_atom_rank = int(mol.GetAtomWithIdx(begin_atom_idx).GetProp('_CIPRank'))
+            end_atom_rank = int(mol.GetAtomWithIdx(end_atom_idx).GetProp('_CIPRank'))
+            try:
+                begin_another_atom_rank = int(mol.GetAtomWithIdx(begin_another_atom_idx).GetProp('_CIPRank'))
+            except:
+                begin_another_atom_rank = np.inf
+            try:
+                end_another_atom_rank = int(mol.GetAtomWithIdx(end_another_atom_idx).GetProp('_CIPRank'))
+            except:
+                end_another_atom_rank = np.inf
+            if begin_atom_rank < begin_another_atom_rank\
+               and end_atom_rank < end_another_atom_rank:
+                pass
+            elif begin_atom_rank < begin_another_atom_rank\
+                 and end_atom_rank > end_another_atom_rank:
+                # (begin_atom_idx +) end_another_atom_idx should be in StereoAtoms
+                if each_bond.GetStereo() == 2:
+                    # set stereo
+                    each_bond.SetStereo(Chem.rdchem.BondStereo.values[3])
+                    # set bond dir
+                    mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 3)
+                    mol = safe_set_bond_dir(mol, begin_another_atom_idx, begin_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, end_another_atom_idx, end_stereo_atom_idx, 3)
+                elif each_bond.GetStereo() == 3:
+                    # set stereo
+                    each_bond.SetStereo(Chem.rdchem.BondStereo.values[2])
+                    # set bond dir
+                    mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 3)
+                    mol = safe_set_bond_dir(mol, begin_another_atom_idx, begin_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, end_another_atom_idx, end_stereo_atom_idx, 4)
+                else:
+                    raise ValueError
+                each_bond.SetStereoAtoms(begin_atom_idx, end_another_atom_idx)
+            elif begin_atom_rank > begin_another_atom_rank\
+                 and end_atom_rank < end_another_atom_rank:
+                # (end_atom_idx +) begin_another_atom_idx should be in StereoAtoms
+                if each_bond.GetStereo() == 2:
+                    # set stereo
+                    each_bond.SetStereo(Chem.rdchem.BondStereo.values[3])
+                    # set bond dir
+                    mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, begin_another_atom_idx, begin_stereo_atom_idx, 4)
+                    mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 4)
+                    mol = safe_set_bond_dir(mol, end_another_atom_idx, end_stereo_atom_idx, 0)
+                elif each_bond.GetStereo() == 3:
+                    # set stereo
+                    each_bond.SetStereo(Chem.rdchem.BondStereo.values[2])
+                    # set bond dir
+                    mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, begin_another_atom_idx, begin_stereo_atom_idx, 4)
+                    mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 3)
+                    mol = safe_set_bond_dir(mol, end_another_atom_idx, end_stereo_atom_idx, 0)
+                else:
+                    raise ValueError
+                each_bond.SetStereoAtoms(begin_another_atom_idx, end_atom_idx)
+            elif begin_atom_rank > begin_another_atom_rank\
+                 and end_atom_rank > end_another_atom_rank:
+                # begin_another_atom_idx + end_another_atom_idx should be in StereoAtoms
+                if each_bond.GetStereo() == 2:
+                    # set bond dir
+                    mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, begin_another_atom_idx, begin_stereo_atom_idx, 4)
+                    mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, end_another_atom_idx, end_stereo_atom_idx, 3)
+                elif each_bond.GetStereo() == 3:
+                    # set bond dir
+                    mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, begin_another_atom_idx, begin_stereo_atom_idx, 4)
+                    mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 0)
+                    mol = safe_set_bond_dir(mol, end_another_atom_idx, end_stereo_atom_idx, 4)
+                else:
+                    raise ValueError
+                each_bond.SetStereoAtoms(begin_another_atom_idx, end_another_atom_idx)
+            else:
+                raise RuntimeError
+    return mol
+
+
+def set_stereo(mol):
+    '''
+ 0: rdkit.Chem.rdchem.BondDir.NONE,
+ 1: rdkit.Chem.rdchem.BondDir.BEGINWEDGE,
+ 2: rdkit.Chem.rdchem.BondDir.BEGINDASH,
+ 3: rdkit.Chem.rdchem.BondDir.ENDDOWNRIGHT,
+ 4: rdkit.Chem.rdchem.BondDir.ENDUPRIGHT,
+    '''
+    _mol = Chem.MolFromSmiles(Chem.MolToSmiles(mol))
+    Chem.Kekulize(_mol, True)
+    substruct_match = mol.GetSubstructMatch(_mol)
+    if not substruct_match:
+        ''' mol and _mol are kekulized.
+        sometimes, the order of '=' and '-' changes, which causes mol and _mol not matched.
+        '''
+        Chem.SetAromaticity(mol)
+        Chem.SetAromaticity(_mol)
+        substruct_match = mol.GetSubstructMatch(_mol)
+    try:
+        atom_match = {substruct_match[_mol_atom_idx]: _mol_atom_idx for _mol_atom_idx in range(_mol.GetNumAtoms())} # mol to _mol
+    except:
+        raise ValueError('two molecules obtained from the same data do not match.')
+        
+    for each_bond in mol.GetBonds():
+        begin_atom_idx = each_bond.GetBeginAtomIdx()
+        end_atom_idx = each_bond.GetEndAtomIdx()
+        _bond = _mol.GetBondBetweenAtoms(atom_match[begin_atom_idx], atom_match[end_atom_idx])
+        _bond.SetStereo(each_bond.GetStereo())
+
+    mol = _mol
+    for each_bond in mol.GetBonds():
+        if int(each_bond.GetStereo()) in [2, 3]: #2=Z (same side), 3=E
+            begin_stereo_atom_idx = each_bond.GetBeginAtomIdx()
+            end_stereo_atom_idx = each_bond.GetEndAtomIdx()
+            begin_atom_idx_set = set([each_neighbor.GetIdx()
+                                      for each_neighbor
+                                      in mol.GetAtomWithIdx(begin_stereo_atom_idx).GetNeighbors()
+                                      if each_neighbor.GetIdx() != end_stereo_atom_idx])
+            end_atom_idx_set = set([each_neighbor.GetIdx()
+                                    for each_neighbor
+                                    in mol.GetAtomWithIdx(end_stereo_atom_idx).GetNeighbors()
+                                    if each_neighbor.GetIdx() != begin_stereo_atom_idx])
+            if not begin_atom_idx_set:
+                each_bond.SetStereo(Chem.rdchem.BondStereo(0))
+                continue
+            if not end_atom_idx_set:
+                each_bond.SetStereo(Chem.rdchem.BondStereo(0))
+                continue
+            if len(begin_atom_idx_set) == 1:
+                begin_atom_idx = begin_atom_idx_set.pop()
+                begin_another_atom_idx = None
+            if len(end_atom_idx_set) == 1:
+                end_atom_idx = end_atom_idx_set.pop()
+                end_another_atom_idx = None
+            if len(begin_atom_idx_set) == 2:
+                atom_idx_1 = begin_atom_idx_set.pop()
+                atom_idx_2 = begin_atom_idx_set.pop()
+                if int(mol.GetAtomWithIdx(atom_idx_1).GetProp('_CIPRank')) < int(mol.GetAtomWithIdx(atom_idx_2).GetProp('_CIPRank')):
+                    begin_atom_idx = atom_idx_1
+                    begin_another_atom_idx = atom_idx_2
+                else:
+                    begin_atom_idx = atom_idx_2
+                    begin_another_atom_idx = atom_idx_1
+            if len(end_atom_idx_set) == 2:
+                atom_idx_1 = end_atom_idx_set.pop()
+                atom_idx_2 = end_atom_idx_set.pop()
+                if int(mol.GetAtomWithIdx(atom_idx_1).GetProp('_CIPRank')) < int(mol.GetAtomWithIdx(atom_idx_2).GetProp('_CIPRank')):
+                    end_atom_idx = atom_idx_1
+                    end_another_atom_idx = atom_idx_2
+                else:
+                    end_atom_idx = atom_idx_2
+                    end_another_atom_idx = atom_idx_1
+
+            if each_bond.GetStereo() == 2: # same side
+                mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 3)
+                mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 4)
+                each_bond.SetStereoAtoms(begin_atom_idx, end_atom_idx)
+            elif each_bond.GetStereo() == 3: # opposite side
+                mol = safe_set_bond_dir(mol, begin_atom_idx, begin_stereo_atom_idx, 3)
+                mol = safe_set_bond_dir(mol, end_atom_idx, end_stereo_atom_idx, 3)
+                each_bond.SetStereoAtoms(begin_atom_idx, end_atom_idx)
+            else:
+                raise ValueError
+    return mol
+
+
+def safe_set_bond_dir(mol, atom_idx_1, atom_idx_2, bond_dir_val):
+    if atom_idx_1 is None or atom_idx_2 is None:
+        return mol
+    else:
+        mol.GetBondBetweenAtoms(atom_idx_1, atom_idx_2).SetBondDir(Chem.rdchem.BondDir.values[bond_dir_val])
+        return mol
+        

models/mhg_model/graph_grammar/nn/__init__.py

@@ -0,0 +1,11 @@
+# -*- coding:utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""

models/mhg_model/graph_grammar/nn/dataset.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Apr 18 2018"
+
+from torch.utils.data import Dataset, DataLoader
+import torch
+import numpy as np
+
+
+def left_padding(sentence_list, max_len, pad_idx=-1, inverse=False):
+    ''' pad left
+
+    Parameters
+    ----------
+    sentence_list : list of sequences of integers
+    max_len : int
+        maximum length of sentences.
+        if a sentence is shorter than `max_len`, its left part is padded.
+    pad_idx : int
+        integer for padding
+    inverse : bool
+        if True, the sequence is inversed.
+
+    Returns
+    -------
+    List of torch.LongTensor
+        each sentence is left-padded.
+    '''
+    max_in_list = max([len(each_sen) for each_sen in sentence_list])
+        
+    if max_in_list > max_len:
+        raise ValueError('`max_len` should be larger than the maximum length of input sequences, {}.'.format(max_in_list))
+
+    if inverse:
+        return [torch.LongTensor([pad_idx] * (max_len - len(each_sen)) + each_sen[::-1]) for each_sen in sentence_list]
+    else:
+        return [torch.LongTensor([pad_idx] * (max_len - len(each_sen)) + each_sen) for each_sen in sentence_list]
+
+
+def right_padding(sentence_list, max_len, pad_idx=-1):
+    ''' pad right
+
+    Parameters
+    ----------
+    sentence_list : list of sequences of integers
+    max_len : int
+        maximum length of sentences.
+        if a sentence is shorter than `max_len`, its right part is padded.
+    pad_idx : int
+        integer for padding
+
+    Returns
+    -------
+    List of torch.LongTensor
+        each sentence is right-padded.
+    '''
+    max_in_list = max([len(each_sen) for each_sen in sentence_list])
+    if max_in_list > max_len:
+        raise ValueError('`max_len` should be larger than the maximum length of input sequences, {}.'.format(max_in_list))
+
+    return [torch.LongTensor(each_sen + [pad_idx] * (max_len - len(each_sen))) for each_sen in sentence_list]
+
+
+class HRGDataset(Dataset):
+
+    '''
+    A class of HRG data
+    '''
+
+    def __init__(self, hrg, prod_rule_seq_list, max_len, target_val_list=None, inversed_input=False):
+        self.hrg = hrg
+        self.left_prod_rule_seq_list = left_padding(prod_rule_seq_list,
+                                                    max_len,
+                                                    inverse=inversed_input)
+            
+        self.right_prod_rule_seq_list = right_padding(prod_rule_seq_list, max_len)
+        self.inserved_input = inversed_input
+        self.target_val_list = target_val_list
+        if target_val_list is not None:
+            if len(prod_rule_seq_list) != len(target_val_list):
+                raise ValueError(f'prod_rule_seq_list and target_val_list have inconsistent lengths: {len(prod_rule_seq_list)}, {len(target_val_list)}')
+
+    def __len__(self):
+        return len(self.left_prod_rule_seq_list)
+
+    def __getitem__(self, idx):
+        if self.target_val_list is not None:
+            return self.left_prod_rule_seq_list[idx], self.right_prod_rule_seq_list[idx], np.float32(self.target_val_list[idx])
+        else:
+            return self.left_prod_rule_seq_list[idx], self.right_prod_rule_seq_list[idx]
+
+    @property
+    def vocab_size(self):
+        return self.hrg.num_prod_rule
+
+def batch_padding(each_batch, batch_size, padding_idx):
+    num_pad = batch_size - len(each_batch[0])
+    if num_pad:
+        each_batch[0] = torch.cat([each_batch[0],
+                                   padding_idx * torch.ones((batch_size - len(each_batch[0]),
+                                                             len(each_batch[0][0])), dtype=torch.int64)], dim=0)
+        each_batch[1] = torch.cat([each_batch[1],
+                                   padding_idx * torch.ones((batch_size - len(each_batch[1]),
+                                                             len(each_batch[1][0])), dtype=torch.int64)], dim=0)
+    return each_batch, num_pad

models/mhg_model/graph_grammar/nn/decoder.py

@@ -0,0 +1,158 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Aug 9 2018"
+
+
+import abc
+import numpy as np
+import torch
+from torch import nn
+
+
+class DecoderBase(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.hidden_dict = {}
+
+    @abc.abstractmethod
+    def forward_one_step(self, tgt_emb_in):
+        ''' one-step forward model
+
+        Parameters
+        ----------
+        tgt_emb_in : Tensor, shape (batch_size, input_dim)
+
+        Returns
+        -------
+        Tensor, shape (batch_size, hidden_dim)
+        '''
+        tgt_emb_out = None
+        return tgt_emb_out
+
+    @abc.abstractmethod
+    def init_hidden(self):
+        ''' initialize the hidden states
+        '''
+        pass
+
+    @abc.abstractmethod
+    def feed_hidden(self, hidden_dict_0):
+        for each_hidden in self.hidden_dict.keys():
+            self.hidden_dict[each_hidden][0] = hidden_dict_0[each_hidden]
+
+
+class GRUDecoder(DecoderBase):
+
+    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int,
+                 dropout: float, batch_size: int, use_gpu: bool,
+                 no_dropout=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.batch_size = batch_size
+        self.use_gpu = use_gpu
+        self.model = nn.GRU(input_size=self.input_dim,
+                            hidden_size=self.hidden_dim,
+                            num_layers=self.num_layers,
+                            batch_first=True,
+                            bidirectional=False,
+                            dropout=self.dropout if not no_dropout else 0
+        )
+        if self.use_gpu:
+            self.model.cuda()
+        self.init_hidden()
+
+    def init_hidden(self):
+        self.hidden_dict['h'] = torch.zeros((self.num_layers,
+                                             self.batch_size,
+                                             self.hidden_dim),
+                                            requires_grad=False)
+        if self.use_gpu:
+            self.hidden_dict['h'] = self.hidden_dict['h'].cuda()
+
+    def forward_one_step(self, tgt_emb_in):
+        ''' one-step forward model
+
+        Parameters
+        ----------
+        tgt_emb_in : Tensor, shape (batch_size, input_dim)
+
+        Returns
+        -------
+        Tensor, shape (batch_size, hidden_dim)
+        '''
+        tgt_emb_out, self.hidden_dict['h'] \
+            = self.model(tgt_emb_in.view(self.batch_size, 1, -1),
+                         self.hidden_dict['h'])
+        return tgt_emb_out
+
+
+class LSTMDecoder(DecoderBase):
+
+    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int,
+                 dropout: float, batch_size: int, use_gpu: bool,
+                 no_dropout=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.batch_size = batch_size
+        self.use_gpu = use_gpu
+        self.model = nn.LSTM(input_size=self.input_dim,
+                             hidden_size=self.hidden_dim,
+                             num_layers=self.num_layers,
+                             batch_first=True,
+                             bidirectional=False,
+                             dropout=self.dropout if not no_dropout else 0)
+        if self.use_gpu:
+            self.model.cuda()
+        self.init_hidden()
+
+    def init_hidden(self):
+        self.hidden_dict['h'] = torch.zeros((self.num_layers,
+                                             self.batch_size,
+                                             self.hidden_dim),
+                                            requires_grad=False)
+        self.hidden_dict['c'] = torch.zeros((self.num_layers,
+                                             self.batch_size,
+                                             self.hidden_dim),
+                                            requires_grad=False)
+        if self.use_gpu:
+            for each_hidden in self.hidden_dict.keys():
+                self.hidden_dict[each_hidden] = self.hidden_dict[each_hidden].cuda()
+
+    def forward_one_step(self, tgt_emb_in):
+        ''' one-step forward model
+
+        Parameters
+        ----------
+        tgt_emb_in : Tensor, shape (batch_size, input_dim)
+
+        Returns
+        -------
+        Tensor, shape (batch_size, hidden_dim)
+        '''
+        tgt_hidden_out, self.hidden_dict['h'], self.hidden_dict['c'] \
+            = self.model(tgt_emb_in.view(self.batch_size, 1, -1),
+                         self.hidden_dict['h'], self.hidden_dict['c'])
+        return tgt_hidden_out

models/mhg_model/graph_grammar/nn/encoder.py

@@ -0,0 +1,199 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Aug 9 2018"
+
+
+import abc
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from typing import List
+
+
+class EncoderBase(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @abc.abstractmethod
+    def forward(self, in_seq):
+        ''' forward model
+
+        Parameters
+        ----------
+        in_seq_emb : Variable, shape (batch_size, max_len, input_dim)
+
+        Returns
+        -------
+        hidden_seq_emb : Tensor, shape (batch_size, max_len, 1 + bidirectional, hidden_dim)
+        '''
+        pass
+
+    @abc.abstractmethod
+    def init_hidden(self):
+        ''' initialize the hidden states
+        '''
+        pass
+
+
+class GRUEncoder(EncoderBase):
+
+    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int,
+                 bidirectional: bool, dropout: float, batch_size: int, use_gpu: bool,
+                 no_dropout=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.dropout = dropout
+        self.batch_size = batch_size
+        self.use_gpu = use_gpu
+        self.model = nn.GRU(input_size=self.input_dim,
+                            hidden_size=self.hidden_dim,
+                            num_layers=self.num_layers,
+                            batch_first=True,
+                            bidirectional=self.bidirectional,
+                            dropout=self.dropout if not no_dropout else 0)
+        if self.use_gpu:
+            self.model.cuda()
+        self.init_hidden()
+
+
+    def init_hidden(self):
+        self.h0 = torch.zeros(((self.bidirectional + 1) * self.num_layers,
+                               self.batch_size,
+                               self.hidden_dim),
+                              requires_grad=False)
+        if self.use_gpu:        
+            self.h0 = self.h0.cuda()
+
+    def forward(self, in_seq_emb):
+        ''' forward model
+
+        Parameters
+        ----------
+        in_seq_emb : Tensor, shape (batch_size, max_len, input_dim)
+
+        Returns
+        -------
+        hidden_seq_emb : Tensor, shape (batch_size, max_len, 1 + bidirectional, hidden_dim)
+        '''
+        max_len = in_seq_emb.size(1)
+        hidden_seq_emb, self.h0 = self.model(
+            in_seq_emb, self.h0)
+        hidden_seq_emb = hidden_seq_emb.view(self.batch_size,
+                                             max_len,
+                                             1 + self.bidirectional,
+                                             self.hidden_dim)
+        return hidden_seq_emb
+
+
+class LSTMEncoder(EncoderBase):
+
+    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int,
+                 bidirectional: bool, dropout: float, batch_size: int, use_gpu: bool,
+                 no_dropout=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.dropout = dropout
+        self.batch_size = batch_size
+        self.use_gpu = use_gpu
+        self.model = nn.LSTM(input_size=self.input_dim,
+                             hidden_size=self.hidden_dim,
+                             num_layers=self.num_layers,
+                             batch_first=True,
+                             bidirectional=self.bidirectional,
+                             dropout=self.dropout if not no_dropout else 0)
+        if self.use_gpu:
+            self.model.cuda()
+        self.init_hidden()
+
+    def init_hidden(self):
+        self.h0 = torch.zeros(((self.bidirectional + 1) * self.num_layers,
+                               self.batch_size,
+                               self.hidden_dim),
+                              requires_grad=False)
+        self.c0 = torch.zeros(((self.bidirectional + 1) * self.num_layers,
+                               self.batch_size,
+                               self.hidden_dim),
+                              requires_grad=False)
+        if self.use_gpu:        
+            self.h0 = self.h0.cuda()
+            self.c0 = self.c0.cuda()
+
+    def forward(self, in_seq_emb):
+        ''' forward model
+
+        Parameters
+        ----------
+        in_seq_emb : Tensor, shape (batch_size, max_len, input_dim)
+
+        Returns
+        -------
+        hidden_seq_emb : Tensor, shape (batch_size, max_len, 1 + bidirectional, hidden_dim)
+        '''
+        max_len = in_seq_emb.size(1)
+        hidden_seq_emb, (self.h0, self.c0) = self.model(
+            in_seq_emb, (self.h0, self.c0))
+        hidden_seq_emb = hidden_seq_emb.view(self.batch_size,
+                                             max_len,
+                                             1 + self.bidirectional,
+                                             self.hidden_dim)
+        return hidden_seq_emb
+
+
+class FullConnectedEncoder(EncoderBase):
+
+    def __init__(self, input_dim: int, hidden_dim: int, max_len: int, hidden_dim_list: List[int],
+                 batch_size: int, use_gpu: bool):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.max_len = max_len
+        self.hidden_dim_list = hidden_dim_list
+        self.use_gpu = use_gpu
+        in_out_dim_list = [input_dim * max_len] + list(hidden_dim_list) + [hidden_dim]
+        self.linear_list = nn.ModuleList(
+            [nn.Linear(in_out_dim_list[each_idx], in_out_dim_list[each_idx + 1])\
+             for each_idx in range(len(in_out_dim_list) - 1)])
+
+    def forward(self, in_seq_emb):
+        ''' forward model
+
+        Parameters
+        ----------
+        in_seq_emb : Tensor, shape (batch_size, max_len, input_dim)
+
+        Returns
+        -------
+        hidden_seq_emb : Tensor, shape (batch_size, max_len, 1 + bidirectional, hidden_dim)
+        '''
+        batch_size = in_seq_emb.size(0)
+        x = in_seq_emb.view(batch_size, -1)
+        for each_linear in self.linear_list:
+            x = F.relu(each_linear(x))
+        return x.view(batch_size, 1, -1)
+
+    def init_hidden(self):
+        pass

models/mhg_model/graph_grammar/nn/graph.py

@@ -0,0 +1,313 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES THE ORIGINAL SOURCE CODE (AND SOME ADAPTATIONS)
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE. 
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+""" Title """
+
+__author__ = "Hiroshi Kajino <[email protected]>"
+__copyright__ = "(c) Copyright IBM Corp. 2018"
+__version__ = "0.1"
+__date__ = "Jan 1 2018"
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from graph_grammar.graph_grammar.hrg import ProductionRuleCorpus
+from torch import nn
+from torch.autograd import Variable
+
+class MolecularProdRuleEmbedding(nn.Module):
+    
+    ''' molecular fingerprint layer
+    '''
+
+    def __init__(self, prod_rule_corpus, layer2layer_activation, layer2out_activation,
+                 out_dim=32, element_embed_dim=32,
+                 num_layers=3, padding_idx=None, use_gpu=False):
+        super().__init__()
+        if padding_idx is not None:
+            assert padding_idx == -1, 'padding_idx must be -1.'
+        self.prod_rule_corpus = prod_rule_corpus
+        self.layer2layer_activation = layer2layer_activation
+        self.layer2out_activation = layer2out_activation
+        self.out_dim = out_dim
+        self.element_embed_dim = element_embed_dim
+        self.num_layers = num_layers
+        self.padding_idx = padding_idx
+        self.use_gpu = use_gpu
+
+        self.layer2layer_list = []
+        self.layer2out_list = []
+
+        if self.use_gpu:
+            self.atom_embed = torch.randn(self.prod_rule_corpus.num_edge_symbol,
+                                          self.element_embed_dim, requires_grad=True).cuda()
+            self.bond_embed = torch.randn(self.prod_rule_corpus.num_node_symbol,
+                                          self.element_embed_dim, requires_grad=True).cuda()
+            self.ext_id_embed = torch.randn(self.prod_rule_corpus.num_ext_id,
+                                            self.element_embed_dim, requires_grad=True).cuda()
+            for _ in range(num_layers):
+                self.layer2layer_list.append(nn.Linear(self.element_embed_dim, self.element_embed_dim).cuda())
+                self.layer2out_list.append(nn.Linear(self.element_embed_dim, self.out_dim).cuda())
+        else:
+            self.atom_embed = torch.randn(self.prod_rule_corpus.num_edge_symbol,
+                                          self.element_embed_dim, requires_grad=True)
+            self.bond_embed = torch.randn(self.prod_rule_corpus.num_node_symbol,
+                                          self.element_embed_dim, requires_grad=True)
+            self.ext_id_embed = torch.randn(self.prod_rule_corpus.num_ext_id,
+                                            self.element_embed_dim, requires_grad=True)
+            for _ in range(num_layers):
+                self.layer2layer_list.append(nn.Linear(self.element_embed_dim, self.element_embed_dim))
+                self.layer2out_list.append(nn.Linear(self.element_embed_dim, self.out_dim))
+
+
+    def forward(self, prod_rule_idx_seq):
+        ''' forward model for mini-batch
+
+        Parameters
+        ----------
+        prod_rule_idx_seq : (batch_size, length)
+
+        Returns
+        -------
+        Variable, shape (batch_size, length, out_dim)
+        '''
+        batch_size, length = prod_rule_idx_seq.shape
+        if self.use_gpu:
+            out = Variable(torch.zeros((batch_size, length, self.out_dim))).cuda()
+        else:
+            out = Variable(torch.zeros((batch_size, length, self.out_dim)))
+        for each_batch_idx in range(batch_size):
+            for each_idx in range(length):
+                if int(prod_rule_idx_seq[each_batch_idx, each_idx]) == len(self.prod_rule_corpus.prod_rule_list):
+                    continue
+                else:
+                    each_prod_rule = self.prod_rule_corpus.prod_rule_list[int(prod_rule_idx_seq[each_batch_idx, each_idx])]
+                    layer_wise_embed_dict = {each_edge: self.atom_embed[
+                        each_prod_rule.rhs.edge_attr(each_edge)['symbol_idx']]
+                                             for each_edge in each_prod_rule.rhs.edges}
+                    layer_wise_embed_dict.update({each_node: self.bond_embed[
+                        each_prod_rule.rhs.node_attr(each_node)['symbol_idx']]
+                                                  for each_node in each_prod_rule.rhs.nodes})
+                    for each_node in each_prod_rule.rhs.nodes:
+                        if 'ext_id' in each_prod_rule.rhs.node_attr(each_node):
+                            layer_wise_embed_dict[each_node] \
+                                = layer_wise_embed_dict[each_node] \
+                                + self.ext_id_embed[each_prod_rule.rhs.node_attr(each_node)['ext_id']]
+
+                    for each_layer in range(self.num_layers):
+                        next_layer_embed_dict = {}
+                        for each_edge in each_prod_rule.rhs.edges:
+                            v = layer_wise_embed_dict[each_edge]
+                            for each_node in each_prod_rule.rhs.nodes_in_edge(each_edge):
+                                v = v + layer_wise_embed_dict[each_node]
+                            next_layer_embed_dict[each_edge] = self.layer2layer_activation(self.layer2layer_list[each_layer](v))
+                            out[each_batch_idx, each_idx, :] \
+                                = out[each_batch_idx, each_idx, :] + self.layer2out_activation(self.layer2out_list[each_layer](v))
+                        for each_node in each_prod_rule.rhs.nodes:
+                            v = layer_wise_embed_dict[each_node]
+                            for each_edge in each_prod_rule.rhs.adj_edges(each_node):
+                                v = v + layer_wise_embed_dict[each_edge]
+                            next_layer_embed_dict[each_node] = self.layer2layer_activation(self.layer2layer_list[each_layer](v))
+                            out[each_batch_idx, each_idx, :]\
+                                = out[each_batch_idx, each_idx, :] + self.layer2out_activation(self.layer2out_list[each_layer](v))
+                        layer_wise_embed_dict = next_layer_embed_dict
+                        
+        return out
+
+
+class MolecularProdRuleEmbeddingLastLayer(nn.Module):
+    
+    ''' molecular fingerprint layer
+    '''
+
+    def __init__(self, prod_rule_corpus, layer2layer_activation, layer2out_activation,
+                 out_dim=32, element_embed_dim=32,
+                 num_layers=3, padding_idx=None, use_gpu=False):
+        super().__init__()
+        if padding_idx is not None:
+            assert padding_idx == -1, 'padding_idx must be -1.'
+        self.prod_rule_corpus = prod_rule_corpus
+        self.layer2layer_activation = layer2layer_activation
+        self.layer2out_activation = layer2out_activation
+        self.out_dim = out_dim
+        self.element_embed_dim = element_embed_dim
+        self.num_layers = num_layers
+        self.padding_idx = padding_idx
+        self.use_gpu = use_gpu
+
+        self.layer2layer_list = []
+        self.layer2out_list = []
+
+        if self.use_gpu:
+            self.atom_embed = nn.Embedding(self.prod_rule_corpus.num_edge_symbol, self.element_embed_dim).cuda()
+            self.bond_embed = nn.Embedding(self.prod_rule_corpus.num_node_symbol, self.element_embed_dim).cuda()
+            for _ in range(num_layers+1):
+                self.layer2layer_list.append(nn.Linear(self.element_embed_dim, self.element_embed_dim).cuda())
+                self.layer2out_list.append(nn.Linear(self.element_embed_dim, self.out_dim).cuda())
+        else:
+            self.atom_embed = nn.Embedding(self.prod_rule_corpus.num_edge_symbol, self.element_embed_dim)
+            self.bond_embed = nn.Embedding(self.prod_rule_corpus.num_node_symbol, self.element_embed_dim)
+            for _ in range(num_layers+1):
+                self.layer2layer_list.append(nn.Linear(self.element_embed_dim, self.element_embed_dim))
+                self.layer2out_list.append(nn.Linear(self.element_embed_dim, self.out_dim))
+
+
+    def forward(self, prod_rule_idx_seq):
+        ''' forward model for mini-batch
+
+        Parameters
+        ----------
+        prod_rule_idx_seq : (batch_size, length)
+
+        Returns
+        -------
+        Variable, shape (batch_size, length, out_dim)
+        '''
+        batch_size, length = prod_rule_idx_seq.shape
+        if self.use_gpu:
+            out = Variable(torch.zeros((batch_size, length, self.out_dim))).cuda()
+        else:
+            out = Variable(torch.zeros((batch_size, length, self.out_dim)))
+        for each_batch_idx in range(batch_size):
+            for each_idx in range(length):
+                if int(prod_rule_idx_seq[each_batch_idx, each_idx]) == len(self.prod_rule_corpus.prod_rule_list):
+                    continue
+                else:
+                    each_prod_rule = self.prod_rule_corpus.prod_rule_list[int(prod_rule_idx_seq[each_batch_idx, each_idx])]
+
+                    if self.use_gpu:
+                        layer_wise_embed_dict = {each_edge: self.atom_embed(
+                            Variable(torch.LongTensor(
+                                [each_prod_rule.rhs.edge_attr(each_edge)['symbol_idx']]
+                            ), requires_grad=False).cuda())
+                                                 for each_edge in each_prod_rule.rhs.edges}
+                        layer_wise_embed_dict.update({each_node: self.bond_embed(
+                                                         Variable(
+                                                             torch.LongTensor([
+                                                                     each_prod_rule.rhs.node_attr(each_node)['symbol_idx']]),
+                                                             requires_grad=False).cuda()
+                                                     ) for each_node in each_prod_rule.rhs.nodes})
+                    else:
+                        layer_wise_embed_dict = {each_edge: self.atom_embed(
+                            Variable(torch.LongTensor(
+                                [each_prod_rule.rhs.edge_attr(each_edge)['symbol_idx']]
+                            ), requires_grad=False))
+                                                 for each_edge in each_prod_rule.rhs.edges}
+                        layer_wise_embed_dict.update({each_node: self.bond_embed(
+                                                         Variable(
+                                                             torch.LongTensor([
+                                                                     each_prod_rule.rhs.node_attr(each_node)['symbol_idx']]), 
+                                                             requires_grad=False)
+                                                     ) for each_node in each_prod_rule.rhs.nodes})
+
+                    for each_layer in range(self.num_layers):
+                        next_layer_embed_dict = {}
+                        for each_edge in each_prod_rule.rhs.edges:
+                            v = layer_wise_embed_dict[each_edge]
+                            for each_node in each_prod_rule.rhs.nodes_in_edge(each_edge):
+                                v += layer_wise_embed_dict[each_node]
+                            next_layer_embed_dict[each_edge] = self.layer2layer_activation(self.layer2layer_list[each_layer](v))
+                        for each_node in each_prod_rule.rhs.nodes:
+                            v = layer_wise_embed_dict[each_node]
+                            for each_edge in each_prod_rule.rhs.adj_edges(each_node):
+                                v += layer_wise_embed_dict[each_edge]
+                            next_layer_embed_dict[each_node] = self.layer2layer_activation(self.layer2layer_list[each_layer](v))
+                        layer_wise_embed_dict = next_layer_embed_dict
+                    for each_edge in each_prod_rule.rhs.edges:
+                        out[each_batch_idx, each_idx, :] = self.layer2out_activation(self.layer2out_list[self.num_layers](v))
+                    for each_edge in each_prod_rule.rhs.edges:
+                        out[each_batch_idx, each_idx, :] = self.layer2out_activation(self.layer2out_list[self.num_layers](v))
+                        
+        return out
+
+
+class MolecularProdRuleEmbeddingUsingFeatures(nn.Module):
+    
+    ''' molecular fingerprint layer
+    '''
+
+    def __init__(self, prod_rule_corpus, layer2layer_activation, layer2out_activation,
+                 out_dim=32, num_layers=3, padding_idx=None, use_gpu=False):
+        super().__init__()
+        if padding_idx is not None:
+            assert padding_idx == -1, 'padding_idx must be -1.'
+        self.feature_dict, self.feature_dim = prod_rule_corpus.construct_feature_vectors()
+        self.prod_rule_corpus = prod_rule_corpus
+        self.layer2layer_activation = layer2layer_activation
+        self.layer2out_activation = layer2out_activation
+        self.out_dim = out_dim
+        self.num_layers = num_layers
+        self.padding_idx = padding_idx
+        self.use_gpu = use_gpu
+
+        self.layer2layer_list = []
+        self.layer2out_list = []
+
+        if self.use_gpu:
+            for each_key in self.feature_dict:
+                self.feature_dict[each_key] = self.feature_dict[each_key].to_dense().cuda()
+            for _ in range(num_layers):
+                self.layer2layer_list.append(nn.Linear(self.feature_dim, self.feature_dim).cuda())
+                self.layer2out_list.append(nn.Linear(self.feature_dim, self.out_dim).cuda())
+        else:
+            for _ in range(num_layers):
+                self.layer2layer_list.append(nn.Linear(self.feature_dim, self.feature_dim))
+                self.layer2out_list.append(nn.Linear(self.feature_dim, self.out_dim))
+
+
+    def forward(self, prod_rule_idx_seq):
+        ''' forward model for mini-batch
+
+        Parameters
+        ----------
+        prod_rule_idx_seq : (batch_size, length)
+
+        Returns
+        -------
+        Variable, shape (batch_size, length, out_dim)
+        '''
+        batch_size, length = prod_rule_idx_seq.shape
+        if self.use_gpu:
+            out = Variable(torch.zeros((batch_size, length, self.out_dim))).cuda()
+        else:
+            out = Variable(torch.zeros((batch_size, length, self.out_dim)))
+        for each_batch_idx in range(batch_size):
+            for each_idx in range(length):
+                if int(prod_rule_idx_seq[each_batch_idx, each_idx]) == len(self.prod_rule_corpus.prod_rule_list):
+                    continue
+                else:
+                    each_prod_rule = self.prod_rule_corpus.prod_rule_list[int(prod_rule_idx_seq[each_batch_idx, each_idx])]
+                    edge_list = sorted(list(each_prod_rule.rhs.edges))
+                    node_list = sorted(list(each_prod_rule.rhs.nodes))
+                    adj_mat = torch.FloatTensor(each_prod_rule.rhs_adj_mat(edge_list + node_list).todense() + np.identity(len(edge_list)+len(node_list)))
+                    if self.use_gpu:
+                        adj_mat = adj_mat.cuda()
+                    layer_wise_embed = [
+                        self.feature_dict[each_prod_rule.rhs.edge_attr(each_edge)['symbol']]
+                        for each_edge in edge_list]\
+                            + [self.feature_dict[each_prod_rule.rhs.node_attr(each_node)['symbol']]
+                               for each_node in node_list]
+                    for each_node in each_prod_rule.ext_node.values():
+                        layer_wise_embed[each_prod_rule.rhs.num_edges + node_list.index(each_node)] \
+                                = layer_wise_embed[each_prod_rule.rhs.num_edges + node_list.index(each_node)] \
+                                + self.feature_dict[('ext_id', each_prod_rule.rhs.node_attr(each_node)['ext_id'])]
+                    layer_wise_embed = torch.stack(layer_wise_embed)
+
+                    for each_layer in range(self.num_layers):
+                        message = adj_mat @ layer_wise_embed
+                        next_layer_embed = self.layer2layer_activation(self.layer2layer_list[each_layer](message))
+                        out[each_batch_idx, each_idx, :] \
+                                = out[each_batch_idx, each_idx, :] \
+                                + self.layer2out_activation(self.layer2out_list[each_layer](message)).sum(dim=0)
+                        layer_wise_embed = next_layer_embed
+        return out

models/mhg_model/load.py

@@ -0,0 +1,103 @@
+# -*- coding:utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+import os
+import pickle
+import sys
+
+from rdkit import Chem
+import torch
+from torch_geometric.utils.smiles import from_smiles
+
+from typing import Any, Dict, List, Optional, Union
+from typing_extensions import Self
+
+from .graph_grammar.io.smi import hg_to_mol
+from .models.mhgvae import GrammarGINVAE
+
+from huggingface_hub import hf_hub_download
+
+
+class PretrainedModelWrapper:
+    model: GrammarGINVAE
+
+    def __init__(self, model_dict: Dict[str, Any]) -> None:
+        json_params = model_dict['gnn_params']
+        encoder_params = json_params['encoder_params']
+        encoder_params['node_feature_size'] = model_dict['num_features']
+        encoder_params['edge_feature_size'] = model_dict['num_edge_features']
+        self.model = GrammarGINVAE(model_dict['hrg'], rank=-1, encoder_params=encoder_params,
+                                   decoder_params=json_params['decoder_params'],
+                                   prod_rule_embed_params=json_params["prod_rule_embed_params"],
+                                   batch_size=512, max_len=model_dict['max_length'])
+        self.model.load_state_dict(model_dict['model_state_dict'])
+
+        self.model.eval()
+
+    def to(self, device: Union[str, int, torch.device]) -> Self:
+        dev_type = type(device)
+        if dev_type != torch.device:
+            if dev_type == str or torch.cuda.is_available():
+                device = torch.device(device)
+            else:
+                device = torch.device("mps", device)
+
+        self.model = self.model.to(device)
+        return self
+
+    def encode(self, data: List[str]) -> List[torch.tensor]:
+        # Need to encode them into a graph nn
+        output = []
+        for d in data:
+            params = next(self.model.parameters())
+            g = from_smiles(d)
+            if (g.cpu() and params != 'cpu') or (not g.cpu() and params == 'cpu'):
+                g.to(params.device)
+            ltvec = self.model.graph_embed(g.x, g.edge_index, g.edge_attr, g.batch)
+            output.append(ltvec[0])
+        return output
+
+    def decode(self, data: List[torch.tensor]) -> List[str]:
+        output = []
+        for d in data:
+            mu, logvar = self.model.get_mean_var(d.unsqueeze(0))
+            z = self.model.reparameterize(mu, logvar)
+            flags, _, hgs = self.model.decode(z)
+            if flags[0]:
+                reconstructed_mol, _ = hg_to_mol(hgs[0], True)
+                output.append(Chem.MolToSmiles(reconstructed_mol))
+            else:
+                output.append(None)
+        return output
+
+
+def load(model_name: str = "mhg_model/pickles/mhggnn_pretrained_model_0724_2023.pickle") -> Optional[
+    PretrainedModelWrapper]:
+
+    repo_id = "ibm/materials.mhg-ged"
+    filename = "pytorch_model.bin" #"mhggnn_pretrained_model_0724_2023.pickle"
+    file_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    with open(file_path, "rb") as f:
+        model_dict = torch.load(f, weights_only=False)
+        return PretrainedModelWrapper(model_dict)
+
+
+    """try:
+        if os.path.isfile(model_name):
+            with open(model_name, "rb") as f:
+                model_dict = pickle.load(f)
+                print("MHG Model Loaded")
+                return PretrainedModelWrapper(model_dict)
+
+    except:
+
+        for p in sys.path:
+            file = p + "/" + model_name
+            if os.path.isfile(file):
+                with open(file, "rb") as f:
+                    model_dict = pickle.load(f)
+                    return PretrainedModelWrapper(model_dict)"""
+    return None

models/mhg_model/mhg_gnn.egg-info/PKG-INFO

@@ -0,0 +1,102 @@
+Metadata-Version: 2.1
+Name: mhg-gnn
+Version: 0.0
+Summary: Package for mhg-gnn
+Author: team
+License: TBD
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Description-Content-Type: text/markdown
+Requires-Dist: networkx>=2.8
+Requires-Dist: numpy<2.0.0,>=1.23.5
+Requires-Dist: pandas>=1.5.3
+Requires-Dist: rdkit-pypi<2023.9.6,>=2022.9.4
+Requires-Dist: torch>=2.0.0
+Requires-Dist: torchinfo>=1.8.0
+Requires-Dist: torch-geometric>=2.3.1
+
+# mhg-gnn
+
+This repository provides PyTorch source code assosiated with our publication, "MHG-GNN: Combination of Molecular Hypergraph Grammar with Graph Neural Network"
+
+**Paper:** [Arxiv Link](https://arxiv.org/pdf/2309.16374)
+
+For more information contact: [email protected]
+
+![mhg-gnn](images/mhg_example1.png)
+
+## Introduction
+
+We present MHG-GNN, an autoencoder architecture
+that has an encoder based on GNN and a decoder based on a sequential model with MHG.
+Since the encoder is a GNN variant, MHG-GNN can accept any molecule as input, and  
+demonstrate high predictive performance on molecular graph data.
+In addition, the decoder inherits the theoretical guarantee of MHG on always generating a structurally valid molecule as output.
+
+## Table of Contents
+
+1. [Getting Started](#getting-started)
+    1. [Pretrained Models and Training Logs](#pretrained-models-and-training-logs)
+    2. [Replicating Conda Environment](#replicating-conda-environment)
+2. [Feature Extraction](#feature-extraction)
+
+## Getting Started
+
+**This code and environment have been tested on Intel E5-2667 CPUs at 3.30GHz and NVIDIA A100 Tensor Core GPUs.**
+
+### Pretrained Models and Training Logs
+
+We provide checkpoints of the MHG-GNN model pre-trained on a dataset of ~1.34M molecules curated from PubChem. (later) For model weights: [HuggingFace Link]()
+
+Add the MHG-GNN `pre-trained weights.pt` to the `models/` directory according to your needs. 
+
+### Replacicating Conda Environment
+
+Follow these steps to replicate our Conda environment and install the necessary libraries:
+
+```
+conda create --name mhg-gnn-env python=3.8.18
+conda activate mhg-gnn-env
+```
+
+#### Install Packages with Conda
+
+```
+conda install -c conda-forge networkx=2.8
+conda install numpy=1.23.5
+# conda install -c conda-forge rdkit=2022.9.4
+conda install pytorch=2.0.0 torchvision torchaudio -c pytorch
+conda install -c conda-forge torchinfo=1.8.0
+conda install pyg -c pyg
+```
+
+#### Install Packages with pip
+```
+pip install rdkit torch-nl==0.3 torch-scatter torch-sparse
+```
+
+## Feature Extraction
+
+The example notebook [mhg-gnn_encoder_decoder_example.ipynb](notebooks/mhg-gnn_encoder_decoder_example.ipynb) contains code to load checkpoint files and use the pre-trained model for encoder and decoder tasks.
+
+To load mhg-gnn, you can simply use:
+
+```python
+import torch
+import load
+
+model = load.load()
+```
+
+To encode SMILES into embeddings, you can use:
+
+```python
+with torch.no_grad():
+    repr = model.encode(["CCO", "O=C=O", "OC(=O)c1ccccc1C(=O)O"])
+```
+
+For decoder, you can use the function, so you can return from embeddings to SMILES strings:
+
+```python
+orig = model.decode(repr)
+```

models/mhg_model/mhg_gnn.egg-info/SOURCES.txt

@@ -0,0 +1,46 @@
+README.md
+setup.cfg
+setup.py
+./graph_grammar/__init__.py
+./graph_grammar/hypergraph.py
+./graph_grammar/algo/__init__.py
+./graph_grammar/algo/tree_decomposition.py
+./graph_grammar/graph_grammar/__init__.py
+./graph_grammar/graph_grammar/base.py
+./graph_grammar/graph_grammar/corpus.py
+./graph_grammar/graph_grammar/hrg.py
+./graph_grammar/graph_grammar/symbols.py
+./graph_grammar/graph_grammar/utils.py
+./graph_grammar/io/__init__.py
+./graph_grammar/io/smi.py
+./graph_grammar/nn/__init__.py
+./graph_grammar/nn/dataset.py
+./graph_grammar/nn/decoder.py
+./graph_grammar/nn/encoder.py
+./graph_grammar/nn/graph.py
+./models/__init__.py
+./models/mhgvae.py
+graph_grammar/__init__.py
+graph_grammar/hypergraph.py
+graph_grammar/algo/__init__.py
+graph_grammar/algo/tree_decomposition.py
+graph_grammar/graph_grammar/__init__.py
+graph_grammar/graph_grammar/base.py
+graph_grammar/graph_grammar/corpus.py
+graph_grammar/graph_grammar/hrg.py
+graph_grammar/graph_grammar/symbols.py
+graph_grammar/graph_grammar/utils.py
+graph_grammar/io/__init__.py
+graph_grammar/io/smi.py
+graph_grammar/nn/__init__.py
+graph_grammar/nn/dataset.py
+graph_grammar/nn/decoder.py
+graph_grammar/nn/encoder.py
+graph_grammar/nn/graph.py
+mhg_gnn.egg-info/PKG-INFO
+mhg_gnn.egg-info/SOURCES.txt
+mhg_gnn.egg-info/dependency_links.txt
+mhg_gnn.egg-info/requires.txt
+mhg_gnn.egg-info/top_level.txt
+models/__init__.py
+models/mhgvae.py

models/mhg_model/mhg_gnn.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

models/mhg_model/mhg_gnn.egg-info/requires.txt

@@ -0,0 +1,7 @@
+networkx>=2.8
+numpy<2.0.0,>=1.23.5
+pandas>=1.5.3
+rdkit-pypi<2023.9.6,>=2022.9.4
+torch>=2.0.0
+torchinfo>=1.8.0
+torch-geometric>=2.3.1

models/mhg_model/mhg_gnn.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+graph_grammar
+models

models/mhg_model/models/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding:utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#

models/mhg_model/models/mhgvae.py

@@ -0,0 +1,956 @@
+# -*- coding:utf-8 -*-
+# Rhizome
+# Version beta 0.0, August 2023
+# Property of IBM Research, Accelerated Discovery
+#
+
+"""
+PLEASE NOTE THIS IMPLEMENTATION INCLUDES ADAPTED SOURCE CODE
+OF THE MHG IMPLEMENTATION OF HIROSHI KAJINO AT IBM TRL ALREADY PUBLICLY AVAILABLE, 
+E.G., GRUEncoder/GRUDecoder, GrammarSeq2SeqVAE AND EVEN SOME METHODS OF GrammarGINVAE.
+THIS MIGHT INFLUENCE THE DECISION OF THE FINAL LICENSE SO CAREFUL CHECK NEEDS BE DONE. 
+"""
+
+import numpy as np
+import logging
+
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.modules.loss import _Loss
+
+from torch_geometric.nn import MessagePassing
+from torch_geometric.nn import global_add_pool
+
+
+from ..graph_grammar.graph_grammar.symbols import NTSymbol
+from ..graph_grammar.nn.encoder import EncoderBase
+from ..graph_grammar.nn.decoder import DecoderBase
+
+def get_atom_edge_feature_dims():
+    from torch_geometric.utils.smiles import x_map, e_map
+    func = lambda x: len(x[1])
+    return list(map(func, x_map.items())), list(map(func, e_map.items()))
+
+
+class FeatureEmbedding(nn.Module):
+    def __init__(self, input_dims, embedded_dim):
+        super().__init__()
+        self.embedding_list = nn.ModuleList()
+        for dim in input_dims:
+            embedding = nn.Embedding(dim, embedded_dim)
+            self.embedding_list.append(embedding)
+    
+    def forward(self, x):
+        output = 0
+        for i in range(x.shape[1]):
+            input = x[:, i].to(torch.int)
+            device = next(self.parameters()).device
+            if device != input.device:
+                input = input.to(device)
+            emb = self.embedding_list[i](input)
+            output += emb
+        return output
+    
+
+class GRUEncoder(EncoderBase):
+
+    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int,
+                 bidirectional: bool, dropout: float, batch_size: int, rank: int=-1,
+                 no_dropout: bool=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.dropout = dropout
+        self.batch_size = batch_size
+        self.rank = rank
+        self.model = nn.GRU(input_size=self.input_dim,
+                            hidden_size=self.hidden_dim,
+                            num_layers=self.num_layers,
+                            batch_first=True,
+                            bidirectional=self.bidirectional,
+                            dropout=self.dropout if not no_dropout else 0)
+        if self.rank >= 0:
+            if torch.cuda.is_available():
+                self.model = self.model.to(rank)
+            else:
+                # support mac mps
+                self.model = self.model.to(torch.device("mps", rank))
+        self.init_hidden(self.batch_size)
+        
+    def init_hidden(self, bsize):
+        self.h0 = torch.zeros(((self.bidirectional + 1) * self.num_layers,
+                               min(self.batch_size, bsize),
+                               self.hidden_dim),
+                              requires_grad=False)
+        if self.rank >= 0:        
+            if torch.cuda.is_available():
+                self.h0 = self.h0.to(self.rank)
+            else:
+                # support mac mps
+                self.h0 = self.h0.to(torch.device("mps", self.rank))
+
+    def to(self, device):
+        newself = super().to(device)
+        newself.model = newself.model.to(device)
+        newself.h0 = newself.h0.to(device)
+        newself.rank = next(newself.parameters()).get_device()
+        return newself
+
+    def forward(self, in_seq_emb):
+        ''' forward model
+
+        Parameters
+        ----------
+        in_seq_emb : Tensor, shape (batch_size, max_len, input_dim)
+
+        Returns
+        -------
+        hidden_seq_emb : Tensor, shape (batch_size, max_len, 1 + bidirectional, hidden_dim)
+        '''
+        # Kishi: I think original MHG had this init_hidden()
+        self.init_hidden(in_seq_emb.size(0))
+        max_len = in_seq_emb.size(1)
+        hidden_seq_emb, self.h0 = self.model(
+            in_seq_emb, self.h0)
+        # As shown as returns, convert hidden_seq_emb: (batch_size, seq_len, (1 or 2) * hidden_size) -->
+        # (batch_size, seq_len, 1 or 2, hidden_size)
+        # In the original input the original GRU/LSTM with bidirectional encoding 
+        # has contactinated tensors 
+        # (first half for forward RNN, latter half for backward RNN)
+        # so convert them in a more friendly format packed for each RNN
+        hidden_seq_emb = hidden_seq_emb.view(-1,
+                                             max_len,
+                                             1 + self.bidirectional,
+                                             self.hidden_dim)
+        return hidden_seq_emb
+
+
+class GRUDecoder(DecoderBase):
+
+    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int,
+                 dropout: float, batch_size: int, rank: int=-1,
+                 no_dropout: bool=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.batch_size = batch_size
+        self.rank = rank
+        self.model = nn.GRU(input_size=self.input_dim,
+                            hidden_size=self.hidden_dim,
+                            num_layers=self.num_layers,
+                            batch_first=True,
+                            bidirectional=False,
+                            dropout=self.dropout if not no_dropout else 0
+        )
+        if self.rank >= 0:
+            if torch.cuda.is_available():
+                self.model = self.model.to(self.rank)
+            else:
+                # support mac mps
+                self.model = self.model.to(torch.device("mps", self.rank))
+        self.init_hidden(self.batch_size)
+
+    def init_hidden(self, bsize):
+        self.hidden_dict['h'] = torch.zeros((self.num_layers,
+                                             min(self.batch_size, bsize),
+                                             self.hidden_dim),
+                                            requires_grad=False)
+        if self.rank >= 0:
+            if torch.cuda.is_available():
+                self.hidden_dict['h'] = self.hidden_dict['h'].to(self.rank)
+            else:
+                self.hidden_dict['h'] = self.hidden_dict['h'].to(torch.device("mps", self.rank))
+
+    def to(self, device):
+        newself = super().to(device)
+        newself.model = newself.model.to(device)
+        for k in self.hidden_dict.keys():
+            newself.hidden_dict[k] = newself.hidden_dict[k].to(device)
+        newself.rank = next(newself.parameters()).get_device()
+        return newself
+
+    def forward_one_step(self, tgt_emb_in):
+        ''' one-step forward model
+
+        Parameters
+        ----------
+        tgt_emb_in : Tensor, shape (batch_size, input_dim)
+
+        Returns
+        -------
+        Tensor, shape (batch_size, hidden_dim)
+        '''
+        bsize = tgt_emb_in.size(0)
+        tgt_emb_out, self.hidden_dict['h'] \
+            = self.model(tgt_emb_in.view(bsize, 1, -1),
+                         self.hidden_dict['h'])
+        return tgt_emb_out
+
+
+class NodeMLP(nn.Module):
+    def __init__(self, input_size, output_size, hidden_size):
+        super().__init__()
+        self.lin1 = nn.Linear(input_size, hidden_size)
+        self.nbat = nn.BatchNorm1d(hidden_size)
+        self.lin2 = nn.Linear(hidden_size, output_size)
+        
+    def forward(self, x):
+        x = self.lin1(x)
+        x = self.nbat(x)
+        x = x.relu()
+        x = self.lin2(x)
+        return x
+    
+    
+class GINLayer(MessagePassing):
+    def __init__(self, node_input_size, node_output_size, node_hidden_size, edge_input_size):
+        super().__init__()
+        self.node_mlp = NodeMLP(node_input_size, node_output_size, node_hidden_size)
+        self.edge_mlp = FeatureEmbedding(edge_input_size, node_output_size)
+        self.eps = nn.Parameter(torch.tensor([0.0]))
+        
+    def forward(self, x, edge_index, edge_attr):
+        msg = self.propagate(edge_index, x=x ,edge_attr=edge_attr)
+        x = (1.0 + self.eps) * x + msg
+        x = x.relu()
+        x = self.node_mlp(x)
+        return x
+    
+    def message(self, x_j, edge_attr):
+        edge_attr = self.edge_mlp(edge_attr)
+        x_j = x_j + edge_attr
+        x_j = x_j.relu()
+        return x_j
+    
+    def update(self, aggr_out):
+        return aggr_out
+
+#TODO implement the case where features of atoms and edges are considered
+# Check GraphMVP and ogb (open graph benchmark) to realize this
+class GIN(torch.nn.Module):
+    def __init__(self, node_feature_size, edge_feature_size, hidden_channels=64, 
+                 proximity_size=3, dropout=0.1):
+        super().__init__()
+        #print("(num node features, num edge features)=", (node_feature_size, edge_feature_size))
+        hsize = hidden_channels * 2
+        atom_dim, edge_dim = get_atom_edge_feature_dims()
+        self.trans = FeatureEmbedding(atom_dim, hidden_channels)
+        ml = []
+        for _ in range(proximity_size):
+            ml.append(GINLayer(hidden_channels, hidden_channels, hsize, edge_dim))
+        self.mlist = nn.ModuleList(ml)
+        #It is possible to calculate relu with x.relu() where x is an output
+        #self.activations = nn.ModuleList(actl)
+        self.dropout = dropout
+        self.proximity_size = proximity_size
+        
+    def forward(self, x, edge_index, edge_attr, batch_size):
+        x = x.to(torch.float)
+        #print("before: edge_weight.shape=", edge_attr.shape)
+        edge_attr = edge_attr.to(torch.float)
+        #print("after: edge_weight.shape=", edge_attr.shape)
+        x = self.trans(x)
+        # TODO Check if this x is consistent with global_add_pool
+        hlist = [global_add_pool(x, batch_size)]
+        for id, m in enumerate(self.mlist):
+            x = m(x, edge_index=edge_index, edge_attr=edge_attr)
+            #print("Done with one layer") 
+            ###if id != self.proximity_size - 1:
+            x = x.relu()
+            x = F.dropout(x, p=self.dropout, training=self.training)
+            #h = global_mean_pool(x, batch_size)
+            h = global_add_pool(x, batch_size)
+            hlist.append(h)
+            #print("Done with one relu call: x.shape=", x.shape)
+        #print("calling golbal mean pool")
+        #print("calling dropout x.shape=", x.shape)
+        #print("x=", x)
+        #print("hlist[0].shape=", hlist[0].shape)
+        x = torch.cat(hlist, dim=1)
+        #print("x.shape=", x.shape)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+
+        return x
+
+
+# TODO copied from MHG implementation and adapted here. 
+class GrammarSeq2SeqVAE(nn.Module):
+
+    '''
+    Variational seq2seq with grammar.
+    TODO: rewrite this class using mixin
+    '''
+
+    def __init__(self, hrg, rank=-1, latent_dim=64, max_len=80,
+                 batch_size=64, padding_idx=-1, 
+                 encoder_params={'hidden_dim': 384, 'num_layers': 3, 'bidirectional': True,
+                                 'dropout': 0.1},
+                 decoder_params={'hidden_dim': 384, #'num_layers': 2,
+                                 'num_layers': 3,
+                                 'dropout': 0.1},
+                 prod_rule_embed_params={'out_dim': 128},
+                 no_dropout=False):
+
+        super().__init__()
+        # TODO USE GRU FOR ENCODING AND DECODING
+        self.hrg = hrg
+        self.rank = rank
+        self.prod_rule_corpus = hrg.prod_rule_corpus
+        self.prod_rule_embed_params = prod_rule_embed_params
+
+        self.vocab_size = hrg.num_prod_rule + 1
+        self.batch_size = batch_size
+        self.padding_idx = np.mod(padding_idx, self.vocab_size)
+        self.no_dropout = no_dropout
+
+        self.latent_dim = latent_dim
+        self.max_len = max_len
+        self.encoder_params = encoder_params
+        self.decoder_params = decoder_params
+        
+        # TODO Simple embedding is used. Check if a domain-dependent embedding works or not. 
+        embed_out_dim = self.prod_rule_embed_params['out_dim']        
+        #use MolecularProdRuleEmbedding later on
+        self.src_embedding = nn.Embedding(self.vocab_size, embed_out_dim,
+                                          padding_idx=self.padding_idx)
+        self.tgt_embedding = nn.Embedding(self.vocab_size, embed_out_dim,
+                                          padding_idx=self.padding_idx)
+
+        # USE a GRU-based encoder in MHG
+        self.encoder = GRUEncoder(input_dim=embed_out_dim, batch_size=self.batch_size,
+                                      rank=self.rank, no_dropout=self.no_dropout,
+                                  **self.encoder_params)
+
+        lin_dim = (self.encoder_params.get('bidirectional', False) + 1) * self.encoder_params['hidden_dim']
+        lin_out_dim = self.latent_dim
+        self.hidden2mean = nn.Linear(lin_dim, lin_out_dim, bias=False)
+        self.hidden2logvar = nn.Linear(lin_dim, lin_out_dim)
+   
+        # USE a GRU-based decoder in MHG
+        self.decoder = GRUDecoder(input_dim=embed_out_dim, batch_size=self.batch_size,
+                                  rank=self.rank, no_dropout=self.no_dropout, **self.decoder_params)
+        self.latent2tgt_emb = nn.Linear(self.latent_dim, embed_out_dim)
+        self.latent2hidden_dict = nn.ModuleDict()
+        dec_lin_out_dim = self.decoder_params['hidden_dim']
+        for each_hidden in self.decoder.hidden_dict.keys():
+            self.latent2hidden_dict[each_hidden] = nn.Linear(self.latent_dim, dec_lin_out_dim)
+            if self.rank >= 0:
+                if torch.cuda.is_available():
+                    self.latent2hidden_dict[each_hidden] = self.latent2hidden_dict[each_hidden].to(self.rank)
+                else:
+                    # support mac mps
+                    self.latent2hidden_dict[each_hidden] = self.latent2hidden_dict[each_hidden].to(torch.device("mps", self.rank))
+
+        self.dec2vocab = nn.Linear(dec_lin_out_dim, self.vocab_size)
+        self.encoder.init_hidden(self.batch_size)
+        self.decoder.init_hidden(self.batch_size)
+
+        # TODO Do we need this?
+        if hasattr(self.src_embedding, 'weight'):
+            self.src_embedding.weight.data.uniform_(-0.1, 0.1)
+        if hasattr(self.tgt_embedding, 'weight'):
+            self.tgt_embedding.weight.data.uniform_(-0.1, 0.1)
+            
+        self.encoder.init_hidden(self.batch_size)
+        self.decoder.init_hidden(self.batch_size)
+    
+    def to(self, device):
+        newself = super().to(device)
+        newself.src_embedding = newself.src_embedding.to(device)
+        newself.tgt_embedding = newself.tgt_embedding.to(device)
+        newself.encoder = newself.encoder.to(device)
+        newself.decoder = newself.decoder.to(device)
+        newself.dec2vocab = newself.dec2vocab.to(device)
+        newself.hidden2mean = newself.hidden2mean.to(device)
+        newself.hidden2logvar = newself.hidden2logvar.to(device)
+        newself.latent2tgt_emb = newself.latent2tgt_emb.to(device)
+        newself.latent2hidden_dict = newself.latent2hidden_dict.to(device)
+        return newself
+    
+    def forward(self, in_seq, out_seq):
+        ''' forward model
+
+        Parameters
+        ----------
+        in_seq : Variable, shape (batch_size, length)
+            each element corresponds to word index.
+            where the index should be less than `vocab_size`
+
+        Returns
+        -------
+        Variable, shape (batch_size, length, vocab_size)
+            logit of each word (applying softmax yields the probability)
+        '''
+        mu, logvar = self.encode(in_seq)
+        z = self.reparameterize(mu, logvar)
+        return self.decode(z, out_seq), mu, logvar
+
+    def encode(self, in_seq):
+        src_emb = self.src_embedding(in_seq)
+        src_h = self.encoder.forward(src_emb)
+        if self.encoder_params.get('bidirectional', False):
+            concat_src_h = torch.cat((src_h[:, -1, 0, :], src_h[:, 0, 1, :]), dim=1)
+            return self.hidden2mean(concat_src_h), self.hidden2logvar(concat_src_h)
+        else:
+            return self.hidden2mean(src_h[:, -1, :]), self.hidden2logvar(src_h[:, -1, :])
+
+    def reparameterize(self, mu, logvar, training=True):
+        if training:
+            std = logvar.mul(0.5).exp_()
+            device = next(self.parameters()).device
+            eps = Variable(std.data.new(std.size()).normal_())
+            if device != eps.get_device():
+                eps.to(device)
+            return eps.mul(std).add_(mu)
+        else:
+            return mu
+
+    #TODO Not tested. Need to implement this in case of molecular structure generation
+    def sample(self, sample_size=-1, deterministic=True, return_z=False):
+        self.eval()
+        self.init_hidden()
+        if sample_size == -1:
+            sample_size = self.batch_size
+
+        num_iter = int(np.ceil(sample_size / self.batch_size))
+        hg_list = []
+        z_list = []
+        for _ in range(num_iter):
+            z = Variable(torch.normal(
+                torch.zeros(self.batch_size, self.latent_dim),
+                torch.ones(self.batch_size * self.latent_dim))).cuda()
+            _, each_hg_list = self.decode(z, deterministic=deterministic)
+            z_list.append(z)
+            hg_list += each_hg_list
+        z = torch.cat(z_list)[:sample_size]
+        hg_list = hg_list[:sample_size]
+        if return_z:
+            return hg_list, z.cpu().detach().numpy()
+        else:
+            return hg_list
+
+    def decode(self, z=None, out_seq=None, deterministic=True):
+        if z is None:
+            z = Variable(torch.normal(
+                torch.zeros(self.batch_size, self.latent_dim),
+                torch.ones(self.batch_size * self.latent_dim)))
+        if self.rank >= 0:
+            z = z.to(next(self.parameters()).device)
+
+        hidden_dict_0 = {}
+        for each_hidden in self.latent2hidden_dict.keys():
+            hidden_dict_0[each_hidden] = self.latent2hidden_dict[each_hidden](z)
+        bsize = z.size(0)
+        self.decoder.init_hidden(bsize) 
+        self.decoder.feed_hidden(hidden_dict_0)
+
+        if out_seq is not None:
+            tgt_emb0 = self.latent2tgt_emb(z)
+            tgt_emb0 = tgt_emb0.view(tgt_emb0.shape[0], 1, tgt_emb0.shape[1])
+            out_seq_emb = self.tgt_embedding(out_seq)
+            tgt_emb = torch.cat((tgt_emb0, out_seq_emb), dim=1)[:, :-1, :]
+            tgt_emb_pred_list = []
+            for each_idx in range(self.max_len):
+                tgt_emb_pred = self.decoder.forward_one_step(tgt_emb[:, each_idx, :].view(bsize, 1, -1))
+                tgt_emb_pred_list.append(tgt_emb_pred)
+            vocab_logit = self.dec2vocab(torch.cat(tgt_emb_pred_list, dim=1))
+            return vocab_logit
+        else:
+            with torch.no_grad():
+                tgt_emb = self.latent2tgt_emb(z)
+                tgt_emb = tgt_emb.view(tgt_emb.shape[0], 1, tgt_emb.shape[1])
+                tgt_emb_pred_list = []
+                stack_list = []
+                hg_list = []
+                nt_symbol_list = []
+                nt_edge_list = []
+                gen_finish_list = []
+                for _ in range(bsize):
+                    stack_list.append([])
+                    hg_list.append(None)
+                    nt_symbol_list.append(NTSymbol(degree=0,
+                                                   is_aromatic=False,
+                                                   bond_symbol_list=[]))
+                    nt_edge_list.append(None)
+                    gen_finish_list.append(False)
+
+                for idx in range(self.max_len):
+                    tgt_emb_pred = self.decoder.forward_one_step(tgt_emb)
+                    tgt_emb_pred_list.append(tgt_emb_pred)
+                    vocab_logit = self.dec2vocab(tgt_emb_pred)
+                    for each_batch_idx in range(bsize):
+                        if not gen_finish_list[each_batch_idx]: # if generation has not finished
+                            # get production rule greedily
+                            prod_rule = self.hrg.prod_rule_corpus.sample(vocab_logit[each_batch_idx, :, :-1].squeeze().cpu().numpy(),
+                                                                         nt_symbol_list[each_batch_idx],
+                                                                         deterministic=deterministic)
+                            # convert production rule into an index
+                            tgt_id = self.hrg.prod_rule_list.index(prod_rule)
+                            # apply the production rule
+                            hg_list[each_batch_idx], nt_edges = prod_rule.applied_to(hg_list[each_batch_idx], nt_edge_list[each_batch_idx])
+                            # add non-terminals to the stack
+                            stack_list[each_batch_idx].extend(nt_edges[::-1])
+                            # if the stack size is 0, generation has finished!
+                            if len(stack_list[each_batch_idx]) == 0:
+                                gen_finish_list[each_batch_idx] = True
+                            else:
+                                nt_edge_list[each_batch_idx] = stack_list[each_batch_idx].pop()
+                                nt_symbol_list[each_batch_idx] = hg_list[each_batch_idx].edge_attr(nt_edge_list[each_batch_idx])['symbol']
+                        else:
+                            tgt_id = np.mod(self.padding_idx, self.vocab_size)
+                        indice_tensor = torch.LongTensor([tgt_id])
+                        device = next(self.parameters()).device
+                        if indice_tensor.device != device: 
+                            indice_tensor = indice_tensor.to(device)
+                        tgt_emb[each_batch_idx, :] = self.tgt_embedding(indice_tensor)
+                vocab_logit = self.dec2vocab(torch.cat(tgt_emb_pred_list, dim=1))
+                #for id, v in enumerate(gen_finish_list):
+                #if not v:
+                #    print("bacth id={} not finished generating a sequence: ".format(id))
+                return gen_finish_list, vocab_logit, hg_list
+
+
+# TODO A lot of duplicates with GrammarVAE. Clean up it if necessary
+class GrammarGINVAE(nn.Module):
+
+    '''
+    Variational autoencoder based on GIN and grammar
+    '''
+
+    def __init__(self, hrg, rank=-1, max_len=80,
+                 batch_size=64, padding_idx=-1, 
+                 encoder_params={'node_feature_size': 4, 'edge_feature_size': 3, 
+                                 'hidden_channels': 64, 'proximity_size': 3,
+                                 'dropout': 0.1},
+                 decoder_params={'hidden_dim': 384, 'num_layers': 3,
+                                 'dropout': 0.1},
+                 prod_rule_embed_params={'out_dim': 128},
+                 no_dropout=False):
+
+        super().__init__()
+        # TODO USE GRU FOR ENCODING AND DECODING
+        self.hrg = hrg
+        self.rank = rank
+        self.prod_rule_corpus = hrg.prod_rule_corpus
+        self.prod_rule_embed_params = prod_rule_embed_params
+
+        self.vocab_size = hrg.num_prod_rule + 1
+        self.batch_size = batch_size
+        self.padding_idx = np.mod(padding_idx, self.vocab_size)
+        self.no_dropout = no_dropout
+        self.max_len = max_len
+        self.encoder_params = encoder_params
+        self.decoder_params = decoder_params
+        
+        # TODO Simple embedding is used. Check if a domain-dependent embedding works or not. 
+        embed_out_dim = self.prod_rule_embed_params['out_dim']        
+        #use MolecularProdRuleEmbedding later on
+        self.tgt_embedding = nn.Embedding(self.vocab_size, embed_out_dim,
+                                          padding_idx=self.padding_idx)
+
+        self.encoder = GIN(**self.encoder_params)
+        self.latent_dim = self.encoder_params['hidden_channels']
+        self.proximity_size = self.encoder_params['proximity_size']
+        hidden_dim = self.decoder_params['hidden_dim']
+        self.hidden2mean = nn.Linear(self.latent_dim * (1 + self.proximity_size), self.latent_dim, bias=False)
+        self.hidden2logvar = nn.Linear(self.latent_dim * (1 + self.proximity_size), self.latent_dim)
+   
+        self.decoder = GRUDecoder(input_dim=embed_out_dim, batch_size=self.batch_size,
+                                  rank=self.rank, no_dropout=self.no_dropout, **self.decoder_params)
+        self.latent2tgt_emb = nn.Linear(self.latent_dim, embed_out_dim)
+        self.latent2hidden_dict = nn.ModuleDict()
+        for each_hidden in self.decoder.hidden_dict.keys():
+            self.latent2hidden_dict[each_hidden] = nn.Linear(self.latent_dim, hidden_dim)
+            if self.rank >= 0:
+                if torch.cuda.is_available():
+                    self.latent2hidden_dict[each_hidden] = self.latent2hidden_dict[each_hidden].to(self.rank)
+                else:
+                    # support mac mps
+                    self.latent2hidden_dict[each_hidden] = self.latent2hidden_dict[each_hidden].to(torch.device("mps", self.rank))
+
+        self.dec2vocab = nn.Linear(hidden_dim, self.vocab_size)
+        self.decoder.init_hidden(self.batch_size)
+
+        # TODO Do we need this?
+        if hasattr(self.tgt_embedding, 'weight'):
+            self.tgt_embedding.weight.data.uniform_(-0.1, 0.1)
+        self.decoder.init_hidden(self.batch_size)
+    
+    def to(self, device):
+        newself = super().to(device)
+        newself.encoder = newself.encoder.to(device)
+        newself.decoder = newself.decoder.to(device)
+        newself.rank = next(newself.encoder.parameters()).get_device()
+        return newself
+    
+    def forward(self, x, edge_index, edge_attr, batch_size, out_seq=None, sched_prob = None):
+        mu, logvar = self.encode(x, edge_index, edge_attr, batch_size)
+        z = self.reparameterize(mu, logvar)
+        return self.decode(z, out_seq, sched_prob=sched_prob), mu, logvar
+
+    #TODO Not tested. Need to implement this in case of molecular structure generation
+    def sample(self, sample_size=-1, deterministic=True, return_z=False):
+        self.eval()
+        self.init_hidden()
+        if sample_size == -1:
+            sample_size = self.batch_size
+
+        num_iter = int(np.ceil(sample_size / self.batch_size))
+        hg_list = []
+        z_list = []
+        for _ in range(num_iter):
+            z = Variable(torch.normal(
+                torch.zeros(self.batch_size, self.latent_dim),
+                torch.ones(self.batch_size * self.latent_dim))).cuda()
+            _, each_hg_list = self.decode(z, deterministic=deterministic)
+            z_list.append(z)
+            hg_list += each_hg_list
+        z = torch.cat(z_list)[:sample_size]
+        hg_list = hg_list[:sample_size]
+        if return_z:
+            return hg_list, z.cpu().detach().numpy()
+        else:
+            return hg_list
+
+    def decode(self, z=None, out_seq=None, deterministic=True, sched_prob=None):
+        if z is None:
+            z = Variable(torch.normal(
+                torch.zeros(self.batch_size, self.latent_dim),
+                torch.ones(self.batch_size * self.latent_dim)))
+        if self.rank >= 0:
+            z = z.to(next(self.parameters()).device)
+
+        hidden_dict_0 = {}
+        for each_hidden in self.latent2hidden_dict.keys():
+            hidden_dict_0[each_hidden] = self.latent2hidden_dict[each_hidden](z)
+        bsize = z.size(0)
+        self.decoder.init_hidden(bsize) 
+        self.decoder.feed_hidden(hidden_dict_0)
+
+        if out_seq is not None:
+            tgt_emb0 = self.latent2tgt_emb(z)
+            tgt_emb0 = tgt_emb0.view(tgt_emb0.shape[0], 1, tgt_emb0.shape[1])
+            out_seq_emb = self.tgt_embedding(out_seq)
+            tgt_emb = torch.cat((tgt_emb0, out_seq_emb), dim=1)[:, :-1, :]
+            tgt_emb_pred_list = []
+            tgt_emb_pred = None
+            for each_idx in range(self.max_len):
+                if tgt_emb_pred is None or sched_prob is None or torch.rand(1)[0] <= sched_prob:
+                    inp = tgt_emb[:, each_idx, :].view(bsize, 1, -1)
+                else:
+                    cur_logit = self.dec2vocab(tgt_emb_pred)
+                    yi = torch.argmax(cur_logit, dim=2)
+                    inp = self.tgt_embedding(yi)
+                tgt_emb_pred = self.decoder.forward_one_step(inp)
+                tgt_emb_pred_list.append(tgt_emb_pred)
+            vocab_logit = self.dec2vocab(torch.cat(tgt_emb_pred_list, dim=1))
+            return vocab_logit
+        else:
+            with torch.no_grad():
+                tgt_emb = self.latent2tgt_emb(z)
+                tgt_emb = tgt_emb.view(tgt_emb.shape[0], 1, tgt_emb.shape[1])
+                tgt_emb_pred_list = []
+                stack_list = []
+                hg_list = []
+                nt_symbol_list = []
+                nt_edge_list = []
+                gen_finish_list = []
+                for _ in range(bsize):
+                    stack_list.append([])
+                    hg_list.append(None)
+                    nt_symbol_list.append(NTSymbol(degree=0,
+                                                   is_aromatic=False,
+                                                   bond_symbol_list=[]))
+                    nt_edge_list.append(None)
+                    gen_finish_list.append(False)
+
+                for _ in range(self.max_len):
+                    tgt_emb_pred = self.decoder.forward_one_step(tgt_emb)
+                    tgt_emb_pred_list.append(tgt_emb_pred)
+                    vocab_logit = self.dec2vocab(tgt_emb_pred)
+                    for each_batch_idx in range(bsize):
+                        if not gen_finish_list[each_batch_idx]: # if generation has not finished
+                            # get production rule greedily
+                            prod_rule = self.hrg.prod_rule_corpus.sample(vocab_logit[each_batch_idx, :, :-1].squeeze().cpu().numpy(),
+                                                                         nt_symbol_list[each_batch_idx],
+                                                                         deterministic=deterministic)
+                            # convert production rule into an index
+                            tgt_id = self.hrg.prod_rule_list.index(prod_rule)
+                            # apply the production rule
+                            hg_list[each_batch_idx], nt_edges = prod_rule.applied_to(hg_list[each_batch_idx], nt_edge_list[each_batch_idx])
+                            # add non-terminals to the stack
+                            stack_list[each_batch_idx].extend(nt_edges[::-1])
+                            # if the stack size is 0, generation has finished!
+                            if len(stack_list[each_batch_idx]) == 0:
+                                gen_finish_list[each_batch_idx] = True
+                            else:
+                                nt_edge_list[each_batch_idx] = stack_list[each_batch_idx].pop()
+                                nt_symbol_list[each_batch_idx] = hg_list[each_batch_idx].edge_attr(nt_edge_list[each_batch_idx])['symbol']
+                        else:
+                            tgt_id = np.mod(self.padding_idx, self.vocab_size)
+                        indice_tensor = torch.LongTensor([tgt_id])
+                        if self.rank >= 0:
+                            indice_tensor = indice_tensor.to(next(self.parameters()).device)
+                        tgt_emb[each_batch_idx, :] = self.tgt_embedding(indice_tensor)
+                vocab_logit = self.dec2vocab(torch.cat(tgt_emb_pred_list, dim=1))
+                return gen_finish_list, vocab_logit, hg_list
+
+    #TODO Not tested. Need to implement this in case of molecular structure generation
+    def conditional_distribution(self, z, tgt_id_list):
+        self.eval()
+        self.init_hidden()
+        z = z.cuda()
+
+        hidden_dict_0 = {}
+        for each_hidden in self.latent2hidden_dict.keys():
+            hidden_dict_0[each_hidden] = self.latent2hidden_dict[each_hidden](z)
+        self.decoder.feed_hidden(hidden_dict_0)
+
+        with torch.no_grad():
+            tgt_emb = self.latent2tgt_emb(z)
+            tgt_emb = tgt_emb.view(tgt_emb.shape[0], 1, tgt_emb.shape[1])
+            nt_symbol_list = []
+            stack_list = []
+            hg_list = []
+            nt_edge_list = []
+            gen_finish_list = []
+            for _ in range(self.batch_size):
+                nt_symbol_list.append(NTSymbol(degree=0,
+                                               is_aromatic=False,
+                                               bond_symbol_list=[]))
+                stack_list.append([])
+                hg_list.append(None)
+                nt_edge_list.append(None)
+                gen_finish_list.append(False)
+
+            for each_position in range(len(tgt_id_list[0])):
+                tgt_emb_pred = self.decoder.forward_one_step(tgt_emb)
+                for each_batch_idx in range(self.batch_size):
+                    if not gen_finish_list[each_batch_idx]: # if generation has not finished
+                        # use the prespecified target ids
+                        tgt_id = tgt_id_list[each_batch_idx][each_position]
+                        prod_rule = self.hrg.prod_rule_list[tgt_id]
+                        # apply the production rule
+                        hg_list[each_batch_idx], nt_edges = prod_rule.applied_to(hg_list[each_batch_idx], nt_edge_list[each_batch_idx])
+                        # add non-terminals to the stack
+                        stack_list[each_batch_idx].extend(nt_edges[::-1])
+                        # if the stack size is 0, generation has finished!
+                        if len(stack_list[each_batch_idx]) == 0:
+                            gen_finish_list[each_batch_idx] = True
+                        else:
+                            nt_edge_list[each_batch_idx] = stack_list[each_batch_idx].pop()
+                            nt_symbol_list[each_batch_idx] = hg_list[each_batch_idx].edge_attr(nt_edge_list[each_batch_idx])['symbol']
+                    else:
+                        tgt_id = np.mod(self.padding_idx, self.vocab_size)
+                    indice_tensor = torch.LongTensor([tgt_id])
+                    indice_tensor = indice_tensor.cuda()
+                    tgt_emb[each_batch_idx, :] = self.tgt_embedding(indice_tensor)
+
+            # last one step
+            conditional_logprob_list = []
+            tgt_emb_pred = self.decoder.forward_one_step(tgt_emb)
+            vocab_logit = self.dec2vocab(tgt_emb_pred)
+            for each_batch_idx in range(self.batch_size):
+                if not gen_finish_list[each_batch_idx]: # if generation has not finished
+                    # get production rule greedily
+                    masked_logprob = self.hrg.prod_rule_corpus.masked_logprob(
+                        vocab_logit[each_batch_idx, :, :-1].squeeze().cpu().numpy(),
+                        nt_symbol_list[each_batch_idx])
+                    conditional_logprob_list.append(masked_logprob)
+                else:
+                    conditional_logprob_list.append(None)
+        return conditional_logprob_list
+
+    #TODO Not tested. Need to implement this in case of molecular structure generation
+    def decode_with_beam_search(self, z, beam_width=1):
+        ''' Decode a latent vector using beam search.
+
+        Parameters
+        ----------
+        z
+            latent vector
+        beam_width : int
+            parameter for beam search
+
+        Returns
+        -------
+        List of Hypergraphs
+        '''
+        if self.batch_size != 1:
+            raise ValueError('this method works only under batch_size=1')
+        if self.padding_idx != -1:
+            raise ValueError('this method works only under padding_idx=-1')
+        top_k_tgt_id_list = [[]] * beam_width
+        logprob_list = [0.] * beam_width
+
+        for each_len in range(self.max_len):
+            expanded_logprob_list = np.repeat(logprob_list, self.vocab_size) # including padding_idx
+            expanded_length_list = np.array([0] * (beam_width * self.vocab_size))
+            for each_beam_idx, each_candidate in enumerate(top_k_tgt_id_list):
+                conditional_logprob = self.conditional_distribution(z, [each_candidate])[0]
+                if conditional_logprob is None:
+                    expanded_logprob_list[(each_beam_idx + 1) * self.vocab_size - 1]\
+                        = logprob_list[each_beam_idx]
+                    expanded_logprob_list[each_beam_idx * self.vocab_size : (each_beam_idx + 1) * self.vocab_size - 1]\
+                        = -np.inf
+                    expanded_length_list[each_beam_idx * self.vocab_size : (each_beam_idx + 1) * self.vocab_size]\
+                        = len(each_candidate)
+                else:
+                    expanded_logprob_list[each_beam_idx * self.vocab_size : (each_beam_idx + 1) * self.vocab_size - 1]\
+                        = logprob_list[each_beam_idx] + conditional_logprob
+                    expanded_logprob_list[(each_beam_idx + 1) * self.vocab_size - 1]\
+                        = -np.inf
+                    expanded_length_list[each_beam_idx * self.vocab_size : (each_beam_idx + 1) * self.vocab_size]\
+                        = len(each_candidate) + 1
+            score_list = np.array(expanded_logprob_list) / np.array(expanded_length_list)
+            if each_len == 0:
+                top_k_list = np.argsort(score_list[:self.vocab_size])[::-1][:beam_width]
+            else:
+                top_k_list = np.argsort(score_list)[::-1][:beam_width]
+            next_top_k_tgt_id_list = []
+            next_logprob_list = []
+            for each_top_k in top_k_list:
+                beam_idx = each_top_k // self.vocab_size
+                vocab_idx = each_top_k % self.vocab_size
+                if vocab_idx == self.vocab_size - 1:
+                    next_top_k_tgt_id_list.append(top_k_tgt_id_list[beam_idx])
+                    next_logprob_list.append(expanded_logprob_list[each_top_k])
+                else:
+                    next_top_k_tgt_id_list.append(top_k_tgt_id_list[beam_idx] + [vocab_idx])
+                    next_logprob_list.append(expanded_logprob_list[each_top_k])
+            top_k_tgt_id_list = next_top_k_tgt_id_list
+            logprob_list = next_logprob_list
+
+        # construct hypergraphs
+        hg_list = []
+        for each_tgt_id_list in top_k_tgt_id_list:
+            hg = None
+            stack = []
+            nt_edge = None
+            for each_idx, each_prod_rule_id in enumerate(each_tgt_id_list):
+                prod_rule = self.hrg.prod_rule_list[each_prod_rule_id]
+                hg, nt_edges = prod_rule.applied_to(hg, nt_edge)
+                stack.extend(nt_edges[::-1])
+                try:
+                    nt_edge = stack.pop()
+                except IndexError:
+                    if each_idx == len(each_tgt_id_list) - 1:
+                        break
+                    else:
+                        raise ValueError('some bugs')
+            hg_list.append(hg)
+        return hg_list
+
+    def graph_embed(self, x, edge_index, edge_attr, batch_size):
+        src_h = self.encoder.forward(x, edge_index, edge_attr, batch_size)
+        return src_h        
+
+    def encode(self, x, edge_index, edge_attr, batch_size):
+        #print("device for src_emb=", src_emb.get_device())
+        #print("device for self.encoder=", next(self.encoder.parameters()).get_device())
+        src_h = self.graph_embed(x, edge_index, edge_attr, batch_size)
+        mu, lv = self.get_mean_var(src_h)
+        return mu, lv
+    
+    def get_mean_var(self, src_h):
+        #src_h = torch.tanh(src_h)
+        mu = self.hidden2mean(src_h)
+        lv = self.hidden2logvar(src_h)
+        mu = torch.tanh(mu)
+        lv = torch.tanh(lv)
+        return mu, lv
+
+    def reparameterize(self, mu, logvar, training=True):
+        if training:
+            std = logvar.mul(0.5).exp_()
+            eps = Variable(std.data.new(std.size()).normal_())
+            if self.rank >= 0:
+                eps = eps.to(next(self.parameters()).device)
+            return eps.mul(std).add_(mu)
+        else:
+            return mu
+
+# Copied from the MHG implementation and adapted
+class GrammarVAELoss(_Loss):
+
+    '''
+    a loss function for Grammar VAE
+
+    Attributes
+    ----------
+    hrg : HyperedgeReplacementGrammar
+    beta : float
+        coefficient of KL divergence
+    '''
+
+    def __init__(self, rank, hrg, beta=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.hrg = hrg
+        self.beta = beta
+        self.rank = rank
+
+    def forward(self, mu, logvar, in_seq_pred, in_seq):
+        ''' compute VAE loss
+
+        Parameters
+        ----------
+        in_seq_pred : torch.Tensor, shape (batch_size, max_len, vocab_size)
+            logit
+        in_seq : torch.Tensor, shape (batch_size, max_len)
+            each element corresponds to a word id in vocabulary.
+        mu : torch.Tensor, shape (batch_size, hidden_dim)
+        logvar : torch.Tensor, shape (batch_size, hidden_dim)
+            mean and log variance of the normal distribution
+        '''
+        batch_size = in_seq_pred.shape[0]
+        max_len = in_seq_pred.shape[1]
+        vocab_size = in_seq_pred.shape[2]
+        mask = torch.zeros(in_seq_pred.shape)
+        
+        for each_batch in range(batch_size):
+            flag = True
+            for each_idx in range(max_len):
+                prod_rule_idx = in_seq[each_batch, each_idx]
+                if prod_rule_idx == vocab_size - 1:
+                    #### DETERMINE WHETHER THIS SHOULD BE SKIPPED OR NOT
+                    mask[each_batch, each_idx, prod_rule_idx] = 1
+                    #break
+                    continue
+                lhs = self.hrg.prod_rule_corpus.prod_rule_list[prod_rule_idx].lhs_nt_symbol
+                lhs_idx = self.hrg.prod_rule_corpus.nt_symbol_list.index(lhs)
+                mask[each_batch, each_idx, :-1] = torch.FloatTensor(self.hrg.prod_rule_corpus.lhs_in_prod_rule[lhs_idx])
+            if self.rank >= 0:
+                mask = mask.to(next(self.parameters()).device)
+        in_seq_pred = mask * in_seq_pred
+
+        cross_entropy = F.cross_entropy(
+            in_seq_pred.view(-1, vocab_size),
+            in_seq.view(-1),
+            reduction='sum',
+            #ignore_index=self.ignore_index if self.ignore_index is not None else -100
+            )
+        kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
+        return cross_entropy + self.beta * kl_div
+
+
+class VAELoss(_Loss):
+    def __init__(self, beta=0.01):
+        super().__init__()
+        self.beta = beta
+
+    def forward(self, mean, log_var, dec_outputs, targets):
+
+        device = mean.get_device()
+        if device >= 0:
+            targets = targets.to(mean.get_device())
+        reconstruction = F.cross_entropy(dec_outputs.view(-1, dec_outputs.size(2)), targets.view(-1), reduction='sum')
+        
+        KL = 0.5 * torch.sum(1 + log_var - mean ** 2 - torch.exp(log_var))
+        loss = - self.beta * KL + reconstruction
+        return loss

models/mhg_model/notebooks/mhg-gnn_encoder_decoder_example.ipynb

@@ -0,0 +1,114 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "829ddc03",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.append('..')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ea820e23",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import load"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b9a51fa8",
+   "metadata": {},
+   "source": [
+    "# Load MHG-GNN"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6ea1fc8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_ckp = \"models/model_checkpoints/mhg_model/pickles/mhggnn_pretrained_model_radius7_1116_2023.pickle\"\n",
+    "\n",
+    "model = load.load(model_name = model_ckp)\n",
+    "if model is None:\n",
+    "    print(\"Model not loaded, please check you have MHG pickle file\")\n",
+    "else:\n",
+    "    print(\"MHG model loaded\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b4a0b557",
+   "metadata": {},
+   "source": [
+    "# Embeddings\n",
+    "\n",
+    "※ replace the smiles exaple list with your dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c63a6be6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with torch.no_grad():\n",
+    "    repr = model.encode([\"CCO\", \"O=C=O\", \"OC(=O)c1ccccc1C(=O)O\"])\n",
+    "    \n",
+    "# Print the latent vectors\n",
+    "print(repr)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a59f9442",
+   "metadata": {},
+   "source": [
+    "# Decoding"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a0d8a41",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "orig = model.decode(repr)\n",
+    "print(orig)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models/mhg_model/paper/MHG-GNN_Combination of Molecular Hypergraph Grammar with Graph Neural Network.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa61ca0975afd93bf1f3c43f3b720594018f57ac9df3cb740def2eb28af6529d
+size 342570

models/mhg_model/setup.cfg

@@ -0,0 +1,37 @@
+[metadata]
+name = mhg-gnn
+version = attr: .__version__
+description = Package for mhg-gnn
+author= team
+long_description_content_type=text/markdown
+long_description = file: README.md
+python_requires = >= 3.9.7 
+license = TBD
+
+classifiers =
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3.9
+
+[options]
+install_requires =
+	networkx>=2.8
+	numpy>=1.23.5, <2.0.0
+	pandas>=1.5.3
+	rdkit-pypi>=2022.9.4, <2023.9.6
+	torch>=2.0.0
+	torchinfo>=1.8.0
+	torch-geometric>=2.3.1
+	requests>=2.32.2
+	scikit-learn>=1.5.0
+	urllib3>=2.2.2
+
+
+setup_requires =
+    setuptools
+package_dir =
+    = .
+packages=find:
+include_package_data = True
+
+[options.packages.find]
+where = .

models/mhg_model/setup.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+
+import setuptools
+
+if __name__ == "__main__":
+    setuptools.setup()

models/selfies_model/README.md

@@ -0,0 +1,87 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: feature-extraction
+tags:
+- chemistry
+---
+
+# selfies-ted
+
+selfies-ted is a project for encoding SMILES (Simplified Molecular Input Line Entry System) into SELFIES (SELF-referencing Embedded Strings) and generating embeddings for molecular representations.
+
+![selfies-ted](selfies-ted.png)
+## Model Architecture
+
+Configuration details
+
+Encoder and Decoder FFN dimensions: 256
+Number of attention heads: 4
+Number of encoder and decoder layers: 2
+Total number of hidden layers: 6
+Maximum position embeddings: 128
+Model dimension (d_model): 256
+
+## Pretrained Models and Training Logs
+We provide checkpoints of the selfies-ted model pre-trained on a dataset of molecules curated from PubChem. The pre-trained model shows competitive performance on molecular representation tasks. For model weights: "HuggingFace link".
+
+To install and use the pre-trained model:
+
+Download the selfies_ted_model.pkl file from the "HuggingFace link".
+Add the selfies-ted selfies_ted_model.pkl to the models/ directory. The directory structure should look like the following:
+
+```
+models/
+└── selfies_ted_model.pkl
+```
+
+## Installation
+
+To use this project, you'll need to install the required dependencies. We recommend using a virtual environment:
+
+```bash
+python -m venv venv
+source venv/bin/activate  # On Windows use `venv\Scripts\activate`
+```
+
+Install the required dependencies
+
+```
+pip install -r requirements.txt
+```
+
+
+## Usage
+
+### Import
+
+```
+import load
+```
+### Training the Model
+
+To train the model, use the train.py script:
+
+```
+python train.py -f <path_to_your_data_file>
+```
+
+
+Note: The actual usage may depend on the specific implementation in load.py. Please refer to the source code for detailed functionality.
+
+### Load the model and tokenizer
+```
+load.load("path/to/checkpoint.pkl")
+```
+### Encode SMILES strings
+```
+smiles_list = ["COC", "CCO"]
+```
+```
+embeddings = load.encode(smiles_list)
+```
+
+
+## Example Notebook
+
+Example notebook of this project is `selfies-ted-example.ipynb`. 

models/selfies_model/load.py

@@ -0,0 +1,102 @@
+import torch
+import selfies as sf
+import numpy as np
+import pandas as pd
+from rdkit import Chem
+from transformers import AutoTokenizer, AutoModel
+import gc
+from torch.utils.data import DataLoader, Dataset
+from multiprocessing import Pool, cpu_count
+from tqdm import tqdm
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+class SELFIESDataset(Dataset):
+    def __init__(self, selfies_list):
+        self.selfies = selfies_list
+
+    def __len__(self):
+        return len(self.selfies)
+
+    def __getitem__(self, idx):
+        return self.selfies[idx]
+
+class SELFIES(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = None
+        self.tokenizer = None
+        self.invalid = []
+
+    def smiles_to_selfies(self, smiles):
+        try:
+            return sf.encoder(smiles.strip()).replace('][', '] [')
+        except:
+            try:
+                smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles.strip()))
+                return sf.encoder(smiles).replace('][', '] [')
+            except:
+                return None
+
+    def get_selfies(self, smiles_list):
+        with Pool(cpu_count()) as pool:
+            selfies = list(pool.map(self.smiles_to_selfies, smiles_list))
+
+        self.invalid = [i for i, s in enumerate(selfies) if s is None]
+        selfies = [s if s is not None else '[nop]' for s in selfies]
+        return selfies
+
+    @torch.no_grad()
+    def get_embedding_batch(self, selfies_batch):
+        encodings = self.tokenizer(
+            selfies_batch,
+            return_tensors='pt',
+            max_length=128,
+            truncation=True,
+            padding='max_length'
+        )
+        encodings = {k: v.to(self.model.device) for k, v in encodings.items()}
+
+        outputs = self.model.encoder(
+            input_ids=encodings['input_ids'],
+            attention_mask=encodings['attention_mask']
+        )
+
+        model_output = outputs.last_hidden_state
+        input_mask_expanded = encodings['attention_mask'].unsqueeze(-1).expand(model_output.size()).float()
+        sum_embeddings = torch.sum(model_output * input_mask_expanded, 1)
+        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+        pooled_output = sum_embeddings / sum_mask
+
+        return pooled_output.cpu().numpy()
+
+    def load(self, checkpoint=None):
+        self.tokenizer = AutoTokenizer.from_pretrained("ibm/materials.selfies-ted")
+        self.model = AutoModel.from_pretrained("ibm/materials.selfies-ted")
+        self.model.eval()
+
+    def encode(self, smiles_list=[], use_gpu=False, return_tensor=False, batch_size=128, num_workers=4):
+        selfies = self.get_selfies(smiles_list)
+        dataset = SELFIESDataset(selfies)
+
+        device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
+        self.model.to(device)
+
+        loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
+
+        embeddings = []
+        for batch in tqdm(loader, desc="Encoding"):
+            emb = self.get_embedding_batch(batch)
+            embeddings.append(emb)
+            del emb
+            gc.collect()
+
+        emb = np.vstack(embeddings)
+
+        for idx in self.invalid:
+            emb[idx] = np.nan
+            print(f"Cannot encode {smiles_list[idx]} to selfies. Embedding replaced by NaN.")
+
+        return torch.tensor(emb) if return_tensor else pd.DataFrame(emb)

models/selfies_model/requirements.txt

@@ -0,0 +1,12 @@
+torch>=2.1.0
+transformers>=4.38
+numpy>=1.26.1
+datasets>=2.13.1
+evaluate>=0.4.0
+selfies>=2.1.0
+scikit-learn>=1.2.1
+pyarrow>=14.0.1
+requests>=2.31.0
+urllib3>=2.0.7
+aiohttp>=3.9.0
+zipp>=3.17.0