Skip to content

refactor(models): remove unused x_init #330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
marcuscollins merged 2 commits into from
Jul 21, 2026
Merged

refactor(models): remove unused x_init #330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1b7fa53
refactor(models): remove unused x_init
xraymemory Jul 20, 2026
0f53cca
test(models): cover batched prior steps
xraymemory Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -397,8 +397,7 @@ class MyModelWrapper:
atom_array = structure["asym_unit"]
# ... model-specific featurization ...
conditioning = my_model_features(atom_array)
x_init = torch.zeros(n_atoms, 3, device=self.device)
return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
return GenerativeModelInput(conditioning=conditioning)

def step(
self,
Expand All @@ -418,7 +417,7 @@ class MyModelWrapper:
shape: tuple[int, ...] | None = None,
) -> Tensor:
"""Sample from the prior (typically Gaussian noise)."""
n_atoms = features.x_init.shape[-2] if features else shape[-2]
n_atoms = features.conditioning.num_atoms if features else shape[-2]
return torch.randn(batch_size, n_atoms, 3, device=self.device)
```

Expand Down
86 changes: 10 additions & 76 deletions src/sampleworks/models/boltz/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,16 +314,13 @@ class BoltzConfig:
Output directory for processed Boltz files.
num_workers : int
Number of parallel workers for preprocessing.
ensemble_size : int
Number of samples to generate (batch dimension of x_init).
recycling_steps : int
Number of recycling steps to perform during featurization Pairformer pass.

"""

out_dir: str | Path | None = None
num_workers: int = 0
ensemble_size: int = 1
recycling_steps: int = 3


Expand All @@ -332,7 +329,6 @@ def process_structure_for_boltz(
*,
out_dir: str | Path | None = None,
num_workers: int = 0,
ensemble_size: int = 1,
recycling_steps: int | None = 3,
) -> dict:
"""Annotate an Atomworks structure with Boltz-specific configuration.
Expand All @@ -346,8 +342,6 @@ def process_structure_for_boltz(
Defaults to structure metadata ID or "boltz_output".
num_workers : int
Number of parallel workers for preprocessing.
ensemble_size : int
Number of samples to generate (batch dimension of x_init).
recycling_steps : int | None
Number of recycling steps to perform during featurization Pairformer pass.
Will set to 3 if None.
Expand All @@ -372,7 +366,6 @@ def process_structure_for_boltz(
config = BoltzConfig(
out_dir=out_dir or structure.get("metadata", {}).get("id", "boltz_output"),
num_workers=num_workers,
ensemble_size=ensemble_size,
recycling_steps=recycling_steps,
)
return {**structure, "_boltz_config": config}
Expand Down Expand Up @@ -649,7 +642,7 @@ def _setup_data_module(
def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
"""From an Atomworks structure, calculate Boltz-2 input features.

Runs Pairformer pass and initializes x_init from prior distribution.
Runs the Pairformer pass to produce conditioning features.

NOTE: Has side effect of creating Boltz input YAML and initial processed files with
create_boltz_input_from_structure() and the data module setup.
Expand All @@ -665,7 +658,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
Returns
-------
GenerativeModelInput[BoltzConditioning]
Model input with x_init and Pairformer conditioning.
Model input with Pairformer conditioning.
"""
self.cached_representations.clear()

Expand All @@ -681,8 +674,6 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
else base_dir
)
num_workers = config.num_workers
ensemble_size = config.ensemble_size

input_path = create_boltz_input_from_structure( # side effect of creating Boltz input YAML
structure,
out_dir,
Expand Down Expand Up @@ -721,32 +712,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
model_atom_array=model_atom_array,
)

# Use atom count from Boltz featurization in atom_pad_mask (via conditioning) to ensure
# consistency between x_init shape and atom_pad_mask used in step(). Note: x_init=None is
# a temporary placeholder; initialize_from_prior derives shape from conditioning.
feats = pairformer_out["feats"]
atom_mask = cast(Tensor, feats.get("atom_pad_mask"))
num_atoms = int(atom_mask.sum())

# x_init should be the reference coordinates for alignment purposes.
if true_atom_array is not None and len(true_atom_array) == num_atoms:
x_init = torch.tensor(true_atom_array.coord, device=self.device, dtype=torch.float32)
x_init = match_batch(x_init.unsqueeze(0), target_batch_size=ensemble_size).clone()
else:
# TODO: The temp features situation is not ideal and I think we can do better probably
# not sure exactly what the best way to handle x_init is, should define it a bit better
# most likely
logger.warning(
"True structure not available or atom count mismatch; initializing "
"x_init from prior. This means align_to_input will not work properly,"
" and reward functions dependent on this won't be accurate."
)
temp_features = GenerativeModelInput[BoltzConditioning](
x_init=cast(Any, None), conditioning=conditioning
)
x_init = self.initialize_from_prior(batch_size=ensemble_size, features=temp_features)

return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
return GenerativeModelInput(conditioning=conditioning)

def _pairformer_pass(
self, features: dict[str, Any], recycling_steps: int = 3
Expand Down Expand Up @@ -951,8 +917,7 @@ def initialize_from_prior(
if shape is not None:
if len(shape) != 2 or shape[1] != 3:
raise ValueError("shape must be of the form (num_atoms, 3)")
x_init = torch.randn((batch_size, *shape), device=self.device)
return x_init
return torch.randn((batch_size, *shape), device=self.device)
if features is None or features.conditioning is None:
raise ValueError("Either features or shape must be provided to initialize_from_prior()")

Expand All @@ -963,9 +928,7 @@ def initialize_from_prior(

num_atoms = int(atom_mask.sum())

x_init = torch.randn((batch_size, num_atoms, 3), device=self.device)

return x_init
return torch.randn((batch_size, num_atoms, 3), device=self.device)


class Boltz1Wrapper:
Expand Down Expand Up @@ -1109,7 +1072,7 @@ def _setup_data_module(
def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
"""From an Atomworks structure, calculate Boltz-1 input features.

Runs Pairformer pass and initializes x_init from prior distribution.
Runs the Pairformer pass to produce conditioning features.

NOTE: Has side effect of creating Boltz input YAML and initial processed files with
create_boltz_input_from_structure() and the data module setup.
Expand All @@ -1125,7 +1088,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
Returns
-------
GenerativeModelInput[BoltzConditioning]
Model input with x_init and Pairformer conditioning.
Model input with Pairformer conditioning.
"""
self.cached_representations.clear()

Expand All @@ -1141,8 +1104,6 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
else base_dir
)
num_workers = config.num_workers
ensemble_size = config.ensemble_size

input_path = create_boltz_input_from_structure(
structure,
out_dir,
Expand Down Expand Up @@ -1180,31 +1141,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[BoltzConditioning]:
model_atom_array=model_atom_array,
)

# x_init should be the reference coordinates for alignment purposes.
# The guidance scalers call initialize_from_prior() separately for starting noise.
# Use atom count from Boltz featurization to ensure shape consistency with model.
feats = pairformer_out["feats"]
atom_mask = cast(Tensor, feats.get("atom_pad_mask"))
num_atoms = int(atom_mask.sum())

if true_atom_array is not None and len(true_atom_array) == num_atoms:
x_init = torch.tensor(true_atom_array.coord, device=self.device, dtype=torch.float32)
x_init = match_batch(x_init.unsqueeze(0), target_batch_size=ensemble_size).clone()
else:
# TODO: The temp features situation is not ideal and I think we can do better probably
# not sure exactly what the best way to handle x_init is, should define it a bit better
# most likely
logger.warning(
"True structure not available or atom count mismatch; initializing "
"x_init from prior. This means align_to_input will not work properly,"
" and reward functions dependent on this won't be accurate."
)
temp_features = GenerativeModelInput[BoltzConditioning](
x_init=cast(Any, None), conditioning=conditioning
)
x_init = self.initialize_from_prior(batch_size=ensemble_size, features=temp_features)

return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
return GenerativeModelInput(conditioning=conditioning)

def step(
self,
Expand Down Expand Up @@ -1312,8 +1249,7 @@ def initialize_from_prior(
if shape is not None:
if len(shape) != 2 or shape[1] != 3:
raise ValueError("shape must be of the form (num_atoms, 3)")
x_init = torch.randn((batch_size, *shape), device=self.device)
return x_init
return torch.randn((batch_size, *shape), device=self.device)
if features is None or features.conditioning is None:
raise ValueError("Either features or shape must be provided to initialize_from_prior()")

Expand All @@ -1324,9 +1260,7 @@ def initialize_from_prior(

num_atoms = int(atom_mask.sum())

x_init = torch.randn((batch_size, num_atoms, 3), device=self.device)

return x_init
return torch.randn((batch_size, num_atoms, 3), device=self.device)

def _pairformer_pass(
self, features: dict[str, Any], recycling_steps: int = 3
Expand Down
44 changes: 4 additions & 40 deletions src/sampleworks/models/protenix/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ class ProtenixConditioning:
Raw feature dict for diffusion module.
num_atoms : int
Number of atoms in the Protenix model's representation. This is the authoritative
atom count for the diffusion module and should be used for x_init shape.
atom count for the diffusion module and prior initialization.
pair_z : Tensor | None
Cached pair representation for diffusion conditioning.
p_lm : Tensor | None
Expand Down Expand Up @@ -93,8 +93,6 @@ class ProtenixConfig:
----------
out_dir : str | Path | None
Output directory for intermediate JSON file.
ensemble_size : int
Number of ensemble members to generate.
recycling_steps : int | None
Number of recycling steps to perform. If None, uses model default.
use_msa : bool
Expand All @@ -104,7 +102,6 @@ class ProtenixConfig:
"""

out_dir: str | Path | None = None
ensemble_size: int = 1
recycling_steps: int | None = None
use_msa: bool = True
enable_diffusion_shared_vars_cache: bool = True
Expand All @@ -114,7 +111,6 @@ def annotate_structure_for_protenix(
structure: dict,
*,
out_dir: str | Path | None = None,
ensemble_size: int = 1,
recycling_steps: int | None = None,
use_msa: bool = True,
enable_diffusion_shared_vars_cache: bool = True,
Expand All @@ -128,8 +124,6 @@ def annotate_structure_for_protenix(
out_dir : str | Path | None
Output directory for intermediate files.
Defaults to structure metadata ID or "protenix_output".
ensemble_size : int
Number of ensemble members to generate.
recycling_steps : int | None
Number of recycling steps to perform. If None, uses model default.
use_msa : bool
Expand All @@ -144,7 +138,6 @@ def annotate_structure_for_protenix(
"""
config = ProtenixConfig(
out_dir=out_dir or structure.get("metadata", {}).get("id", "protenix_output"),
ensemble_size=ensemble_size,
recycling_steps=recycling_steps,
use_msa=use_msa,
enable_diffusion_shared_vars_cache=enable_diffusion_shared_vars_cache,
Expand Down Expand Up @@ -327,7 +320,7 @@ def __init__(
def featurize(self, structure: dict) -> GenerativeModelInput[ProtenixConditioning]:
"""From an Atomworks structure, calculate Protenix input features.

Runs Pairformer pass and initializes x_init from prior distribution.
Runs the Pairformer pass to produce conditioning features.

Parameters
----------
Expand All @@ -340,14 +333,13 @@ def featurize(self, structure: dict) -> GenerativeModelInput[ProtenixConditionin
Returns
-------
GenerativeModelInput[ProtenixConditioning]
Model input with x_init and Pairformer conditioning.
Model input with Pairformer conditioning.
"""
config = structure.get("_protenix_config", ProtenixConfig())
if isinstance(config, dict):
config = ProtenixConfig(**config)

out_dir = config.out_dir or structure.get("metadata", {}).get("id", "protenix_output")
ensemble_size = config.ensemble_size
use_msa = config.use_msa
recycling_steps = config.recycling_steps
enable_diffusion_shared_vars_cache = config.enable_diffusion_shared_vars_cache
Expand Down Expand Up @@ -505,35 +497,7 @@ def featurize(self, structure: dict) -> GenerativeModelInput[ProtenixConditionin
model_atom_array=model_aa,
)

# x_init should be the reference coordinates for alignment purposes.
# Must match num_atoms_protenix so downstream samplers/scalers see
# consistent shapes with the model representation.
if "asym_unit" in structure:
n_input = len(atom_array)
if n_input != num_atoms_protenix:
logger.warning(
f"Atom-count mismatch: atom_array has {n_input} atoms, "
f"atom_array_protenix has {num_atoms_protenix} atoms. "
"Using atom_array_protenix coords for x_init to match model "
"atom count.",
)
x_init = torch.as_tensor(
atom_array_protenix.coord, device=self.device, dtype=torch.float32
)
else:
x_init = torch.as_tensor(atom_array.coord, device=self.device, dtype=torch.float32)
x_init = match_batch(x_init.unsqueeze(0), target_batch_size=ensemble_size).clone()
else:
logger.warning(
"True structure not available, so initializing "
"x_init from prior. This means align_to_input will not work properly,"
" and reward functions dependent on this won't be accurate."
)
x_init = self.initialize_from_prior(
batch_size=ensemble_size, shape=(num_atoms_protenix, 3)
)

return GenerativeModelInput(x_init=x_init, conditioning=conditioning)
return GenerativeModelInput(conditioning=conditioning)

def _pairformer_pass(
self, features: dict[str, Any], grad_needed: bool = False, **kwargs
Expand Down
13 changes: 0 additions & 13 deletions src/sampleworks/models/protocol.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,25 +26,12 @@
class GenerativeModelInput(Generic[C]): # noqa: UP046 (for Python 3.11 compatibility)
"""Container for inputs to generative models.

The x_init tensor is typically sampled from a prior distribution,
with shape determined by the input data (e.g., sequence length
determines atom count).

Attributes
----------
x_init : Float[Array, "*batch atoms 3"]
Initial structure coordinates.
This can be a reference structure (e.g. for alignment during sampling)
or a noisy sample from a prior distribution. This should have the proper
shape expected for the given ensemble being sampled, e.g. (4, atoms, 3) for
ensemble size 4.
conditioning : C | None
Model-specific conditioning features, or None.
"""

# TODO: make x_init more general (not just Float),
# relate this to StateT in Sampler protocol?
x_init: Float[Array, "*batch atoms 3"]
conditioning: C | None


Expand Down
Loading
Loading