amplfi.train.models module

class amplfi.train.models.flow.FlowModel(*args, arch, filter_params=True, samples_per_event=10000, min_samples_per_pix_dist=5, max_samples_per_pixel=20, target_prior=None, crossmatch=True, create_pp_plot=True, **kwargs)

Bases: AmplfiModel

A LightningModule for training normalizing flows

Parameters:

• *args – See arguments in amplfi.train.models.base.AmplfiModel

• arch (FlowArchitecture) – Neural network architecture to train. This should be a subclass of FlowArchitecture.

• filter_params (bool) – If True, filter the samples produced by the flow to keep only valid samples within the prior boundaries.

• samples_per_event (int) – Number of samples to draw per event for testing

• min_samples_per_pix_dist (int)

• max_samples_per_pixel (int)

• crossmatch (bool) – If True, run ligo.skymap.postprocess.crossmatch on result objects at the end of testing epoch and produce searched area and volume cdfs

• target_prior (Optional[Path]) – Path to a bilby prior file for reweighting posterior samples to a new prior.

forward(context, parameters)

Return type:

Tensor

training_step(batch, _)

validation_step(batch, _)

on_predict_epoch_start()

on_test_epoch_start()

on_test_batch_end(outputs, *_)

analyze_event(strain, asds, parameters=None, snr=None, gpstime=None)

Return type:

tuple[AmplfiResult, Optional[AmplfiResult]]

test_step(batch, _)

Return type:

tuple[AmplfiResult, Optional[AmplfiResult]]

predict_step(batch, _)

Return type:

tuple[AmplfiResult, Optional[AmplfiResult]]

cast_as_bilby_result(samples, log_probs, log_prior_probs, injection_parameters=None)

Cast posterior samples as Bilby Result object for ease of producing corner and pp plots

Parameters:

• samples (ndarray) – An array of posterior samples of shape (1, num_samples, num_params)

• log_probs (ndarray) – An array of log probabilities of posterior samples as predicted under the normalizing flow model

• log_prior_probs (ndarray) – An array of log prior probabilities of posterior samples as predicted under the training prior

• injection_parameters (Optional[dict[str, float]]) – For injections, a dictionary mapping from parameter string to the true injection value

Return type:

AmplfiResult

filter_parameters(parameters)

Filter the descaled parameters to keep only valid samples within their boundaries.

Parameters:

descaled (torch.Tensor) – The descaled parameters tensor.

Returns:

The filtered descaled parameters.

Return type:

torch.Tensor

configure_callbacks()

class amplfi.train.models.similarity.SimilarityModel(*args, arch, similarity_loss, **kwargs)

Bases: AmplfiModel

A LightningModule for training similarity embeddings

Parameters:

arch (SimilarityEmbedding) – A neural network architecture that maps waveforms to lower dimensional embedded space

forward(ref, aug)

validation_step(batch, _)

training_step(batch, _)

configure_callbacks()