Module curlew.fields.geoinr

TODO - implement neural field based on the GeoINR approach.

Classes

class GeoINR (name: str,
H: HSet,
C: CSet = None,
input_dim: int = None,
output_dim: int = 1,
transform=None,
seed=42,
vloss=MSELoss(),
scale=100.0,
**kwargs)

Expand source code

class GeoINR(BaseNF):
    """
    GeoINR-inspired neural field for interpolation of geological structures. 
    
    See Hillier et al., 2023 for further details: 
    
    `Hillier, Michael, et al. "GeoINR 1.0: an implicit neural network approach to three-dimensional geological modelling." Geoscientific Model Development 16.23 (2023): 6987-7012.`
    """

    def initField(self, 
                  hidden_layers: list = [],
                  activation: nn.Module = None,
                  rff_features: int = 8,
                  length_scales: list = [1e2, 2e2, 3e2],
                  stochastic_scales : bool = True,
                  learning_rate: float = 1e-1):
        """
            Initialise and build this neural field.
            
            hidden_layers : list of int, optional
                A list of integer sizes for the hidden layers of the MLP. Default is [,], which indicates the input encoding is directly translated to the output (i.e. no hidden layers).
            activation : nn.Module, optional
                The activation function to use for each hidden layer. Default is None, though `nn.SiLU()` can be useful for some fields.
            learning_rate : float
                The learning rate of the optimizer used to train this NF.
        """
        # -------------------- Random Fourier Features -------------------- #
        self.activation = activation
            
        # -------------------- MLP Construction -------------------- #
        # Determine input dimension for the MLP
        mlp_input_dim = self.input_dim

        # Define layer shapes
        self.dims = [mlp_input_dim] + hidden_layers + [self.output_dim]

        # Build layers in nn.Sequential
        layers = []
        for i in range(len(self.dims) - 2):
            layers.append(nn.Linear(self.dims[i], self.dims[i + 1],
                                    device=curlew.device, dtype=curlew.dtype))
            if self.activation is not None:
                layers.append(self.activation)

        # Final layer
        layers.append(nn.Linear(self.dims[-2], self.dims[-1],
                                device=curlew.device, dtype=curlew.dtype))
        self.mlp = nn.Sequential(*layers) # Combine layers into nn.Sequential

        # Xavier initialization
        for layer in self.mlp:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_normal_(layer.weight)
                
        # push onto device
        self.to(curlew.device)

        # Initialise optimiser used for this MLP.
        self.init_optim(lr=learning_rate)
        
    def evaluate(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the network to create a scalar value or property estimate.

        If random Fourier features are enabled, the input is first encoded accordingly.

        Parameters
        ----------
        x : torch.Tensor
            A tensor of shape (N, input_dim), where N is the batch size.

        Returns
        -------
        torch.Tensor
            A tensor of shape (N, output_dim), representing the scalar potential.
        """
        # Pass through all layers and return
        out = self.scale * self.mlp(x)
        return out
    
    def loss(self, transform=True) -> torch.Tensor:
        """
        Compute the loss associated with this neural field given its current state.
        """
        C = self.C # curlew-style constraints
        return super().loss(transform) # todo some funky loss 
    
    def fit(self, epochs, C=None, **kwargs):
        """
        Train this neural field using the specified constraints.
        """
        return super().fit(epochs, C=C, **kwargs)

GeoINR-inspired neural field for interpolation of geological structures.

See Hillier et al., 2023 for further details:

Hillier, Michael, et al. "GeoINR 1.0: an implicit neural network approach to three-dimensional geological modelling." Geoscientific Model Development 16.23 (2023): 6987-7012.

Parameters

name : str

A (ideally unique) name for this neural field. Should typically match the name of the GeoEvent instance that uses this field.

H : HSet

Hyperparameters used to tune the loss function for this NF.

C : CSet, optinoal

Constraint sent used when learning this implicit field. Default is None (can be set using field.bind(…)).

input_dim : int, optional

The dimensionality of the input space (e.g., 3 for (x, y, z)). If None (default), then default_dim will be used.

output_dim : int, optional

Dimensionality of the output (usually 1 for a scalar potential).

transform : callable

A function that transforms input coordinates prior to predictions. Must take exactly one argument as input (a tensor of positions) and return the transformed positions.

seed : callable, optional

The random seed to use for any random operations.

vloss : callable, optional

The loss function to use for value fitting. Default is mean squared error (nn.MSELoss()).

scale : float, optional

A scaling factor to apply to outputs of the neural field, as often these struggle to learn functions with a large (>1) amplitude. Default is 1e2.

This value should be approximately equal to the expected range (max - min) of the scalar field that is being learned. It can be especially important when using a drift (trend), as it determines the extent to which the model initialisation is determined by the drift. Larger values should allow the model to deviate farther from the trend. Also note that this term also tends to control the magnitude of residuals (to value or (in)equality constraints), so will also interact with the learning rate.

N.B. The actual implementation of this scale depends on the neural field method being used.

Keywords

All keywords are passed to the initField(…) function of the child class, to build the relevant neural architecture.

Ancestors

• BaseNF
• BaseSF
• LearnableBase
• torch.nn.modules.module.Module

Methods

def evaluate(self, x: torch.Tensor) ‑> torch.Tensor

Expand source code

def evaluate(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass of the network to create a scalar value or property estimate.

    If random Fourier features are enabled, the input is first encoded accordingly.

    Parameters
    ----------
    x : torch.Tensor
        A tensor of shape (N, input_dim), where N is the batch size.

    Returns
    -------
    torch.Tensor
        A tensor of shape (N, output_dim), representing the scalar potential.
    """
    # Pass through all layers and return
    out = self.scale * self.mlp(x)
    return out

Forward pass of the network to create a scalar value or property estimate.

If random Fourier features are enabled, the input is first encoded accordingly.

Parameters

x : torch.Tensor

A tensor of shape (N, input_dim), where N is the batch size.

Returns

torch.Tensor

A tensor of shape (N, output_dim), representing the scalar potential.

def fit(self, epochs, C=None, **kwargs)

Expand source code

def fit(self, epochs, C=None, **kwargs):
    """
    Train this neural field using the specified constraints.
    """
    return super().fit(epochs, C=C, **kwargs)

Train this neural field using the specified constraints.

def initField(self,
hidden_layers: list = [],
activation: torch.nn.modules.module.Module = None,
rff_features: int = 8,
length_scales: list = [100.0, 200.0, 300.0],
stochastic_scales: bool = True,
learning_rate: float = 0.1)

Expand source code

def initField(self, 
              hidden_layers: list = [],
              activation: nn.Module = None,
              rff_features: int = 8,
              length_scales: list = [1e2, 2e2, 3e2],
              stochastic_scales : bool = True,
              learning_rate: float = 1e-1):
    """
        Initialise and build this neural field.
        
        hidden_layers : list of int, optional
            A list of integer sizes for the hidden layers of the MLP. Default is [,], which indicates the input encoding is directly translated to the output (i.e. no hidden layers).
        activation : nn.Module, optional
            The activation function to use for each hidden layer. Default is None, though `nn.SiLU()` can be useful for some fields.
        learning_rate : float
            The learning rate of the optimizer used to train this NF.
    """
    # -------------------- Random Fourier Features -------------------- #
    self.activation = activation
        
    # -------------------- MLP Construction -------------------- #
    # Determine input dimension for the MLP
    mlp_input_dim = self.input_dim

    # Define layer shapes
    self.dims = [mlp_input_dim] + hidden_layers + [self.output_dim]

    # Build layers in nn.Sequential
    layers = []
    for i in range(len(self.dims) - 2):
        layers.append(nn.Linear(self.dims[i], self.dims[i + 1],
                                device=curlew.device, dtype=curlew.dtype))
        if self.activation is not None:
            layers.append(self.activation)

    # Final layer
    layers.append(nn.Linear(self.dims[-2], self.dims[-1],
                            device=curlew.device, dtype=curlew.dtype))
    self.mlp = nn.Sequential(*layers) # Combine layers into nn.Sequential

    # Xavier initialization
    for layer in self.mlp:
        if isinstance(layer, nn.Linear):
            nn.init.xavier_normal_(layer.weight)
            
    # push onto device
    self.to(curlew.device)

    # Initialise optimiser used for this MLP.
    self.init_optim(lr=learning_rate)

Initialise and build this neural field.

hidden_layers : list of int, optional A list of integer sizes for the hidden layers of the MLP. Default is [,], which indicates the input encoding is directly translated to the output (i.e. no hidden layers). activation : nn.Module, optional The activation function to use for each hidden layer. Default is None, though nn.SiLU() can be useful for some fields. learning_rate : float The learning rate of the optimizer used to train this NF.

def loss(self, transform=True) ‑> torch.Tensor

Expand source code

def loss(self, transform=True) -> torch.Tensor:
    """
    Compute the loss associated with this neural field given its current state.
    """
    C = self.C # curlew-style constraints
    return super().loss(transform) # todo some funky loss 

Compute the loss associated with this neural field given its current state.

Inherited members

• BaseNF:
  • bind
  • forward
  • gradient
  • init_optim
  • level
  • reset_mnorm
  • set_rate