Gradient-Based Guidance
=======================

The surrogate model enables gradient-based optimization of molecular latent vectors
toward desired properties.

Basic Guidance
--------------

.. code-block:: python

   from moltenflow.guidance.guidance import compute_guidance
   import torch.nn as nn

   # Define target properties
   target = torch.tensor([[0.5, 2.0]])  # Desired property values
   loss_fn = nn.MSELoss()

   # Compute gradient
   grad = compute_guidance(z, target, model, loss_fn)

   # Update latent
   z_new = z - 0.1 * grad  # Gradient descent step

Iterative Optimization
----------------------

Optimize latents over multiple steps:

.. code-block:: python

   z_opt = z.clone()
   n_steps = 50
   step_size = 0.1

   for step in range(n_steps):
       grad = compute_guidance(z_opt, target, model, loss_fn)
       z_opt = z_opt - step_size * grad

       # Monitor progress
       with torch.no_grad():
           pred = model(z_opt)
           print(f"Step {step}: Predicted = {pred.numpy()}")

With Conditions
---------------

Include fixed conditions during guidance:

.. code-block:: python

   # Fixed conditions (e.g., 300K, 1 bar)
   c_fixed = torch.tensor([[300.0, 1.0]])
   c_scaled = torch.from_numpy(cond_scaler.transform(c_fixed.numpy()))

   grad = compute_guidance(z_opt, target, model, loss_fn, c=c_scaled)

Multi-Objective Optimization
----------------------------

For joint models predicting multiple properties:

.. code-block:: python

   # Target both properties simultaneously
   target = torch.tensor([[0.3, 1.5]])  # [CO2_solubility, ln_viscosity]

   # Optimization will balance both objectives
   grad = compute_guidance(z_opt, target, model_joint, loss_fn)

Limitations
-----------

- **Fingerprint Decoding**: RDKit fingerprints cannot be decoded back to SMILES.
  Use a VAE for end-to-end molecule generation.
- **Local Optima**: Gradient descent may converge to local minima.
- **Valid Chemistry**: No guarantee optimized latents correspond to valid molecules.
