manimol 0.2.0__py3-none-any.whl

args_parse.py

@@ -0,0 +1,593 @@
+import argparse
+
+
+def args_parser():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Probability-manifold molecular conformation generation: "
+            "Stage 1 manifold learning and Stage 2 multimode conformer generation."
+        )
+    )
+
+    # =========================================================
+    # Experiment
+    # =========================================================
+    parser.add_argument("--exp_name", type=str, default="run")
+    parser.add_argument("--dump_path", type=str, default="dump/")
+    parser.add_argument("--exp_id", type=str, default="")
+    parser.add_argument(
+        "--fixed_dump_name",
+        action="store_true",
+        help="Store runs under dump/exp_name/exp_id instead of dump/MMDD-exp_name/exp_id.",
+    )
+    parser.add_argument("--gpu", type=str, default="0")
+    parser.add_argument("--auto_device_config", dest="auto_device_config", action="store_true", default=True)
+    parser.add_argument("--no_auto_device_config", dest="auto_device_config", action="store_false")
+    parser.add_argument("--cpu_bs", type=int, default=32)
+    parser.add_argument("--cuda_bs", type=int, default=96)
+    parser.add_argument("--num_workers", type=int, default=None)
+    parser.add_argument("--cpu_num_workers", type=int, default=0)
+    parser.add_argument("--cuda_num_workers", type=int, default=0)
+    parser.add_argument("--prefetch_factor", type=int, default=None)
+    parser.add_argument("--cuda_prefetch_factor", type=int, default=2)
+    parser.add_argument("--pin_memory", dest="pin_memory", action="store_true", default=None)
+    parser.add_argument("--no_pin_memory", dest="pin_memory", action="store_false")
+    parser.add_argument("--random_seed", type=int, default=0)
+    parser.add_argument("--train_mol_limit", type=int, default=0)
+    parser.add_argument("--valid_mol_limit", type=int, default=0)
+    parser.add_argument("--test_mol_limit", type=int, default=0)
+
+    # checkpoint / resume
+    parser.add_argument("--load_path", type=str, default=None)
+    parser.add_argument("--checkpoint_path", type=str, default=None)
+    parser.add_argument(
+        "--resume_progress_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Resume a MANIMOL training progress checkpoint, restoring model, "
+            "optimizer, epoch, batch position, and global step."
+        ),
+    )
+    parser.add_argument(
+        "--resume_next_epoch",
+        action="store_true",
+        default=False,
+        help="Resume from the next epoch without iterating over the unfinished part of a saved epoch.",
+    )
+    parser.add_argument("--checkpoint_interval_steps", type=int, default=1000)
+
+    # =========================================================
+    # Dataset
+    # =========================================================
+    parser.add_argument("--data_root", type=str, default="data")
+    parser.add_argument("--config_path", type=str, default="configs")
+    parser.add_argument("--dataset", type=str, default="QM9")
+    parser.add_argument(
+        "--raw_prefix",
+        type=str,
+        default="converted",
+        choices=["converted", "geom_qm9", "geom_qm9_full", "geom_drugs", "geom_drugs75k_subset"],
+        help="Raw data file set under data_root. geom_drugs75k_subset uses geom_drugs75k_subset_train/val/test.",
+    )
+    parser.add_argument("--train_raw_filename", type=str, default=None)
+    parser.add_argument("--val_raw_filename", type=str, default=None)
+    parser.add_argument("--test_raw_filename", type=str, default=None)
+
+    # 0 opt-in auto-detects the train/validation conformer capacity.
+    parser.add_argument("--max_confs", type=int, default=5)
+    parser.add_argument("--full_test_conformers", action="store_true", default=False,
+                        help="Opt-in: retain every test conformer in a separate v2 cache.")
+    parser.add_argument("--sharded_data_root", type=str, default="",
+                        help="Optional isolated tensor-only shard cache; legacy InMemoryDataset remains default.")
+    parser.add_argument("--shard_cache_size", type=int, default=2)
+    parser.add_argument(
+        "--stage2_assignment_cache",
+        type=str,
+        default="",
+        help="Optional directory containing fixed per-split Stage-II cluster assignments.",
+    )
+    parser.add_argument(
+        "--distributed_stage2",
+        action="store_true",
+        default=False,
+        help="Enable synchronous multi-GPU Stage-II training under torchrun.",
+    )
+
+    # =========================================================
+    # Model: GNN encoder
+    # =========================================================
+    parser.add_argument("--emb_dim", type=int, default=128)
+    parser.add_argument("--layer", type=int, default=4)
+    parser.add_argument("--dropout", type=float, default=0.5)
+
+    parser.add_argument(
+        "--gnn_type",
+        type=str,
+        default="gin",
+        choices=["gcn", "gin"],
+    )
+
+    parser.add_argument(
+        "--pooling_type",
+        type=str,
+        default="mean",
+        choices=["mean", "add", "max"],
+    )
+
+    # Input feature dimensions. These must match dataset / gnnconv.py.
+    parser.add_argument("--node_dim", type=int, default=39)
+    parser.add_argument("--edge_dim", type=int, default=13)
+
+    # =========================================================
+    # Model: probability map / MLP hidden dims
+    # =========================================================
+    parser.add_argument("--mlp_hidden", type=int, default=128)
+    parser.add_argument("--mlp_layer", type=int, default=4)
+
+    # =========================================================
+    # Model: probability mode generation
+    # =========================================================
+    # Number of generated probability modes/conformers during multimode training/inference.
+    parser.add_argument("--num_modes", type=int, default=10)
+
+    # Latent code dimension for P^(k) = f(G, z_k).
+    # Kept as mode_dim to avoid CVAE wording.
+    parser.add_argument("--mode_dim", type=int, default=32)
+
+    # Compatibility alias only. train.py should map latent_dim -> mode_dim if needed.
+    parser.add_argument("--latent_dim", type=int, default=None)
+
+    # =========================================================
+    # Model: coordinate generation / refinement
+    # =========================================================
+    parser.add_argument("--coarse_hidden", type=int, default=128)
+    parser.add_argument("--refine_hidden", type=int, default=128)
+    parser.add_argument("--num_refine_steps", type=int, default=4)
+
+    # Compatibility alias only. train.py should map refine_steps -> num_refine_steps if needed.
+    parser.add_argument("--refine_steps", type=int, default=None)
+
+    # =========================================================
+    # Training
+    # =========================================================
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--bs", type=int, default=32)
+    parser.add_argument("--epoch", type=int, default=300)
+    parser.add_argument("--eta_min", type=float, default=1e-6)
+    parser.add_argument("--patience", type=int, default=20)
+    parser.add_argument("--lambda_torsion_supervision", type=float, default=0.0)
+
+    parser.add_argument(
+        "--metric",
+        type=str,
+        default="RMSD",
+        choices=["CE", "RMSD"],
+    )
+
+    # Default True, with explicit opt-out.
+    parser.add_argument("--early_stop", dest="early_stop", action="store_true", default=True)
+    parser.add_argument("--no_early_stop", dest="early_stop", action="store_false")
+
+    parser.add_argument(
+        "--train_stage",
+        type=str,
+        default="gnn",
+        choices=[
+            "gnn",
+            "multimode",
+        ],
+    )
+
+    # Compatibility. Current code always uses teacher-style supervision.
+    # Kept to avoid breaking old commands, but should not affect logic.
+    parser.add_argument("--train_model", type=str, default="teacher")
+
+    # =========================================================
+    # Losses: Stage 1 probability manifold alignment
+    # =========================================================
+    # Main CE weight for CE(P_global, Q_global).
+    parser.add_argument("--lambda_global", type=float, default=1.0)
+    parser.add_argument(
+        "--mean_dispersion_prior",
+        action="store_true",
+        help=(
+            "Train the optional two-head Stage-I prior: a mean pairwise proximity map "
+            "and an ensemble-derived pairwise dispersion map."
+        ),
+    )
+    parser.add_argument(
+        "--lambda_stage1_dispersion",
+        type=float,
+        default=1.0,
+        help="Weight of the pairwise dispersion regression loss in the mean--dispersion Stage-I objective.",
+    )
+    parser.add_argument(
+        "--stage1_target_strategy",
+        type=str,
+        default="conformer",
+        choices=["fusion", "lowest_energy", "first", "conformer"],
+        help=(
+            "How Stage 1 builds the teacher probability graph. "
+            "fusion keeps the old learnable weighted fusion; lowest_energy uses the lowest-energy conformer; "
+            "first uses the first conformer; conformer averages per-conformer CE without constructing a fused Q_global."
+        ),
+    )
+    parser.add_argument(
+        "--stage1_pair_relevance_weight",
+        type=float,
+        default=0.0,
+        help=(
+            "Optional isolated Route-2 extension: upweight atom pairs whose teacher proximity varies "
+            "across the reference conformer ensemble. Zero preserves the legacy Stage-I objective."
+        ),
+    )
+    parser.add_argument(
+        "--stage1_pair_relevance_clip",
+        type=float,
+        default=4.0,
+        help="Maximum normalized conformer-variation weight before applying stage1_pair_relevance_weight.",
+    )
+    parser.add_argument(
+        "--stage1_nonlocal_heavy_only",
+        action="store_true",
+        help=(
+            "Restrict Stage-I probability supervision to heavy-atom pairs separated by at least "
+            "three bonds. This removes bonded and angle-local pairs that cannot distinguish conformers."
+        ),
+    )
+    parser.add_argument(
+        "--functional_teacher_mode",
+        type=str,
+        default="geometry",
+        choices=["geometry", "functional", "dual"],
+        help=(
+            "Optional Stage-I teacher extension. 'geometry' uses the original proximity teacher; "
+            "'functional' uses a geometry-preserving teacher enriched for intramolecular functional pairs; "
+            "'dual' jointly aligns to both teachers. This never uses protein structures."
+        ),
+    )
+    parser.add_argument(
+        "--functional_teacher_alpha",
+        type=float,
+        default=0.80,
+        help="Geometry retention in the functional-aware teacher Q_mix = alpha*Q_geom + (1-alpha)*Q_func.",
+    )
+    parser.add_argument(
+        "--lambda_stage1_geometry",
+        type=float,
+        default=1.0,
+        help="Geometry-teacher BCE weight when --functional_teacher_mode=dual.",
+    )
+    parser.add_argument(
+        "--lambda_stage1_functional",
+        type=float,
+        default=0.25,
+        help="Functional-aware teacher BCE weight when --functional_teacher_mode=dual.",
+    )
+    parser.add_argument(
+        "--functional_hbond_weight",
+        type=float,
+        default=0.35,
+        help="Relative teacher enrichment for intramolecular donor--acceptor atom pairs.",
+    )
+    parser.add_argument(
+        "--functional_aromatic_weight",
+        type=float,
+        default=0.20,
+        help="Relative teacher enrichment for aromatic atom pairs.",
+    )
+    parser.add_argument(
+        "--functional_hydrophobic_weight",
+        type=float,
+        default=0.10,
+        help="Relative teacher enrichment for hydrophobic atom pairs.",
+    )
+    parser.add_argument(
+        "--stage1_kernel",
+        type=str,
+        default="umap",
+        choices=[
+            "umap",
+            "gaussian",
+            "laplacian",
+            "student",
+            "cauchy",
+            "rational_quadratic",
+            "inverse_multiquadric",
+            "matern32",
+            "matern52",
+            "triweight",
+            "logistic",
+        ],
+        help=(
+            "Kernel used to build teacher Q(Y) in Stage 1. "
+            "umap uses learnable a,b; gaussian/laplacian/cauchy/matern/triweight/logistic use learnable bandwidth from a; "
+            "student and rational_quadratic use b as a learnable tail/shape parameter."
+        ),
+    )
+    parser.add_argument(
+        "--stage1_fixed_kernel_a",
+        type=float,
+        default=0.0,
+        help="Positive value fixes kernel parameter a during Stage-I and inference; zero keeps it learnable.",
+    )
+    parser.add_argument(
+        "--stage1_fixed_kernel_b",
+        type=float,
+        default=0.0,
+        help="Positive value fixes kernel parameter b during Stage-I and inference; zero keeps it learnable.",
+    )
+
+    # Learnable a,b regularization. Keep this: a,b are part of teacher Q construction.
+    parser.add_argument("--lambda_ab", type=float, default=1e-4)
+
+    # =========================================================
+    # Losses: Stage 2 coarse coordinate generation
+    # =========================================================
+    parser.add_argument("--lambda_coarse_rmsd", type=float, default=1.0)
+
+    parser.add_argument("--lambda_refine_rmsd", type=float, default=1.0)
+    parser.add_argument(
+        "--lambda_bond",
+        type=float,
+        default=0.0,
+        help="Stage 2 bonded-distance matching loss weight, using each reference conformer's bond distances.",
+    )
+    parser.add_argument(
+        "--lambda_angle",
+        type=float,
+        default=0.0,
+        help="Stage 2 lightweight angle regularization weight.",
+    )
+    parser.add_argument(
+        "--lambda_pair_dist",
+        type=float,
+        default=0.0,
+        help="Stage 2 all-pair heavy-atom distance matching weight for refined coordinates.",
+    )
+    parser.add_argument(
+        "--lambda_coarse_pair_dist",
+        type=float,
+        default=0.0,
+        help="Stage 2 all-pair heavy-atom distance matching weight for coarse decoder coordinates.",
+    )
+    parser.add_argument(
+        "--pair_dist_max_pairs",
+        type=int,
+        default=2048,
+        help="Maximum atom pairs used by the pairwise distance loss per conformer.",
+    )
+    parser.add_argument(
+        "--pair_dist_min_ref",
+        type=float,
+        default=0.0,
+        help="Minimum reference pair distance included in the pairwise distance loss.",
+    )
+    parser.add_argument(
+        "--pair_dist_max_ref",
+        type=float,
+        default=12.0,
+        help="Maximum reference pair distance included in the pairwise distance loss.",
+    )
+    parser.add_argument(
+        "--lambda_clash",
+        type=float,
+        default=0.0,
+        help="Stage 2 non-bonded steric clash penalty weight.",
+    )
+    parser.add_argument(
+        "--lambda_steric",
+        type=float,
+        default=0.0,
+        help="Backward-compatible alias for --lambda_clash.",
+    )
+
+    # Backward-compatible alias.
+    parser.add_argument("--lambda_rmsd", type=float, default=None)
+
+    # =========================================================
+    # Losses: Stage 2 manifold preservation
+    # =========================================================
+    parser.add_argument("--lambda_ce_manifold", type=float, default=0.1)
+    parser.add_argument(
+        "--lambda_set_recall",
+        type=float,
+        default=0.0,
+        help="Stage 2 set-level recall loss weight: each reference conformer should be covered by some generated mode.",
+    )
+    parser.add_argument(
+        "--lambda_set_precision",
+        type=float,
+        default=0.0,
+        help="Stage 2 set-level precision loss weight: each generated mode should stay near some reference conformer.",
+    )
+    parser.add_argument(
+        "--set_loss_clamp",
+        type=float,
+        default=10.0,
+        help="Clamp value for RMSD entries used by set-level coverage loss.",
+    )
+    parser.add_argument(
+        "--detach_refiner_input",
+        dest="detach_refiner_input",
+        action="store_true",
+        default=True,
+        help="Detach coarse coordinates before refiner, preserving the legacy Stage 2 behavior.",
+    )
+    parser.add_argument(
+        "--no_detach_refiner_input",
+        dest="detach_refiner_input",
+        action="store_false",
+        help="Allow refiner losses to backpropagate into the coarse coordinate decoder.",
+    )
+    parser.add_argument(
+        "--lambda_mode_diversity",
+        type=float,
+        default=0.0,
+        help="Stage 2 mode diversity weight. Encourages generated modes for the same molecule to spread without changing P_global.",
+    )
+    parser.add_argument(
+        "--use_end2end_pglobal",
+        action="store_true",
+        help="Stage 2: unfreeze GNN/P_global branch and fine-tune it with a smaller learning rate.",
+    )
+    parser.add_argument(
+        "--pglobal_lr",
+        type=float,
+        default=3e-5,
+        help="Learning rate for GNN/P_global branch when --use_end2end_pglobal is enabled.",
+    )
+    parser.add_argument(
+        "--use_mode_specific_p",
+        action="store_true",
+        help=(
+            "Stage 2: generate a mode-specific probability graph P_k for each latent mode. "
+            "This keeps Stage 1 P_global frozen and learns P_k = f(P_global, z_k) in the decoder stage."
+        ),
+    )
+    parser.add_argument(
+        "--lambda_mode_p_ce",
+        type=float,
+        default=0.0,
+        help="Stage 2 weight for CE(P_k, Q_i) when a generated mode is matched to a reference conformer.",
+    )
+    parser.add_argument(
+        "--mode_p_delta_scale",
+        type=float,
+        default=1.0,
+        help="Maximum logit-space residual scale used by the mode-specific probability generator.",
+    )
+    parser.add_argument(
+        "--mode_assignment_strategy",
+        type=str,
+        default="min",
+        choices=["min", "conf_id", "cluster"],
+        help=(
+            "Stage 2 mode-reference assignment. min keeps the old min-over-modes objective; "
+            "conf_id assigns each reference conformer to a stable mode by conf_id modulo num_modes; "
+            "cluster groups reference conformers by heavy-atom RMSD with farthest-point anchors."
+        ),
+    )
+    parser.add_argument(
+        "--deterministic_modes",
+        action="store_true",
+        help="Use deterministic one-hot mode codes z_k instead of random z during Stage 2 training and inference.",
+    )
+    parser.add_argument(
+        "--use_mode_pair_attention",
+        action="store_true",
+        help="Stage 2: enable mode-conditioned atom-pair attention in the coordinate decoder.",
+    )
+    parser.add_argument(
+        "--mode_attention_heads",
+        type=int,
+        default=4,
+        help="Number of heads for mode-conditioned pair attention.",
+    )
+    parser.add_argument(
+        "--mode_attention_dropout",
+        type=float,
+        default=0.0,
+        help="Dropout used inside mode-conditioned pair attention.",
+    )
+    parser.add_argument(
+        "--mode_diversity_margin",
+        type=float,
+        default=0.6,
+        help="Fallback heavy-atom RMSD margin for mode diversity when reference diversity is unavailable.",
+    )
+    parser.add_argument(
+        "--mode_diversity_ref_scale",
+        type=float,
+        default=0.7,
+        help="Scale applied to average reference-reference RMSD to set the diversity margin.",
+    )
+    parser.add_argument(
+        "--mode_diversity_max_pairs",
+        type=int,
+        default=15,
+        help="Maximum generated mode pairs used by the diversity loss per molecule.",
+    )
+    parser.add_argument(
+        "--lambda_signed_torsion",
+        type=float,
+        default=0.0,
+        help="Stage 2 signed torsion supervision weight. Penalizes mirror-like torsion flips.",
+    )
+    parser.add_argument(
+        "--use_torsion_residual",
+        action="store_true",
+        help="Stage 2: predict and apply learned residual rotations on rotatable torsions after coordinate refinement.",
+    )
+    parser.add_argument(
+        "--lambda_torsion_residual_rmsd",
+        type=float,
+        default=0.0,
+        help="Additional RMSD weight for the torsion-residual refined coordinates.",
+    )
+    parser.add_argument(
+        "--lambda_torsion_delta",
+        type=float,
+        default=0.0,
+        help="Supervise predicted torsion residuals against reference torsion angle deltas.",
+    )
+    parser.add_argument(
+        "--torsion_residual_max",
+        type=int,
+        default=32,
+        help="Maximum number of torsions updated by the learned torsion residual module.",
+    )
+    parser.add_argument(
+        "--torsion_residual_scale",
+        type=float,
+        default=1.0,
+        help="Maximum absolute learned torsion residual in radians.",
+    )
+    parser.add_argument(
+        "--lambda_chiral_volume",
+        type=float,
+        default=0.0,
+        help="Stage 2 signed local volume supervision weight. Penalizes mirror-like local handedness flips.",
+    )
+    parser.add_argument("--signed_torsion_max", type=int, default=64)
+    parser.add_argument("--chiral_volume_max_centers", type=int, default=64)
+    parser.add_argument("--chiral_volume_ref_abs_min", type=float, default=0.05)
+
+    # =========================================================
+    # Evaluation / inference helpers
+    # =========================================================
+    parser.add_argument("--eval_num_samples", type=int, default=100)
+
+    # =========================================================
+    # Visualization
+    # =========================================================
+    parser.add_argument("--smiles", type=str, default=None)
+    parser.add_argument("--get_image", action="store_true")
+
+    args = parser.parse_args()
+
+    # =========================================================
+    # Compatibility normalization
+    # =========================================================
+    if args.latent_dim is not None:
+        args.mode_dim = int(args.latent_dim)
+    else:
+        args.latent_dim = int(args.mode_dim)
+
+    if args.refine_steps is not None:
+        args.num_refine_steps = int(args.refine_steps)
+    else:
+        args.refine_steps = int(args.num_refine_steps)
+
+    if args.lambda_rmsd is not None:
+        args.lambda_refine_rmsd = float(args.lambda_rmsd)
+    else:
+        args.lambda_rmsd = float(args.lambda_refine_rmsd)
+
+    args.pos_w = float(args.lambda_refine_rmsd)
+    args.inv_alpha = 0.0
+    args.model_type = "GNNEncoder"
+    args.use_pos_loss = True
+
+    return args