diff --git a/README.md b/README.md index 3092118..9ecad5c 100644 --- a/README.md +++ b/README.md @@ -434,7 +434,6 @@ uv run python -m scripts.inference \ | `--save_gifs` | `false` | Save trajectory GIFs (slower) | | `--threshold` | `1.0` | Distance threshold for precision/recall (A) | | `--water_ratio` | `None` | Sample `num_residues * ratio` waters (if not set, uses ground truth count) | -| `--use_sc` | `false` | Use self-conditioning during integration | > **`--water_ratio` counts mate residues too.** `num_residues` covers ASU *and* mate > residues, so `--include_mates` emits ~1.7x more waters at the same ratio (~440 vs diff --git a/diffuse.yaml b/diffuse.yaml index 6a8aaf6..c5ab137 100644 --- a/diffuse.yaml +++ b/diffuse.yaml @@ -51,9 +51,6 @@ profiles: - key: num_workers type: number default: 4 - - key: use_self_cond - type: boolean - default: true - key: run_name type: text - key: train_list @@ -92,8 +89,7 @@ profiles: save_dir: --save_dir run_name: --run_name num_workers: --num_workers - boolean_args: - use_self_cond: --use_self_cond + boolean_args: {} static_args: [--base_pdb_dir, /data/pdb] - slug: waterflow-train-esm @@ -132,9 +128,6 @@ profiles: - key: num_workers type: number default: 4 - - key: use_self_cond - type: boolean - default: true - key: run_name type: text - key: train_list @@ -173,8 +166,7 @@ profiles: save_dir: --save_dir run_name: --run_name num_workers: --num_workers - boolean_args: - use_self_cond: --use_self_cond + boolean_args: {} static_args: [--base_pdb_dir, /data/pdb] - slug: waterflow-train-slae @@ -213,9 +205,6 @@ profiles: - key: num_workers type: number default: 4 - - key: use_self_cond - type: boolean - default: true - key: run_name type: text - key: train_list @@ -254,8 +243,7 @@ profiles: save_dir: --save_dir run_name: --run_name num_workers: --num_workers - boolean_args: - use_self_cond: --use_self_cond + boolean_args: {} static_args: [--base_pdb_dir, /data/pdb] - slug: waterflow-inference @@ -280,9 +268,6 @@ profiles: - key: method type: text default: rk4 - - key: use_sc - type: boolean - default: true run_config_defaults: shared_memory_size: "16Gi" runtime_class: nvidia @@ -306,8 +291,7 @@ profiles: pdb_list: --pdb_list method: --method num_steps: --num_steps - boolean_args: - use_sc: --use_sc + boolean_args: {} static_args: [--output_dir, /data/outputs, --base_pdb_dir, /data/pdb] - slug: waterflow-generate-esm diff --git a/scripts/inference.py b/scripts/inference.py index c3d5017..8f18ecb 100644 --- a/scripts/inference.py +++ b/scripts/inference.py @@ -132,11 +132,6 @@ def parse_args(): default=100, help="Number of integration steps (default: 100)", ) - p.add_argument( - "--use_sc", - action="store_true", - help="Use self-conditioning during integration", - ) p.add_argument( "--save_gifs", @@ -321,7 +316,6 @@ def run_inference_batch( graphs: list, method: str, num_steps: int, - use_sc: bool, device: str, water_ratio: float = None, ) -> list: @@ -333,7 +327,6 @@ def run_inference_batch( graphs: List of HeteroData graphs method: Integration method ('euler' or 'rk4') num_steps: Number of integration steps - use_sc: Whether to use self-conditioning device: Device to run on water_ratio: If provided, sample num_residues * water_ratio waters @@ -347,7 +340,6 @@ def run_inference_batch( results = flow_matcher.rk4_integrate( graphs, num_steps=num_steps, - use_sc=use_sc, device=device, return_trajectory=True, water_ratio=water_ratio, @@ -356,7 +348,6 @@ def run_inference_batch( results = flow_matcher.euler_integrate( graphs, num_steps=num_steps, - use_sc=use_sc, device=device, water_ratio=water_ratio, ) @@ -441,7 +432,6 @@ def main(): # Create FlowMatcher flow_matcher = FlowMatcher( model=model, - p_self_cond=config.get("p_self_cond", 0.5), sampling_strategy=config.get("sampling_strategy", "uniform_ball"), ) @@ -476,7 +466,6 @@ def main(): # run inference logger.info(f"Running inference with method={args.method}, steps={args.num_steps}") - logger.info(f"Self-conditioning: {args.use_sc}") logger.info(f"Threshold for metrics: {args.threshold}Å") logger.info(f"Batch size: {args.batch_size}") @@ -526,7 +515,6 @@ def main(): batch_graphs, method=args.method, num_steps=args.num_steps, - use_sc=args.use_sc, device=args.device, water_ratio=args.water_ratio, ) @@ -627,7 +615,6 @@ def main(): "checkpoint": args.checkpoint, "method": args.method, "num_steps": args.num_steps, - "use_sc": args.use_sc, "threshold": args.threshold, "include_mates": include_mates, "water_ratio": args.water_ratio, diff --git a/scripts/train.py b/scripts/train.py index ec7b0dc..be4ffae 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -385,14 +385,6 @@ def parse_args(): ) p.add_argument("--step_gamma", type=float, default=0.5, help="StepLR gamma") - # flow matching - p.add_argument("--use_self_cond", action="store_true") - p.add_argument("--p_self_cond", type=float, default=0.5) - p.add_argument("--use_distortion", action="store_true") - p.add_argument("--p_distort", type=float, default=0.2) - p.add_argument("--t_distort", type=float, default=0.5) - p.add_argument("--sigma_distort", type=float, default=0.5) - # checkpointing p.add_argument("--save_dir", type=str, default="/home/srivasv/flow_checkpoints") p.add_argument( @@ -703,7 +695,6 @@ def run_eval_sampling( out = flow_matcher.rk4_integrate( graph, num_steps=args.rk4_steps, - use_sc=args.use_self_cond, device=device, return_trajectory=True, )[0] # rk4_integrate returns a list, get the single result @@ -817,7 +808,6 @@ def train_epoch( with flow_matcher.model.no_sync() if no_sync else contextlib.nullcontext(): metrics = flow_matcher.training_step( batch, - use_self_conditioning=args.use_self_cond, accumulation_steps=args.grad_accum_steps, ) @@ -1294,12 +1284,7 @@ def main(): flow_matcher = FlowMatcher( model=model, - p_self_cond=args.p_self_cond, sampling_strategy=args.sampling_strategy, - use_distortion=args.use_distortion, - p_distort=args.p_distort, - t_distort=args.t_distort, - sigma_distort=args.sigma_distort, ) optimizer = AdamW( diff --git a/src/flow.py b/src/flow.py index 7253ec5..b5250ad 100644 --- a/src/flow.py +++ b/src/flow.py @@ -810,19 +810,6 @@ def __init__( k_wp=k_wp, ) - self.sc_vec_encoder = GVP( - in_dims=(0, 1), - out_dims=(0, v_h), - activations=(None, None), - vector_gate=True, - ) - - self.sc_sca_encoder = nn.Sequential( - nn.Linear(1, s_h), - nn.GELU(), - nn.LayerNorm(s_h), - ) - # Water vector field head: project (s_h, v_h) -> (s_h // 4, 1) -> single vector channel # NOTE: vector_gate=True requires scalar input features. GVP gating works by # computing gate values from scalars via a learned linear map, then applying @@ -837,7 +824,6 @@ def forward( self, data: HeteroData, t: torch.Tensor, - self_cond: dict[str, torch.Tensor] | None = None, ) -> torch.Tensor: """ Predict velocity field for water nodes given protein context and time. @@ -851,8 +837,6 @@ def forward( positions: (N_w, 3) Cartesian coordinates features: (N_w, 16) element one-hot encoding t: (B,) flow time per complex in batch, values in [0, 1] - self_cond: Optional self-conditioning dict with 'x1_pred' key containing - previous prediction (N_w, 3) for iterative refinement Returns: (N_w, 3) predicted velocity vector field at each water node @@ -886,25 +870,6 @@ def forward( device=device, ) - # self conditioning - if ( - self_cond is not None - and ("x1_pred" in self_cond) - and self_cond["x1_pred"] is not None - ): - delta = self_cond["x1_pred"] - data["water"].pos - delta_vec = delta.unsqueeze(1) - - # vector conditioning (equivariant) - s_empty = torch.empty(data["water"].num_nodes, 0, device=device) - _, v_sc = self.sc_vec_encoder((s_empty, delta_vec)) - v_w = v_w + v_sc - - # scalar conditioning (invariant) on ||delta|| - d_mag = delta.norm(dim=-1, keepdim=True) - s_sc = self.sc_sca_encoder(d_mag) - s_w = s_w + s_sc - # build hetero feature dict for GVP multi-edge updates x_dict = { "protein": (s_p, v_p_latent), @@ -934,11 +899,6 @@ class FlowMatcher: def __init__( self, model, - p_self_cond: float = 0.5, - use_distortion: bool = False, - p_distort: float = 0.2, - t_distort: float = 0.5, - sigma_distort: float = 0.5, loss_eps: float = 1e-3, sampling_strategy: str = "uniform_ball", ): @@ -947,11 +907,6 @@ def __init__( Args: model: FlowWaterGVP model instance - p_self_cond: Probability of using self-conditioning during training - use_distortion: Whether to apply late-stage path distortion - p_distort: Probability of applying distortion per sample - t_distort: Time threshold after which distortion may be applied - sigma_distort: Standard deviation of distortion noise loss_eps: Small constant for numerical stability in loss weighting sampling_strategy: Source distribution for flow matching noise. "uniform_ball" samples uniformly in balls around protein atoms. @@ -967,11 +922,6 @@ def __init__( f"got '{sampling_strategy}'" ) self.model = model - self.p_self_cond = p_self_cond - self.use_distortion = use_distortion - self.p_distort = p_distort - self.t_distort = t_distort - self.sigma_distort = sigma_distort self.loss_eps = loss_eps # Read the graph cutoff off the flow model. Under DDP the attribute lives # on the wrapped `.module` (DDP does not forward attribute lookups), so @@ -1121,7 +1071,6 @@ def compute_sigma_per_graph( def training_step( self, batch: HeteroData, - use_self_conditioning: bool = True, accumulation_steps: int = 1, ) -> dict[str, object]: """ @@ -1131,7 +1080,6 @@ def training_step( Args: batch: HeteroData batch - use_self_conditioning: Whether to use self-conditioning accumulation_steps: Number of gradient accumulation steps (loss is scaled by 1/accumulation_steps) Returns: @@ -1174,26 +1122,9 @@ def training_step( x_t = (1.0 - t_per_atom) * x0_star + t_per_atom * x1_star - # late stage path distortion - if self.use_distortion: - indicator = (t_per_atom >= self.t_distort).float() - if indicator.any(): - mask = (torch.rand_like(t_per_atom) < self.p_distort).float() - eps = torch.randn_like(x_t) * self.sigma_distort - x_t = x_t + indicator * mask * eps - - # self-conditioning - self_cond = None - if use_self_conditioning and torch.rand(1).item() < self.p_self_cond: - with torch.no_grad(): - batch["water"].pos = x_t - v_pred_sc = self.model(batch, t, self_cond=None) - x1_pred_sc = x_t + (1.0 - t_per_atom) * v_pred_sc - self_cond = {"x1_pred": x1_pred_sc} - # forward pass batch["water"].pos = x_t - v_pred = self.model(batch, t, self_cond=self_cond) + v_pred = self.model(batch, t) # target velocity v_target = x1_star - x0_star @@ -1267,7 +1198,7 @@ def validation_step(self, batch: HeteroData) -> dict[str, float]: x_t = (1.0 - t_per_atom) * x0_star + t_per_atom * x1_star batch["water"].pos = x_t - v_pred = self.model(batch, t, self_cond=None) + v_pred = self.model(batch, t) v_target = x1_star - x0_star @@ -1416,8 +1347,6 @@ def euler_integrate( self, graphs: HeteroData | list[HeteroData], num_steps: int = 100, - use_sc: bool = True, - sc_ema_alpha: float = 0.2, device: str | torch.device = "cuda", water_ratio: float | None = None, water_count: int | None = None, @@ -1428,8 +1357,6 @@ def euler_integrate( Args: graphs: Single HeteroData or list of HeteroData graphs to process num_steps: Number of integration steps - use_sc: Whether to use self-conditioning - sc_ema_alpha: EMA decay for self-conditioning device: Device to run on water_ratio: If provided, sample num_residues * water_ratio waters instead of using ground truth water count. Ignored when @@ -1468,8 +1395,6 @@ def euler_integrate( x, batch_w = self._setup_water_nodes(g, water_ratio, water_count, device) - x1_pred_ema = x.clone() - ts = torch.linspace(0, 1, num_steps, device=device) dt = ts[1] - ts[0] @@ -1478,20 +1403,9 @@ def euler_integrate( t = t_scalar.expand(num_graphs) # (num_graphs,) all same value g["water"].pos = x - self_cond = {"x1_pred": x1_pred_ema} if use_sc else None - v = self.model(g, t, self_cond=self_cond) + v = self.model(g, t) x = x + dt * v - if use_sc: - t_next_scalar = ts[i + 1] - t_next = t_next_scalar.expand(num_graphs) - g["water"].pos = x - v_next = self.model(g, t_next, self_cond={"x1_pred": x1_pred_ema}) - x1_pred_now = x + (1.0 - t_next_scalar) * v_next - x1_pred_ema = ( - 1.0 - sc_ema_alpha - ) * x1_pred_ema + sc_ema_alpha * x1_pred_now - # split results by graph x_cpu = x.detach().cpu() protein_pos_cpu = g["protein"].pos.detach().cpu() @@ -1521,8 +1435,6 @@ def rk4_integrate( self, graphs: HeteroData | list[HeteroData], num_steps: int = 500, - use_sc: bool = True, - sc_ema_alpha: float = 0.2, device: str | torch.device = "cuda", return_trajectory: bool = True, water_ratio: float | None = None, @@ -1534,8 +1446,6 @@ def rk4_integrate( Args: graphs: Single HeteroData or list of HeteroData graphs to process num_steps: Number of integration steps - use_sc: Whether to use self-conditioning - sc_ema_alpha: EMA decay for self-conditioning device: Device to run on return_trajectory: Whether to return full trajectory and metrics water_ratio: If provided, sample num_residues * water_ratio waters @@ -1575,8 +1485,6 @@ def rk4_integrate( x, batch_w = self._setup_water_nodes(g, water_ratio, water_count, device) - x1_pred_ema = x.clone() - ts = torch.linspace(0, 1, num_steps, device=device) dt = ts[1] - ts[0] @@ -1596,8 +1504,7 @@ def rk4_integrate( def f(xpos, t_tensor): g["water"].pos = xpos - self_cond = {"x1_pred": x1_pred_ema} if use_sc else None - return self.model(g, t_tensor, self_cond=self_cond) + return self.model(g, t_tensor) k1 = f(x, t0) k2 = f(x + 0.5 * dt * k1, (t0_scalar + 0.5 * dt).expand(num_graphs)) @@ -1606,16 +1513,6 @@ def f(xpos, t_tensor): x = x + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4) - if use_sc: - t1_scalar = ts[step + 1] - t1 = t1_scalar.expand(num_graphs) - g["water"].pos = x - v_next = self.model(g, t1, self_cond={"x1_pred": x1_pred_ema}) - x1_pred_now = x + (1.0 - t1_scalar) * v_next - x1_pred_ema = ( - 1.0 - sc_ema_alpha - ) * x1_pred_ema + sc_ema_alpha * x1_pred_now - if return_trajectory: x_cpu = x.detach().cpu() for i in range(num_graphs): @@ -1655,7 +1552,6 @@ def sample( graphs: HeteroData | list[HeteroData], num_steps: int = 100, method: str = "euler", - use_sc: bool = True, device: str = "cuda", ) -> np.ndarray | list[np.ndarray]: """ @@ -1665,7 +1561,6 @@ def sample( graphs: Single HeteroData or list of HeteroData graphs num_steps: Number of integration steps method: 'euler' or 'rk4' - use_sc: Whether to use self-conditioning device: Device to run on Returns: @@ -1675,11 +1570,11 @@ def sample( single_input = isinstance(graphs, HeteroData) if method == "euler": - results = self.euler_integrate(graphs, num_steps, use_sc, device=device) + results = self.euler_integrate(graphs, num_steps, device=device) results = [r["water_pred"] for r in results] elif method == "rk4": results = self.rk4_integrate( - graphs, num_steps, use_sc, device=device, return_trajectory=False + graphs, num_steps, device=device, return_trajectory=False ) results = [r["water_pred"] for r in results] else: diff --git a/tests/test_flow.py b/tests/test_flow.py index c40f45d..0202721 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -621,21 +621,6 @@ def test_forward_no_water(self, device): assert v_pred.shape == (0, 3) - def test_self_conditioning(self, simple_hetero_data, device, gvp_encoder): - model = FlowWaterGVP( - encoder=gvp_encoder, - hidden_dims=(64, 8), - layers=1, - ).to(device) - - n_water = simple_hetero_data["water"].num_nodes - sc = {"x1_pred": torch.randn(n_water, 3, device=device)} - t = torch.tensor([0.5], device=device) - - v_pred = model(simple_hetero_data, t, self_cond=sc) - - assert v_pred.shape == (n_water, 3) - # ============== Tests for FlowMatcher ============== @@ -650,7 +635,7 @@ def flow_matcher(self, device, gvp_encoder): layers=1, ).to(device) - return FlowMatcher(model, p_self_cond=0.5) + return FlowMatcher(model) def test_compute_sigma(self, simple_hetero_data): sigma = FlowMatcher.compute_sigma(simple_hetero_data) @@ -672,9 +657,7 @@ def test_training_step(self, flow_matcher, simple_hetero_data, device): optimizer = torch.optim.Adam(flow_matcher.model.parameters(), lr=1e-4) optimizer.zero_grad() - result = flow_matcher.training_step( - simple_hetero_data, use_self_conditioning=False - ) + result = flow_matcher.training_step(simple_hetero_data) optimizer.step() assert "loss" in result @@ -682,21 +665,6 @@ def test_training_step(self, flow_matcher, simple_hetero_data, device): assert "sigma" in result assert result["loss"] >= 0 - def test_training_step_with_self_cond( - self, flow_matcher, simple_hetero_data, device - ): - optimizer = torch.optim.Adam(flow_matcher.model.parameters(), lr=1e-4) - - # Force self-conditioning - flow_matcher.p_self_cond = 1.0 - optimizer.zero_grad() - result = flow_matcher.training_step( - simple_hetero_data, use_self_conditioning=True - ) - optimizer.step() - - assert "loss" in result - def test_validation_step(self, flow_matcher, simple_hetero_data): result = flow_matcher.validation_step(simple_hetero_data) @@ -726,7 +694,7 @@ def test_edge_config_propagates_to_updater(self, device, gvp_encoder): @pytest.mark.slow def test_euler_integrate(self, flow_matcher, simple_hetero_data, device): results = flow_matcher.euler_integrate( - simple_hetero_data, num_steps=5, use_sc=False, device=str(device) + simple_hetero_data, num_steps=5, device=str(device) ) # euler_integrate returns List[Dict], one per input graph result = results[0] @@ -745,7 +713,6 @@ def test_rk4_integrate(self, flow_matcher, simple_hetero_data, device): results = flow_matcher.rk4_integrate( simple_hetero_data, num_steps=5, - use_sc=False, device=str(device), return_trajectory=True, ) @@ -1300,38 +1267,6 @@ def test_negative_water_count_raises(self, device): fm._setup_water_nodes_from_count(g, -1, device) -# ============== Tests for distortion ============== - - -@pytest.mark.unit -class TestDistortion: - def test_distortion_enabled(self, device): - base_encoder = ProteinGVPEncoder( - node_scalar_in=16, - hidden_dims=(64, 8), - n_edge_scalar_in=16, - pool_residue=False, - ).to(device) - encoder = GVPEncoder(encoder=base_encoder, freeze=False) - - model = FlowWaterGVP( - encoder=encoder, - hidden_dims=(64, 8), - layers=1, - ).to(device) - - fm = FlowMatcher( - model, - use_distortion=True, - p_distort=1.0, # Always apply - t_distort=0.0, # Apply at all times - sigma_distort=0.5, - ) - - assert fm.use_distortion is True - assert fm.p_distort == 1.0 - - # ============== Edge case tests ============== diff --git a/tests/test_forward.py b/tests/test_forward.py index 52b3598..e8236d6 100644 --- a/tests/test_forward.py +++ b/tests/test_forward.py @@ -267,8 +267,6 @@ def test_training_step_no_nan_tripwire(device): fm = FlowMatcher( model=model, - p_self_cond=0.0, # simpler/cleaner for tripwire - use_distortion=False, loss_eps=1e-3, ) @@ -282,7 +280,7 @@ def test_training_step_no_nan_tripwire(device): for step in range(5): opt.zero_grad() - out = fm.training_step(data, use_self_conditioning=False) + out = fm.training_step(data) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() @@ -543,13 +541,12 @@ def test_integration_trajectory_length(self, device): k_ww=8, ).to(device) - fm = FlowMatcher(model, p_self_cond=0.0) + fm = FlowMatcher(model) num_steps = 20 results = fm.rk4_integrate( data, num_steps=num_steps, - use_sc=False, device=str(device), return_trajectory=True, ) @@ -610,7 +607,7 @@ def test_velocity_field_finite(self, device, gvp_encoder): # Test at multiple t values for t_val in [0.0, 0.25, 0.5, 0.75, 1.0]: t = torch.tensor([t_val], device=device) - v_pred = model(data, t, self_cond=None) + v_pred = model(data, t) assert torch.isfinite(v_pred).all(), f"Velocity has NaN/Inf at t={t_val}" @@ -637,8 +634,8 @@ def test_velocity_field_changes_with_t(self, device, gvp_encoder): t0 = torch.tensor([0.1], device=device) t1 = torch.tensor([0.9], device=device) - v0 = model(data, t0, self_cond=None) - v1 = model(data, t1, self_cond=None) + v0 = model(data, t0) + v1 = model(data, t1) # Velocities should be different - check that MOST individual waters show change per_water_diff = torch.norm(v0 - v1, dim=-1) # per-water norm