nessai.flowmodel.base ===================== .. py:module:: nessai.flowmodel.base .. autoapi-nested-parse:: Object and functions to handle training the normalising flow. .. !! processed by numpydoc !! Classes ------- .. autoapisummary:: nessai.flowmodel.base.FlowModel Module Contents --------------- .. py:class:: FlowModel(flow_config: Union[dict, None] = None, training_config: Union[dict, None] = None, output: Union[str, None] = None, rng: Optional[numpy.random.Generator] = None) Object that contains the normalising flows and handles training and data pre-processing. Does NOT use structured arrays for live points, :obj:`~nessai.proposal.base.Proposal` object should act as the interface between structured used by the sampler and unstructured arrays of live points used for training. :Parameters: **config** : dict, optional Configuration used for the normalising flow. If None, default values are used. **output** : str, optional Path for output, this includes weights files and the loss plot. If not specified, the current working directory is used. .. !! processed by numpydoc !! .. py:method:: setup_from_input_dict(flow_config, training_config) Setup the trainer from a dictionary, all keys in the dictionary are added as methods to the object. Input is automatically saved. :Parameters: **config** : dict Dictionary with parameters that are used to update the defaults. .. !! processed by numpydoc !! .. py:method:: update_mask() Method to update the ask upon calling ``initialise`` By default the mask is left unchanged. .. !! processed by numpydoc !! .. py:method:: get_optimiser(optimiser=None, **kwargs) Get the optimiser and ensure it is always correctly initialised. By default kwargs will include any optimiser specific values, the learning rate and any values included in :code:`optimiser_kwargs`. These will then be updated with any additional keyword arguments that are passed. :Returns: :obj:`torch.optim.Adam` Instance of the Adam optimiser from torch.optim .. !! processed by numpydoc !! .. py:method:: initialise() Initialise the model and optimiser. This includes: - Updating the model configuration - Initialising the normalising flow - Initialising the optimiser - Configuring the inference device .. !! processed by numpydoc !! .. py:method:: move_to(device, update_default=False) Move the flow to a different device. :Parameters: **device** : str Device to move flow to. **update_default** : bool, optional If True, the default device for the flow (and data) will be updated. .. !! processed by numpydoc !! .. py:method:: check_batch_size(x, batch_size, min_fraction=0.1) :staticmethod: Check that the batch size is valid. Tries to ensure that the last batch is at least a minimum fraction of the size of the batch size. :Parameters: **x** : numpy.ndarray Training data. **batch_size** : int The user-specified batch size **min_fraction** : float Fraction of user-specified batch size to use a lower limit. .. !! processed by numpydoc !! .. py:method:: prep_data(samples, val_size, batch_size, weights=None, use_dataloader=False, conditional=None) Prep data and return dataloaders for training :Parameters: **samples** : array_like Array of samples to split in to training and validation. **val_size** : float Float between 0 and 1 that defines the fraction of data used for validation. **batch_size** : int Batch size used when constructing dataloaders. **use_dataloader** : bool, optional If True data is returned in a dataloader else a tensor is returned. **weights** : array_like, optional Array of weights for each sample, weights will used when computing the loss. :Returns: train_data, val_data Training and validation data as either a tensor or dataloader .. !! processed by numpydoc !! .. py:method:: end_iteration() Calls any functions that should be applied at the end of the iteration. This functions is called after the flow has been updated on all batches of data but before the validation step. Calls :py:meth:`nessai.flows.base.BaseFlow.end_iteration` .. !! processed by numpydoc !! .. py:method:: finalise() Method to finalise the flow before using it for inference. .. !! processed by numpydoc !! .. py:method:: train(samples, weights=None, conditional=None, max_epochs=None, patience=None, output=None, val_size=None, plot=True) Train the flow on a set of samples. Allows for training parameters to specified instead of those given in initial config. :Parameters: **samples** : ndarray Unstructured numpy array containing data to train on **weights** : array_like Array of weights to use with the weight KL when computing the loss. **max_epochs** : int, optional Maximum number of epochs that is used instead of value in the configuration. **patience** : int, optional Patience in number of epochs that is used instead of value in the configuration. **val_size** : float, optional Fraction of the samples to use for validation **output** : str, optional Path to output directory that is used instead of the path specified when the object is initialised **plot** : bool, optional Boolean to enable or disable plotting the loss function :Returns: **history** : dict Dictionary that contains the training and validation losses. .. !! processed by numpydoc !! .. py:method:: save_weights(weights_file) Save the weights file. If the file already exists move it to ``.old`` and then save the file. :Parameters: **weights_file** : str Path to to file to save weights. Recommended file type is ``.pt``. .. !! processed by numpydoc !! .. py:method:: load_weights(weights_file) Load weights for the model and initialises the model if it is not initialised. The weights_file attribute is also updated. Model is loaded in evaluation mode (``model.eval()``) :Parameters: **weights_files** : str Path to weights file .. !! processed by numpydoc !! .. py:method:: reload_weights(weights_file) Tries to the load the weights file and if not, tries to load the weights file stored internally. :Parameters: **weights_files** : str Path to weights file .. !! processed by numpydoc !! .. py:method:: reset_model(weights=True, permutations=False) Reset the weights of the model and optimiser. :Parameters: **weights** : bool, optional If true the model weights are reset. **permutations** : bool, optional If true any permutations (linear transforms) are reset. .. !! processed by numpydoc !! .. py:method:: numpy_array_to_tensor(array: numpy.ndarray, /) -> torch.Tensor Convert a numpy array to a tensor and move it to the device .. !! processed by numpydoc !! .. py:method:: forward_and_log_prob(x: numpy.ndarray, conditional: Optional[numpy.ndarray] = None) -> Tuple[numpy.ndarray, numpy.ndarray] Forward pass through the model and return the samples in the latent space with their log probabilities :Parameters: **x** : ndarray Array of samples :Returns: **z** : ndarray Samples in the latent space **log_prob** : ndarray Log probabilities for each samples .. !! processed by numpydoc !! .. py:method:: inverse(z: numpy.ndarray, conditional: Optional[numpy.ndarray] = None) -> tuple[numpy.ndarray, numpy.ndarray] Apply the inverse transform. :Parameters: **z** Array of samples in the latent space. **conditional** Conditional inputs. :Returns: x Array of samples in the X-prime space. log_j Log-Jacobian determinant for each sample. .. !! processed by numpydoc !! .. py:method:: log_prob(x: numpy.ndarray, conditional: Optional[numpy.ndarray] = None) -> numpy.ndarray Compute the log-probability of a sample. :Parameters: **x** : ndarray Array of samples in the X-prime space. :Returns: ndarray Array of log-probabilities. .. !! processed by numpydoc !! .. py:method:: sample(n: int = 1, conditional: Optional[numpy.ndarray] = None) -> numpy.ndarray Sample from the flow. :Parameters: **n** : int Number of samples to draw **conditional** : numpy.ndarray Array of conditional inputs :Returns: numpy.ndarray Array of samples .. !! processed by numpydoc !! .. py:method:: sample_latent_distribution(n: int = 1) -> numpy.ndarray Sample from the latent distribution. :Parameters: **n** : int Number of samples to draw :Returns: numpy.ndarray Array of samples .. !! processed by numpydoc !! .. py:method:: sample_and_log_prob(N=1, z=None, alt_dist=None, conditional=None) Generate samples from samples drawn from the base distribution or and alternative distribution from provided latent samples :Parameters: **N** : int, optional Number of samples to draw if z is not specified **z** : ndarray, optional Array of latent samples to map the the data space, if ``alt_dist`` is not specified they are assumed to be drawn from the base distribution of the flow. **alt_dist** : :obj:`glasflow.nflows.distribution.Distribution` Distribution object from which the latent samples z were drawn from. Must have a ``log_prob`` method that accepts an instance of ``torch.Tensor`` :Returns: **samples** : ndarray Array containing samples in the latent space. **log_prob** : ndarray Array containing the log probability that corresponds to each sample. .. !! processed by numpydoc !!