model.parameterization.neural_parameterization

Learning-based parameterization module for ADELM.

This module provides two trainable parameterization components:

NNGlobalParameter Holds one nn.Parameter per nn_global parameter. Parameters are constant across locations and optimised directly via gradient descent.

NNFeatureParameter Predicts location-varying nn_feature_based parameters using one independent shared MLP per physical parameter. The module inspects the target parameter shape at runtime: site-level parameters use site attributes only, while layer-wise parameters use site attributes plus soil-profile inputs.

Module Contents

Classes

NNGlobalParameter	Trainable scalar parameterization for nn_global parameters.
NNFeatureParameter	Location-varying parameterization for nn_feature_based parameters.
NNBasedParams	Combined container for all learning-based ADELM parameters.

API

class model.parameterization.neural_parameterization.NNGlobalParameter(param_names, bounds, seed=42)

Bases: torch.nn.Module

Trainable scalar parameterization for nn_global parameters.

Each declared nn_global parameter is stored as a raw nn.Parameter in unconstrained logit space. forward applies sigmoid rescaling to return physical-space tensors bounded within the registered parameter bounds::

physical = lower + (upper - lower) * sigmoid(logit)

This mapping is smooth, differentiable, and never reaches the exact boundary, allowing gradient-based optimisation without clamping artefacts.

Parameters

param_names : list[str] Names of the nn_global parameters, as returned by RuntimeConfig.global_nn_params. bounds : dict[str, tuple[float, float]] Physical bounds {name: (lower, upper)}, as returned by RuntimeConfig.parameter_bounds. seed : int RNG seed for reproducible initialisation of logit values.

Attributes

logits : nn.ParameterDict Unconstrained logit scalars, one per parameter.

Initialization

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward()

Return physical-space parameter values from current logit weights.

Returns

dict[str, torch.Tensor] Mapping of parameter name to a 0-D tensor in physical units, bounded within (lower, upper). Every tensor participates in the autograd graph.

class model.parameterization.neural_parameterization.NNFeatureParameter(param_names, bounds, attri_features, hidden_dims, seed=42, dropout_rate=0.2)

Bases: torch.nn.Module

Location-varying parameterization for nn_feature_based parameters.

Predicts one physical-space value per location for each declared nn_feature_based parameter using a fully-connected MLP. Each physical parameter owns an independent network, and that network is shared across all sites. Runtime behavior depends on the target parameter shape already present in the ADELM parameter container:

• site-level targets [n_entities] use site attributes only

• layer-wise targets [n_entities, n_layers] use site attributes and soil texture

Runtime layout

For every declared nn_feature_based parameter, this module builds shared MLPs as needed:

• one network per physical parameter

• shared weights across all sites

• one scalar output per site or per [site, layer] pair

Each logit is mapped to physical space using the same bounded sigmoid transform as NNGlobalParameter::

physical = lower + (upper - lower) * sigmoid(logit)

Architecture example for hidden_dims = [5, 5] and one output parameter::

[n_sites, n_features] → Linear(n_features, 5) → ReLU
                      → Linear(5, 5)          → ReLU
                      → Linear(5, 1)          → sigmoid → physical bounds

Parameters

param_names : list[str] Names of the nn_feature_based parameters. bounds : dict[str, tuple[float, float]] Physical bounds {name: (lower, upper)}. attri_features : list[str] Ordered attribute names (keys in attris) used as input features. Each must be a 1-D entity-level tensor [n_entities]. hidden_dims : list[int] Hidden layer widths, e.g. [5, 5] for two hidden layers of width 5. seed : int RNG seed for reproducible weight initialisation.

Initialization

Initialize internal Module state, shared by both nn.Module and ScriptModule.

initialize(n_entities, params=None, device=None)

Validate runtime usage and optionally move the shared MLPs to device.

Parameters

n_entities : int Number of sites / entities in the currently loaded dataset. device : torch.device or str, optional Device on which to place the shared parameter networks.

forward(attris, params=None, structure=None)

Predict physical-space parameter values from location attributes.

Parameters

attris : dict ADELM attribute container. Each key listed in attri_features must map to a 1-D tensor of shape [n_entities].

Returns

dict[str, torch.Tensor] Mapping of parameter name to a tensor of shape [n_entities], bounded within (lower, upper).

class model.parameterization.neural_parameterization.NNBasedParams(global_params=None, feature_params=None)

Bases: torch.nn.Module

Combined container for all learning-based ADELM parameters.

Merges the outputs of NNGlobalParameter and NNFeatureParameter into a single dict that can be injected into the ADELM parameter container via _inject_nn_params.

Parameters

global_params : NNGlobalParameter or None Handles all nn_global parameters. feature_params : NNFeatureParameter or None Handles all nn_feature_based parameters. Pass None when no nn_feature_based parameters are declared.

Initialization

Initialize internal Module state, shared by both nn.Module and ScriptModule.

parameter_names()

static checkpoint_parameter_names(state_dict)

filtered_state_dict(state_dict, parameter_names)

compatible_checkpoint_parameter_names(state_dict, parameter_names)

freeze_parameters(parameter_names)

parameter_tensors(name)

classmethod from_runtime_config(config)

Construct an NNBasedParams instance from an already-loaded RuntimeConfig.

Parameters

config : RuntimeConfig Validated runtime configuration.

Returns

NNBasedParams

forward(attris=None, params=None, structure=None)

Return physical-space values for all declared NN-based parameters.

Parameters

attris : dict or None ADELM attribute container. Required when feature_params is set.

Returns

dict[str, torch.Tensor] Merged mapping of parameter name to physical-space tensor.