{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":14313,"databundleVersionId":862019,"sourceType":"competition"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q swifter\n!pip install -U -q ipywidgets","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:06.115672Z","iopub.execute_input":"2024-11-05T02:43:06.116364Z","iopub.status.idle":"2024-11-05T02:43:40.001723Z","shell.execute_reply.started":"2024-11-05T02:43:06.116298Z","shell.execute_reply":"2024-11-05T02:43:40.000139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\n","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:40.005224Z","iopub.execute_input":"2024-11-05T02:43:40.005669Z","iopub.status.idle":"2024-11-05T02:43:40.465304Z","shell.execute_reply.started":"2024-11-05T02:43:40.005622Z","shell.execute_reply":"2024-11-05T02:43:40.464290Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Configuration","metadata":{}},{"cell_type":"code","source":"class Configuration:\n    seed = 42\n    train_path = '/kaggle/input/champs-scalar-coupling/train.csv'\n    test_path = '/kaggle/input/champs-scalar-coupling/test.csv'\n    structures_path = '/kaggle/input/champs-scalar-coupling/structures.csv'\n    structures_dir = '/kaggle/input/champs-scalar-coupling/structures'\n    xyz_output_dir = '/kaggle/working/xyz'\n    picture_output_dir = '/kaggle/working/picture'\ncfg = Configuration()","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:40.466804Z","iopub.execute_input":"2024-11-05T02:43:40.467448Z","iopub.status.idle":"2024-11-05T02:43:40.473518Z","shell.execute_reply.started":"2024-11-05T02:43:40.467395Z","shell.execute_reply":"2024-11-05T02:43:40.472293Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Setup Random Seed","metadata":{}},{"cell_type":"code","source":"import random\nimport numpy as np \nimport torch \n\nseed = cfg.seed\nrandom.seed(seed)\nnp.random.seed(seed)\ntorch.manual_seed(seed)\ntorch.cuda.manual_seed_all(seed)","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:40.475214Z","iopub.execute_input":"2024-11-05T02:43:40.475925Z","iopub.status.idle":"2024-11-05T02:43:43.718655Z","shell.execute_reply.started":"2024-11-05T02:43:40.475883Z","shell.execute_reply":"2024-11-05T02:43:43.717639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# EDA","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(cfg.train_path)\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:43.721670Z","iopub.execute_input":"2024-11-05T02:43:43.722217Z","iopub.status.idle":"2024-11-05T02:43:50.789981Z","shell.execute_reply.started":"2024-11-05T02:43:43.722175Z","shell.execute_reply":"2024-11-05T02:43:50.788712Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"n = np.random.randint(0, len(df))\nprint(df['id'][n])\nprint(df['molecule_name'][n])\nprint(df['atom_index_0'][n])\nprint(df['atom_index_1'][n])\nprint(df['type'][n])\nprint(df['scalar_coupling_constant'][n]) ","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:50.791435Z","iopub.execute_input":"2024-11-05T02:43:50.791854Z","iopub.status.idle":"2024-11-05T02:43:50.800370Z","shell.execute_reply.started":"2024-11-05T02:43:50.791788Z","shell.execute_reply":"2024-11-05T02:43:50.799072Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nxyz_file = os.path.join(cfg.structures_dir, f\"{df['molecule_name'][n]}.xyz\")","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:50.802173Z","iopub.execute_input":"2024-11-05T02:43:50.802535Z","iopub.status.idle":"2024-11-05T02:43:50.837069Z","shell.execute_reply.started":"2024-11-05T02:43:50.802498Z","shell.execute_reply":"2024-11-05T02:43:50.835750Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n# open xyzfile and print it\nwith open(xyz_file) as f:\n    print(f.read())","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:50.838761Z","iopub.execute_input":"2024-11-05T02:43:50.839372Z","iopub.status.idle":"2024-11-05T02:43:50.856851Z","shell.execute_reply.started":"2024-11-05T02:43:50.839272Z","shell.execute_reply":"2024-11-05T02:43:50.855441Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"structures = pd.read_csv(cfg.structures_path)\nstructures.head()","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:50.858630Z","iopub.execute_input":"2024-11-05T02:43:50.859058Z","iopub.status.idle":"2024-11-05T02:43:55.058426Z","shell.execute_reply.started":"2024-11-05T02:43:50.858982Z","shell.execute_reply":"2024-11-05T02:43:55.057237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# filter out other rows and get rows with molecule_name = df['molecule_name'][n]\nstructure = structures[structures['molecule_name'] == df['molecule_name'][n]]\nstructure","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:55.059930Z","iopub.execute_input":"2024-11-05T02:43:55.060388Z","iopub.status.idle":"2024-11-05T02:43:55.550371Z","shell.execute_reply.started":"2024-11-05T02:43:55.060343Z","shell.execute_reply":"2024-11-05T02:43:55.549262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Merge df with structures twice to add coordinates for both atom_index_0 and atom_index_1\ndf = df.merge(\n    structures.rename(columns={'atom_index': 'atom_index_0', 'x': 'x0', 'y': 'y0', 'z': 'z0'}),\n    on=['molecule_name', 'atom_index_0']\n).merge(\n    structures.rename(columns={'atom_index': 'atom_index_1', 'x': 'x1', 'y': 'y1', 'z': 'z1'}),\n    on=['molecule_name', 'atom_index_1']\n)\n\n# Calculate the distance using vectorized operations\ndf['distance'] = np.sqrt((df['x0'] - df['x1'])**2 + (df['y0'] - df['y1'])**2 + (df['z0'] - df['z1'])**2)\ndf = df.drop(columns=['x0', 'y0', 'z0', 'x1', 'y1', 'z1'])\ndf.head()\n","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:43:55.551699Z","iopub.execute_input":"2024-11-05T02:43:55.552095Z","iopub.status.idle":"2024-11-05T02:44:02.541149Z","shell.execute_reply.started":"2024-11-05T02:43:55.552040Z","shell.execute_reply":"2024-11-05T02:44:02.540052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for _ in df['type'].unique():\n    print(_)\n    _df = df[df['type'] == _]\n    plt.figure() \n    plt.scatter(_df['distance'], _df['scalar_coupling_constant'], c=np.random.rand(1,3),s=0.1, alpha=0.1)\n    plt.xlabel('Distance [angstrom]')\n    plt.ylabel('Coupling Constant [Hz]')\n    plt.title(_)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:02.542552Z","iopub.execute_input":"2024-11-05T02:44:02.542944Z","iopub.status.idle":"2024-11-05T02:44:15.271232Z","shell.execute_reply.started":"2024-11-05T02:44:02.542904Z","shell.execute_reply":"2024-11-05T02:44:15.269966Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(df)","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.273094Z","iopub.execute_input":"2024-11-05T02:44:15.273563Z","iopub.status.idle":"2024-11-05T02:44:15.282401Z","shell.execute_reply.started":"2024-11-05T02:44:15.273512Z","shell.execute_reply":"2024-11-05T02:44:15.281029Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Utilities","metadata":{}},{"cell_type":"code","source":"structure['molecule_name'].unique()","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.288038Z","iopub.execute_input":"2024-11-05T02:44:15.288550Z","iopub.status.idle":"2024-11-05T02:44:15.297396Z","shell.execute_reply.started":"2024-11-05T02:44:15.288501Z","shell.execute_reply":"2024-11-05T02:44:15.296065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def df_to_xyz(df=None):\n    molecule_name = df['molecule_name'].unique()[0]\n    assert df['molecule_name'].nunique() == 1\n    \n    sorted_df = df.sort_values(by='atom_index')\n    sorted_df = sorted_df.drop(columns=['molecule_name', 'atom_index'])\n    arr = sorted_df.to_numpy()\n    xyz_str = \"\"\n    xyz_str += f\"{len(arr)}\\n\\n\" \n    for row in arr:\n        xyz_str += f\"{row[0]} {row[1]} {row[2]} {row[3]}\\n\"\n    return xyz_str, molecule_name\nxyz_str, molecule_name = df_to_xyz(structure) \n\nprint(molecule_name)\nprint(\"-\"*60)\nprint(xyz_str) ","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.299283Z","iopub.execute_input":"2024-11-05T02:44:15.299733Z","iopub.status.idle":"2024-11-05T02:44:15.316178Z","shell.execute_reply.started":"2024-11-05T02:44:15.299679Z","shell.execute_reply":"2024-11-05T02:44:15.314764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"row = df.loc[n] \n\nr1 = structure[structure['atom_index'] == row['atom_index_0']][['x', 'y', 'z']].to_numpy()\nr2 = structure[structure['atom_index'] == row['atom_index_1']][['x', 'y', 'z']].to_numpy()\n\nc = (r1 + r2) / 2\nd = (r1 - r2) / 2\n#load_and_set_view(xyz_file, camera_position=\"side\")\nc,d","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.318207Z","iopub.execute_input":"2024-11-05T02:44:15.318601Z","iopub.status.idle":"2024-11-05T02:44:15.334919Z","shell.execute_reply.started":"2024-11-05T02:44:15.318561Z","shell.execute_reply":"2024-11-05T02:44:15.333611Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\ndef get_matrix(c=None, d=None):\n    # Assuming `c` is the center point and `d` is the direction vector\n    c = np.array(c).flatten()  # The center point\n    d = np.array(d).flatten()  # The direction vector\n    \n    # Step 1: Translation matrix\n    T = np.array([[1, 0, 0, -c[0]],\n                  [0, 1, 0, -c[1]],\n                  [0, 0, 1, -c[2]],\n                  [0, 0, 0, 1]], dtype=np.float32)\n\n    # Step 2: Normalize the direction vector d\n    d_norm = d / np.linalg.norm(d)\n\n    # Step 3: Create the rotation matrix\n    # Define the z-axis\n    z_axis = np.array([0, 0, 1])\n    # Calculate the rotation angle and axis\n    angle = -np.arccos(np.dot(z_axis, d_norm))\n    rotation_axis = np.cross(z_axis, d_norm)\n\n    # Normalize the rotation axis\n    if np.linalg.norm(rotation_axis) > 1e-6:  # Avoid division by zero\n        rotation_axis /= np.linalg.norm(rotation_axis)\n    else:\n        rotation_axis = np.array([1, 0, 0])  # Arbitrarily choose x-axis\n\n    # Using Rodrigues' rotation formula to create rotation matrix\n    K = np.array([[0, -rotation_axis[2], rotation_axis[1]],\n                  [rotation_axis[2], 0, -rotation_axis[0]],\n                  [-rotation_axis[1], rotation_axis[0], 0]])\n\n    R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K@K\n\n    # Convert to a 4x4 matrix\n    R_full = np.eye(4)\n    R_full[:3, :3] = R\n\n    # Step 4: Combine transformations\n    M = R_full @ T\n    \n    # Now `M` can be used to transform your molecule's coordinates\n    return M\n\ndef rotate(M, arr):\n    '''arr is shaped (natom, 3)'''\n    natom = arr.shape[0]\n    before_trans = np.concatenate((arr.T, np.ones((1,natom))), axis=0)\n    #print(before_trans.shape)\n    after_trans = M @ before_trans\n    return after_trans[:3, :].T","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.336456Z","iopub.execute_input":"2024-11-05T02:44:15.336930Z","iopub.status.idle":"2024-11-05T02:44:15.354219Z","shell.execute_reply.started":"2024-11-05T02:44:15.336879Z","shell.execute_reply":"2024-11-05T02:44:15.352998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"M = get_matrix(c=c, d=d)\nstructure_rotated = structure.copy()\nstructure_rotated[['x', 'y', 'z']] = rotate(M, structure[['x', 'y', 'z']].to_numpy())\nstructure","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.355898Z","iopub.execute_input":"2024-11-05T02:44:15.356437Z","iopub.status.idle":"2024-11-05T02:44:15.384988Z","shell.execute_reply.started":"2024-11-05T02:44:15.356385Z","shell.execute_reply":"2024-11-05T02:44:15.383491Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"structure_rotated","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.386906Z","iopub.execute_input":"2024-11-05T02:44:15.387581Z","iopub.status.idle":"2024-11-05T02:44:15.408342Z","shell.execute_reply.started":"2024-11-05T02:44:15.387533Z","shell.execute_reply":"2024-11-05T02:44:15.407084Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"xyz_str_rotated, molecule_name = df_to_xyz(structure_rotated) \nprint(xyz_str_rotated)","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.409982Z","iopub.execute_input":"2024-11-05T02:44:15.410497Z","iopub.status.idle":"2024-11-05T02:44:15.420883Z","shell.execute_reply.started":"2024-11-05T02:44:15.410445Z","shell.execute_reply":"2024-11-05T02:44:15.419615Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"row","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.424305Z","iopub.execute_input":"2024-11-05T02:44:15.425204Z","iopub.status.idle":"2024-11-05T02:44:15.443186Z","shell.execute_reply.started":"2024-11-05T02:44:15.425146Z","shell.execute_reply":"2024-11-05T02:44:15.441900Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n# create cfg.xyz_output_dir if it doesn't exist\nos.makedirs(cfg.xyz_output_dir, exist_ok=True)\nmolecule_name = row['molecule_name']\nrow_id = row['id']\nfname_xyz = os.path.join(cfg.xyz_output_dir, f\"{molecule_name}_{row_id}.xyz\")\nwith open(fname_xyz, \"w\") as file:\n    file.write(xyz_str_rotated)\n","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.444684Z","iopub.execute_input":"2024-11-05T02:44:15.445334Z","iopub.status.idle":"2024-11-05T02:44:15.456156Z","shell.execute_reply.started":"2024-11-05T02:44:15.445268Z","shell.execute_reply":"2024-11-05T02:44:15.454913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nfrom matplotlib import colors as mcolors\n\n# Define colors for different atoms\natom_colors = {\n    'H': 'white',\n    'C': 'gray',\n    'N': 'blue',\n    'O': 'red',\n    'S': 'yellow',\n    # Add other elements as needed\n}\n\ndef load_xyz(filename):\n    \"\"\"Load XYZ file and return atom types and coordinates.\"\"\"\n    atom_types = []\n    coordinates = []\n    with open(filename, 'r') as file:\n        # Skip the first two lines (comment and atom count)\n        file.readline()\n        file.readline()\n        for line in file:\n            parts = line.split()\n            atom_types.append(parts[0])\n            coordinates.append([float(parts[1]), float(parts[2]), float(parts[3])])\n    return atom_types, np.array(coordinates)\n\ndef plot_molecule(filename, output_filename=\"molecule_view.jpg\", loud=True):\n    atom_types, coordinates = load_xyz(filename)\n\n    # Set up 3D plot\n    fig = plt.figure(figsize=(3, 3), dpi=300)\n    ax = fig.add_subplot(111, projection='3d')\n    #fig.subplots_adjust(left=0, right=1, top=1, bottom=0)\n\n    plt.tight_layout(pad=0, w_pad=0, h_pad=0)\n    ax.set_position([0, 0, 1, 1])\n    # Set camera to point at +x direction\n    ax.view_init(elev=0, azim=0)  # Camera pointing along +x-axis\n\n    # Plot each atom with the corresponding color\n    for atom, coord in zip(atom_types, coordinates):\n        color = atom_colors.get(atom, 'black')  # Use black if atom type is not specified\n        r = (coord[0]**2 + coord[1]**2 + coord[2]**2)**0.5\n        alpha = np.exp(-r**2/5**2)\n        ax.scatter(coord[0], coord[1], coord[2], color=color, alpha=alpha, s=60, edgecolor='k', label=atom if atom not in ax.get_legend_handles_labels()[1] else \"\")\n    ax.plot([0, 0], [0, 0], [-5, 5], 'k--', linewidth=1)\n    # Customize plot appearance\n    #ax.set_xlabel(\"X\")\n    #ax.set_ylabel(\"Y\")\n    #ax.set_zlabel(\"Z\")\n    ax.set_xlim(-2, 2)\n    ax.set_ylim(-2, 2)\n    ax.set_zlim(-2, 2)\n    #ax.legend(loc=\"upper right\", bbox_to_anchor=(1.15, 1))\n    ax.grid(False)\n    ax.set_axis_off()\n    ax.set_box_aspect([1, 1, 1])  # Equal scaling for all axes\n \n    plt.savefig(output_filename) \n    if loud:\n        plt.show()\n    plt.close()\n# Usage\nplot_molecule(fname_xyz)\n","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:15.457757Z","iopub.execute_input":"2024-11-05T02:44:15.458190Z","iopub.status.idle":"2024-11-05T02:44:16.098984Z","shell.execute_reply.started":"2024-11-05T02:44:15.458150Z","shell.execute_reply":"2024-11-05T02:44:16.097816Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generate Pictures","metadata":{}},{"cell_type":"code","source":"\nos.makedirs(cfg.picture_output_dir, exist_ok=True)\nos.makedirs(cfg.xyz_output_dir, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:16.100325Z","iopub.execute_input":"2024-11-05T02:44:16.100659Z","iopub.status.idle":"2024-11-05T02:44:16.106609Z","shell.execute_reply.started":"2024-11-05T02:44:16.100624Z","shell.execute_reply":"2024-11-05T02:44:16.105401Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"row['molecule_name']","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:16.108138Z","iopub.execute_input":"2024-11-05T02:44:16.108582Z","iopub.status.idle":"2024-11-05T02:44:16.118823Z","shell.execute_reply.started":"2024-11-05T02:44:16.108533Z","shell.execute_reply":"2024-11-05T02:44:16.117583Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pandas as pd\n\nn = np.random.randint(0, len(df))\nrow = df.loc[n] \n\n# Create directories only once\nos.makedirs(cfg.xyz_output_dir, exist_ok=True)\nos.makedirs(cfg.picture_output_dir, exist_ok=True)\n\n# Convert structures to a dictionary for faster access\nstructures_dict = {name: group for name, group in structures.groupby('molecule_name')}\n\ndef process_row(row, structures_dict=structures_dict, cfg=cfg):\n    try:\n        structure = structures_dict.get(row['molecule_name'], pd.DataFrame())\n        if structure.empty:\n            print(f\"{row['molecule_name']} structure not found.\")\n            return row\n\n        r1 = structure[structure['atom_index'] == row['atom_index_0']][['x', 'y', 'z']].to_numpy()\n        r2 = structure[structure['atom_index'] == row['atom_index_1']][['x', 'y', 'z']].to_numpy()\n\n        c = (r1 + r2) / 2\n        d = (r1 - r2) / 2\n        M = get_matrix(c=c, d=d)\n        structure_rotated = structure.copy()\n        structure_rotated[['x', 'y', 'z']] = rotate(M, structure[['x', 'y', 'z']].to_numpy())\n        xyz_str_rotated, molecule_name = df_to_xyz(structure_rotated) \n\n        fname_xyz = os.path.join(cfg.xyz_output_dir, f\"{molecule_name}_{row['id']}.xyz\")\n        with open(fname_xyz, \"w\") as file:\n            file.write(xyz_str_rotated)\n\n        fname_pic = os.path.join(cfg.picture_output_dir, f\"{molecule_name}_{row['id']}.jpg\")\n        plot_molecule(fname_xyz, output_filename=fname_pic, loud=False)\n        del structure_rotated\n    except Exception as e:\n        print(f\"{row['molecule_name']} failed with atom pairs indexes: {row['atom_index_0']}, {row['atom_index_1']} - {e}\")\n\n    return row\n\nprocess_row(row, structures_dict=structures_dict, cfg=cfg)\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''n = np.random.randint(0, len(df))\nrow = df.loc[n] \n\ndef process_row(row, structures=None, cfg=cfg):\n    try:\n        structure = structures[structures['molecule_name'] == row['molecule_name']]\n        r1 = structure[structure['atom_index'] == row['atom_index_0']][['x', 'y', 'z']].to_numpy()\n        r2 = structure[structure['atom_index'] == row['atom_index_1']][['x', 'y', 'z']].to_numpy()\n\n        c = (r1 + r2) / 2\n        d = (r1 - r2) / 2\n        #load_and_set_view(xyz_file, camera_position=\"side\")\n        c,d\n        M = get_matrix(c=c, d=d)\n        structure_rotated = structure.copy()\n        structure_rotated[['x', 'y', 'z']] = rotate(M, structure[['x', 'y', 'z']].to_numpy())\n        structure\n        xyz_str_rotated, molecule_name = df_to_xyz(structure_rotated) \n\n\n        import os\n        # create cfg.xyz_output_dir if it doesn't exist \n        # os.makedirs(cfg.xyz_output_dir, exist_ok=True)\n        molecule_name = row['molecule_name']\n        row_id = row['id']\n        fname_xyz = os.path.join(cfg.xyz_output_dir, f\"{molecule_name}_{row_id}.xyz\")\n        with open(fname_xyz, \"w\") as file:\n            file.write(xyz_str_rotated)\n\n        fname_pic = os.path.join(cfg.picture_output_dir, f\"{molecule_name}_{row_id}.jpg\")\n        plot_molecule(fname_xyz, output_filename=fname_pic, loud=False)\n    except:\n        print(f\"{row['molecule_name']} failed with atom pairs indexes: {row['atom_index_0']}, {row['atom_index_1']}\" )\n    return row\nprocess_row(row, structures=structures, cfg=cfg)\n'''","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:16.120211Z","iopub.execute_input":"2024-11-05T02:44:16.120621Z","iopub.status.idle":"2024-11-05T02:44:17.014628Z","shell.execute_reply.started":"2024-11-05T02:44:16.120562Z","shell.execute_reply":"2024-11-05T02:44:17.013283Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport swifter\nimport gc\nfrom tqdm import tqdm\ntqdm.pandas()\nchunk_size = 100  # Adjust chunk size based on available resources\nnum_rows=100\nfor cc_type in df['type'].unique():\n    print(f\"Processing type {cc_type}\")\n    df_type = df[df['type']==cc_type]\n    # pick random 1000 rows within df_1JHN as df_1JHN\n    df_type = df_type.sample(n=num_rows, random_state=42)\n    _ = lambda row: process_row(row, structures_dict=structures_dict, cfg=cfg)\n\n    df_chunks = np.array_split(df_type, len(df_type) // chunk_size)\n\n    processed_chunks = []\n    for i, chunk in enumerate(df_chunks):\n        print(f\"Processing chunk {i + 1}/{len(df_chunks)}\")\n        processed_chunk = chunk.progress_apply(_, axis=1)\n        processed_chunks.append(processed_chunk)\n        # Collect garbage after processing each chunk\n        gc.collect()\n    df_type_processed = pd.concat(processed_chunks, ignore_index=True)\n    df_type.to_csv(f'train_{cc_type}.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2024-11-05T02:44:17.016666Z","iopub.execute_input":"2024-11-05T02:44:17.017204Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''# use swift to process df[df['type']=='1JHN']\nimport swifter\nfrom tqdm import tqdm\ntqdm.pandas()\ndf_1JHN = df[df['type']=='1JHN']#[:100]\n_ = lambda row: process_row(row, structures=structures, cfg=cfg)\n\ndf_1JHN = df_1JHN.progress_apply(_, axis=1)\n#df_1JHN.swifter.apply(_, axis=1)\n#df_1JHN = df_1JHN.swifter.progress_bar(True).apply(_, axis=1)'''","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''import matplotlib.pyplot as plt\n\nplt.imshow(plt.imread('molecule_view.jpg'))\nplt.axis('off')\nplt.show()'''","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''len(df[df['type']=='1JHN']), len(df)'''","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''# get file size for molecule_view.jpg\nimport os\nimport shutil\nfile_size = os.path.getsize(\"molecule_view.jpg\")\nprint(file_size/1024,'KB')\nprint(f\"{file_size*len(df)/1024/1024/1024:.3f} GB\")\nprint(f\"{file_size*len(df[df['type']=='1JHN'])/1024/1024/1024:.3f} GB\")\n\n#print(file_size*len(df[df['type']=='1JHC'])/1024/1024/1024,'GB')\n\n#print(file_size*len(df[df['type']=='3JHC'])/1024/1024/1024,'GB')'''","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''df_1JHN'''","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}