diff --git a/gstlal-inspiral/python/stats/inspiral_intrinsics.py b/gstlal-inspiral/python/stats/inspiral_intrinsics.py index 721431d430ac7e9ab1b6bf8eb6246ae3e527718f..40851dfdac659af72f23adeb5e60e1e2f0eee972 100644 --- a/gstlal-inspiral/python/stats/inspiral_intrinsics.py +++ b/gstlal-inspiral/python/stats/inspiral_intrinsics.py @@ -112,32 +112,28 @@ class SourcePopulationModel(object): """ if filename is not None: with h5py.File(filename, 'r') as model: - coefficients = model['coefficients'][()] snr_bp = model['SNR'][()] try: - model_ids = model['template_id'][()] + template_indices = {x:n for n,x in enumerate(numpy.array(model['template_id']))} except KeyError: # FIXME: assume sequential order if model['event_id'] doesn't exist - #model_ids = numpy.arange(numpy.shape(model['coefficients'][()])[-1]) - model_ids = model['event_id'][()] - # PPoly can construct an array of polynomials by just - # feeding it the coefficients array all in one go, but then - # it insists on evaluating all of them at once. we don't - # want to suffer that cost, so we have to make an array of - # PPoly objects ourselves, and index into it to evaluate - # just one. since we have to do this anyway, we use a - # dictionary to also solve the problem of mapping - # template_id to a specific polynomial - template_indices = {} - for template_id in template_ids: - # maps template ID to the right coefficient array, since template IDs - # in the bank may not be in sequential order + template_indices = {x:n for n,x in enumerate(numpy.array(model['event_id']))} + # PPoly can construct an array of polynomials by just + # feeding it the coefficients array all in one go, but then + # it insists on evaluating all of them at once. we don't + # want to suffer that cost, so we have to make an array of + # PPoly objects ourselves, and index into it to evaluate + # just one. since we have to do this anyway, we use a + # dictionary to also solve the problem of mapping + # template_id to a specific polynomial + template_ids = sorted(template_ids) try: - template_indices[template_id] = numpy.where(model_ids==template_id)[0][0] - except IndexError: - raise IndexError("template ID %d is not in this model" % template_id) - self.polys = dict((template_id, PPoly(coefficients[:,:,[template_indices[template_id]]], snr_bp)) for template_id in template_ids) - self.max_snr = snr_bp.max() + these_template_id_indices = numpy.array([template_indices[t] for t in template_ids], dtype=numpy.int) + except KeyError: + raise KeyError("One or more template IDs are not in this model") + coefficients = model['coefficients'][:,:,these_template_id_indices] + self.polys = dict((template_id, PPoly(coefficients[:,:, n], snr_bp)) for n, template_id in enumerate(template_ids)) + self.max_snr = snr_bp.max() else: self.polys = None self.max_snr = None @@ -149,6 +145,5 @@ class SourcePopulationModel(object): lnP_vs_snr = self.polys[template_id] except KeyError: raise KeyError("template ID %d is not in this model" % template_id) - # PPoly's .__call__() returns an array, so we need the - # final [0] to flatten it - return lnP_vs_snr(min(snr, self.max_snr))[0] + # PPoly's .__call__() returns an array, so we need the .item() + return lnP_vs_snr(min(snr, self.max_snr)).item()