sampling/EquiFM_sampling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "qm9_with_h = {\n",
    "    'name':\n",
    "    'qm9',\n",
    "    'atom_encoder': {\n",
    "        'H': 0,\n",
    "        'C': 1,\n",
    "        'N': 2,\n",
    "        'O': 3,\n",
    "        'F': 4\n",
    "    },\n",
    "    'atom_decoder': ['H', 'C', 'N', 'O', 'F'],\n",
    "    'n_nodes': {\n",
    "        22: 3393,\n",
    "        17: 13025,\n",
    "        23: 4848,\n",
    "        21: 9970,\n",
    "        19: 13832,\n",
    "        20: 9482,\n",
    "        16: 10644,\n",
    "        13: 3060,\n",
    "        15: 7796,\n",
    "        25: 1506,\n",
    "        18: 13364,\n",
    "        12: 1689,\n",
    "        11: 807,\n",
    "        24: 539,\n",
    "        14: 5136,\n",
    "        26: 48,\n",
    "        7: 16,\n",
    "        10: 362,\n",
    "        8: 49,\n",
    "        9: 124,\n",
    "        27: 266,\n",
    "        4: 4,\n",
    "        29: 25,\n",
    "        6: 9,\n",
    "        5: 5,\n",
    "        3: 1\n",
    "    },\n",
    "    'max_n_nodes':\n",
    "    29,\n",
    "    'atom_types': {\n",
    "        1: 635559,\n",
    "        2: 101476,\n",
    "        0: 923537,\n",
    "        3: 140202,\n",
    "        4: 2323\n",
    "    },\n",
    "    'distances': [\n",
    "        903054, 307308, 111994, 57474, 40384, 29170, 47152, 414344, 2202212,\n",
    "        573726, 1490786, 2970978, 756818, 969276, 489242, 1265402, 4587994,\n",
    "        3187130, 2454868, 2647422, 2098884, 2001974, 1625206, 1754172, 1620830,\n",
    "        1710042, 2133746, 1852492, 1415318, 1421064, 1223156, 1322256, 1380656,\n",
    "        1239244, 1084358, 981076, 896904, 762008, 659298, 604676, 523580,\n",
    "        437464, 413974, 352372, 291886, 271948, 231328, 188484, 160026, 136322,\n",
    "        117850, 103546, 87192, 76562, 61840, 49666, 43100, 33876, 26686, 22402,\n",
    "        18358, 15518, 13600, 12128, 9480, 7458, 5088, 4726, 3696, 3362, 3396,\n",
    "        2484, 1988, 1490, 984, 734, 600, 456, 482, 378, 362, 168, 124, 94, 88,\n",
    "        52, 44, 40, 18, 16, 8, 6, 2, 0, 0, 0, 0, 0, 0, 0\n",
    "    ],\n",
    "    'colors_dic': ['#FFFFFF99', 'C7', 'C0', 'C3', 'C1'],\n",
    "    'radius_dic': [0.46, 0.77, 0.77, 0.77, 0.77],\n",
    "    'with_h':\n",
    "    True\n",
    "}\n",
    "allowed_bonds = {\n",
    "    'H': 1,\n",
    "    'C': 4,\n",
    "    'N': 3,\n",
    "    'O': 2,\n",
    "    'F': 1,\n",
    "    'B': 3,\n",
    "    'Al': 3,\n",
    "    'Si': 4,\n",
    "    'P': [3, 5],\n",
    "    'S': 4,\n",
    "    'Cl': 1,\n",
    "    'As': 3,\n",
    "    'Br': 1,\n",
    "    'I': 1,\n",
    "    'Hg': [1, 2],\n",
    "    'Bi': [3, 5]\n",
    "}\n",
    "# Bond lengths from:\n",
    "# http://www.wiredchemist.com/chemistry/data/bond_energies_lengths.html\n",
    "# And:\n",
    "# http://chemistry-reference.com/tables/Bond%20Lengths%20and%20Enthalpies.pdf\n",
    "\n",
    "bonds1 = {\n",
    "    'H': {\n",
    "        'H': 74,\n",
    "        'C': 109,\n",
    "        'N': 101,\n",
    "        'O': 96,\n",
    "        'F': 92,\n",
    "        'B': 119,\n",
    "        'Si': 148,\n",
    "        'P': 144,\n",
    "        'As': 152,\n",
    "        'S': 134,\n",
    "        'Cl': 127,\n",
    "        'Br': 141,\n",
    "        'I': 161\n",
    "    },\n",
    "    'C': {\n",
    "        'H': 109,\n",
    "        'C': 154,\n",
    "        'N': 147,\n",
    "        'O': 143,\n",
    "        'F': 135,\n",
    "        'Si': 185,\n",
    "        'P': 184,\n",
    "        'S': 182,\n",
    "        'Cl': 177,\n",
    "        'Br': 194,\n",
    "        'I': 214\n",
    "    },\n",
    "    'N': {\n",
    "        'H': 101,\n",
    "        'C': 147,\n",
    "        'N': 145,\n",
    "        'O': 140,\n",
    "        'F': 136,\n",
    "        'Cl': 175,\n",
    "        'Br': 214,\n",
    "        'S': 168,\n",
    "        'I': 222,\n",
    "        'P': 177\n",
    "    },\n",
    "    'O': {\n",
    "        'H': 96,\n",
    "        'C': 143,\n",
    "        'N': 140,\n",
    "        'O': 148,\n",
    "        'F': 142,\n",
    "        'Br': 172,\n",
    "        'S': 151,\n",
    "        'P': 163,\n",
    "        'Si': 163,\n",
    "        'Cl': 164,\n",
    "        'I': 194\n",
    "    },\n",
    "    'F': {\n",
    "        'H': 92,\n",
    "        'C': 135,\n",
    "        'N': 136,\n",
    "        'O': 142,\n",
    "        'F': 142,\n",
    "        'S': 158,\n",
    "        'Si': 160,\n",
    "        'Cl': 166,\n",
    "        'Br': 178,\n",
    "        'P': 156,\n",
    "        'I': 187\n",
    "    },\n",
    "    'B': {\n",
    "        'H': 119,\n",
    "        'Cl': 175\n",
    "    },\n",
    "    'Si': {\n",
    "        'Si': 233,\n",
    "        'H': 148,\n",
    "        'C': 185,\n",
    "        'O': 163,\n",
    "        'S': 200,\n",
    "        'F': 160,\n",
    "        'Cl': 202,\n",
    "        'Br': 215,\n",
    "        'I': 243\n",
    "    },\n",
    "    'Cl': {\n",
    "        'Cl': 199,\n",
    "        'H': 127,\n",
    "        'C': 177,\n",
    "        'N': 175,\n",
    "        'O': 164,\n",
    "        'P': 203,\n",
    "        'S': 207,\n",
    "        'B': 175,\n",
    "        'Si': 202,\n",
    "        'F': 166,\n",
    "        'Br': 214\n",
    "    },\n",
    "    'S': {\n",
    "        'H': 134,\n",
    "        'C': 182,\n",
    "        'N': 168,\n",
    "        'O': 151,\n",
    "        'S': 204,\n",
    "        'F': 158,\n",
    "        'Cl': 207,\n",
    "        'Br': 225,\n",
    "        'Si': 200,\n",
    "        'P': 210,\n",
    "        'I': 234\n",
    "    },\n",
    "    'Br': {\n",
    "        'Br': 228,\n",
    "        'H': 141,\n",
    "        'C': 194,\n",
    "        'O': 172,\n",
    "        'N': 214,\n",
    "        'Si': 215,\n",
    "        'S': 225,\n",
    "        'F': 178,\n",
    "        'Cl': 214,\n",
    "        'P': 222\n",
    "    },\n",
    "    'P': {\n",
    "        'P': 221,\n",
    "        'H': 144,\n",
    "        'C': 184,\n",
    "        'O': 163,\n",
    "        'Cl': 203,\n",
    "        'S': 210,\n",
    "        'F': 156,\n",
    "        'N': 177,\n",
    "        'Br': 222\n",
    "    },\n",
    "    'I': {\n",
    "        'H': 161,\n",
    "        'C': 214,\n",
    "        'Si': 243,\n",
    "        'N': 222,\n",
    "        'O': 194,\n",
    "        'S': 234,\n",
    "        'F': 187,\n",
    "        'I': 266\n",
    "    },\n",
    "    'As': {\n",
    "        'H': 152\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "bonds2 = {\n",
    "    'C': {\n",
    "        'C': 134,\n",
    "        'N': 129,\n",
    "        'O': 120,\n",
    "        'S': 160\n",
    "    },\n",
    "    'N': {\n",
    "        'C': 129,\n",
    "        'N': 125,\n",
    "        'O': 121\n",
    "    },\n",
    "    'O': {\n",
    "        'C': 120,\n",
    "        'N': 121,\n",
    "        'O': 121,\n",
    "        'P': 150\n",
    "    },\n",
    "    'P': {\n",
    "        'O': 150,\n",
    "        'S': 186\n",
    "    },\n",
    "    'S': {\n",
    "        'P': 186\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "\n",
    "bonds3 = {\n",
    "    'C': {\n",
    "        'C': 120,\n",
    "        'N': 116,\n",
    "        'O': 113\n",
    "    },\n",
    "    'N': {\n",
    "        'C': 116,\n",
    "        'N': 110\n",
    "    },\n",
    "    'O': {\n",
    "        'C': 113\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "margin1, margin2, margin3 = 10, 5, 3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import pickle\n",
    "import torch\n",
    "import numpy as np\n",
    "from cnf_models import Cnflows, EGNN_dynamics_QM9, DistributionNodes, UniformDequantizer, assert_correctly_masked, assert_mean_zero_with_mask\n",
    "\n",
    "\n",
    "\n",
    "# A sample wrapper\n",
    "def sample(\n",
    "    args,\n",
    "    device,\n",
    "    dequantizer,\n",
    "    generative_model,\n",
    "    dataset_info,\n",
    "    prop_dist=None,\n",
    "    nodesxsample=torch.tensor([10]),\n",
    "    context=None,\n",
    "    fix_noise=False,\n",
    "):\n",
    "    max_n_nodes = dataset_info[\"max_n_nodes\"]  # this is the maximum node_size in QM9\n",
    "\n",
    "    assert int(torch.max(nodesxsample)) <= max_n_nodes\n",
    "    batch_size = len(nodesxsample)\n",
    "\n",
    "    node_mask = torch.zeros(batch_size, max_n_nodes)\n",
    "    for i in range(batch_size):\n",
    "        node_mask[i, 0 : nodesxsample[i]] = 1\n",
    "\n",
    "    # Compute edge_mask\n",
    "\n",
    "    edge_mask = node_mask.unsqueeze(1) * node_mask.unsqueeze(2)\n",
    "    diag_mask = ~torch.eye(edge_mask.size(1), dtype=torch.bool).unsqueeze(0)\n",
    "    edge_mask *= diag_mask\n",
    "    edge_mask = edge_mask.view(batch_size * max_n_nodes * max_n_nodes, 1).to(device)\n",
    "    node_mask = node_mask.unsqueeze(2).to(device)\n",
    "\n",
    "    # TODO FIX: This conditioning just zeros.\n",
    "    if args.context_node_nf > 0:\n",
    "        if context is None:\n",
    "            context = prop_dist.sample_batch(nodesxsample)\n",
    "        context = context.unsqueeze(1).repeat(1, max_n_nodes, 1).to(device) * node_mask\n",
    "    else:\n",
    "        context = None\n",
    "\n",
    "    if args.probabilistic_model == \"flow_matching\":\n",
    "        x, h = generative_model.sample(\n",
    "            dequantizer,\n",
    "            batch_size,\n",
    "            max_n_nodes,\n",
    "            node_mask,\n",
    "            edge_mask,\n",
    "            context,\n",
    "            fix_noise=fix_noise,\n",
    "        )\n",
    "\n",
    "        assert_correctly_masked(x, node_mask)\n",
    "        assert_mean_zero_with_mask(x, node_mask)\n",
    "\n",
    "        one_hot = h[\"categorical\"]\n",
    "        charges = h[\"integer\"]\n",
    "        if args.include_charges:\n",
    "            assert_correctly_masked(charges.float(), node_mask)\n",
    "\n",
    "    else:\n",
    "        raise ValueError(args.probabilistic_model)\n",
    "\n",
    "    return one_hot, charges, x, node_mask\n",
    "\n",
    "\n",
    "def get_bond_order(atom1, atom2, distance, check_exists=False):\n",
    "    distance = 100 * distance  # We change the metric\n",
    "\n",
    "    # Check exists for large molecules where some atom pairs do not have a\n",
    "    # typical bond length.\n",
    "    if check_exists:\n",
    "        if atom1 not in bonds1:\n",
    "            return 0\n",
    "        if atom2 not in bonds1[atom1]:\n",
    "            return 0\n",
    "\n",
    "    # margin1, margin2 and margin3 have been tuned to maximize the stability of\n",
    "    # the QM9 true samples.\n",
    "    if distance < bonds1[atom1][atom2] + margin1:\n",
    "\n",
    "        # Check if atoms in bonds2 dictionary.\n",
    "        if atom1 in bonds2 and atom2 in bonds2[atom1]:\n",
    "            thr_bond2 = bonds2[atom1][atom2] + margin2\n",
    "            if distance < thr_bond2:\n",
    "                if atom1 in bonds3 and atom2 in bonds3[atom1]:\n",
    "                    thr_bond3 = bonds3[atom1][atom2] + margin3\n",
    "                    if distance < thr_bond3:\n",
    "                        return 3  # Triple\n",
    "                return 2  # Double\n",
    "        return 1  # Single\n",
    "    return 0  # No bond\n",
    "\n",
    "\n",
    "def check_stability(positions, atom_type, dataset_info, debug=False):\n",
    "    assert len(positions.shape) == 2\n",
    "    assert positions.shape[1] == 3\n",
    "    atom_decoder = dataset_info['atom_decoder']\n",
    "    x = positions[:, 0]\n",
    "    y = positions[:, 1]\n",
    "    z = positions[:, 2]\n",
    "\n",
    "    nr_bonds = np.zeros(len(x), dtype='int')\n",
    "\n",
    "    for i in range(len(x)):\n",
    "        for j in range(i + 1, len(x)):\n",
    "            p1 = np.array([x[i], y[i], z[i]])\n",
    "            p2 = np.array([x[j], y[j], z[j]])\n",
    "            dist = np.sqrt(np.sum((p1 - p2)**2))\n",
    "            atom1, atom2 = atom_decoder[atom_type[i]], atom_decoder[\n",
    "                atom_type[j]]\n",
    "            pair = sorted([atom_type[i], atom_type[j]])\n",
    "            if dataset_info['name'] == 'qm9' or dataset_info[\n",
    "                    'name'] == 'qm9_second_half' or dataset_info[\n",
    "                        'name'] == 'qm9_first_half':\n",
    "                order = get_bond_order(atom1, atom2, dist)\n",
    "\n",
    "            nr_bonds[i] += order\n",
    "            nr_bonds[j] += order\n",
    "    nr_stable_bonds = 0\n",
    "    for atom_type_i, nr_bonds_i in zip(atom_type, nr_bonds):\n",
    "        possible_bonds = allowed_bonds[atom_decoder[atom_type_i]]\n",
    "        if type(possible_bonds) == int:\n",
    "            is_stable = possible_bonds == nr_bonds_i\n",
    "        else:\n",
    "            is_stable = nr_bonds_i in possible_bonds\n",
    "        if not is_stable and debug:\n",
    "            print(\"Invalid bonds for molecule %s with %d bonds\" %\n",
    "                  (atom_decoder[atom_type_i], nr_bonds_i))\n",
    "        nr_stable_bonds += int(is_stable)\n",
    "\n",
    "    molecule_stable = nr_stable_bonds == len(x)\n",
    "    return molecule_stable, nr_stable_bonds, len(x)\n",
    "\n",
    "def analyze_stability_for_molecules(molecule_list, dataset_info):\n",
    "    use_rdkit = True\n",
    "    one_hot = molecule_list['one_hot']\n",
    "    x = molecule_list['x']\n",
    "    node_mask = molecule_list['node_mask']\n",
    "\n",
    "    if isinstance(node_mask, torch.Tensor):\n",
    "        atomsxmol = torch.sum(node_mask, dim=1)\n",
    "    else:\n",
    "        atomsxmol = [torch.sum(m) for m in node_mask]\n",
    "\n",
    "    n_samples = len(x)\n",
    "\n",
    "    molecule_stable = 0\n",
    "    nr_stable_bonds = 0\n",
    "    n_atoms = 0\n",
    "\n",
    "    processed_list = []\n",
    "\n",
    "    for i in range(n_samples):\n",
    "        atom_type = one_hot[i].argmax(1).cpu().detach()\n",
    "        pos = x[i].cpu().detach()\n",
    "\n",
    "        atom_type = atom_type[0:int(atomsxmol[i])]\n",
    "        pos = pos[0:int(atomsxmol[i])]\n",
    "        processed_list.append((pos, atom_type))\n",
    "\n",
    "    for mol in processed_list:\n",
    "        pos, atom_type = mol\n",
    "        validity_results = check_stability(pos, atom_type, dataset_info)\n",
    "\n",
    "        molecule_stable += int(validity_results[0])\n",
    "        nr_stable_bonds += int(validity_results[1])\n",
    "        n_atoms += int(validity_results[2])\n",
    "\n",
    "    # Validity\n",
    "    fraction_mol_stable = molecule_stable / float(n_samples)\n",
    "    fraction_atm_stable = nr_stable_bonds / float(n_atoms)\n",
    "    validity_dict = {\n",
    "        'mol_stable': fraction_mol_stable,\n",
    "        'atm_stable': fraction_atm_stable,\n",
    "    }\n",
    "\n",
    "    \n",
    "    return validity_dict, None\n",
    "\n",
    "\n",
    "def analyze_and_save(args,\n",
    "                     device,\n",
    "                     dequantizer,\n",
    "                     generative_model,\n",
    "                     nodes_dist,\n",
    "                     prop_dist,\n",
    "                     dataset_info,\n",
    "                     n_samples=10,\n",
    "                     batch_size=10,\n",
    "                     save_to_xyz=False):\n",
    "    batch_size = min(batch_size, n_samples)\n",
    "    assert n_samples % batch_size == 0\n",
    "    molecules = {'one_hot': [], 'x': [], 'node_mask': []}\n",
    "    start_time = time.time()\n",
    "    for i in range(int(n_samples / batch_size)):\n",
    "        nodesxsample = nodes_dist.sample(batch_size)\n",
    "        one_hot, charges, x, node_mask = sample(args,\n",
    "                                                device,\n",
    "                                                dequantizer,\n",
    "                                                generative_model,\n",
    "                                                dataset_info,\n",
    "                                                prop_dist=prop_dist,\n",
    "                                                nodesxsample=nodesxsample)\n",
    "\n",
    "        molecules['one_hot'].append(one_hot.detach().cpu())\n",
    "        molecules['x'].append(x.detach().cpu())\n",
    "        molecules['node_mask'].append(node_mask.detach().cpu())\n",
    "\n",
    "        current_num_samples = (i + 1) * batch_size\n",
    "        secs_per_sample = (time.time() - start_time) / current_num_samples\n",
    "        print('\\t %d/%d Molecules generated at %.2f secs/sample' %\n",
    "              (current_num_samples, n_samples, secs_per_sample))\n",
    "\n",
    "    molecules = {key: torch.cat(molecules[key], dim=0) for key in molecules}\n",
    "    stability_dict, rdkit_metrics = analyze_stability_for_molecules(\n",
    "        molecules, dataset_info)\n",
    "\n",
    "    return stability_dict, rdkit_metrics\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "                \n",
    "def get_model(args, device, dataset_info):\n",
    "    histogram = dataset_info[\"n_nodes\"]\n",
    "    in_node_nf = len(dataset_info[\"atom_decoder\"]) + int(args.include_charges)\n",
    "    nodes_dist = DistributionNodes(histogram)\n",
    "\n",
    "    # prop_dist = None\n",
    "    # if len(args.conditioning) > 0:\n",
    "    #     prop_dist = DistributionProperty(dataloader_train, args.conditioning)\n",
    "\n",
    "    if args.condition_time:\n",
    "        dynamics_in_node_nf = in_node_nf + 1\n",
    "    else:\n",
    "        print(\"Warning: dynamics model is _not_ conditioned on time.\")\n",
    "        dynamics_in_node_nf = in_node_nf\n",
    "\n",
    "    net_dynamics = EGNN_dynamics_QM9(\n",
    "        in_node_nf=dynamics_in_node_nf,\n",
    "        context_node_nf=args.context_node_nf,\n",
    "        n_dims=3,\n",
    "        device=device,\n",
    "        hidden_nf=args.nf,\n",
    "        act_fn=torch.nn.SiLU(),\n",
    "        n_layers=args.n_layers,\n",
    "        attention=args.attention,\n",
    "        tanh=args.tanh,\n",
    "        mode=args.model,\n",
    "        norm_constant=args.norm_constant,\n",
    "        inv_sublayers=args.inv_sublayers,\n",
    "        sin_embedding=args.sin_embedding,\n",
    "        normalization_factor=args.normalization_factor,\n",
    "        aggregation_method=args.aggregation_method,\n",
    "    )\n",
    "    dequantizer = UniformDequantizer()\n",
    "    cnf = Cnflows(\n",
    "            dynamics=net_dynamics,\n",
    "            in_node_nf=in_node_nf,\n",
    "            n_dims=3,\n",
    "            timesteps=args.diffusion_steps,\n",
    "            noise_schedule=args.diffusion_noise_schedule,\n",
    "            noise_precision=args.diffusion_noise_precision,\n",
    "            loss_type=args.diffusion_loss_type,\n",
    "            norm_values=args.normalize_factors,\n",
    "            include_charges=args.include_charges,\n",
    "            discrete_path=args.discrete_path,\n",
    "            cat_loss=args.cat_loss,\n",
    "            cat_loss_step=args.cat_loss_step,\n",
    "            on_hold_batch=args.on_hold_batch,\n",
    "            sampling_method=args.sampling_method,\n",
    "            weighted_methods=args.weighted_methods,\n",
    "            ode_method=args.ode_method,\n",
    "            without_cat_loss=args.without_cat_loss,\n",
    "            angle_penalty=args.angle_penalty,\n",
    "        )\n",
    "\n",
    "    return cnf, nodes_dist, None, dequantizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entropy of n_nodes: H[N] -2.475700616836548\n",
      "\t 10/100 Molecules generated at 0.34 secs/sample\n",
      "\t 20/100 Molecules generated at 0.38 secs/sample\n",
      "\t 30/100 Molecules generated at 0.36 secs/sample\n",
      "\t 40/100 Molecules generated at 0.36 secs/sample\n",
      "\t 50/100 Molecules generated at 0.37 secs/sample\n",
      "\t 60/100 Molecules generated at 0.37 secs/sample\n",
      "\t 70/100 Molecules generated at 0.36 secs/sample\n",
      "\t 80/100 Molecules generated at 0.37 secs/sample\n",
      "\t 90/100 Molecules generated at 0.37 secs/sample\n",
      "\t 100/100 Molecules generated at 0.37 secs/sample\n",
      "{'mol_stable': 0.9, 'atm_stable': 0.9916154276131918}\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import pickle\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "with open('args.pickle', 'rb') as f:\n",
    "    args = pickle.load(f)\n",
    "if not hasattr(args, 'normalization_factor'):\n",
    "    args.normalization_factor = 1\n",
    "if not hasattr(args, 'aggregation_method'):\n",
    "    args.aggregation_method = 'sum'\n",
    "\n",
    "model_path = \"generative_model_ema_0.npy\"\n",
    "args.cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if args.cuda else \"cpu\")\n",
    "args.device = device\n",
    "dtype = torch.float32\n",
    "# utils.create_folders(args)\n",
    "# print(args)\n",
    "dataset_info = qm9_with_h\n",
    "generative_model, nodes_dist, prop_dist, deq = get_model(\n",
    "        args, device, dataset_info)\n",
    "generative_model.to(device)\n",
    "deq.to(device)\n",
    "\n",
    "fn = 'generative_model_ema_0.npy'\n",
    "flow_state_dict = torch.load(fn,\n",
    "                                map_location=device)\n",
    "generative_model.load_state_dict(flow_state_dict)\n",
    "\n",
    "stability_dict, rdkit_metrics = analyze_and_save(\n",
    "        args,\n",
    "        device,\n",
    "        deq,\n",
    "        generative_model,\n",
    "        nodes_dist,\n",
    "        prop_dist,\n",
    "        dataset_info,\n",
    "        n_samples=100,\n",
    "        batch_size=10,\n",
    "        save_to_xyz=False)\n",
    "\n",
    "print(stability_dict)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}