nif.model

The code represents a TensorFlow implementation of the Neural Implicit Flow

NIF is a neural network architecture designed for learning and representing implicit flow. The model consists of two sub-networks : a shape network and a parameter network. The shape network takes the input data as input and outputs a point on a manifold, while the parameter network takes the input data as input and outputs the parameters that determine the shape of the manifold.

The NIF class is the main class that represents the NIF model. It takes two configuration dictionaries as inputs, one for the shape network and one for the parameter network, and initializes the shape and parameter networks. The call method of the class takes the input data as input, passes it through the parameter network to obtain the parameters, and then passes it through the shape network to obtain the output.

The NIF class uses several custom layers defined in the layers module, including SIREN (Sinusoidal Representation Networks), SIREN_ResNet (SIREN with residual connections), MLP_ResNet (Multi-Layer Perceptron with residual connections), MLP_SimpleShortCut (Multi-Layer Perceptron with shortcut connections), JacRegLatentLayer (a layer that computes the Jacobian of the output with respect to the latent variables and applies a regularization term), and EinsumLayer (a layer that computes the dot product of two tensors using Einstein summation notation). The NIF class also uses various TensorFlow modules, including tf.keras, tf.keras.layers, tf.keras.regularizers, tf.initializers, and tf.cast.

class nif.model.NIFMultiScale(cfg_shape_net, cfg_parameter_net, mixed_policy='float32')

Bases: NIF

The NIFMultiScale class is a subclass of the NIF class, extending its functionality to support multiscale computations. The class is designed to work with neural implicit flow in a multiscale context.

Variables:

  • cfg_shape_net (dict) – Configuration dictionary for the shape network.

  • cfg_parameter_net (dict) – Configuration dictionary for the parameter network.

  • mixed_policy (str, optional) – The mixed precision policy to be used for TensorFlow computations. Defaults to “float32”.

call(inputs, training=None, mask=None)

Implements the forward pass of the NIFMultiScale model, which takes the input tensors and computes the output using the multiscale architecture.

Parameters:

  • inputs (tensor) – Input tensor with concatenated parameter and shape information.

  • training (bool, optional) – Whether the model is in training mode. Defaults to None, which will use the model’s training mode.

  • mask (tensor, optional) – Mask tensor for masked input. Defaults to None.

Returns:

tensor – Output tensor computed by the multiscale shape network.

model_x_to_u_given_w()

Constructs a Keras model for mapping input tensor x to output tensor u given the weights and biases from the parameter network.

Returns:

tf.keras.Model –

A Keras model that takes two inputs, input_s and

input_pnet, and returns the output tensor u.

class nif.model.NIF(cfg_shape_net, cfg_parameter_net, mixed_policy='float32')

Bases: object

Neural Implicit Flow class represents a network with two sub-networks to reduce the dimensionality of spatial temporal fields

Variables:

  • cfg_shape_net (dict) – Configuration dictionary for the shape network.

  • cfg_parameter_net (dict) – Configuration dictionary for the parameter network.

  • mixed_policy (str) – The data type for mixed precision training (default is ‘float32’).

call(self, inputs, training=None, mask=None)

Forward pass for the NIF model.

build(self)

Builds and returns the NIF model with a Jacobian regularization layer.

model(self)

Builds and returns the NIF model.

model_p_to_w(self)

Builds and returns a model that maps input parameters to weights and biases of the shape net.

model_p_to_lr(self)

Builds and returns a model that maps input parameters to the hidden layer representation.

model_lr_to_w(self)

Builds and returns a model that maps the hidden layer representation to shape net weights and biases.

model_x_to_u_given_w(self)

Builds and returns a model that maps input states to output, given shape net weights and biases.

save_config(self, filename='config.json')

Saves the NIF configuration to a JSON file.

call(inputs, training=None, mask=None)

Performs the forward pass for the NIF model, given input parameters and states.

Parameters:

  • inputs (tf.Tensor) – A tensor containing input parameters and states.

  • training (bool, optional) – Whether the model is in training mode. Defaults to None.

  • mask (tf.Tensor, optional) – A tensor representing masked elements. Defaults to None.

Returns:

tf.Tensor – The output tensor after passing through the shape network.

build()

Builds and returns the NIF model with a Jacobian regularization layer if specified in the configuration. Otherwise it is the same as .model()

Returns:

tf.keras.Model – The NIF model with or without the Jacobian regularization layer.

model()

Builds and returns the NIF model.

Returns:

tf.keras.Model – The NIF model.

model_p_to_w()

Builds and returns a model that maps input parameters to weights and biases of the shape network.

Returns:

tf.keras.Model – The model mapping input parameters to shape network weights and biases.

model_p_to_lr()

Builds and returns a model that maps input parameters to the hidden layer representation.

Returns:

tf.keras.Model – The model mapping input parameters to the hidden layer representation.

model_lr_to_w()

Builds and returns a model that maps the hidden layer representation to shape network weights and biases.

Returns:

tf.keras.Model – The model mapping the hidden layer representation to shape network weights and biases.

model_x_to_u_given_w()

Builds and returns a model that maps input states to output, given shape network weights and biases.

Returns:

tf.keras.Model – The model mapping input states to output, given shape network weights and biases.

save_config(filename='config.json')

Saves the NIF model configuration to a JSON file.

Parameters:

  • filename (str, optional) – The name of the file to save the

  • configuration. Defaults to “config.json”.

class nif.model.NIFMultiScaleLastLayerParameterized(cfg_shape_net, cfg_parameter_net, mixed_policy='float32')

Bases: NIFMultiScale

NIFMultiScaleLastLayerParameterized is a subclass of NIFMultiScale representing a Neural Information Flow (NIF) model with a multi-scale architecture that parameterizes only the last layer of the shape network. This class is designed to have its shape network and parameter network work in conjunction to map the inputs to the outputs.

Variables:

  • cfg_shape_net (dict) – Configuration for the shape network.

  • cfg_parameter_net (dict) – Configuration for the parameter network.

  • mixed_policy (str) – Policy to be used for mixed precision calculations.

call(inputs, training=None, mask=None)

Implements the forward pass of the model.

model_p_to_lr()

Returns a Keras model that maps input parameters to hidden learning rates.

model_x_to_phi()

Returns a Keras model that maps input spatial data to the output phi_x.

model_lr_to_w()

Raises a ValueError as ‘w’ is the same as ‘lr’ in this class.

model_x_to_u_given_w()

Returns a Keras model that maps input spatial data to the :wq:wq :wq output ‘u’ given ‘w’.

call(inputs, training=None, mask=None)

Forward pass of the NIFMultiScaleLastLayerParameterized model.

Parameters:

  • inputs (tf.Tensor) – Input tensor of shape (batch_size, input_dim).

  • training (bool, optional) – Whether the model is in training mode. Defaults to None.

  • mask (tf.Tensor, optional) – Mask tensor for masked inputs. Defaults to None.

Returns:

tf.Tensor – Output tensor of shape (batch_size, output_dim).

model_p_to_lr()

Creates a Keras model for mapping input parameters to hidden layer representation.

Returns:

tf.keras.Model – A Keras model that maps input parameters (t, mu) to a hidden layer representation (LR).

model_x_to_phi()

Creates a Keras model to compute phi_x from input spatial features.

Returns:

tf.keras.Model – A Keras model that takes input spatial features and computes the phi_x values using the shape network.

model_lr_to_w()

Raises an error as ‘w’ and ‘lr’ are the same in NIFMultiScaleLastLayerParameterized.

Raises:

ValueError – Error stating that ‘w’ is the same as ‘lr’ in this class.

model_x_to_u_given_w()

Creates a Keras model that maps input_s and input_pnet to the output of the shape network with the given parameters.

Returns:

Model –

A Keras model with input_s and input_pnet as inputs, and the output

of the shape network as the output.