OCDocker.OCScore.Optimization.StagedOptuna module¶

Current staged Optuna protocol for OCScore PDBbind/DUDEz modeling.

It is imported as:

from OCDocker.OCScore.Optimization.StagedOptuna import StagedProtocol

class OCDocker.OCScore.Optimization.StagedOptuna.DUDEzOptunaConfig(kind_column='kind', target_group_column='receptor', primary_metric='BEDROC', bedroc_alpha=20.0, n_trials=10, epochs=100, storage='auto', study_name='DUDEz_Screening_Optimization', load_if_exists=False, validation_size=0.2, test_size=0.2, direction='maximize', sampler_seed=None, random_seed=None, use_gpu=True, verbose=False, n_jobs=1, allow_scratch=True, use_class_weighting=True, search_space=None, split_config=None, dudez_scaling_config=None, calibration_report_mode='ranking_only')[source]¶

Bases: object

Configuration for the DUDEz screening Optuna stage.

Parameters:

kind_column (str, optional) – Column used to derive active/decoy labels, by default “kind”.
target_group_column (str, optional) – Column used for target-grouped splits, by default “receptor”.
primary_metric (str, optional) – DUDEz validation objective metric, by default “BEDROC”.
bedroc_alpha (float, optional) – Exponential BEDROC weighting factor, by default 20.0.
n_trials (int, optional) – Number of Optuna trials, by default 10.
epochs (int, optional) – Training epochs per trial, by default 100.
storage (str | None, optional) – Optuna storage URL. Use “auto” to create optuna.db in the protocol output directory (shared across stages and replicas), by default “auto”.
study_name (str, optional) – Optuna study name, by default “DUDEz_Screening_Optimization”.
load_if_exists (bool, optional) – Reuse an existing Optuna study with the same name, by default False.
validation_size (float, optional) – Fraction of non-test data held out for validation, by default 0.2.
test_size (float, optional) – Fraction held out for test metrics, by default 0.2.
direction (str, optional) – Optuna direction. Must be “maximize”, by default “maximize”.
sampler_seed (int | None, optional) – Seed for the Optuna sampler, by default None.
random_seed (int | None, optional) – Stage-specific seed override, by default None.
use_gpu (bool, optional) – Use CUDA when available, by default True.
verbose (bool, optional) – Reserved for compatibility with OCScore verbosity patterns, by default False.
n_jobs (int, optional) – Number of parallel Optuna jobs, by default 1.
allow_scratch (bool, optional) – Allow from-scratch DUDEz feature extractors as a tunable option, by default True.
use_class_weighting (bool, optional) – Apply positive-class weighting in BCEWithLogitsLoss, by default True.
search_space (DUDEzSearchSpaceConfig | None, optional) – Centralized Optuna search-space definition. If None, defaults are used.
split_config (DUDEzSplitConfig | None, optional) – Receptor/kind-aware train/validation/test split configuration. When None, defaults to receptor-wise stratified splitting using validation_size, test_size, target_group_column, and kind_column.
dudez_scaling_config (DUDEzScalingConfig | None, optional) – DUDEz feature scaling policy. When None, the PDBbind-fitted scaler is reused strictly for transfer-compatible DUDEz inputs.
calibration_report_mode (str, optional) – Controls diagnostic vs validated calibration reporting in export metrics, by default "ranking_only".

kind_column: str = 'kind'¶

target_group_column: str = 'receptor'¶

primary_metric: str = 'BEDROC'¶

bedroc_alpha: float = 20.0¶

n_trials: int = 10¶

epochs: int = 100¶

storage: str | None = 'auto'¶

study_name: str = 'DUDEz_Screening_Optimization'¶

load_if_exists: bool = False¶

validation_size: float = 0.2¶

test_size: float = 0.2¶

direction: str = 'maximize'¶

sampler_seed: int | None = None¶

random_seed: int | None = None¶

use_gpu: bool = True¶

verbose: bool = False¶

n_jobs: int = 1¶

allow_scratch: bool = True¶

use_class_weighting: bool = True¶

search_space: DUDEzSearchSpaceConfig | None = None¶

split_config: DUDEzSplitConfig | None = None¶

dudez_scaling_config: DUDEzScalingConfig | None = None¶

calibration_report_mode: str = 'ranking_only'¶

class OCDocker.OCScore.Optimization.StagedOptuna.DUDEzOptunaStage(config=None)[source]¶

Bases: object

Optuna stage that optimizes DUDEz screening by one screening metric only.

Parameters:: config (DUDEzOptunaConfig | None, optional) – Stage configuration. If None, default DUDEz settings are used.

name = 'dudez_optuna'¶

__init__(config=None)[source]¶

Initialize the DUDEz Optuna stage.

Parameters:: config (DUDEzOptunaConfig | None, optional) – Stage configuration. If None, default DUDEz settings are used.
Return type:: None

run(context)[source]¶

Run DUDEz screening optimization.

Parameters:: context (ProtocolContext) – Protocol context containing DUDEz data and transferred extractor.
Returns:: Updated context with DUDEz artifacts.
Return type:: ProtocolContext

class OCDocker.OCScore.Optimization.StagedOptuna.DUDEzScreeningModel(*args, **kwargs)[source]¶

Bases: Module

DUDEz classifier/ranking model using a feature extractor and new head.

Parameters:

feature_extractor (FeatureExtractor) – Transferred or newly initialized feature extractor.
classifier_hidden_size (int, optional) – Hidden size for the classifier head, by default 128.
dropout (float, optional) – Classifier dropout probability, by default 0.0.
activation (str, optional) – Activation name, by default “GELU”.

forward(x)[source]¶

Run DUDEz screening prediction.

Parameters:: x (torch.Tensor) – Input feature tensor.
Returns:: One-dimensional classifier logits.
Return type:: torch.Tensor

class OCDocker.OCScore.Optimization.StagedOptuna.FeatureExtractor(*args, **kwargs)[source]¶

Bases: Module

Monotonic MLP feature extractor with optional projection block.

Parameters:

input_size (int) – Input feature dimension.
hidden_sizes (Sequence[int]) – Monotonic hidden layer sizes.
latent_dim (int) – Latent encoder dimension.
activation (str, optional) – Activation name, by default “GELU”.
dropout (float, optional) – Dropout probability, by default 0.0.
projection_dim (int, optional) – Optional projection output dimension. A value of 0 disables projection, by default 0.

projection: nn.Sequential | None¶

forward(x)[source]¶

Compute feature embeddings.

Parameters:: x (torch.Tensor) – Input feature tensor.
Returns:: Latent or projected feature tensor.
Return type:: torch.Tensor

class OCDocker.OCScore.Optimization.StagedOptuna.PDBbindOptunaConfig(target_column='experimental', n_trials=10, epochs=100, storage='auto', study_name='PDBbind_Regression_Optimization', search_phase='full', enable_pruning=True, pruner_n_startup_trials=None, pruner_n_warmup_steps=None, load_if_exists=False, validation_size=0.2, test_size=0.2, objective_metric='RMSE', direction='minimize', sampler_seed=None, random_seed=None, use_gpu=True, verbose=False, n_jobs=1, search_space=None, split_config=None)[source]¶

Bases: object

Configuration for the PDBbind regression Optuna stage.

Parameters:

target_column (str, optional) – Regression target column, by default “experimental”.
n_trials (int, optional) – Number of Optuna trials, by default 10.
epochs (int, optional) – Training epochs per trial, by default 100.
storage (str | None, optional) – Optuna storage URL. Use “auto” to create optuna.db in the protocol output directory (shared across stages and replicas), by default “auto”.
study_name (str, optional) – Optuna study name, by default “PDBbind_Regression_Optimization”.
load_if_exists (bool, optional) – Reuse an existing Optuna study with the same name, by default False.
validation_size (float, optional) – Fraction of non-test data held out for validation, by default 0.2.
test_size (float, optional) – Fraction held out for test metrics, by default 0.2.
objective_metric (str, optional) – Stage objective metric. Must be “RMSE”, by default “RMSE”.
direction (str, optional) – Optuna direction. Must be “minimize”, by default “minimize”.
sampler_seed (int | None, optional) – Seed for the Optuna sampler, by default None.
random_seed (int | None, optional) – Stage-specific seed override, by default None.
use_gpu (bool, optional) – Use CUDA when available, by default True.
verbose (bool, optional) – Reserved for compatibility with OCScore verbosity patterns, by default False.
n_jobs (int, optional) – Number of parallel Optuna jobs, by default 1.
search_space (PDBbindSearchSpaceConfig | None, optional) – Centralized Optuna search-space definition. If None, defaults are used.
search_phase (str, optional) – Staged search phase: full (default) or encoder_regression (Phase 1: encoder + regression only). Ignored when search_space is set explicitly.
enable_pruning (bool, optional) – When False, Optuna uses NopPruner (recommended for Phase 1 experiments).
pruner_n_startup_trials (int | None, optional) – MedianPruner startup trials. When None, max(10% of n_trials, 5).
pruner_n_warmup_steps (int | None, optional) – MedianPruner warmup epochs before pruning. When None, max(10% of epochs, 10).
split_config (PDBbindSplitConfig | None, optional) – Affinity-aware PDBbind split configuration. When None, defaults to quantile-bin stratified splitting using target_column, validation_size, and test_size.

target_column: str = 'experimental'¶

n_trials: int = 10¶

epochs: int = 100¶

storage: str | None = 'auto'¶

study_name: str = 'PDBbind_Regression_Optimization'¶

search_phase: str = 'full'¶

enable_pruning: bool = True¶

pruner_n_startup_trials: int | None = None¶

pruner_n_warmup_steps: int | None = None¶

load_if_exists: bool = False¶

validation_size: float = 0.2¶

test_size: float = 0.2¶

objective_metric: str = 'RMSE'¶

direction: str = 'minimize'¶

sampler_seed: int | None = None¶

random_seed: int | None = None¶

use_gpu: bool = True¶

verbose: bool = False¶

n_jobs: int = 1¶

search_space: PDBbindSearchSpaceConfig | None = None¶

split_config: PDBbindSplitConfig | None = None¶

class OCDocker.OCScore.Optimization.StagedOptuna.PDBbindOptunaStage(config=None)[source]¶

Bases: object

Optuna stage that optimizes PDBbind regression by validation RMSE only.

Parameters:: config (PDBbindOptunaConfig | None, optional) – Stage configuration. If None, default PDBbind settings are used.

name = 'pdbbind_optuna'¶

__init__(config=None)[source]¶

Initialize the PDBbind Optuna stage.

Parameters:: config (PDBbindOptunaConfig | None, optional) – Stage configuration. If None, default PDBbind settings are used.
Return type:: None

run(context)[source]¶

Run PDBbind regression optimization.

Parameters:: context (ProtocolContext) – Protocol context with reduced PDBbind data and selected features.
Returns:: Updated context with PDBbind artifacts.
Return type:: ProtocolContext

class OCDocker.OCScore.Optimization.StagedOptuna.PDBbindRegressionModel(*args, **kwargs)[source]¶

Bases: Module

PDBbind affinity regression model with optional weak decoder.

Parameters:

feature_extractor (FeatureExtractor) – Encoder/projection module used for affinity prediction.
input_size (int) – Original feature dimension used by the optional decoder.
activation (str) – Activation name used by the optional decoder.
decoder_sizes (Sequence[int] | None, optional) – Decoder hidden sizes. If None, reconstruction is disabled.

decoder: nn.Sequential | None¶

forward(x, return_reconstruction=False)[source]¶

Run affinity prediction and optional reconstruction.

Parameters:

x (torch.Tensor) – Input feature tensor.
return_reconstruction (bool, optional) – Return decoder reconstruction when available, by default False.

Returns:

Feature embedding, regression prediction, and optional reconstruction.

Return type:

dict[str, torch.Tensor | None]

class OCDocker.OCScore.Optimization.StagedOptuna.TransferFeatureExtractorStage[source]¶

Bases: object

Transfer reusable PDBbind feature extractor layers to the DUDEz stage.

Notes

The final PDBbind regression head is not transferred to DUDEz. The decoder is also excluded because it is only a PDBbind reconstruction regularizer.

name = 'transfer_feature_extractor'¶

run(context)[source]¶

Transfer the PDBbind feature extractor without the regression head.

Parameters:: context (ProtocolContext) – Protocol context containing a PDBbind model or checkpoint path.
Returns:: Updated context with transferred_feature_extractor.
Return type:: ProtocolContext

OCDocker.OCScore.Optimization.StagedOptuna.apply_fine_tuning_mode(feature_extractor, mode, num_unfrozen_layers=1)[source]¶

Apply frozen, partial, or full fine-tuning to a transferred extractor.

Parameters:

feature_extractor (nn.Module) – Feature extractor to update in-place.
mode (str) – One of "frozen", "partial", or "full".
num_unfrozen_layers (int, optional) – Number of final linear layers unfrozen when mode is "partial".

Return type:

None

OCDocker.OCScore.Optimization.StagedOptuna.build_dudez_model(input_size, params, transferred_extractor=None, feature_extractor_architecture=None)[source]¶

Build a DUDEz screening model from sampled parameters.

Parameters:

input_size (int) – Input feature dimension.
params (dict[str, Any]) – Sampled DUDEz model and training parameters.
transferred_extractor (FeatureExtractor | None, optional) – Transferred PDBbind feature extractor used when transfer is enabled.
feature_extractor_architecture (dict[str, Any] | None, optional) – Resolved extractor architecture for scratch DUDEz models.

Returns:

Initialized DUDEz screening model.

Return type:

DUDEzScreeningModel

OCDocker.OCScore.Optimization.StagedOptuna.build_pdbbind_model(input_size, params)[source]¶

Build a PDBbind regression model from sampled parameters.

Parameters:

input_size (int) – Input feature dimension.
params (dict[str, Any]) – Sampled PDBbind model and training parameters.

Returns:

Initialized PDBbind regression model.

Return type:

PDBbindRegressionModel

OCDocker.OCScore.Optimization.StagedOptuna.compute_regression_reconstruction_loss(prediction, target, reconstruction, features, regression_loss, reconstruction_loss, lambda_rec)[source]¶

Compute regression loss plus optional reconstruction regularization.

Parameters:

prediction (torch.Tensor) – Regression predictions.
target (torch.Tensor) – Regression targets.
reconstruction (torch.Tensor | None) – Reconstructed input features, if decoder regularization is enabled.
features (torch.Tensor) – Original input features.
regression_loss (nn.Module) – Primary regression loss.
reconstruction_loss (nn.Module) – Reconstruction loss.
lambda_rec (float) – Reconstruction loss weight. A value of 0 disables reconstruction.

Returns:

Total loss for backpropagation.

Return type:

torch.Tensor

OCDocker.OCScore.Optimization.StagedOptuna.derive_dudez_labels(df, kind_column='kind')[source]¶

Derive DUDEz labels from kind values.

ligands/ligand map to 1 and decoys/decoy map to 0.

Parameters:

df (pd.DataFrame) – DUDEz dataframe containing active/decoy kind values.
kind_column (str, optional) – Column used to derive labels, by default “kind”.

Returns:

Binary labels as float32 values.

Return type:

np.ndarray

OCDocker.OCScore.Optimization.StagedOptuna.dudez_search_space_summary(search_space=None)[source]¶

Return the DUDEz Optuna search-space summary.

Parameters:: search_space (DUDEzSearchSpaceConfig | None, optional) – Search-space configuration. Defaults to DEFAULT_DUDEZ_SEARCH_SPACE.
Returns:: JSON-serializable DUDEz search-space description.
Return type:: dict[str, Any]

OCDocker.OCScore.Optimization.StagedOptuna.evaluate_regression_metrics(y_true, y_pred)[source]¶

Evaluate PDBbind regression metrics.

Parameters:

y_true (np.ndarray) – Experimental affinity values.
y_pred (np.ndarray) – Predicted affinity values.

Returns:

RMSE, MAE, Pearson r, Spearman rho, and R2 metrics.

Return type:

dict[str, float]

OCDocker.OCScore.Optimization.StagedOptuna.pdbbind_phase1_experiment_config(*, n_trials=40, epochs=100, study_name='PDBbind_EncoderRegression_Phase1', enable_pruning=False, **kwargs)[source]¶

Preset PDBbind Optuna config for Phase 1 encoder-regression experiments.

Parameters:

n_trials (int, optional) – Number of Optuna trials (default 40, within the 30–50 Phase 1 band).
epochs (int, optional) – Training epochs per trial.
study_name (str, optional) – Distinct study name so Phase 1 does not resume the full-search study.
enable_pruning (bool, optional) – When False (default), pruning is disabled for the experiment.
**kwargs – Additional PDBbindOptunaConfig fields.

Returns:

Phase 1 experiment configuration.

Return type:

PDBbindOptunaConfig

OCDocker.OCScore.Optimization.StagedOptuna.pdbbind_search_space_summary(search_space=None)[source]¶

Return the PDBbind Optuna search-space summary.

Parameters:: search_space (PDBbindSearchSpaceConfig | None, optional) – Search-space configuration. Defaults to DEFAULT_PDBBIND_SEARCH_SPACE.
Returns:: JSON-serializable PDBbind search-space description.
Return type:: dict[str, Any]

OCDocker.OCScore.Optimization.StagedOptuna.prepare_dudez_screening_data(df, selected_features, labels, groups, split_config, target_group_column=None, scaling_config=None, pdbbind_scaler=None, *, fixed_train_indices=None, fixed_validation_indices=None, fixed_test_indices=None)[source]¶

Prepare train/validation/test arrays for DUDEz screening.

Parameters:

df (pd.DataFrame) – Reduced DUDEz dataframe.
selected_features (Sequence[str]) – Descriptor columns selected by feature reduction.
labels (np.ndarray) – Binary active/decoy labels.
groups (np.ndarray | None) – Optional target groups aligned with df (used for grouped metrics).
split_config (DUDEzSplitConfig) – Receptor/kind-aware split configuration.
target_group_column (str | None, optional) – Target column name recorded in split diagnostics.
scaling_config (DUDEzScalingConfig | None)
pdbbind_scaler (StandardScaler | None)
fixed_train_indices (Sequence[int] | None)
fixed_validation_indices (Sequence[int] | None)
fixed_test_indices (Sequence[int] | None)

Returns:

Train, validation, and test arrays with source indices.

Return type:

dict[str, Any]

OCDocker.OCScore.Optimization.StagedOptuna.summarize_dudez_split_diagnostics(df, labels, train_idx, val_idx, test_idx, groups=None, target_group_column=None)[source]¶

Summarize DUDEz split composition for logging and reproducibility.

Parameters:

df (pd.DataFrame) – Reduced DUDEz dataframe.
labels (np.ndarray) – Binary active/decoy labels.
train_idx (np.ndarray) – Training row indices.
val_idx (np.ndarray) – Validation row indices.
test_idx (np.ndarray) – Test row indices.
groups (np.ndarray | None, optional) – Target/receptor group labels aligned with labels.
target_group_column (str | None, optional) – Source column name used for grouped metrics.

Returns:

JSON-compatible split diagnostics.

Return type:

dict[str, Any]

OCDocker.OCScore.Optimization.StagedOptuna.prepare_pdbbind_regression_data(df, selected_features, split_config, *, fixed_train_indices=None, fixed_validation_indices=None, fixed_test_indices=None)[source]¶

Prepare scaled train/validation/test arrays for PDBbind regression.

Parameters:

df (pd.DataFrame) – Reduced PDBbind dataframe.
selected_features (Sequence[str]) – Descriptor columns selected by feature reduction.
split_config (PDBbindSplitConfig) – PDBbind split configuration (used for metadata when fixed indices are supplied).
fixed_train_indices (Sequence[int] | None, optional) – When provided, reuse this fixed training partition instead of recomputing a split.
fixed_validation_indices (Sequence[int] | None, optional) – Fixed validation partition indices.
fixed_test_indices (Sequence[int] | None, optional) – Fixed test partition indices.

Returns:

Scaled train, validation, and test arrays with source indices and scaler.

Return type:

dict[str, Any]

OCDocker.OCScore.Optimization.StagedOptuna.resolve_dudez_primary_metric(requested_metric, metrics)[source]¶

Return the effective DUDEz objective metric.

BEDROC is preferred by default. If it is missing or non-finite, PR-AUC is used as the primary objective.

Parameters:

requested_metric (str) – Requested screening objective metric.
metrics (dict[str, float]) – Validation screening metrics.

Returns:

Effective metric name used for optimization.

Return type:

str

OCDocker.OCScore.Optimization.StagedOptuna.run_staged_ocscore_optuna_protocol(context, pdbbind_config=None, dudez_config=None)[source]¶

Run the current staged OCScore Optuna protocol.

Parameters:

context (ProtocolContext) – Initial protocol context.
pdbbind_config (PDBbindOptunaConfig | None, optional) – PDBbind stage configuration, by default None.
dudez_config (DUDEzOptunaConfig | None, optional) – DUDEz stage configuration, by default None.

Returns:

Updated protocol context after all stages complete.

Return type:

ProtocolContext

OCDocker.OCScore.Optimization.StagedOptuna.suggest_decoder_hidden_sizes(trial, search_space, latent_dim, projection_dim, input_dim)[source]¶

Sample optional decoder hidden sizes for PDBbind reconstruction.

Parameters:

trial (optuna.Trial) – Optuna trial object.
search_space (DecoderSearchSpace) – Decoder search-space definition.
latent_dim (int) – Encoder latent dimension.
input_dim (int) – Original input feature dimension.
projection_dim (int)

Returns:

Decoder hidden sizes and reconstruction-loss weight.

Return type:

tuple[list[int] | None, float]

OCDocker.OCScore.Optimization.StagedOptuna.suggest_dudez_trial_params(trial, allow_scratch=True, search_space=None)[source]¶

Sample DUDEz screening hyperparameters for one Optuna trial.

Parameters:

trial (optuna.Trial) – Optuna trial object.
allow_scratch (bool, optional) – Allow sampling from-scratch feature extractors, by default True.
search_space (DUDEzSearchSpaceConfig | None, optional) – Centralized search-space definition.

Returns:

Sampled DUDEz hyperparameters.

Return type:

dict[str, Any]

OCDocker.OCScore.Optimization.StagedOptuna.suggest_encoder_architecture(trial, input_dim, search_space=None)[source]¶

Sample monotonic encoder hidden sizes and latent dimension.

Parameters:

trial (optuna.Trial) – Optuna trial object.
input_dim (int) – Input feature dimension.
search_space (EncoderSearchSpace | None, optional) – Encoder search-space definition.

Returns:

Encoder architecture index, hidden layer sizes, and latent dimension.

Return type:

tuple[int, list[int], int]

OCDocker.OCScore.Optimization.StagedOptuna.suggest_pdbbind_trial_params(trial, input_dim, search_space=None)[source]¶

Sample PDBbind regression hyperparameters for one Optuna trial.

Parameters:

trial (optuna.Trial) – Optuna trial object.
input_dim (int) – Input feature dimension.
search_space (PDBbindSearchSpaceConfig | None, optional) – Centralized search-space definition.

Returns:

Sampled PDBbind hyperparameters.

Return type:

dict[str, Any]

__all__ re-exports several names from other modules for convenience; they’re excluded above to avoid documenting them twice under two different module paths. See their original modules (DUDEzSplit, PDBbindSplit, Protocol, Ranking) for their documentation.