Update defog with global pbar and data preprocessing

Author: liugangcode

tests/generator/defog.py

@@ -4,7 +4,7 @@
 
 from torch_molecule import DeFoGMolecularGenerator
 
-EPOCHS = 2
+EPOCHS = 10
 BATCH_SIZE = 32
 
 def test_defog_generator():
@@ -20,32 +20,43 @@ def test_defog_generator():
         'CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F'
     ]
     smiles_list = smiles_list * 25  # Create 100 molecules for training
-    properties = [0, 0, 1, 1] * 25  # Create 100 properties for training
+    
+    # Multi-dimensional properties: each row is a molecule, each column is a property
+    # Properties range from 0 to 1
+    np.random.seed(42)  # For reproducible results
+    properties = np.random.rand(100, 3)  # 100 molecules, 3 properties each
 
-    # 1. Conditional Model Testing
-    print("\n=== Testing Conditional DeFoG Model ===")
+    # 1. Multi-Conditional Model Testing
+    print("\n=== Testing Multi-Conditional DeFoG Model ===")
     conditional_model = DeFoGMolecularGenerator(
-        task_type=['regression'],
+        task_type=['regression', 'regression', 'regression'],  # 3 regression tasks
         epochs=EPOCHS,
         batch_size=BATCH_SIZE,
         learning_rate=5e-4,
         sample_steps=10,  # Fewer steps for faster testing
         guidance_weight=0.2,
         verbose=True,
     )
-    print("Conditional DeFoG Model initialized successfully.")
+    print("Multi-Conditional DeFoG Model initialized successfully.")
     print(f"Input dim y: {conditional_model.input_dim_y}")
 
-    print("\n--- Fitting conditional model ---")
+    print("\n--- Fitting multi-conditional model ---")
     conditional_model.fit(smiles_list, properties)
-    print("Conditional DeFoG Model fitting completed.")
-
-    print("\n--- Testing conditional generation ---")
-    target_properties = [[0], [0], [1], [1]]
+    print("Multi-Conditional DeFoG Model fitting completed.")
+
+    print("\n--- Testing multi-conditional generation ---")
+    # Generate molecules with specific multi-dimensional properties
+    target_properties = [
+        [0.1, 0.2, 0.3],  # Low values for all properties
+        [0.4, 0.5, 0.6],  # Medium values for all properties
+        [0.7, 0.8, 0.9],  # High values for all properties
+        [0.9, 0.1, 0.5]   # Mixed values
+    ]
     generated_smiles = conditional_model.generate(labels=target_properties)
-    print(f"Conditionally generated {len(generated_smiles)} molecules.")
+    print(f"Multi-conditionally generated {len(generated_smiles)} molecules.")
     assert len(generated_smiles) == len(target_properties)
     print("Example SMILES:", generated_smiles[:2])
+    print("Target properties for first molecule:", target_properties[0])
 
     print("\n--- Testing model saving and loading ---")
     save_path = "conditional_defog_test_model.pt"
@@ -70,7 +81,29 @@ def test_defog_generator():
         os.remove(save_path)
         print(f"Cleaned up {save_path}")
 
-    # 2. Unconditional Model Testing
+    # 2. Single-property conditional testing (backwards compatibility)
+    print("\n=== Testing Single-Property Conditional DeFoG Model ===")
+    single_properties = properties[:, 0:1]  # Use only first property
+    single_conditional_model = DeFoGMolecularGenerator(
+        task_type=['regression'],  # Single regression task
+        epochs=EPOCHS,
+        batch_size=BATCH_SIZE,
+        learning_rate=5e-4,
+        sample_steps=10,
+        guidance_weight=0.2,
+        verbose=True,
+    )
+    print("Single-Property Conditional DeFoG Model initialized successfully.")
+    
+    single_conditional_model.fit(smiles_list, single_properties)
+    print("Single-Property DeFoG Model fitting completed.")
+    
+    single_target_properties = [[0.2], [0.5], [0.8], [0.1]]
+    single_generated_smiles = single_conditional_model.generate(labels=single_target_properties)
+    print(f"Single-conditionally generated {len(single_generated_smiles)} molecules.")
+    assert len(single_generated_smiles) == len(single_target_properties)
+
+    # 3. Unconditional Model Testing
     print("\n=== Testing Unconditional DeFoG Model ===")
     unconditional_model = DeFoGMolecularGenerator(
         task_type=[],  # Empty task_type for unconditional generation

torch_molecule/generator/defog/modeling_defog.py

@@ -33,10 +33,10 @@ class DeFoGMolecularGenerator(BaseMolecularGenerator):
     num_layer : int, default=6
         Number of transformer layers
     hidden_mlp_dims : Dict[str, int], default={'X': 256, 'E': 128, 'y': 128}
-        Hidden dimensions for MLP layers in X, E, and y components
+        Hidden dimensions for MLP layers in X (node dim), E (edge dim), and y (property dim) components
     hidden_dims : Dict[str, Any], default={'dx': 256, 'de': 64, 'dy': 64, 'n_head': 8, 'dim_ffX': 256, 'dim_ffE': 128, 'dim_ffy': 128}
         Hidden dimensions for transformer components including attention heads and feed-forward layers
-        Keys: 'dx', 'de', 'dy', 'n_head', 'dim_ffX', 'dim_ffE', 'dim_ffy'
+        Keys: 'dx' (node dim), 'de' (edge dim), 'dy' (property dim), 'n_head' (number of attention heads), 'dim_ffX' (feed-forward dim for node features), 'dim_ffE' (feed-forward dim for edge features), 'dim_ffy' (feed-forward dim for property features)
     transition : str, default='marginal'
         Transition type for flow matching.
         Options: 'marginal', 'absorbing', 'uniform', 'absorbfirst', 'argmax', 'edge_marginal', 'node_marginal'
@@ -239,8 +239,6 @@ def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
 
         return optimizer, scheduler
 
-
-    
     def _convert_to_pytorch_data(self, X, y=None):
         """Convert numpy arrays to PyTorch Geometric data format."""
         if self.verbose:
@@ -256,7 +254,9 @@ def _convert_to_pytorch_data(self, X, y=None):
             g = Data()
 
             node_type = torch.from_numpy(graph['node_feat'][:, 0] - 1)
-            
+            if node_type.numel() <= 1:
+                continue
+
             valid_mask = node_type >= 0
             if not valid_mask.all():
                 # Get valid nodes and adjust edge indices
@@ -398,24 +398,30 @@ def fit(self, X_train: List[str], y_train: Optional[Union[List, np.ndarray]] = N
         train_dataset = self._convert_to_pytorch_data(X_train, y_train)
         train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
 
+        # Calculate total steps for global progress bar
+        total_steps = self.epochs * len(train_loader)
+        global_progress = tqdm(total=total_steps, desc="Training Progress", leave=True) if self.verbose else None
+
         self.fitting_loss = []
         for epoch in range(self.epochs):
-            train_losses = self._train_epoch(train_loader, optimizer, epoch)
+            train_losses = self._train_epoch(train_loader, optimizer, epoch, global_progress)
             avg_loss = np.mean(train_losses)
             self.fitting_loss.append(avg_loss)
             if scheduler:
                 scheduler.step(avg_loss)
 
+        if global_progress:
+            global_progress.close()
+        
         self.is_fitted_ = True
         return self
 
-    def _train_epoch(self, train_loader, optimizer, epoch):
+    def _train_epoch(self, train_loader, optimizer, epoch, global_progress=None):
         self.model.train()
         losses = []
-        iterator = tqdm(train_loader, desc=f"Epoch {epoch}", leave=False) if self.verbose else train_loader
 
         active_index = self.dataset_info["active_index"]
-        for batched_data in iterator:
+        for batched_data in train_loader:
             batched_data = batched_data.to(self.device)
             optimizer.zero_grad()
 
@@ -467,8 +473,17 @@ def _train_epoch(self, train_loader, optimizer, epoch):
             optimizer.step()
 
             losses.append(loss.item())
-            if self.verbose:
-                iterator.set_postfix({"Loss": f"{loss.item():.4f}"})
+            
+            # Update global progress bar
+            if global_progress:
+                global_progress.set_postfix({
+                    "Epoch": f"{epoch+1}",
+                    "Loss": f"{loss.item():.4f}",
+                    "Loss_X": f"{masked_loss_X.item():.4f}",
+                    "Loss_E": f"{masked_loss_E.item():.4f}",
+                    "Loss_y": f"{loss_y.item():.4f}"
+                })
+                global_progress.update(1)
 
         return losses