dymad.modules.collections

Module Attributes

NN_MAP

Options for preset neural network models.

AE_MAP

Options for preset autoencoder types.

Functions

make_autoencoder(ae_type, input_dim, ...[, ...])

Factory function to create preset autoencoder models.

make_kernel(k_type, input_dim[, output_dim, ...])

Factory function to create preset kernels.

make_krr(type, kernel[, ridge_init, jitter, ...])

Factory function to create preset Kernel Ridge Regression (KRR) models.

make_network(nn_type, input_dim, hidden_dim, ...)

Factory function to create preset neural network models based on NN_MAP.
dymad.modules.collections.AE_MAP = {'gnn_cat': ('gnn_cat', 'gnn_1st'), 'gnn_res': ('gnn_res', 'gnn_res'), 'gnn_seq_rnn': ('seq_rnn', 'seq_gnn_smp'), 'gnn_seq_smp': ('seq_gnn_smp', 'seq_gnn_smp'), 'gnn_seq_std': ('seq_std', 'seq_gnn_smp'), 'gnn_smp': ('gnn_smp', 'gnn_smp'), 'mlp_cat': ('mlp_cat', 'mlp_1st'), 'mlp_res': ('mlp_res', 'mlp_res'), 'mlp_seq_rnn': ('seq_rnn', 'seq_mlp_smp'), 'mlp_seq_smp': ('seq_mlp_smp', 'seq_mlp_smp'), 'mlp_seq_std': ('seq_std', 'seq_mlp_smp'), 'mlp_smp': ('mlp_smp', 'mlp_smp')}

Options for preset autoencoder types.

dymad.modules.collections.NN_MAP = {'gnn_1st': <class 'dymad.modules.misc.TakeFirstGraph'>, 'gnn_cat': <class 'dymad.modules.gnn.IdenCatGNN'>, 'gnn_res': <class 'dymad.modules.gnn.ResBlockGNN'>, 'gnn_smp': <class 'dymad.modules.gnn.GNN'>, 'mlp_1st': <class 'dymad.modules.misc.TakeFirst'>, 'mlp_cat': <class 'dymad.modules.mlp.IdenCatMLP'>, 'mlp_res': <class 'dymad.modules.mlp.ResBlockMLP'>, 'mlp_smp': <class 'dymad.modules.mlp.MLP'>, 'seq_rnn': <class 'dymad.modules.sequential.VanillaRNN'>, 'seq_std': <class 'dymad.modules.sequential.SimpleRNN'>}

Options for preset neural network models.

dymad.modules.collections.make_autoencoder(ae_type, input_dim, hidden_dim, latent_dim, enc_depth, dec_depth, output_dim=None, seq_len=None, **kwargs)

Factory function to create preset autoencoder models. Including:

  • [mlp_smp] Simple version: MLP-in MLP-out

  • [mlp_res] Simple version but with ResBlockMLP

  • [mlp_cat] Concatenation as encoder [x MLP(x)], then TakeFirst as decoder

  • [mlp_seq_rnn] RNN-in MLP-out, using a 1-layer unidirectional RNN

  • [mlp_seq_std] RNN-in MLP-out, using standard RNN from pytorch

  • [mlp_seq_smp] MLP-in MLP-out, applied stepwise to sequences

  • The graph version of the above, e.g., gnn_smp, gnn_seq_rnn, etc.

Parameters:

  • ae_type (str) – Type of autoencoder to create.

  • input_dim (int) – Dimension of the input features.

  • hidden_dim (int) – Width of the hidden layers.

  • latent_dim (int) – Dimension of the latent/encoded space.

  • enc_depth (int) – Number of layers in the encoder.

  • dec_depth (int) – Number of layers in the decoder.

  • output_dim (int, optional) – Dimension of the output features, defaults to input_dim.

  • seq_len (int, optional) – Length of the input sequences (for sequence-based autoencoders).

  • **kwargs – Additional keyword arguments passed to the specific constructors.

Return type:

tuple[Module, Module]

dymad.modules.collections.make_kernel(k_type, input_dim, output_dim=None, kopts=None, dtype=None, **kwargs)

Factory function to create preset kernels. Including:

  • [sc_rbf] Scalar: Radial basis function kernel

  • [sc_dm] Scalar: Diffusion Map kernel

  • [op_sep] Operator-valued: Separable kernel with multiple scalar kernels

Parameters:

  • k_type (str) – Type of kernel to create. One of {‘sc_rbf’, ‘sc_dm’}.

  • input_dim (int) – Dimension of the input features.

  • output_dim (int, optional) – Dimension of the output features.

  • kopts (List, optional) – List of scalar kernel options (for operator-valued kernels).

  • dtype – Data type of the kernel parameters.

  • **kwargs – Additional keyword arguments passed to the kernel constructors.

Return type:

Module

dymad.modules.collections.make_krr(type, kernel, ridge_init=0, jitter=1e-10, dtype=None, device=None)

Factory function to create preset Kernel Ridge Regression (KRR) models. Including:

  • [share] Multi-output KRR with shared scalar kernel

  • [indep] Multi-output KRR with independent scalar kernels

  • [opval] Multi-output KRR with operator-valued kernel

  • [tangent] KRR for vector fields on manifolds

Parameters:

  • type (str) – Type of KRR model to create. One of {‘krr_shared’, ‘krr_indep’, ‘krr_opval’, ‘krr_tangent’}.

  • kernel (Union[Dict, List[Dict]]) – Kernel configuration(s).

  • ridge_init (float, optional) – Initial value for the ridge regularization parameter.

  • jitter (float, optional) – Jitter added to the diagonal for numerical stability.

Return type:

Module

dymad.modules.collections.make_network(nn_type, input_dim, hidden_dim, output_dim, n_layers, seq_len=None, **kwargs)

Factory function to create preset neural network models based on NN_MAP.

Parameters:

  • nn_type (str) – Type of network to create. One of the keys in NN_MAP: {‘mlp_smp’, ‘mlp_res’, ‘mlp_cat’, ‘mlp_1st’, ‘gnn_smp’, ‘gnn_res’, ‘gnn_cat’, ‘gnn_1st’, ‘seq_std’, ‘seq_rnn’}, or ‘seq’ prefixed versions of MLP and GNN types for sequence models.

  • input_dim (int) – Dimension of the input features.

  • hidden_dim (int) – Width of the hidden layers.

  • output_dim (int) – Dimension of the output space.

  • n_layers (int) – Number of layers in the network.

  • seq_len (int, optional) – Length of the input sequences (for sequence-based networks).

  • **kwargs – Additional keyword arguments passed to the specific constructors.

Returns:

The constructed neural network module.

Return type:

nn.Module